refactor(eval): plan resumed rollouts by task content and occurrence

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

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)

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