From 761573efe2397de501224d97118dd303ea4a4351 Mon Sep 17 00:00:00 2001 From: Xeophon <46377542+xeophon@users.noreply.github.com> Date: Sat, 3 Oct 2026 10:25:10 -0700 Subject: [PATCH] refactor(eval): share rollout planning across runners --- verifiers/v1/cli/eval/resume.py | 73 ++++++++++----------------------- verifiers/v1/cli/eval/runner.py | 24 ++++------- verifiers/v1/utils/eval.py | 36 ++++++++++++++++ 3 files changed, 66 insertions(+), 67 deletions(-) create mode 100644 verifiers/v1/utils/eval.py diff --git a/verifiers/v1/cli/eval/resume.py b/verifiers/v1/cli/eval/resume.py index 97221ad93d..940ece9ba5 100644 --- a/verifiers/v1/cli/eval/resume.py +++ b/verifiers/v1/cli/eval/resume.py @@ -1,46 +1,35 @@ """Resume an interrupted eval: reload its finished rollouts and run only the rest. -`--resume` re-enters the run dir the resolved config points at and writes back into -it. `load` keeps the good saved rollouts and re-runs what's owed: missing rollouts -(never written) and errored ones (dropped and redone). - -A saved episode is matched to a selected task by its persisted `episode.task.hash`, -falling back to hashing `episode.task.data` for older rows without one. Tasks with -identical data are interchangeable, a task whose data changed since the interrupted -run re-runs, and nothing depends on `data.idx`. +`load` validates saved episodes, applies the environment's completion verdict, and +uses the shared rollout planner to retain only those matching the current tasks. """ -from collections import Counter, defaultdict from collections.abc import Callable from pathlib import Path from pydantic_core import from_json from verifiers.v1.episode import WireEpisode -from verifiers.v1.task import task_key +from verifiers.v1.task import TaskT +from verifiers.v1.utils.eval import plan_rollouts from verifiers.v1.utils.trace_store import TRACES_FILE def load( resume_dir: Path, - selected_keys: list[str], + tasks: list[TaskT], num_rollouts: int, complete: Callable[[WireEpisode], bool] | None = None, -) -> tuple[list[WireEpisode], dict[str, int]]: - """Load the good saved rollouts and diff them against the run's target: returns - (kept episodes, rollouts owed per task key). `selected_keys` is one key per - selected task (duplicates allowed — a key selected k times is owed up to - `k * num_rollouts`; spread back over the tasks with `distribute`). `complete` - is the keep-verdict (default `episode.ok`). Rewrites `traces.jsonl` to the - kept rows via a temp file + atomic rename; a torn or malformed row is owed - again, never a crash.""" - path = resume_dir / TRACES_FILE - targets = { - key: count * num_rollouts for key, count in Counter(selected_keys).items() - } +) -> list[tuple[TaskT, list[WireEpisode], int]]: + """Plan the remaining rollouts and atomically rewrite the stream to kept rows. + ``complete`` is the keep-verdict (default ``episode.ok``). Torn or malformed + rows are discarded; retained rows preserve their original serialized fields. + """ + path = resume_dir / TRACES_FILE verdict = complete if complete is not None else (lambda episode: episode.ok) - good: dict[str, list[tuple[bytes, WireEpisode]]] = defaultdict(list) + episodes: list[WireEpisode] = [] + lines: dict[str, bytes] = {} if path.exists(): with path.open("rb") as results: for line in results: @@ -50,40 +39,20 @@ def load( row = from_json(line) if "traces" not in row: continue - task = row["task"] - if not isinstance(task, dict): - continue - key = task.get("hash") - if key is None: - key = task_key(task["data"]) - except (ValueError, KeyError, IndexError, TypeError): - # A torn final line (the run died mid-write) or a foreign shape - # is not a keepable rollout — it's owed again, never a crash. - continue - if not isinstance(key, str): - continue - if key not in targets or len(good[key]) >= targets[key]: - continue - try: episode = WireEpisode.model_validate(row) if not verdict(episode): continue # A malformed row from any task/episode plugin is owed again. except Exception: # noqa: BLE001, S112 continue - good[key].append( - (line if line.endswith(b"\n") else line + b"\n", episode) + episodes.append(episode) + lines.setdefault( + episode.id, line if line.endswith(b"\n") else line + b"\n" ) - keep: list[bytes] = [] - episodes: list[WireEpisode] = [] - owed: dict[str, int] = {} - for key, target in targets.items(): - rows = good.get(key, []) - keep.extend(line for line, _ in rows) - episodes.extend(episode for _, episode in rows) - if missing := target - len(rows): - owed[key] = missing + plan = plan_rollouts(tasks, num_rollouts, episodes) tmp = path.with_suffix(".jsonl.tmp") - tmp.write_bytes(b"".join(keep)) + tmp.write_bytes( + b"".join(lines[episode.id] for _, kept, _ in plan for episode in kept) + ) tmp.replace(path) - return episodes, owed + return plan diff --git a/verifiers/v1/cli/eval/runner.py b/verifiers/v1/cli/eval/runner.py index a6598517bb..22f348db35 100644 --- a/verifiers/v1/cli/eval/runner.py +++ b/verifiers/v1/cli/eval/runner.py @@ -18,12 +18,12 @@ from verifiers.v1.cli.eval import resume from verifiers.v1.cli.eval.hint import PRIME_RL_HINT from verifiers.v1.cli.output import output_path, save_config -from verifiers.v1.cli.resume import distribute from verifiers.v1.clients import ModelContext from verifiers.v1.configs.cli.eval import EvalConfig from verifiers.v1.env import Env, RunSlot from verifiers.v1.episode import Episode, EvalRunInfo from verifiers.v1.utils.aio import run_shielded +from verifiers.v1.utils.eval import plan_rollouts from verifiers.v1.utils.platform import ( PushState, abort_run, @@ -89,40 +89,34 @@ async def run_eval(config: EvalConfig) -> list[Episode]: ) tasks = list(taskset) out = output_path(config) - # One (task, rollouts-to-run) pair per selected task; resume shrinks the counts. - plan = [(task, config.num_rollouts) for task in tasks] - # Kept on-disk rollouts rejoin the run as finished episodes; only owed ones re-run. - finished: list[Episode] = [] if config.resume: - keys = [task.hash for task in tasks] # the env's own keep-verdict decides what resumes - loaded, owed = resume.load( + plan = resume.load( out, - keys, + tasks, config.num_rollouts, lambda episode: env.complete(cast(Episode, episode)), ) - finished = [cast(Episode, episode) for episode in loaded] - if not owed: # already complete - report it and exit successfully + if not any(n for _, _, n in plan): print( f"nothing to resume in {out}: all {len(tasks)}x{config.num_rollouts} " "rollouts already completed without error" ) raise SystemExit(0) - counts = distribute(keys, owed, config.num_rollouts) - plan = [(task, n) for task, n in zip(tasks, counts) if n] logger.info( "resuming %s: %d task(s), %d rollout(s) owed", out, - len(plan), - sum(owed.values()), + sum(n > 0 for _, _, n in plan), + sum(n for _, _, n in plan), ) else: + plan = plan_rollouts(tasks, config.num_rollouts) save_config(config, out) logger.info( "running %dx%d rollouts on %s", len(plan), config.num_rollouts, config.model ) logger.info(PRIME_RL_HINT) + finished = [cast(Episode, episode) for _, kept, _ in plan for episode in kept] start = time.time() logger.info("results: %s", out) @@ -148,7 +142,7 @@ async def on_complete(episode: Episode) -> None: try: async with _in_process(env, config, semaphore, on_complete) as run_slot: # the env's own slots: it fills their live traces as the rollouts run - planned = [slot for task, n in plan for slot in env.slots(task, n)] + planned = [slot for task, _, n in plan if n for slot in env.slots(task, n)] slots = [RunSlot.finished(episode) for episode in finished] + planned display = ( dashboard(slots, config, start, push=push_state) diff --git a/verifiers/v1/utils/eval.py b/verifiers/v1/utils/eval.py new file mode 100644 index 0000000000..865f496a54 --- /dev/null +++ b/verifiers/v1/utils/eval.py @@ -0,0 +1,36 @@ +"""Rollout planning shared by in-process and served evaluations.""" + +from collections import defaultdict, deque +from collections.abc import Iterable + +from verifiers.v1.episode import WireEpisode +from verifiers.v1.task import Task, TaskT + + +def plan_rollouts( + tasks: Iterable[TaskT], + num_rollouts: int, + completed: Iterable[WireEpisode] = (), +) -> list[tuple[TaskT, list[WireEpisode], int]]: + """Assign distinct accepted episodes to selected task occurrences by content hash. + + Each occurrence gets up to ``num_rollouts`` saved episodes and the count still + missing. Call separately per evaluation source; callers own acceptance and I/O. + """ + saved: dict[str, deque[WireEpisode]] = defaultdict(deque) + seen: set[str] = set() + for episode in completed: + if episode.id in seen: + continue + seen.add(episode.id) + key = episode.task.hash + if key is None: + key = Task(episode.task.data).hash + saved[key].append(episode) + + plan = [] + for task in tasks: + episodes = saved[task.hash] + kept = [episodes.popleft() for _ in range(min(num_rollouts, len(episodes)))] + plan.append((task, kept, num_rollouts - len(kept))) + return plan