perf: keep record streams as numpy arrays

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
     "gepa>=0.0.6",
     "pyzmq>=27.1.0",
     "msgpack>=1.1.2",
+    "orjson>=3.10.0",
     "httpx>=0.27.0",
     "httpx2>=2.12.0",
     "h11>=0.16.0",

tests/unit/test_trace.py

@@ -2,14 +2,17 @@
 
 import json
 
+import numpy as np
+import orjson
 import pytest
 
 import verifiers.v1 as vf
 from verifiers.v1.configs.agent import resolve_agent
 from verifiers.v1.dialects.chat import ChatDialect
 from verifiers.v1.dialects.responses import ResponsesDialect, fold_assistant
-from verifiers.v1.graph import MessageNode, prepare_turn
+from verifiers.v1.graph import MessageNode, prepare_turn, round_stream
 from verifiers.v1.session import RolloutLimits
+from verifiers.v1.trace import EXCLUDE_FIELDS
 from verifiers.v1.types import AssistantMessage, UserMessage
 from verifiers.v1.utils.trace_store import write_episode
 
@@ -73,6 +76,73 @@ def test_bare_trace_round_trip(sampling):
         assert rt.agent.config.sampling.model_dump() == resolved.sampling.model_dump()
 
 
+def _streamed_trace() -> vf.Trace:
+    return vf.Trace(
+        agent=vf.AgentInfo(config=vf.AgentConfig()),
+        task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="q")),
+        nodes=[
+            MessageNode(
+                message=UserMessage(content="q"),
+                token_ids=[1, 2, 3],
+                is_content=[False, True, True],
+            ),
+            MessageNode(
+                parent=0,
+                message=AssistantMessage(content="a"),
+                sampled=True,
+                token_ids=[4, 5, 6],
+                mask=[False, True, True],
+                logprobs=[-0.123456, -1.00005],
+                advantages=[0.333333, 0.333333],
+                loss_weights={"policy": [0.0, 1 / 3, 1 / 3]},
+            ),
+        ],
+    )
+
+
+@pytest.mark.parametrize("float_decimals", [4, None])
+def test_record_streams_are_arrays_with_json_dump_content(float_decimals):
+    tr = _streamed_trace()
+    episode = vf.Episode(task=tr.task, traces=[tr], ok=True)
+    context = {"float_decimals": float_decimals}
+    cases = [
+        (
+            tr.to_record(float_decimals),
+            tr.model_dump(mode="json", exclude=EXCLUDE_FIELDS, context=context),
+        ),
+        (
+            episode.to_record(float_decimals),
+            episode.model_dump(
+                mode="json",
+                exclude={"traces": {"__all__": EXCLUDE_FIELDS}},
+                exclude_none=True,
+                context=context,
+            ),
+        ),
+    ]
+    for record, json_dump in cases:
+        nodes = (record["traces"][0] if "traces" in record else record)["nodes"]
+        assert isinstance(nodes[1]["token_ids"], np.ndarray)
+        assert nodes[1]["logprobs"].dtype == np.float64
+        assert isinstance(nodes[1]["loss_weights"]["policy"], np.ndarray)
+        assert orjson.loads(
+            orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)
+        ) == orjson.loads(json.dumps(json_dump))
+    assert orjson.loads(tr.to_json(float_decimals)) == orjson.loads(
+        json.dumps(cases[0][1])
+    )
+    assert orjson.loads(episode.to_json(float_decimals)) == orjson.loads(
+        json.dumps(cases[1][1])
+    )
+
+
+def test_round_stream_matches_round():
+    values = np.random.default_rng(0).standard_normal(100_000).astype(np.float32) * 10
+    values = np.concatenate([values, np.float32([0.00005, -0.00015, 1.23445, 2.5e-5])])
+    expected = [round(value, 4) for value in values.tolist()]
+    assert round_stream(values, 4).tolist() == expected
+
+
 def test_custom_task_state_round_trip(tmp_path):
     # Custom data and state round-trip into the same parameterization. Data fields are
     # typed (not just `model_extra`); `state` is runtime-only and never crosses the wire.

verifiers/v1/envs/agentic_judge/env.py

@@ -6,10 +6,10 @@
 """
 
 import copy
-import json
 import re
 from pathlib import Path
 
+import orjson
 from pydantic import FiniteFloat
 
 import verifiers.v1 as vf
@@ -147,7 +147,7 @@ def from_trace(
                 message = node["message"]
                 message.pop("reasoning_content", None)
                 message.pop("provider_state", None)
-        files = {TRACE_FILE: json.dumps(record).encode()}
+        files = {TRACE_FILE: orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)}
         template = config.build_prompt()
         body = _render(template, prompt=solved.prompt_text)
         if "{prompt}" not in template:

verifiers/v1/envs/best_of_n_judge/env.py

@@ -9,12 +9,12 @@
 """
 
 import asyncio
-import json
 import random
 import shlex
 import statistics
 import uuid
 
+import orjson
 from pydantic import BaseModel, ConfigDict, Field
 
 import verifiers.v1 as vf
@@ -136,7 +136,7 @@ def record(attempt: vf.Trace) -> bytes:
         for node in record["nodes"]:
             node["message"].pop("reasoning_content", None)
             node["message"].pop("provider_state", None)
-        return json.dumps(record).encode()
+        return orjson.dumps(record, option=orjson.OPT_SERIALIZE_NUMPY)
 
     async def setup(self, trace: vf.Trace, runtime: vf.Runtime) -> None:
         # Archive bytes are the attempts': confine every member to its declared

verifiers/v1/episode.py

@@ -3,13 +3,21 @@
 import uuid
 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
 from verifiers.v1.graph import RECORD_FLOAT_DECIMALS
 from verifiers.v1.state import State, StateT
 from verifiers.v1.task import DataT, WireTaskData
-from verifiers.v1.trace import EXCLUDE_FIELDS, AgentConfigT, Error, Trace, TraceTask
+from verifiers.v1.trace import (
+    RECORD_DUMP_EXCLUDE,
+    AgentConfigT,
+    Error,
+    Trace,
+    TraceTask,
+    add_record_streams,
+)
 from verifiers.v1.types import Usage
 
 
@@ -157,12 +165,22 @@ def to_record(
         self, float_decimals: int | None = RECORD_FLOAT_DECIMALS
     ) -> dict[str, Any]:
         """JSON record without raw trace tensors, which remain on the msgpack wire.
-        Per-token float streams are rounded to `float_decimals` (`None` keeps every digit)."""
-        return self.model_dump(
+        Per-token node streams are numpy arrays (`graph.RECORD_STREAMS`): encode it with
+        `to_json`, or with orjson and `OPT_SERIALIZE_NUMPY`. Per-token float streams are
+        rounded to `float_decimals` (`None` keeps every digit)."""
+        record = self.model_dump(
             mode="json",
-            exclude={"traces": {"__all__": EXCLUDE_FIELDS}},
+            exclude={"traces": {"__all__": RECORD_DUMP_EXCLUDE}},
             exclude_none=True,
-            context={"float_decimals": float_decimals},
+        )
+        for trace_record, trace in zip(record["traces"], self.traces, strict=True):
+            add_record_streams(trace_record, trace, float_decimals, exclude_none=True)
+        return record
+
+    def to_json(self, float_decimals: int | None = RECORD_FLOAT_DECIMALS) -> bytes:
+        """`to_record` as compact UTF-8 JSON. NaN and infinity become `null`."""
+        return orjson.dumps(
+            self.to_record(float_decimals), option=orjson.OPT_SERIALIZE_NUMPY
         )
 
 

verifiers/v1/graph.py

@@ -88,19 +88,40 @@ def same_tokens(a: list | np.ndarray, b: list | np.ndarray) -> bool:
 """Float streams a JSON record rounds to `float_decimals`. Loss weights stay exact: a 1/T
 weight would round to zero."""
 
+RECORD_STREAMS = ("token_ids", "mask", "is_content", *ROUNDED_STREAMS, "loss_weights")
+"""Per-token node fields a JSON record (`to_record`) carries as numpy arrays — `loss_weights`
+as a dict of them — so an encoder writes each stream in one pass instead of element by
+element."""
+
+
+def round_stream(values: list | np.ndarray, decimals: int) -> np.ndarray:
+    """A float stream rounded to `decimals`, in float64."""
+    return np.round(np.asarray(values, dtype=np.float64), decimals)
+
+
+def _record_array(values: list | np.ndarray, decimals: int | None) -> np.ndarray:
+    """A stream as a record array: C-contiguous, which numpy-aware encoders require. Floats
+    widen to float64 so they print the digits a list of Python floats does — float32 would
+    print its own shortest form."""
+    array = np.asarray(values)
+    if array.dtype.kind == "f":
+        if decimals is not None:
+            return round_stream(array, decimals)
+        return array.astype(np.float64)
+    return np.ascontiguousarray(array)
+
 
 def _serialize_stream(
     values: list | np.ndarray, info: FieldSerializationInfo, rounded: bool
 ) -> list | dict:
-    """JSON records: a list. Python dumps: a list too, unless the context sets `token_arrays`,
+    """JSON dumps: a list. Python dumps: a list too, unless the context sets `token_arrays`,
     when the msgpack wire carries a raw-bytes `__nd__` dict."""
     context = info.context or {}
     if info.mode_is_json():
-        values = as_list(values)
         decimals = context.get("float_decimals", RECORD_FLOAT_DECIMALS)
         if not rounded or decimals is None:
-            return values
-        return [round(value, decimals) for value in values]
+            return as_list(values)
+        return round_stream(values, decimals).tolist()
     if context.get("token_arrays"):
         return encode_ndarray(np.asarray(values))
     return as_list(values)
@@ -207,6 +228,26 @@ def logical_ids(self) -> np.ndarray:
             else self.renderer_token_ids
         )
 
+    def record_streams(
+        self, float_decimals: int | None, exclude_none: bool
+    ) -> dict[str, Any]:
+        """This node's `RECORD_STREAMS` for its JSON record, as numpy arrays. Float streams
+        in `ROUNDED_STREAMS` are rounded to `float_decimals` (`None` keeps every digit)."""
+        streams: dict[str, Any] = {}
+        for name in RECORD_STREAMS:
+            values = getattr(self, name)
+            if values is None:
+                if not exclude_none:
+                    streams[name] = None
+            elif name == "loss_weights":
+                streams[name] = {
+                    key: _record_array(stream, None) for key, stream in values.items()
+                }
+            else:
+                decimals = float_decimals if name in ROUNDED_STREAMS else None
+                streams[name] = _record_array(values, decimals)
+        return streams
+
     @field_serializer("token_ids", "mask", "is_content", *ROUNDED_STREAMS)
     def serialize_token_stream(
         self, values: list | np.ndarray | None, info: FieldSerializationInfo

verifiers/v1/trace.py

@@ -9,6 +9,7 @@
 from typing import TYPE_CHECKING, Any, Generic
 
 import numpy as np
+import orjson
 from pydantic import BaseModel, Field, PrivateAttr, computed_field, field_serializer
 from typing_extensions import TypeVar
 
@@ -52,6 +53,25 @@
 }
 """Raw tensor fields kept on the msgpack wire but excluded from disk serialization."""
 
+RECORD_DUMP_EXCLUDE: dict = {
+    "nodes": {"__all__": EXCLUDE_FIELDS["nodes"]["__all__"] | set(graph.RECORD_STREAMS)}
+}
+"""What `to_record` leaves out of its pydantic dump: the wire-only tensors, and the per-token
+streams, which `add_record_streams` adds as numpy arrays. Pydantic's JSON mode cannot emit
+numpy arrays, and building the streams as lists of Python numbers would be most of the cost
+of a record."""
+
+
+def add_record_streams(
+    record: dict[str, Any],
+    trace: Trace,
+    float_decimals: int | None,
+    exclude_none: bool,
+) -> None:
+    """Join a trace's record (dumped with `RECORD_DUMP_EXCLUDE`) with its node streams."""
+    for node_record, node in zip(record["nodes"], trace.nodes, strict=True):
+        node_record.update(node.record_streams(float_decimals, exclude_none))
+
 
 class TimeSpan(BaseModel):
     """Wall-clock timestamps with a derived, non-serialized duration in seconds."""
@@ -799,11 +819,17 @@ def to_record(
         self, float_decimals: int | None = RECORD_FLOAT_DECIMALS
     ) -> dict[str, Any]:
         """JSON record without raw tensors, which remain available on the msgpack wire.
-        Per-token float streams are rounded to `float_decimals` (`None` keeps every digit)."""
-        return self.model_dump(
-            mode="json",
-            exclude=EXCLUDE_FIELDS,
-            context={"float_decimals": float_decimals},
+        Per-token node streams are numpy arrays (`graph.RECORD_STREAMS`): encode it with
+        `to_json`, or with orjson and `OPT_SERIALIZE_NUMPY`. Per-token float streams are
+        rounded to `float_decimals` (`None` keeps every digit)."""
+        record = self.model_dump(mode="json", exclude=RECORD_DUMP_EXCLUDE)
+        add_record_streams(record, self, float_decimals, exclude_none=False)
+        return record
+
+    def to_json(self, float_decimals: int | None = RECORD_FLOAT_DECIMALS) -> bytes:
+        """`to_record` as compact UTF-8 JSON. NaN and infinity become `null`."""
+        return orjson.dumps(
+            self.to_record(float_decimals), option=orjson.OPT_SERIALIZE_NUMPY
         )