diff --git a/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py b/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py index 4d2c19114e..43dd675cf4 100644 --- a/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py +++ b/benchmarks/harbor-buzz-orchestra/src/harbor_buzz_orchestra/container_runtime.py @@ -381,12 +381,26 @@ async def run( ), timeout=manifest.trial_budget.timeout_seconds, ) - # `DONE:` is published before the turn's usage notification is, so - # teardown has to wait or the trial's tokens are lost. Inside the - # try, and only on this path: a timeout has no usage to flush. - await self._settle_usage(environment, agents) await self._verify_m1_output(environment, manifest) finally: + # `DONE:` is published before the turn's usage notification is, so + # teardown has to wait or the trial's tokens are lost. + # + # In the `finally`, so the TIMEOUT path settles too. That path used + # to fall straight through to the kill on the reasoning that "a turn + # that never completed has no usage to flush", which was true while + # buzz-agent reported once per turn and is not any more: it now + # reports after every provider round, so an unfinished turn has + # reported everything but its in-flight request. Skipping the settle + # here is what made `continue_until_timeout` runs uncostable — + # every phase but the last ends on this path, and 97% of one + # measured run's receipt rows came back all zeros. + # + # Cheap in the common case: this returns on the first poll once a + # usage line exists, which after the first round it does. Only a + # phase that never completed a single round can spend the full + # budget, and that phase has nothing to report anyway. + await self._settle_usage(environment, agents) await self._stop_agents(environment, agents + infra) # Logs first, and ahead of anything that touches the network. The # verifier pre-install below used to run first and, when the proxy @@ -1021,10 +1035,10 @@ async def _settle_usage( ) -> None: """Give each agent the moment it needs to report what it spent. - buzz-agent emits its `_goose/unstable/session/update` usage notification - once per turn, immediately *before* returning the `session/prompt` - response (buzz-agent/src/lib.rs:708). A solo agent gets exactly one turn - per trial, so that single notification is the only record of the trial's + buzz-agent emits a `_goose/unstable/session/update` usage notification + after every provider round, and once more immediately *before* returning + the `session/prompt` response. A solo agent gets exactly one turn per + trial, so the final notification is the complete record of the trial's tokens — and it is written after the agent has already published `DONE:` as a tool call. @@ -1039,10 +1053,12 @@ async def _settle_usage( `DONE:` reaches the channel from inside a tool call, so the turn may still owe one model round-trip before it ends and reports, and for a thinking model that is tens of seconds, not milliseconds. - Waiting is pointless on the timeout path — a turn that never completed - has no usage to flush — so callers only invoke this once `DONE:` is seen. - A miss is not fatal; the accounting note already reports an unpriced - trial, and losing the tokens is better than hanging the sweep. + Called on the timeout path too, not just after `DONE:`. An interrupted + turn has still reported every round it finished, and under + `continue_until_timeout` the interrupted turn is the normal case, not + the exception. A miss is not fatal; the accounting note already reports + an unpriced trial, and losing the tokens is better than hanging the + sweep. """ deadline = asyncio.get_running_loop().time() + self.usage_settle_seconds pending = {agent.credential.agent_id: agent for agent in agents} diff --git a/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py b/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py index b7a42c542b..1f1b0717d3 100644 --- a/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py +++ b/benchmarks/harbor-buzz-orchestra/tests/test_container_runtime.py @@ -1,5 +1,6 @@ """The container runtime must launch the production stack, unmodified.""" +import asyncio import hashlib import json import re @@ -783,6 +784,62 @@ async def done(*args, **kwargs): assert "kill" in order +async def test_usage_is_settled_when_the_trial_times_out(tmp_path, monkeypatch): + """The timeout path has to settle too, and it is the path that matters most. + + It used to fall straight through to the kill, on the reasoning that a turn + which never completed had nothing to flush. buzz-agent now reports after + every provider round, so an interrupted turn HAS reported — and under + `continue_until_timeout` every phase but the last ends here. Skipping the + settle on this path is what left 97% of one measured run's receipt rows at + all zeros while the provider billed it in full. + """ + manifest = write_manifest(tmp_path) + trial = trial_handle((credential("orch-1", "orchestrator", "orch-model", "lead"),)) + rt = runtime(tmp_path, poll_seconds=0, usage_settle_seconds=5) + order = [] + + class TimingOutEnvironment(Environment): + async def exec(self, command, env=None, **kwargs): + if "buzz-acp" in command: + return ExecResult(stdout="99\n", stderr="", return_code=0) + if command.startswith("cat "): + order.append("usage") + return ExecResult( + stdout="goose usage update input=1 output=2\n", + stderr="", + return_code=0, + ) + if "/proc/[0-9]*" in command: + order.append("kill") + return ExecResult(stdout="", stderr="", return_code=0) + + async def never_done(*args, **kwargs): + await asyncio.sleep(3600) + + monkeypatch.setattr(rt, "_install_stack", lambda env: _noop()) + monkeypatch.setattr(rt, "_wait_for_agents_ready", lambda *a, **k: _noop()) + monkeypatch.setattr(rt, "_wait_for_done", never_done) + monkeypatch.setattr(rt, "_buzz_json", lambda *a, **k: _value([])) + # TrialBudget is frozen, so swap in a zero-budget copy rather than mutating. + budget = manifest.trial_budget.model_copy(update={"timeout_seconds": 0}) + manifest = manifest.model_copy(update={"trial_budget": budget}) + + with pytest.raises(asyncio.TimeoutError): + await rt.run( + instruction="do the thing", + environment=TimingOutEnvironment(), + manifest=manifest, + trial=trial, + ) + + assert "usage" in order, "the timeout path never settled usage" + assert order.index("usage") < order.index("kill"), ( + "usage was settled after teardown killed the agent, which is the same " + "as not settling at all" + ) + + async def test_wait_for_agents_ready_requires_every_channel_subscription(tmp_path): rt = runtime(tmp_path, poll_seconds=0) logs = {"orch-1": "", "worker-1": ""}