feat: encode episodes to JSON one node at a time and record runs on traces

tests/unit/test_trace.py

@@ -133,6 +133,11 @@ def test_record_streams_are_arrays_with_json_dump_content(float_decimals):
     assert orjson.loads(episode.to_json(float_decimals)) == orjson.loads(
         json.dumps(cases[1][1])
     )
+    for record, encoded in (
+        (tr.to_record(float_decimals), tr.to_json(float_decimals)),
+        (episode.to_record(float_decimals), episode.to_json(float_decimals)),
+    ):
+        assert encoded == orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)
 
 
 def test_record_exclude_merges_with_the_record_excludes():
@@ -145,9 +150,41 @@ def test_record_exclude_merges_with_the_record_excludes():
     assert record["nodes"][1]["message"]["content"] == "a"
     assert "top_logprobs" not in record["nodes"][1]
     assert record["nodes"][1]["token_ids"].tolist() == [4, 5, 6]
-    assert orjson.loads(tr.to_json(exclude=exclude)) == orjson.loads(
-        orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)
+    assert tr.to_json(exclude=exclude) == orjson.dumps(
+        record, option=orjson.OPT_SERIALIZE_NUMPY
+    )
+    episode = vf.Episode(task=tr.task, traces=[tr, _streamed_trace()], ok=True)
+    exclude = {"ok": True, "traces": {"__all__": exclude, 1: {"tools"}}}
+    assert episode.to_json(exclude=exclude) == orjson.dumps(
+        episode.to_record(exclude=exclude), option=orjson.OPT_SERIALIZE_NUMPY
+    )
+
+
+def test_episode_json_keeps_task_subclass_fields_like_its_record():
+    tr = _streamed_trace()
+    tr.task.data = MyTask(idx=0, prompt="q", answer="42")
+    episode = vf.Episode(task=tr.task, traces=[tr], ok=True)
+    assert episode.to_json() == orjson.dumps(
+        episode.to_record(), option=orjson.OPT_SERIALIZE_NUMPY
+    )
+
+
+def test_record_run_stamps_the_episode_and_every_trace():
+    work = vf.TrainWorkInfo(step=3, policy=vf.PolicySpan(start=0, end=1))
+    episode = vf.Episode(
+        task=_streamed_trace().task, traces=[_streamed_trace(), _streamed_trace()]
+    )
+    run = vf.TrainRunInfo(id="run-1", work=work)
+    episode.record_run(run)
+
+    record = orjson.loads(episode.to_json())
+    expected = run.model_dump(mode="json", exclude_none=True)
+    assert record["run"] == expected
+    assert [trace["run"] for trace in record["traces"]] == [expected, expected]
+    assert episode.to_json() == orjson.dumps(
+        episode.to_record(), option=orjson.OPT_SERIALIZE_NUMPY
     )
+    assert vf.Episode.model_validate(record).traces[1].run == run
 
 
 def test_round_stream_matches_round():

verifiers/v1/__init__.py

@@ -65,19 +65,7 @@
     from verifiers.v1.configs.taskset import SharedTasksetConfig, TasksetConfig
     from verifiers.v1.env import Env
     from verifiers.v1.envs.single_agent import SingleAgentEnv
-    from verifiers.v1.episode import (
-        EnvInfo,
-        Episode,
-        EvalRunInfo,
-        EvalWorkInfo,
-        GroupInfo,
-        PolicySpan,
-        RunInfo,
-        TrainRunInfo,
-        TrainWorkInfo,
-        WireEpisode,
-        WorkInfo,
-    )
+    from verifiers.v1.episode import EnvInfo, Episode, GroupInfo, WireEpisode
     from verifiers.v1.errors import (
         EnvError,
         HarnessError,
@@ -133,17 +121,24 @@
         AgentSpan,
         Branch,
         Error,
+        EvalRunInfo,
+        EvalWorkInfo,
         InterceptRecord,
         ModelCall,
         PolicyEvent,
+        PolicySpan,
         Reward,
+        RunInfo,
         TimeSpan,
         TimeSplit,
         Timing,
         Trace,
         TraceTask,
+        TrainRunInfo,
+        TrainWorkInfo,
         VersionInfo,
         WireTrace,
+        WorkInfo,
     )
     from verifiers.v1.types import (
         AssistantMessage,
@@ -235,8 +230,8 @@
     "Episode": "verifiers.v1.episode",
     "Error": "verifiers.v1.trace",
     "EvalClientConfig": "verifiers.v1.configs.client",
-    "EvalRunInfo": "verifiers.v1.episode",
-    "EvalWorkInfo": "verifiers.v1.episode",
+    "EvalRunInfo": "verifiers.v1.trace",
+    "EvalWorkInfo": "verifiers.v1.trace",
     "GroupInfo": "verifiers.v1.episode",
     "Harness": "verifiers.v1.harness",
     "HarnessConfig": "verifiers.v1.configs.harness",
@@ -263,7 +258,7 @@
     "PATCH_CAP_BYTES": "verifiers.v1.utils.git",
     "ParentLink": "verifiers.v1.semantic",
     "PolicyEvent": "verifiers.v1.trace",
-    "PolicySpan": "verifiers.v1.episode",
+    "PolicySpan": "verifiers.v1.trace",
     "PrimeConfig": "verifiers.v1.configs.runtime",
     "PrimeVolume": "verifiers.v1.configs.runtime",
     "ProgramResult": "verifiers.v1.runtimes",
@@ -279,7 +274,7 @@
     "RolloutError": "verifiers.v1.errors",
     "RubricJudge": "verifiers.v1.judges",
     "RubricJudgeConfig": "verifiers.v1.judges",
-    "RunInfo": "verifiers.v1.episode",
+    "RunInfo": "verifiers.v1.trace",
     "Runtime": "verifiers.v1.runtimes",
     "RuntimeConfig": "verifiers.v1.configs.runtime",
     "RuntimeInfo": "verifiers.v1.runtimes",
@@ -331,8 +326,8 @@
     "Trace": "verifiers.v1.trace",
     "TraceTask": "verifiers.v1.trace",
     "TrainClientConfig": "verifiers.v1.configs.client",
-    "TrainRunInfo": "verifiers.v1.episode",
-    "TrainWorkInfo": "verifiers.v1.episode",
+    "TrainRunInfo": "verifiers.v1.trace",
+    "TrainWorkInfo": "verifiers.v1.trace",
     "TunnelError": "verifiers.v1.errors",
     "TurnTokens": "verifiers.v1.types",
     "Usage": "verifiers.v1.types",
@@ -341,7 +336,7 @@
     "WireEpisode": "verifiers.v1.episode",
     "WireTaskData": "verifiers.v1.task",
     "WireTrace": "verifiers.v1.trace",
-    "WorkInfo": "verifiers.v1.episode",
+    "WorkInfo": "verifiers.v1.trace",
     "capture_patch": "verifiers.v1.utils.git",
     "collect": "verifiers.v1.utils.artifacts",
     "compare_stdout_results": "verifiers.v1.utils.score",

verifiers/v1/cli/eval/runner.py

@@ -21,7 +21,8 @@
 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.episode import Episode
+from verifiers.v1.trace import EvalRunInfo
 from verifiers.v1.utils.aio import run_shielded
 from verifiers.v1.utils.eval import plan_rollouts
 from verifiers.v1.utils.platform import (

verifiers/v1/episode.py

@@ -1,10 +1,10 @@
 """The episode — one run's traces plus their shared standing, whole."""
 
+import functools
 import uuid
-from typing import Annotated, Any, Generic, Literal, Self
+from typing import Any, Generic
 
-import orjson
-from pydantic import BaseModel, Field, computed_field, model_validator
+from pydantic import BaseModel, Field, computed_field
 
 from verifiers.v1.configs.agent import WireAgentConfig
 from verifiers.v1.graph import RECORD_FLOAT_DECIMALS
@@ -14,10 +14,15 @@
     RECORD_DUMP_EXCLUDE,
     AgentConfigT,
     Error,
+    RunInfo,
     Trace,
     TraceTask,
     add_record_streams,
+    item_adapter,
+    item_exclude,
     merge_exclude,
+    splice_json,
+    trace_json_parts,
 )
 from verifiers.v1.types import Usage
 
@@ -37,62 +42,6 @@ class GroupInfo(BaseModel):
     id: str
 
 
-class PolicySpan(BaseModel):
-    """Live policy versions spanned while generating an episode."""
-
-    start: int
-    end: int
-
-    @model_validator(mode="after")
-    def ordered(self) -> Self:
-        if self.end < self.start:
-            raise ValueError("policy span end must be at least its start")
-        return self
-
-    @property
-    def drift(self) -> int:
-        return self.end - self.start
-
-
-class TrainWorkInfo(BaseModel):
-    type: Literal["train"] = "train"
-    step: int
-    """The training step at which the episode was dispatched."""
-    policy: PolicySpan | None = None
-    """The live policy used for generation; None for frozen-policy work."""
-
-
-class EvalWorkInfo(BaseModel):
-    type: Literal["eval"] = "eval"
-    step: int
-    """The training step at which the episode was dispatched."""
-    policy: PolicySpan | None = None
-    """The live policy used for generation; None for frozen-policy work."""
-
-
-WorkInfo = Annotated[TrainWorkInfo | EvalWorkInfo, Field(discriminator="type")]
-
-
-class TrainRunInfo(BaseModel):
-    """A training run and the kind of work this episode performed for it."""
-
-    type: Literal["train"] = "train"
-    id: str
-    name: str | None = None
-    work: WorkInfo
-
-
-class EvalRunInfo(BaseModel):
-    """A standalone evaluation run."""
-
-    type: Literal["eval"] = "eval"
-    id: str
-    name: str | None = None
-
-
-RunInfo = Annotated[TrainRunInfo | EvalRunInfo, Field(discriminator="type")]
-
-
 class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
     """The artifact Env.run produces. Contains multiple agents' traces."""
 
@@ -117,8 +66,10 @@ class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
     """Every agent's trace, in completion order."""
 
     def record_run(self, run: RunInfo) -> None:
-        """Record the run identity on the dispatched episode."""
+        """Record the run identity on the dispatched episode and each of its traces."""
         self.run = run
+        for trace in self.traces:
+            trace.run = run
 
     @property
     def last_error(self) -> Error | None:
@@ -188,10 +139,34 @@ def to_json(
         float_decimals: int | None = RECORD_FLOAT_DECIMALS,
         exclude: dict | set | None = None,
     ) -> bytes:
-        """`to_record` as compact UTF-8 JSON. NaN and infinity become `null`."""
-        return orjson.dumps(
-            self.to_record(float_decimals, exclude), option=orjson.OPT_SERIALIZE_NUMPY
+        """`to_record` as compact UTF-8 JSON, byte for byte. NaN and infinity become
+        `null`.
+
+        Pydantic's dump and orjson each hold the GIL for a whole call, so encoding a
+        large episode in one call stalls every other thread of the process, such as an
+        asyncio event loop, for as long as it takes. This encodes the episode one node at
+        a time instead, and releases the GIL about every millisecond, so other threads
+        run in between and the stall no longer grows with the episode. The total work
+        stays the same."""
+        exclude = merge_exclude({"traces": {"__all__": RECORD_DUMP_EXCLUDE}}, exclude)
+        head = self.model_dump(
+            mode="json",
+            exclude=merge_exclude(exclude, {"traces": {"__all__": True}}),
+            exclude_none=True,
+        )
+        adapter = item_adapter(type(self), "traces")
+        traces = (
+            trace_json_parts(
+                trace,
+                functools.partial(adapter.dump_python, trace),
+                item_exclude(exclude, "traces", index) or {},
+                float_decimals,
+                exclude_none=True,
+            )
+            for index, trace in enumerate(self.traces)
         )
+        # One join at the end: joining per level would copy the whole record each time.
+        return b"".join(splice_json(head, "traces", traces))
 
 
 WireEpisode = Episode[WireTaskData, State, WireAgentConfig]

verifiers/v1/serve/delta.py

@@ -39,7 +39,7 @@
 
 logger = logging.getLogger(__name__)
 
-HEADER_FIELDS = ("version", "id", "verifiers", "task")
+HEADER_FIELDS = ("version", "id", "verifiers", "task", "run")
 """Sent once, when a trace first appears."""
 
 LIST_FIELDS = (