[mikasenghaas]

Summary

• Episode.to_record() and Trace.to_record() return the per-token node streams as C-contiguous numpy arrays. That covers token_ids, mask, is_content, the rounded float streams and loss_weights.
  • Floats are float64, and the rounded streams use np.round.
  • orjson with OPT_SERIALIZE_NUMPY writes each stream in one pass.
  • Before, each stream became a Python list, each float went through round(), and pydantic serialized every element.
  • to_record() dumps the model without the streams, then adds them as arrays. Pydantic's JSON mode cannot emit numpy arrays, and its own encoder (model_dump_json) is 3.5x slower on real episodes (see Verification).
  • The streams now come last in each node record.
• Add Episode.to_json() and Trace.to_json(): to_record() encoded with orjson, as compact UTF-8 bytes. The pair mirrors pydantic's model_dump/model_dump_json, which stay unchanged.
• to_record() and to_json() take exclude, a pydantic exclude spec merged with the record's own excludes. The per-token streams are always included.
• JSON-mode dumps (model_dump(mode="json"), model_dump_json) round with np.round instead of a per-value round(). The output is the same.
• The agentic and best-of-n judge envs write the trace file with to_json(exclude=...), which drops ok and the hidden reasoning fields. Before, they edited the record dict. Add orjson as a dependency.
• vf-eval --push keeps passing the episodes to the prime SDK. The SDK encodes them with orjson from perf: encode json with orjson in prime-runs and prime-traces prime#1109 on.

prime-rl counterpart: PrimeIntellect-ai/prime-rl#4257. In a replayed batch with PrimeIntellect-ai/prime#1109, the ship under monitor contention drops from 1016 s to 114 s.

Pinned until release: PrimeIntellect-ai/prime#1109 is merged but not released yet, so prime-runs (and prime-traces through override-dependencies) are pinned to its merge commit on prime main (adf8d25), with a TODO to replace them with the PyPI release asap and raise the prime-runs minimum. The pins only apply to this repo's lock; consumers need the release.

Breaking

• to_record() node streams are numpy arrays, not lists. json.dumps(record) with the stdlib json raises TypeError.
  • Migration: encode with orjson and OPT_SERIALIZE_NUMPY.
• Prime SDK releases before perf: encode json with orjson in prime-runs and prime-traces prime#1109 cannot encode to_record(), so their uploads fail with TypeError.
  • Migration: upgrade prime-runs.

Verification

Real data: one 2,048-episode training batch with 193M tokens. Each step below ran single-threaded over the whole batch.

Per batch	Before	After
to_record() + orjson (trace-file write)	103.3 s (50.4 ms/episode)	18.1 s (8.8 ms/episode)

A second check ran on 256 episodes of the same batch:

Per episode	Time
to_record() + orjson (this branch)	13.1 ms
model_dump_json with the same excludes	47.9 ms
model_dump(mode="json") + orjson	46.3 ms

• All three paths produce the same parsed content for each episode.
• New tests:
  • the record arrays and to_json() have the content of the plain JSON dump, for traces and episodes, with and without rounding;
  • exclude merges with the record's own excludes and keeps the streams;
  • round_stream matches round() on 100k float32 values.
• The judges' trace files parse to the same content as the old edit-then-encode path, for 32 real episodes.
• tests/unit: all passed (Docker, PRIME_API_KEY and harbor tests skipped).

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes the trace/episode serialization contract (NumPy in records, bytes from to_json) and pins Prime packages from git, which affects eval uploads and any code that assumed list-shaped streams or stdlib JSON.

Overview
Trace and episode JSON records now attach per-token fields (token_ids, mask, is_content, float training streams, loss_weights) as NumPy arrays instead of Python lists. to_record() dumps the model without those streams, then merges them via add_record_streams / MessageNode.record_streams; float streams use np.round (round_stream) at the configured precision.

Adds Trace.to_json() and Episode.to_json() — compact UTF-8 bytes from orjson with OPT_SERIALIZE_NUMPY. Both to_record and to_json accept an optional exclude merged with RECORD_DUMP_EXCLUDE via merge_exclude (judge envs use this to drop ok and hidden reasoning instead of mutating dicts).

Dependencies: runtime orjson, and temporary git pins for prime-runs / prime-traces until a PyPI release ships prime#1109-compatible encoding.

Breaking: to_record() dicts are no longer json.dumps-safe; callers must use to_json or orjson with NumPy support.

^{Reviewed by Cursor Bugbot for commit 4ae941f. Bugbot is set up for automated code reviews on this repo. Configure here.}

Note

Keep record streams as NumPy arrays and encode records with orjson

• Trace.to_record and Episode.to_record now attach node streams (token IDs, mask, content flags, rounded float streams, loss weights) as NumPy arrays, with float streams rounded per float_decimals via round_stream in graph.py
• Trace.to_json and Episode.to_json encode records with orjson's NumPy option; NaN and infinity become null
• to_record accepts an exclude argument merged with built-in record exclusions via merge_exclude in trace.py; JudgeTask.from_trace and RankTask.record now use Trace.to_json with exclusions instead of mutating records and using stdlib JSON
• Adds orjson>=3.10.0 as a dependency in pyproject.toml
• Behavioral Change: JSON output format changes — NaN/inf encode as null, and records contain arrays instead of lists for stream fields; check consumers of to_json/to_record output

^{Macroscope summarized 4ae941f.}

[macroscopeapp]

Approvability

Verdict: Needs Human Approval

Macroscope's review found this PR requires human approval — This changes the default Trace/Episode record representation from Python lists to NumPy arrays and routes production judge snapshots through a new orjson-based serializer. It also pins runtime dependencies to an external Git revision, so the public serialization contract and installation impact need human review.

^{You can add or adjust custom eligibility rules. Learn more.}