From e3986389febd1b79b02bc6fe97d5cd8006e71421 Mon Sep 17 00:00:00 2001 From: CK0607 Date: Wed, 22 Jul 2026 12:46:31 +0530 Subject: [PATCH 1/2] fix(eval): hard rollout wall-clock timeout for HUDRuntime training Bound grade/cancel with a short teardown grace so a wedged tunnel cannot stall Taskset.run. Wire HUDRuntime.run_timeout into rollout, return a terminal failed Run on HostedRuntime timeout, and document GRPO keep- failed-members policy (HUD-2099). Co-authored-by: Cursor --- docs/v6/guides/training-agents.mdx | 5 +- docs/v6/reference/runtime.mdx | 7 ++- docs/v6/reference/training.mdx | 10 ++++ hud/eval/run.py | 87 ++++++++++++++++++++++-------- hud/eval/runtime.py | 16 ++++-- hud/eval/taskset.py | 11 ++-- hud/eval/tests/test_hosted.py | 37 +++++++++++-- hud/eval/tests/test_rollout.py | 55 +++++++++++++++++++ hud/train/tests/__init__.py | 0 hud/train/tests/test_groups.py | 36 +++++++++++++ 10 files changed, 231 insertions(+), 33 deletions(-) create mode 100644 hud/train/tests/__init__.py create mode 100644 hud/train/tests/test_groups.py diff --git a/docs/v6/guides/training-agents.mdx b/docs/v6/guides/training-agents.mdx index 11f1abf13..3c5a494ad 100644 --- a/docs/v6/guides/training-agents.mdx +++ b/docs/v6/guides/training-agents.mdx @@ -93,7 +93,10 @@ One job spans the session; each step appends a batch and trains on it: - **Open the job** with `group=8` - 8 rollouts per task, so the rewards are comparable (next). - **Roll out** the batch, the same eval as the [previous guide](/v6/guides/running-an-eval). The - [runtime](/v6/reference/runtime) sets where it runs; swap `LocalRuntime` for `HUDRuntime()` unchanged. + [runtime](/v6/reference/runtime) sets where it runs; swap `LocalRuntime` for + `HUDRuntime(run_timeout=600)` unchanged. That `run_timeout` is a hard per-rollout wall clock + (agent + grade/cancel) so a disconnected tunnel cannot stall the batch before `trainer.step`. + Keep timed-out runs in the batch — GRPO needs full groups. - **Nudge** with `trainer.step` - the one line that learns. It scores each rollout against its group, shifts the weights, then **promotes** them so the gateway serves the new ones at once. diff --git a/docs/v6/reference/runtime.mdx b/docs/v6/reference/runtime.mdx index 102760d83..beffece90 100644 --- a/docs/v6/reference/runtime.mdx +++ b/docs/v6/reference/runtime.mdx @@ -154,7 +154,10 @@ DaytonaRuntime(snapshot_name=None, *, image=None, command=None, workdir="/app", HUDRuntime(*, run_timeout=3600.0, runtime_url=None) ``` -- **`run_timeout`** - bound on one rollout end to end, including instance startup. +- **`run_timeout`** - hard wall-clock bound on one rollout end to end (session + ready, agent loop, and grade/cancel with a short teardown grace). Forwarded + as `rollout_timeout` into the local rollout driver. Session hello still caps + at 300s via `ready_timeout = min(run_timeout, 300)`. - **`runtime_url`** - override the runtime endpoint the tunnel connects to. The SDK leases your deployed env by name and tunnels to its control channel; the agent loop runs local. @@ -172,6 +175,8 @@ Where `HUDRuntime` runs the agent loop locally against a tunneled env, `HostedRu **whole rollout** off-box: the platform leases an instance, brings the env's container up on it, and runs the agent right next to it. This process only submits the rollout and polls its trace to completion. It requires a gateway agent that can serialize its identity (Claude/OpenAI/Gemini). +A timeout cancels the platform rollout and returns a terminal failed `Run` (it does not raise), +so one wedged hosted rollout cannot collapse the rest of a `Taskset.run` batch. ### `Runtime` diff --git a/docs/v6/reference/training.mdx b/docs/v6/reference/training.mdx index d9a11877f..a12b19ae1 100644 --- a/docs/v6/reference/training.mdx +++ b/docs/v6/reference/training.mdx @@ -76,6 +76,16 @@ policy-gradient objective applied on top, defaulting to `importance_sampling` (n batch must divide evenly into groups - `forward_backward` rejects a partial final group before spending a forward pass. `num_substeps` splits the batch for gradient accumulation. +### Timed-out and failed rollouts + +Keep timed-out / errored runs in the batch. `Taskset.run` always returns a terminal `Run` per slot +(including HUDRuntime / HostedRuntime timeouts), so group cardinality stays intact for +`group_size`. A failed member typically contributes reward `0` (or a salvage grade if the agent +answered before the deadline). Do **not** drop failed members before `trainer.step` — that leaves an +incomplete GRPO group and raises. Prefer bounding each rollout with +`HUDRuntime(run_timeout=…)` or `Taskset.run(..., rollout_timeout=…)` so one wedged tunnel cannot +stall the whole batch. + ## Inputs A training input is a recorded trajectory by id, or an inline one: diff --git a/hud/eval/run.py b/hud/eval/run.py index f947efd6c..2f2f3593a 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -50,6 +50,11 @@ logger = logging.getLogger("hud.eval.run") +#: Extra seconds allowed for grade/cancel after the hard ``rollout_timeout`` +#: deadline. Keeps salvage grading usable without letting a wedged tunnel hang +#: the batch indefinitely. +_TEARDOWN_GRACE_S = 30.0 + def _prompt_message(item: Any) -> mcp_types.PromptMessage: """Coerce one wire prompt turn onto MCP's ``PromptMessage`` vocabulary. @@ -114,10 +119,21 @@ class Run: half-working. """ - def __init__(self, client: HudClient | None, task_id: str, args: dict[str, Any]) -> None: + def __init__( + self, + client: HudClient | None, + task_id: str, + args: dict[str, Any], + *, + deadline: float | None = None, + ) -> None: self._client = client self._task_id = task_id self._args = args + #: Shared wall-clock end (``loop.time()``) for this rollout, if any. + #: Provision/agent honor it strictly; grade/cancel get a short grace + #: after it so a wedged JSON-RPC read cannot stall the batch forever. + self._deadline = deadline #: The task's opening prompt as ``tasks.start`` returned it: plain #: text, or a list of message dicts (``{"role", "content"}``) for #: chat-style / multi-turn prompts. Agents consume the normalized @@ -218,6 +234,18 @@ async def __aenter__(self) -> Self: self.record(Step(source="user", messages=self.prompt_messages)) return self + async def _await_teardown(self, awaitable: Any) -> Any: + """Await teardown RPC (grade/cancel) under the rollout deadline + grace. + + Provision and the agent loop use the hard ``deadline``. Grade and cancel + may run briefly past it so a partial trajectory can still be scored, but + never unbounded — a wedged tunnel during grade used to stall the batch. + """ + if self._deadline is None: + return await awaitable + remaining = self._deadline + _TEARDOWN_GRACE_S - asyncio.get_running_loop().time() + return await asyncio.wait_for(awaitable, max(remaining, 0.0)) + async def __aexit__( self, exc_type: type[BaseException] | None, @@ -229,7 +257,8 @@ async def __aexit__( exc_type, asyncio.CancelledError | KeyboardInterrupt ): self.trace.status = "cancelled" - await self.client.cancel() + with contextlib.suppress(Exception, TimeoutError): + await self._await_teardown(self.client.cancel()) return False answer: dict[str, Any] = {"answer": self.trace.content} @@ -237,17 +266,27 @@ async def __aexit__( # A mid-run error grades best-effort (capture a salvageable reward, keep # status=error), but a grade failure must not mask the original error. A - # clean run grades normally — a grader fault propagates. grade() also - # blocks on an unbounded JSON-RPC read (not bound by rollout_timeout). - if exc_type is not None: + # clean run grades normally — a grader fault propagates. Grade/cancel + # are bound by the rollout deadline plus a short teardown grace. + try: + if exc_type is not None: + self.trace.status = "error" + try: + evaluation = await self._await_teardown(self.client.grade(answer)) + except Exception as grade_exc: + logger.warning("grade failed after mid-run error: %s", grade_exc) + return False + else: + evaluation = await self._await_teardown(self.client.grade(answer)) + except TimeoutError: + logger.warning("rollout grading timed out; leaving run terminal without a grade") self.trace.status = "error" - try: - evaluation = await self.client.grade(answer) - except Exception as grade_exc: - logger.warning("grade failed after mid-run error: %s", grade_exc) - return False - else: - evaluation = await self.client.grade(answer) + if self.trace.stop_reason is None: + self.trace.stop_reason = "timeout" + self.record(Step(source="system", error="grading timed out")) + with contextlib.suppress(Exception, TimeoutError): + await self._await_teardown(self.client.cancel()) + return False self.grade = Grade.from_dict(evaluation) self.record( @@ -311,9 +350,12 @@ async def rollout( start) yields a synthesized :meth:`Run.failed`; a failure *mid-run* keeps the real run — prompt, placement record, and the partial trace the agent built — marked as errored, and still graded best-effort so a salvageable - reward is captured. Either way the logged warning names the lifecycle - phase (``provisioning``, ``starting task``, ``agent loop``, ``grading``) so - callers can tell where the failure landed without reading the trace. + reward is captured. Grade and cancel share the same wall-clock deadline + plus a short teardown grace, so a disconnected tunnel during grading + cannot leave the batch waiting indefinitely. Either way the logged warning + names the lifecycle phase (``provisioning``, ``starting task``, + ``agent loop``, ``grading``) so callers can tell where the failure landed + without reading the trace. """ if job_id is None: # no standalone traces: a lone rollout is a job of one job_id = uuid.uuid4().hex @@ -339,11 +381,14 @@ async def _bounded(awaitable: Any) -> Any: One shared deadline across provision, connect, and the agent loop — not a fresh budget per phase — so the bounded work cannot exceed - ``rollout_timeout`` in total. A client read-timeout is not enough on - its own: a wedged upstream that trickles bytes resets the read timer - forever, so a single stuck sampling call could otherwise hang the - rollout — and the batch waits on it — indefinitely. A breach surfaces - as ``TimeoutError`` (distinct from a Ctrl-C ``CancelledError``). + ``rollout_timeout`` in total. Grade and cancel get a short grace + past that deadline (see :data:`_TEARDOWN_GRACE_S`) so salvage + scoring still works, but a wedged JSON-RPC read cannot hang the + batch. A client read-timeout is not enough on its own: a wedged + upstream that trickles bytes resets the read timer forever, so a + single stuck sampling call could otherwise hang the rollout — and + the batch waits on it — indefinitely. A breach surfaces as + ``TimeoutError`` (distinct from a Ctrl-C ``CancelledError``). """ if deadline is None: return await awaitable @@ -359,7 +404,7 @@ async def _drive() -> None: addr = cast("Runtime", await _bounded(stack.enter_async_context(runtime(task)))) _phase = "starting task" client = cast("HudClient", await _bounded(stack.enter_async_context(connect(addr)))) - live = Run(client, task.id, task.args) + live = Run(client, task.id, task.args, deadline=deadline) live._runtime = addr.url # the placement record for the receipt async with live: # start on enter; grade on exit run = live # bound only once live: an earlier failure synthesizes diff --git a/hud/eval/runtime.py b/hud/eval/runtime.py index a05e16ab1..9806dbd81 100644 --- a/hud/eval/runtime.py +++ b/hud/eval/runtime.py @@ -896,6 +896,9 @@ class HUDRuntime: """ def __init__(self, *, run_timeout: float = 3600.0, runtime_url: str | None = None) -> None: + #: Hard wall-clock bound for one rollout (provision + agent + grade/ + #: cancel with a short teardown grace). Forwarded as + #: ``rollout_timeout`` into :func:`~hud.eval.run.rollout`. self.run_timeout = run_timeout self.runtime_url = runtime_url @@ -915,6 +918,7 @@ async def run( trace_id=trace_id, job_id=job_id, group_id=group_id, + rollout_timeout=self.run_timeout, ) def __call__(self, task: Task) -> AbstractAsyncContextManager[Runtime]: @@ -1054,17 +1058,21 @@ async def run( The platform owns the trace lifecycle (the instance-side driver reports enter/exit and streams telemetry), so this never double-reports. Failures isolating one rollout from its batch (submit rejected, the - env/model unresolved) surface as :meth:`Run.failed`; a timeout or a - local cancel propagate, having first asked the platform to release the - lease. + env/model unresolved, or ``run_timeout``) surface as :meth:`Run.failed` + so ``Taskset.run``'s gather cannot drop the rest of the batch. A local + cancel (Ctrl-C) still propagates after requesting a platform cancel. """ trace_id = trace_id or uuid.uuid4().hex try: state = await self._submit_and_await( task, agent, job_id=job_id, group_id=group_id, trace_id=trace_id ) - except (TimeoutError, asyncio.CancelledError): + except asyncio.CancelledError: raise + except TimeoutError as exc: + logger.warning("hosted rollout timed out: %s", exc) + run = Run.failed(str(exc)) + run.trace.stop_reason = "timeout" except Exception as exc: logger.warning("hosted rollout failed to launch: %s", exc) run = Run.failed(str(exc)) diff --git a/hud/eval/taskset.py b/hud/eval/taskset.py index 2ef0d223a..dae302548 100644 --- a/hud/eval/taskset.py +++ b/hud/eval/taskset.py @@ -271,8 +271,10 @@ async def run( ``rollout_timeout`` is a hard per-rollout wall-clock cap (seconds) for the local (Provider) path: a rollout that exceeds it is cancelled and recorded as a failed/errored run so one wedged rollout (e.g. a stuck sampling - stream) cannot stall the whole batch. ``HUDRuntime`` carries its own - ``run_timeout`` instead. + stream or a disconnected tunnel during grading) cannot stall the whole + batch. When unset and ``runtime`` is a :class:`~hud.eval.runtime.HUDRuntime`, + its ``run_timeout`` is used. :class:`~hud.eval.runtime.HostedRuntime` + always uses its own ``run_timeout``. """ group = group or (job.group if job else 1) if group < 1: @@ -308,6 +310,9 @@ async def run( # An empty taskset schedules nothing, so it needs no placement. placement = runtime if runtime is not None or not task_list else self._resolve_placement() sem = asyncio.Semaphore(max_concurrent) if max_concurrent else None + effective_timeout = rollout_timeout + if effective_timeout is None and isinstance(placement, HUDRuntime): + effective_timeout = placement.run_timeout async def _run(task: Task, group_id: str) -> Run: assert placement is not None # only reached when tasks were expanded @@ -319,7 +324,7 @@ async def _run(task: Task, group_id: str) -> Run: runtime=placement, job_id=job_id, group_id=group_id, - rollout_timeout=rollout_timeout, + rollout_timeout=effective_timeout, ) async def _one(task: Task, group_id: str) -> Run: diff --git a/hud/eval/tests/test_hosted.py b/hud/eval/tests/test_hosted.py index 374aa5481..e83d595fc 100644 --- a/hud/eval/tests/test_hosted.py +++ b/hud/eval/tests/test_hosted.py @@ -193,11 +193,15 @@ async def test_run_timeout_requests_platform_cancel(monkeypatch: pytest.MonkeyPa hosted = HostedRuntime(poll_interval=0.0, run_timeout=0.0) task = Task(env="sums", id="add", args={}) - with pytest.raises(TimeoutError, match="hosted rollout"): - await hosted.run(task, _agent(), job_id=uuid.uuid4().hex) + run = await hosted.run(task, _agent(), job_id=uuid.uuid4().hex) cancel_posts = [(p, b) for p, b in platform.posts if p == "/rollouts/cancel"] assert len(cancel_posts) == 1 + # Timeout must return a terminal failed Run so Taskset.gather cannot drop + # the rest of the batch (HUD-2099). + assert run.trace.is_error + assert run.trace.stop_reason == "timeout" + assert "did not finish" in (run.trace.error or "") @pytest.mark.asyncio @@ -293,7 +297,7 @@ async def fake_rollout(task: Task, agent: Any, **kwargs: Any) -> Run: monkeypatch.setattr("hud.eval.runtime.rollout", fake_rollout) - runtime = HUDRuntime() + runtime = HUDRuntime(run_timeout=90.0) job_id = uuid.uuid4().hex trace_id = uuid.uuid4().hex run = await runtime.run( @@ -309,6 +313,33 @@ async def fake_rollout(task: Task, agent: Any, **kwargs: Any) -> Run: assert seen["job_id"] == job_id assert seen["group_id"] == "g1" assert seen["trace_id"] == trace_id + assert seen["rollout_timeout"] == 90.0 + + +@pytest.mark.asyncio +async def test_taskset_uses_hud_runtime_run_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + import hud.eval.taskset as taskset_mod + from hud.eval.taskset import Taskset + + seen: dict[str, Any] = {} + + async def fake_rollout(task: Task, agent: Any, **kwargs: Any) -> Run: + seen.update(kwargs) + run = Run(None, task.id, {}) + run.trace.status = "completed" + return run + + monkeypatch.setattr(taskset_mod, "rollout", fake_rollout) + monkeypatch.setattr(taskset_mod, "job_enter", lambda *a, **k: asyncio.sleep(0)) + monkeypatch.setattr(taskset_mod, "flush", lambda timeout=120.0: True) + + job = await Taskset("t", [Task(env="e", id="x")]).run( + _agent(), + runtime=HUDRuntime(run_timeout=42.0), + ) + + assert len(job.runs) == 1 + assert seen["rollout_timeout"] == 42.0 @pytest.mark.asyncio diff --git a/hud/eval/tests/test_rollout.py b/hud/eval/tests/test_rollout.py index 4995a9505..a35e85f06 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -338,6 +338,61 @@ async def add(a: int, b: int): assert run.trace_id is not None +async def test_wedged_grade_cannot_stall_the_rollout() -> None: + """A disconnected tunnel during grade must hit the teardown grace, not hang.""" + from hud.eval import run as run_mod + + env = Environment("sums") + + @env.template() + async def add(a: int, b: int): + answer = yield f"add:{a}:{b}" + yield 1.0 if answer == str(a + b) else 0.0 + + class _HangOnGrade: + """Wraps a live client but never returns from grade().""" + + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + async def grade(self, *_a: Any, **_k: Any) -> dict[str, Any]: + await asyncio.sleep(3600) + return {} + + async def cancel(self) -> None: + return None + + real_connect = run_mod.connect + + @asynccontextmanager + async def _connect_hanging(addr: Any) -> AsyncIterator[Any]: + async with real_connect(addr) as client: + yield _HangOnGrade(client) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(run_mod, "connect", _connect_hanging) + monkeypatch.setattr(run_mod, "_TEARDOWN_GRACE_S", 0.2) + try: + started = asyncio.get_running_loop().time() + run = await rollout( + _add_task(2, 3), + _FnAgent(_solve_add), + runtime=lambda _row: _local(env), + rollout_timeout=0.3, + ) + elapsed = asyncio.get_running_loop().time() - started + finally: + monkeypatch.undo() + + assert elapsed < 5.0 # must not wait on the hung grade + assert run.trace.is_error + assert run.trace.stop_reason == "timeout" + assert "grading timed out" in (run.trace.error or "") + + async def test_pre_launch_failure_yields_a_synthesized_failed_run() -> None: @asynccontextmanager async def broken_provider(task: TaskRow) -> AsyncIterator[Runtime]: diff --git a/hud/train/tests/__init__.py b/hud/train/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/hud/train/tests/test_groups.py b/hud/train/tests/test_groups.py new file mode 100644 index 000000000..ee5864598 --- /dev/null +++ b/hud/train/tests/test_groups.py @@ -0,0 +1,36 @@ +"""GRPO group cardinality with timed-out / failed rollouts (HUD-2099).""" + +from __future__ import annotations + +import uuid + +import pytest + +from hud.eval.run import Run +from hud.train.client import _check_groups, _to_inputs + + +def _terminal_run(*, error: str | None = None, stop_reason: str | None = None) -> Run: + run = Run.failed(error) if error else Run(None, "t", {}) + if error is None: + run.trace.status = "completed" + run.trace.trace_id = uuid.uuid4().hex + if stop_reason is not None: + run.trace.stop_reason = stop_reason # type: ignore[assignment] + return run + + +def test_check_groups_rejects_incomplete_cardinality() -> None: + with pytest.raises(ValueError, match="do not divide evenly"): + _check_groups(7, 8) + + +def test_timed_out_runs_keep_grpo_group_full() -> None: + """Keep timed-out members in the batch — dropping them breaks group_size.""" + batch = [ + _terminal_run(error="agent loop timed out", stop_reason="timeout"), + *(_terminal_run() for _ in range(7)), + ] + inputs = _to_inputs(batch) + assert len(inputs) == 8 + _check_groups(len(inputs), 8) # must not raise From 04a8f3d0db5d59ce114e9ca81d0b993434b70a65 Mon Sep 17 00:00:00 2001 From: CK0607 Date: Fri, 24 Jul 2026 15:16:02 +0530 Subject: [PATCH 2/2] fix(eval): keep timeout policy in rollout lifecycle Keep Run focused on task state and result recording. Let rollout own the hard deadline plus independent grading, cancellation, and provider cleanup budgets so cleanup remains usable after a grading timeout. Co-authored-by: Cursor --- hud/eval/run.py | 214 ++++++++++++++++++--------------- hud/eval/tests/test_rollout.py | 48 +++++++- 2 files changed, 165 insertions(+), 97 deletions(-) diff --git a/hud/eval/run.py b/hud/eval/run.py index 2f2f3593a..168c2df67 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -50,10 +50,12 @@ logger = logging.getLogger("hud.eval.run") -#: Extra seconds allowed for grade/cancel after the hard ``rollout_timeout`` -#: deadline. Keeps salvage grading usable without letting a wedged tunnel hang -#: the batch indefinitely. -_TEARDOWN_GRACE_S = 30.0 +#: Lifecycle budgets owned by :func:`rollout`. Grading gets a short salvage +#: window after the hard rollout deadline; cancel and provider teardown each +#: get an independent cleanup window. +_GRADE_GRACE_S = 30.0 +_CANCEL_TIMEOUT_S = 5.0 +_CLEANUP_TIMEOUT_S = 10.0 def _prompt_message(item: Any) -> mcp_types.PromptMessage: @@ -119,21 +121,10 @@ class Run: half-working. """ - def __init__( - self, - client: HudClient | None, - task_id: str, - args: dict[str, Any], - *, - deadline: float | None = None, - ) -> None: + def __init__(self, client: HudClient | None, task_id: str, args: dict[str, Any]) -> None: self._client = client self._task_id = task_id self._args = args - #: Shared wall-clock end (``loop.time()``) for this rollout, if any. - #: Provision/agent honor it strictly; grade/cancel get a short grace - #: after it so a wedged JSON-RPC read cannot stall the batch forever. - self._deadline = deadline #: The task's opening prompt as ``tasks.start`` returned it: plain #: text, or a list of message dicts (``{"role", "content"}``) for #: chat-style / multi-turn prompts. Agents consume the normalized @@ -234,18 +225,6 @@ async def __aenter__(self) -> Self: self.record(Step(source="user", messages=self.prompt_messages)) return self - async def _await_teardown(self, awaitable: Any) -> Any: - """Await teardown RPC (grade/cancel) under the rollout deadline + grace. - - Provision and the agent loop use the hard ``deadline``. Grade and cancel - may run briefly past it so a partial trajectory can still be scored, but - never unbounded — a wedged tunnel during grade used to stall the batch. - """ - if self._deadline is None: - return await awaitable - remaining = self._deadline + _TEARDOWN_GRACE_S - asyncio.get_running_loop().time() - return await asyncio.wait_for(awaitable, max(remaining, 0.0)) - async def __aexit__( self, exc_type: type[BaseException] | None, @@ -256,37 +235,34 @@ async def __aexit__( if exc_type is not None and issubclass( exc_type, asyncio.CancelledError | KeyboardInterrupt ): - self.trace.status = "cancelled" - with contextlib.suppress(Exception, TimeoutError): - await self._await_teardown(self.client.cancel()) + await self.cancel() return False + await self.finish(exc) + return False + + async def cancel(self) -> None: + """Mark the run cancelled and ask the live task to stop.""" + self.trace.status = "cancelled" + await self.client.cancel() + + async def finish(self, exc: BaseException | None = None) -> None: + """Grade and record the final task result, preserving a prior error.""" answer: dict[str, Any] = {"answer": self.trace.content} started_at = now_iso() # A mid-run error grades best-effort (capture a salvageable reward, keep # status=error), but a grade failure must not mask the original error. A - # clean run grades normally — a grader fault propagates. Grade/cancel - # are bound by the rollout deadline plus a short teardown grace. - try: - if exc_type is not None: - self.trace.status = "error" - try: - evaluation = await self._await_teardown(self.client.grade(answer)) - except Exception as grade_exc: - logger.warning("grade failed after mid-run error: %s", grade_exc) - return False - else: - evaluation = await self._await_teardown(self.client.grade(answer)) - except TimeoutError: - logger.warning("rollout grading timed out; leaving run terminal without a grade") + # clean run grades normally — a grader fault propagates. + if exc is not None: self.trace.status = "error" - if self.trace.stop_reason is None: - self.trace.stop_reason = "timeout" - self.record(Step(source="system", error="grading timed out")) - with contextlib.suppress(Exception, TimeoutError): - await self._await_teardown(self.client.cancel()) - return False + try: + evaluation = await self.client.grade(answer) + except Exception as grade_exc: + logger.warning("grade failed after mid-run error: %s", grade_exc) + return + else: + evaluation = await self.client.grade(answer) self.grade = Grade.from_dict(evaluation) self.record( @@ -304,7 +280,6 @@ async def __aexit__( ) if self.trace.status is None: self.trace.status = "completed" - return False @classmethod def failed(cls, error: str) -> Run: @@ -350,12 +325,13 @@ async def rollout( start) yields a synthesized :meth:`Run.failed`; a failure *mid-run* keeps the real run — prompt, placement record, and the partial trace the agent built — marked as errored, and still graded best-effort so a salvageable - reward is captured. Grade and cancel share the same wall-clock deadline - plus a short teardown grace, so a disconnected tunnel during grading - cannot leave the batch waiting indefinitely. Either way the logged warning - names the lifecycle phase (``provisioning``, ``starting task``, - ``agent loop``, ``grading``) so callers can tell where the failure landed - without reading the trace. + reward is captured. ``rollout()`` owns every timeout: provision, connect, + start, and the agent share the hard deadline; grading gets one salvage + grace; cancellation and provider cleanup get independent bounded budgets. + A disconnected tunnel therefore cannot leave the batch waiting + indefinitely. Either way the logged warning names the lifecycle phase + (``provisioning``, ``starting task``, ``agent loop``, ``grading``, + ``cleanup``) so callers can tell where the failure landed. """ if job_id is None: # no standalone traces: a lone rollout is a job of one job_id = uuid.uuid4().hex @@ -379,24 +355,47 @@ async def rollout( async def _bounded(awaitable: Any) -> Any: """Bound work by the rollout's single wall-clock ``deadline``. - One shared deadline across provision, connect, and the agent loop — + One shared deadline across provision, connect, start, and agent — not a fresh budget per phase — so the bounded work cannot exceed - ``rollout_timeout`` in total. Grade and cancel get a short grace - past that deadline (see :data:`_TEARDOWN_GRACE_S`) so salvage - scoring still works, but a wedged JSON-RPC read cannot hang the - batch. A client read-timeout is not enough on its own: a wedged - upstream that trickles bytes resets the read timer forever, so a - single stuck sampling call could otherwise hang the rollout — and - the batch waits on it — indefinitely. A breach surfaces as + ``rollout_timeout`` in total. Grade, cancel, and provider cleanup + use their own explicit budgets below. A breach surfaces as ``TimeoutError`` (distinct from a Ctrl-C ``CancelledError``). """ if deadline is None: return await awaitable return await asyncio.wait_for(awaitable, max(deadline - loop.time(), 0.0)) + async def _cleanup_bounded(awaitable: Any, budget_s: float) -> Any: + """Apply cleanup budgets only when this rollout has a hard deadline.""" + if deadline is None: + return await awaitable + return await asyncio.wait_for(awaitable, budget_s) + + async def _cancel(live: Run) -> None: + """Best-effort task cancellation under its own cleanup budget.""" + try: + await _cleanup_bounded(live.client.cancel(), _CANCEL_TIMEOUT_S) + except TimeoutError: + logger.warning("rollout cancel did not finish within %.0fs", _CANCEL_TIMEOUT_S) + except Exception as exc: + logger.warning("rollout cancel failed: %s", exc) + + async def _grade(live: Run, exc: BaseException | None) -> None: + """Finalize the Run under one bounded salvage-grading window.""" + try: + await _cleanup_bounded(live.finish(exc), _GRADE_GRACE_S) + except TimeoutError: + logger.warning("rollout grading did not finish within %.0fs", _GRADE_GRACE_S) + live.trace.status = "error" + if live.trace.stop_reason is None: + live.trace.stop_reason = "timeout" + live.record(Step(source="system", error="grading timed out")) + await _cancel(live) + async def _drive() -> None: nonlocal run, _phase - async with contextlib.AsyncExitStack() as stack: + stack = contextlib.AsyncExitStack() + try: # Setup (provision + connect) is bounded but not gradable: a # timeout fires before the run is live, so it surfaces as a # pre-launch failure. A cancelled enter still tears the @@ -404,35 +403,62 @@ async def _drive() -> None: addr = cast("Runtime", await _bounded(stack.enter_async_context(runtime(task)))) _phase = "starting task" client = cast("HudClient", await _bounded(stack.enter_async_context(connect(addr)))) - live = Run(client, task.id, task.args, deadline=deadline) + live = Run(client, task.id, task.args) live._runtime = addr.url # the placement record for the receipt - async with live: # start on enter; grade on exit - run = live # bound only once live: an earlier failure synthesizes - _phase = "agent loop" - # File tracking (when enabled) emits setup separately, then - # streams workspace diffs for the duration of the agent loop. + await _bounded(live.__aenter__()) + run = live # bound only once live: an earlier failure synthesizes + _phase = "agent loop" + + async def _agent_loop() -> None: async with file_tracking_observer(client): - try: - await _bounded(agent(run)) - except TimeoutError: - # The run is live with a partial trajectory worth - # grading, so record the truncation and fall through - # to the normal grade-on-exit path. A bare cancel - # (Ctrl-C) does not land here — it is a CancelledError, - # which this does not catch, so it keeps the non-graded - # cancel path in ``__aexit__``. - logger.warning( - "rollout agent loop timed out after %.0fs; grading partial", - rollout_timeout, - ) - run.trace.stop_reason = "timeout" - run.record( - Step( - source="system", - error=f"agent loop timed out after {rollout_timeout:.0f}s", - ) - ) - _phase = "grading" + await agent(live) + + agent_error: Exception | None = None + try: + await _bounded(_agent_loop()) + except TimeoutError: + # Preserve the partial trajectory and spend one independent + # grace window grading it. + logger.warning( + "rollout agent loop timed out after %.0fs; grading partial", + rollout_timeout, + ) + live.trace.stop_reason = "timeout" + live.record( + Step( + source="system", + error=f"agent loop timed out after {rollout_timeout:.0f}s", + ) + ) + except (asyncio.CancelledError, KeyboardInterrupt): + live.trace.status = "cancelled" + await _cancel(live) + raise + except Exception as exc: + agent_error = exc + + _phase = "grading" + await _grade(live, agent_error) + if agent_error is not None: + raise agent_error + finally: + try: + await _cleanup_bounded(stack.aclose(), _CLEANUP_TIMEOUT_S) + except TimeoutError: + _phase = "cleanup" + logger.warning( + "rollout provider cleanup did not finish within %.0fs", + _CLEANUP_TIMEOUT_S, + ) + if run is not None: + run.trace.status = "error" + run.record(Step(source="system", error="provider cleanup timed out")) + except Exception as exc: + _phase = "cleanup" + logger.warning("rollout provider cleanup failed: %s", exc) + if run is not None: + run.trace.status = "error" + run.record(Step(source="system", error=f"provider cleanup failed: {exc}")) try: await _drive() diff --git a/hud/eval/tests/test_rollout.py b/hud/eval/tests/test_rollout.py index a35e85f06..1fd059d97 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -354,6 +354,7 @@ class _HangOnGrade: def __init__(self, inner: Any) -> None: self._inner = inner + self.cancelled = False def __getattr__(self, name: str) -> Any: return getattr(self._inner, name) @@ -363,18 +364,22 @@ async def grade(self, *_a: Any, **_k: Any) -> dict[str, Any]: return {} async def cancel(self) -> None: - return None + self.cancelled = True real_connect = run_mod.connect + hanging_client: _HangOnGrade | None = None @asynccontextmanager async def _connect_hanging(addr: Any) -> AsyncIterator[Any]: + nonlocal hanging_client async with real_connect(addr) as client: - yield _HangOnGrade(client) + hanging_client = _HangOnGrade(client) + yield hanging_client monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr(run_mod, "connect", _connect_hanging) - monkeypatch.setattr(run_mod, "_TEARDOWN_GRACE_S", 0.2) + monkeypatch.setattr(run_mod, "_GRADE_GRACE_S", 0.2) + monkeypatch.setattr(run_mod, "_CANCEL_TIMEOUT_S", 0.2) try: started = asyncio.get_running_loop().time() run = await rollout( @@ -391,6 +396,43 @@ async def _connect_hanging(addr: Any) -> AsyncIterator[Any]: assert run.trace.is_error assert run.trace.stop_reason == "timeout" assert "grading timed out" in (run.trace.error or "") + assert hanging_client is not None and hanging_client.cancelled + + +async def test_wedged_provider_cleanup_cannot_stall_the_rollout() -> None: + """Provider teardown gets its own budget after grading completes.""" + from hud.eval import run as run_mod + + env = Environment("sums") + + @env.template() + async def add(a: int, b: int): + answer = yield f"add:{a}:{b}" + yield 1.0 if answer == str(a + b) else 0.0 + + @asynccontextmanager + async def _wedged_cleanup(_task: TaskRow) -> AsyncIterator[Runtime]: + async with _local(env) as runtime: + yield runtime + await asyncio.sleep(3600) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(run_mod, "_CLEANUP_TIMEOUT_S", 0.2) + try: + started = asyncio.get_running_loop().time() + run = await rollout( + _add_task(2, 3), + _FnAgent(_solve_add), + runtime=_wedged_cleanup, + rollout_timeout=1.0, + ) + elapsed = asyncio.get_running_loop().time() - started + finally: + monkeypatch.undo() + + assert elapsed < 5.0 + assert run.trace.is_error + assert "provider cleanup timed out" in (run.trace.error or "") async def test_pre_launch_failure_yields_a_synthesized_failed_run() -> None: