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..168c2df67 100644 --- a/hud/eval/run.py +++ b/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() 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..1fd059d97 100644 --- a/hud/eval/tests/test_rollout.py +++ b/hud/eval/tests/test_rollout.py @@ -338,6 +338,103 @@ 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 + self.cancelled = False + + 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: + 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: + hanging_client = _HangOnGrade(client) + yield hanging_client + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(run_mod, "connect", _connect_hanging) + 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( + _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 "") + 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: @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