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,37 @@ 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_episode_json_puts_the_given_run_on_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="local", work=work),
     )
+    record = orjson.loads(episode.to_json(run={"id": "platform", "type": "train"}))
+    run = {**episode.run.model_dump(mode="json", exclude_none=True), "id": "platform"}
+    assert record["run"] == run
+    assert [trace["run"] for trace in record["traces"]] == [run, run]
+    assert "run" not in orjson.loads(episode.to_json())["traces"][0]
 
 
 def test_round_stream_matches_round():

verifiers/v1/episode.py

@@ -1,9 +1,10 @@
 """The episode — one run's traces plus their shared standing, whole."""
 
+import functools
 import uuid
+from collections.abc import Mapping
 from typing import Annotated, Any, Generic, Literal, Self
 
-import orjson
 from pydantic import BaseModel, Field, computed_field, model_validator
 
 from verifiers.v1.configs.agent import WireAgentConfig
@@ -17,7 +18,11 @@
     Trace,
     TraceTask,
     add_record_streams,
+    item_adapter,
+    item_exclude,
     merge_exclude,
+    splice_json,
+    trace_json_parts,
 )
 from verifiers.v1.types import Usage
 
@@ -187,11 +192,41 @@ def to_json(
         self,
         float_decimals: int | None = RECORD_FLOAT_DECIMALS,
         exclude: dict | set | None = None,
+        run: Mapping[str, Any] | 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`. `run` sets fields of the episode's `run`, and puts that `run` on every
+        member trace too, which Prime Traces reads per trace.
+
+        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,
+        )
+        member: dict[str, Any] = {}
+        if run is not None:
+            head["run"] = member["run"] = {**head.get("run", {}), **run}
+        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,
+                extra=member,
+            )
+            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/trace.py

@@ -1,16 +1,24 @@
 from __future__ import annotations
 
 import copy
+import functools
 import math
 import time
 import traceback
 import uuid
 from collections.abc import Callable, Iterable, Iterator, Mapping
-from typing import TYPE_CHECKING, Any, Generic
+from typing import TYPE_CHECKING, Any, Generic, get_args
 
 import numpy as np
 import orjson
-from pydantic import BaseModel, Field, PrivateAttr, computed_field, field_serializer
+from pydantic import (
+    BaseModel,
+    Field,
+    PrivateAttr,
+    TypeAdapter,
+    computed_field,
+    field_serializer,
+)
 from typing_extensions import TypeVar
 
 if TYPE_CHECKING:
@@ -90,6 +98,83 @@ def add_record_streams(
         node_record.update(node.record_streams(float_decimals, exclude_none))
 
 
+_SPLICE = f"\x00splice-{uuid.uuid4().hex}"
+"""Holds a list field's place in a record head until its encoded items replace it."""
+_SPLICE_JSON = orjson.dumps(_SPLICE)
+YIELD_INTERVAL_S = 0.001
+"""How long `to_json` encodes before it releases the GIL to other threads."""
+
+
+@functools.cache
+def item_adapter(model: type[BaseModel], field: str) -> TypeAdapter:
+    """Dumps one item of `model`'s list `field` as `model` itself does: as the declared
+    item type, not the item's own class."""
+    return TypeAdapter(get_args(model.model_fields[field].annotation)[0])
+
+
+def item_exclude(exclude: dict, field: str, index: int) -> dict | None:
+    """The part of the pydantic `exclude` spec that applies to item `index` of `field`."""
+    spec = exclude.get(field)
+    if not isinstance(spec, dict):
+        return None
+    return merge_exclude(spec.get("__all__"), spec.get(index)) or None
+
+
+def splice_json(head: dict, field: str, items: Iterable[list[bytes]]) -> list[bytes]:
+    """`head` as JSON parts, with the encoded `items` as the array at its `field`. A head
+    without `field` (excluded whole) is encoded alone."""
+    if field not in head:
+        return [orjson.dumps(head)]
+    head[field] = _SPLICE
+    before, after = orjson.dumps(head).split(_SPLICE_JSON, 1)
+    parts = [before, b"["]
+    for index, item in enumerate(items):
+        if index:
+            parts.append(b",")
+        parts.extend(item)
+    parts += [b"]", after]
+    return parts
+
+
+def trace_json_parts(
+    trace: Trace,
+    dump: Callable[..., dict[str, Any]],
+    exclude: dict,
+    float_decimals: int | None,
+    exclude_none: bool,
+    extra: Mapping[str, Any] | None = None,
+) -> list[bytes]:
+    """A trace's record as JSON parts, one per node: `dump` (the trace's pydantic dump
+    with `exclude`) gives the head, `extra` adds fields to it."""
+    head = dump(
+        mode="json",
+        exclude=merge_exclude(exclude, {"nodes": {"__all__": True}}),
+        exclude_none=exclude_none,
+    )
+    head.update(extra or {})
+    adapter = item_adapter(type(trace), "nodes")
+
+    def nodes() -> Iterator[list[bytes]]:
+        last_yield = time.perf_counter()
+        for index, node in enumerate(trace.nodes):
+            # A busy thread keeps the GIL until CPython's switch interval forces a
+            # handoff, so another thread (an event loop) waits that long each time it
+            # wakes. Releasing the GIL every millisecond lets it run at once.
+            if time.perf_counter() - last_yield > YIELD_INTERVAL_S:
+                time.sleep(0)
+                last_yield = time.perf_counter()
+            record = adapter.dump_python(
+                node,
+                mode="json",
+                exclude=item_exclude(exclude, "nodes", index),
+                exclude_none=exclude_none,
+            )
+            record.update(node.record_streams(float_decimals, exclude_none))
+            yield [orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)]
+
+    return splice_json(head, "nodes", nodes())
+
+
 class TimeSpan(BaseModel):
     """Wall-clock timestamps with a derived, non-serialized duration in seconds."""
 
@@ -853,10 +938,17 @@ 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`. Encoded one node at a time, like `Episode.to_json`, so a long trace does
+        not stall other threads."""
+        parts = trace_json_parts(
+            self,
+            self.model_dump,
+            merge_exclude(RECORD_DUMP_EXCLUDE, exclude),
+            float_decimals,
+            exclude_none=False,
         )
+        return b"".join(parts)
 
 
 WireTrace = Trace[WireTaskData, State, WireAgentConfig]