Repository navigation
feat(v1): seat replay — pin any env agent to a finished run's traces #2118
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
9f38cd1
1327cfb
9b9d495
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,9 +5,11 @@ | |
| entered agent (`async with`) owns one server; un-entered, each run brings its own.""" | ||
|
|
||
| import asyncio | ||
| import json | ||
| import logging | ||
| from collections.abc import Callable, Iterator, Mapping | ||
| from contextlib import asynccontextmanager, nullcontext | ||
| from pathlib import Path | ||
| from typing import AsyncIterator | ||
|
|
||
| from pydantic import SerializeAsAny, model_validator | ||
|
|
@@ -34,7 +36,7 @@ | |
| ) | ||
| from verifiers.v1.session import RolloutLimits | ||
| from verifiers.v1.task import Task | ||
| from verifiers.v1.trace import Trace | ||
| from verifiers.v1.trace import Trace, WireTrace | ||
| from verifiers.v1.types import Sampling, SamplingConfig | ||
| from verifiers.v1.utils.compile import ( | ||
| cap_remote_harness_timeout, | ||
|
|
@@ -68,6 +70,13 @@ class AgentConfig(BaseConfig): | |
| """Endpoint override (None = the run's client).""" | ||
| sampling: SamplingConfig | None = None | ||
| """Sampling override (None = the run's sampling).""" | ||
| replay: Path | None = None | ||
| """Play this seat from a finished run instead of sampling: a run dir (or its | ||
| traces.jsonl). Each task returns that run's saved trace for this seat (a | ||
| single-seat run matches regardless of name), re-stamped with this seat's | ||
| standing — so the other seats iterate against fixed work. A replayed seat | ||
| runs no model and leaves no runtime state: an env that inspects this seat's | ||
| box can't replay it.""" | ||
| timeout: TimeoutConfig = TimeoutConfig() | ||
| retries: RetryConfig = RetryConfig() | ||
| """Whole-run retries: rerun this agent's rollout while its trace ends with a | ||
|
|
@@ -423,6 +432,88 @@ def watch(trace: Trace) -> None: | |
| return trace | ||
|
|
||
|
|
||
| class ReplayStore: | ||
| """A finished run's traces, keyed for seat playback: seat name -> task name | ||
| (or `idx:<idx>`) -> trace record. Episode lines are unwrapped; a run with a | ||
| single seat serves any requested seat, so a plain single-agent run replays | ||
| into a multi-agent env's differently-named seat.""" | ||
|
|
||
| def __init__(self, path: Path) -> None: | ||
| self._file = path / "traces.jsonl" if path.is_dir() else path | ||
| by_seat: dict[str, dict[str, dict]] = {} | ||
| for line in self._file.read_text(encoding="utf-8").splitlines(): | ||
| record = json.loads(line) | ||
| traces = ( | ||
| record["traces"] | ||
| if "traces" in record and "nodes" not in record | ||
| else [record] | ||
| ) | ||
| for t in traces: | ||
| seat = ((t.get("agent") or {}).get("name")) or "agent" | ||
| data = (t.get("task") or {}).get("data") or {} | ||
| key = data.get("name") or f"idx:{data.get('idx')}" | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🟡 Medium When two tasks in the dataset share the same non-empty 🚀 Reply "fix it for me" or copy this AI Prompt for your agent: |
||
| by_seat.setdefault(seat, {})[key] = t | ||
| if not by_seat: | ||
| raise ValueError(f"no traces to replay in {self._file}") | ||
| self._by_seat = by_seat | ||
|
|
||
| def lookup(self, seat: str, task: Task) -> dict: | ||
| records = self._by_seat.get(seat) | ||
| if records is None: | ||
| if len(self._by_seat) != 1: | ||
| raise ValueError( | ||
| f"{self._file} has seats {sorted(self._by_seat)}, none named " | ||
| f"{seat!r} and more than one to fall back to" | ||
| ) | ||
| records = next(iter(self._by_seat.values())) | ||
| key = task.data.name or f"idx:{task.data.idx}" | ||
| record = records.get(key) | ||
| if record is None: | ||
| raise ValueError( | ||
| f"no saved trace for task {key!r} in {self._file}; select the " | ||
| "same tasks the replayed run played" | ||
| ) | ||
| return record | ||
|
|
||
|
|
||
| class _ReplayEpisodeAgent(_EpisodeAgent): | ||
| """A seat pinned to a finished run's traces (`AgentConfig.replay`): `run` | ||
| revalidates the saved trace for the task and re-stamps it with this seat's | ||
| standing instead of sampling. No model call, no runtime — anything | ||
| downstream must live off the trace.""" | ||
|
|
||
| def __init__(self, *args, store: ReplayStore, **kwargs) -> None: | ||
| super().__init__(*args, **kwargs) | ||
| self._store = store | ||
|
|
||
| async def run( | ||
| self, | ||
| task: Task, | ||
| *, | ||
| runtime: Runtime | None = None, | ||
| tools: Mapping[str, SharedToolServer] | None = None, | ||
| on_trace: Callable[[Trace], None] | None = None, | ||
| ) -> Trace: | ||
| trace = WireTrace.model_validate(self._store.lookup(self._name, task)) | ||
| if trace.agent is not None: | ||
| trace.agent.name = self._name | ||
| trace.agent.trainable = self.trainable | ||
| if self._on_trace is not None: | ||
| self._on_trace(trace) | ||
| if on_trace is not None: | ||
| on_trace(trace) | ||
| self._completed.append(trace) | ||
| return trace | ||
|
|
||
| @asynccontextmanager | ||
| async def provision(self, task: Task | None = None) -> AsyncIterator[Runtime]: | ||
| raise RuntimeError( | ||
| f"agent {self._name!r} replays saved traces and has no runtime to " | ||
| "provision; an env that needs this seat's box can't replay it" | ||
| ) | ||
| yield # pragma: no cover - unreachable, keeps the context-manager shape | ||
|
|
||
|
|
||
| def make_agent( | ||
| config: AgentConfig, | ||
| *, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🟡 Medium
v1/agent.py:435ReplayStorekeys each seat's saved traces by task name/idxonly, so when the source run produced multiple rollouts for the same task (e.g.num_rollouts > 1), line 455 silently overwrites every earlier trace for that key with the last one. Each replay episode for that task then receives the same single last trace regardless of which source rollout it should correspond to, duplicating one sample and discarding the rest — biasing re-grading and rollout-group statistics. Consider keying by a per-rollout discriminator (e.g. an episode index) in addition to task name/idx, or preserving all traces per key and selecting by episode.🚀 Reply "fix it for me" or copy this AI Prompt for your agent: