Persist large V1 traces off the event loop

verifiers/v1/cli/eval/runner.py

@@ -38,8 +38,8 @@ async def run_eval(env: Environment, config: EvalConfig) -> list[Trace]:
     out = output_path(config)
     # Write config.toml up front, then persist each trace as it completes (so the results are
     # durable mid-run, not only at the end). On resume, keep the saved config + good traces and
-    # run only the owed rollouts. `append_trace` is a sync single-line append, safe to call from
-    # concurrent rollouts in the one event loop.
+    # run only the owed rollouts. One lock serializes worker-thread appends from concurrent
+    # rollouts while keeping large trace serialization off the event loop.
     owed: dict[str, int] | None = None
     if config.resume is not None:
         group = bool(discover_decorated(env.taskset, "group_reward"))
@@ -68,8 +68,10 @@ async def run_eval(env: Environment, config: EvalConfig) -> list[Trace]:
     start = time.time()
     logger.info("results: %s", out)
 
-    def on_complete(trace: Trace) -> None:
-        append_trace(out, trace)
+    write_lock = asyncio.Lock()
+
+    async def on_complete(trace: Trace) -> None:
+        await append_trace(out, trace, write_lock)
 
     # Shared tool servers (if any) come up once here and their URLs flow into every rollout
     # (non-shared ones start per rollout inside the episodes); the interception pool comes up
@@ -183,6 +185,7 @@ async def run_eval_server(config: EvalConfig) -> list[Trace]:
         semaphore = (
             asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None
         )
+        write_lock = asyncio.Lock()
 
         async def run_group_unit(idx: int) -> list[Trace]:
             async with semaphore or contextlib.nullcontext():
@@ -194,7 +197,7 @@ async def run_group_unit(idx: int) -> list[Trace]:
                     sampling=config.sampling,
                 )
             for trace in traces:
-                append_trace(out, trace)
+                await append_trace(out, trace, write_lock)
             return traces
 
         async def run_rollout_unit(idx: int) -> list[Trace]:
@@ -205,7 +208,7 @@ async def run_rollout_unit(idx: int) -> list[Trace]:
                     model=config.model,
                     sampling=config.sampling,
                 )
-            append_trace(out, trace)
+            await append_trace(out, trace, write_lock)
             return [trace]
 
         # A group-scored taskset must run each task's rollouts together (cross-rollout

verifiers/v1/cli/output.py

@@ -11,9 +11,11 @@
 durable as they land rather than only at the end.
 """
 
+import asyncio
 from pathlib import Path
 
 import tomli_w
+from pydantic import TypeAdapter
 
 from verifiers.v1.configs.eval import EvalConfig
 from verifiers.v1.trace import Trace
@@ -48,9 +50,25 @@ def save_config(config: EvalConfig, results_dir: Path) -> None:
     )  # fresh; appended to as traces complete
 
 
-def append_trace(results_dir: Path, trace: Trace) -> None:
-    """Append one finished trace to `results.jsonl` (one full trace per line). Called per
-    trace as it completes — a synchronous, single-line append, so concurrent rollouts in
-    one event loop never interleave."""
-    with (results_dir / "results.jsonl").open("a") as f:
-        f.write(trace.model_dump_json(exclude_none=True) + "\n")
+def write_trace(results_dir: Path, trace: Trace) -> None:
+    """Serialize and append one trace in the worker thread."""
+    data = TypeAdapter(type(trace)).dump_json(trace, exclude_none=True)
+    with (results_dir / "results.jsonl").open("ab") as f:
+        f.write(data + b"\n")
+
+
+async def append_trace(results_dir: Path, trace: Trace, lock: asyncio.Lock) -> None:
+    """Append one finished trace without blocking the event loop. The run's shared lock
+    preserves whole-line ordering, and awaiting the worker preserves per-trace durability."""
+
+    async def persist() -> None:
+        async with lock:
+            await asyncio.to_thread(write_trace, results_dir, trace)
+
+    # Shield lock acquisition and the worker so finalized traces survive cancellation.
+    persist_task = asyncio.create_task(persist())
+    try:
+        await asyncio.shield(persist_task)
+    except asyncio.CancelledError:
+        await persist_task
+        raise

verifiers/v1/episode.py

@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 from contextlib import nullcontext
 from typing import TYPE_CHECKING
 
@@ -44,7 +44,7 @@ def __init__(
     async def run(
         self,
         semaphore: asyncio.Semaphore | None = None,
-        on_complete: Callable[[Trace], None] = lambda _trace: None,
+        on_complete: Callable[[Trace], Awaitable[None]] | None = None,
     ) -> list[Trace]:
         """Run all rollouts (each under `semaphore`), then group-score across their
         traces. Without `@group_reward`s a rollout's reward is final the moment its own
@@ -62,7 +62,8 @@ async def run_one(rollout: Rollout) -> Trace:
                 trace = await run_with_retry(rollout, self.retry)
             if not group_scored:  # reward already final → don't wait for the group
                 rollout.phase = Phase.DONE
-                on_complete(trace)
+                if on_complete is not None:
+                    await on_complete(trace)
             # hand freed per-turn request bodies (base64 images) back to the OS
             await trim_memory_periodically()
             return trace
@@ -73,5 +74,6 @@ async def run_one(rollout: Rollout) -> Trace:
             for rollout in self.rollouts:
                 rollout.phase = Phase.DONE
             for trace in traces:
-                on_complete(trace)
+                if on_complete is not None:
+                    await on_complete(trace)
         return traces

verifiers/v1/legacy.py

@@ -450,6 +450,7 @@ async def run_legacy_eval(config) -> list[Trace]:
     )
 
     sem = asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None
+    write_lock = asyncio.Lock()
 
     async def run_one(task_idx: int) -> Trace:
         async def go() -> Trace:
@@ -461,7 +462,7 @@ async def go() -> Trace:
                 state_columns=["trajectory"],
             )
             trace = rollout_output_to_trace(out, task_idx)
-            append_trace(out_dir, trace)
+            await append_trace(out_dir, trace, write_lock)
             return trace
 
         if sem is None: