fix(eval): hard rollout wall-clock timeout for HUDRuntime training (HUD-2099)

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.
 

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`
 

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:

hud/eval/run.py

@@ -50,6 +50,13 @@
 
 logger = logging.getLogger("hud.eval.run")
 
+#: 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:
     """Coerce one wire prompt turn onto MCP's ``PromptMessage`` vocabulary.
@@ -228,24 +235,32 @@ async def __aexit__(
         if exc_type is not None and issubclass(
             exc_type, asyncio.CancelledError | KeyboardInterrupt
         ):
-            self.trace.status = "cancelled"
-            await 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() 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.
+        if exc is not None:
             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
+                return
         else:
             evaluation = await self.client.grade(answer)
 
@@ -265,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:
@@ -311,9 +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. 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
@@ -337,21 +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. 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, 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
@@ -361,33 +405,60 @@ async def _drive() -> None:
                 client = cast("HudClient", await _bounded(stack.enter_async_context(connect(addr))))
                 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()

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))

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: