diff --git a/verifiers/v1/agent.py b/verifiers/v1/agent.py index d7ffda028f..052c46ea18 100644 --- a/verifiers/v1/agent.py +++ b/verifiers/v1/agent.py @@ -8,10 +8,12 @@ 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 dataclasses import dataclass +from pathlib import Path from typing import AsyncIterator @@ -37,7 +39,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_VERSION, Trace, WireTrace from verifiers.v1.types import ( AssistantMessage, Messages, @@ -693,6 +695,96 @@ def remember(current: Trace) -> None: self._completed.append(trace) +class ReplayStore: + """A finished run's traces, keyed for seat playback: seat name -> task name + (or `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')}" + 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" + ) + if (version := record.get("version")) != TRACE_VERSION: + # Records don't migrate (same stance as resume/replay CLIs): an old + # run replays only on the build that produced it. + raise ValueError( + f"saved trace for {key!r} is record version {version}, this " + f"build reads {TRACE_VERSION}; re-run the source eval on this " + "build to replay it" + ) + 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, *, diff --git a/verifiers/v1/configs/agent.py b/verifiers/v1/configs/agent.py index 1836fbd484..9302ec3f1e 100644 --- a/verifiers/v1/configs/agent.py +++ b/verifiers/v1/configs/agent.py @@ -1,5 +1,7 @@ """One env agent's config: who plays the seat, and its per-run caps.""" +from pathlib import Path + from pydantic import SerializeAsAny, model_validator from pydantic_config import BaseConfig @@ -38,6 +40,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 diff --git a/verifiers/v1/env.py b/verifiers/v1/env.py index 0564533e96..ab4354dacf 100644 --- a/verifiers/v1/env.py +++ b/verifiers/v1/env.py @@ -13,7 +13,13 @@ ) -from verifiers.v1.agent import Agent, Agents, _EpisodeAgent +from verifiers.v1.agent import ( + Agent, + Agents, + ReplayStore, + _EpisodeAgent, + _ReplayEpisodeAgent, +) from verifiers.v1.configs.agent import AgentConfig from verifiers.v1.configs.env import ( EnvConfig, @@ -142,6 +148,8 @@ def __init__(self, config: ConfigT) -> None: self._interception: Interception | None = None # Clients for endpoint-pinning roles, cached by config, closed with serving(). self._agent_clients: dict[str, Client] = {} + # Saved runs backing replayed seats, loaded once per env. + self._replay_stores: dict[str, ReplayStore] = {} # Resource warnings dedupe env-wide (agents are per-episode). self._warned_resources: set = set() @@ -211,7 +219,12 @@ def make(name: str, spec: AgentConfig) -> Agent: else ctx.sampling, } ) - return _EpisodeAgent( + agent_cls, extra = ( + (_ReplayEpisodeAgent, {"store": self._replay_store(spec.replay)}) + if spec.replay is not None + else (_EpisodeAgent, {}) + ) + return agent_cls( resolved, client=self._client_for(spec.client) if spec.client is not None @@ -225,6 +238,7 @@ def make(name: str, spec: AgentConfig) -> Agent: on_trace=on_trace, on_discard=on_discard, warned_resources=self._warned_resources, + **extra, ) agents = Agents(self.config, make) @@ -237,6 +251,13 @@ def _client_for(self, config: ClientConfig) -> Client: self._agent_clients[key] = resolve_client(config) return self._agent_clients[key] + def _replay_store(self, path) -> ReplayStore: + """Load (and cache) a replayed seat's saved run once per env.""" + key = str(path) + if key not in self._replay_stores: + self._replay_stores[key] = ReplayStore(path) + return self._replay_stores[key] + async def run_episode( self, task: Task, diff --git a/verifiers/v1/envs/agentic_judge/env.py b/verifiers/v1/envs/agentic_judge/env.py index b1ce109f0d..7a002afe6e 100644 --- a/verifiers/v1/envs/agentic_judge/env.py +++ b/verifiers/v1/envs/agentic_judge/env.py @@ -276,6 +276,11 @@ def _check_agents(self) -> None: "that needs no execution is a plugged judge " "(--env.taskset.task.judges), not an agent." ) + if self.config.solver.replay is not None: + raise ValueError( + "agentic-judge judges in the solver's live box, and a replayed " + "solver has no box; replay another seat, or another env" + ) if isinstance(self.config.solver.runtime, vf.SubprocessConfig): raise ValueError( "agentic-judge plays its judge in the solver's box, but the solver "