Repository navigation
feat: encode episodes to JSON one node at a time and record runs on traces - #3042
Conversation
Episode.to_json and Trace.to_json made one pydantic dump and one orjson call over the whole record. Both hold the GIL for the whole call, so encoding a large episode stalled every other thread of the process, such as an asyncio event loop, for as long as it took. They now encode the record one node at a time and join the parts once at the end, so the stall no longer grows with the record. The output is byte for byte the same as orjson.dumps(to_record()). Episode.to_json also takes a run: its fields are set on the episode's run, and that run goes on every member trace too, as Prime Traces reads it per trace. This lets the Prime SDK upload the encoded bytes as they are. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| spec = exclude.get(field) | ||
| if not isinstance(spec, dict): | ||
| return None | ||
| return merge_exclude(spec.get("__all__"), spec.get(index)) or None |
There was a problem hiding this comment.
🟢 Low v1/trace.py:118
Negative-index exclusions are ignored, so trace.to_json(exclude={"nodes": {-1: {"message"}}}) serializes the last node's message instead of excluding it; the same affects Episode's traces exclusions. item_exclude receives nonnegative enumerate indices and looks up only that index, so normalize negative indices against the list length before combining the per-item exclusions.
🚀 Reply "fix it for me" or copy this AI Prompt for your agent:
In file @verifiers/v1/trace.py around line 118:
Negative-index exclusions are ignored, so `trace.to_json(exclude={"nodes": {-1: {"message"}}})` serializes the last node's `message` instead of excluding it; the same affects `Episode`'s `traces` exclusions. `item_exclude` receives nonnegative `enumerate` indices and looks up only that index, so normalize negative indices against the list length before combining the per-item exclusions.
ApprovabilityVerdict: Approved at Macroscope's review found this PR approvable — JSON encoding now proceeds node by node to reduce thread stalls, and traces receive their episode’s run metadata. The reported negative-index exclusion limitation is below the blocking threshold, with no demonstrated consequential use in the inspected callers. No code changes detected at You can add or adjust custom eligibility rules. Learn more. |
Short calls alone did not keep a busy event loop responsive. A thread that keeps working keeps the GIL until CPython's switch interval (5 ms) forces a handoff, and the event loop needs the GIL several times each time it wakes, so its stalls added up to several intervals. to_json now releases the GIL after each millisecond of encoding. Releasing it after every node instead made encoding about 3x slower. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Trace stores such as Prime Traces key each trace by its own run, and derive an episode's run from its members. Trace gets a run field, and Episode.record_run sets it on the episode and on each of its traces, so to_record and to_json carry it without a run parameter. The run types move from episode.py to trace.py, which Trace needs; the package exports them as before. The delta stream sends run with the header. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* gzip trace uploads at level 1 The upload thread spent ~70 ms of its ~78 ms per 2 MB episode in gzip at the default level 9, so it fell behind a fast producer and the upload queue dropped episodes. Level 1 is ~9x faster for ~13% more bytes on the wire. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * send up to eight trace upload requests at once upload_lines sent one batch at a time, so network time and gzip ran in series with encoding. It now keeps up to max_in_flight (default 8, the most the service contract allows per producer) batches in flight on a thread pool. gzip and hashing release the GIL, so the upload rate is bound by record encoding alone. on_batch still runs on the calling thread, and the receipts come back in batch order. A failed batch stops new requests, and its error is raised once the requests in flight have finished. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * never block the producer on a full upload queue The upload queue held 256 submit calls, and a full queue blocked the producer for up to 5 s before it dropped the batch. An async caller such as the prime-rl orchestrator waited that long on its arrival path. The queue is now bounded by records (max_queued_records, default 10,000), and submit never blocks: a batch that does not fit is dropped and counted at once. A batch always fits an empty queue. The traces sink also prepared every merged record as a wire mapping before the upload started, so a backlog of large episodes was held twice in memory. It now prepares each record as the client encodes it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * upload a record's own json encoding when it has one TracesSink built each record's wire mapping with to_record() and stamp_run, and prime-traces encoded it with orjson. Both calls hold the GIL for the whole record, so a large episode stalled the producer's other threads, such as an asyncio event loop. A record whose to_json takes a run (verifiers Episode, from PrimeIntellect-ai/verifiers#3042) now encodes itself with the run stamped on, in short steps, and the sink uploads those bytes as they are. The line is byte for byte what the sink built before. A record whose to_json takes no run, as in released verifiers, keeps the old path. prime-traces accepts an encoded line (bytes) as a record. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * upload to_json only for records that name this run verifiers now records the run on an episode and each of its traces (Episode.record_run), so its to_json() already carries what the sink stamped. The sink uploads to_json() as it is when the record and every member trace name this run, id and type. Any other record is stamped as before: a producer that recorded a local run id, or an episode from a verifiers whose traces carry no run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Cooper Miller <kcoopermiller9@gmail.com>
…node # Conflicts: # verifiers/v1/cli/eval/runner.py
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit dabe65d. Configure here.
## Summary Prime uploads keep up with the orchestrator, and logging to Prime never blocks it. Companion to PrimeIntellect-ai/prime#1136. - **SDK pin.** `prime-runs[train]>=0.1.6`, which pulls `prime-traces` 0.0.12 from PyPI. These releases contain PrimeIntellect-ai/prime#1136: gzip level 1, up to 8 upload requests in flight, an upload queue bounded by records whose `submit` never blocks, and lazily built upload records. - `uv.lock` is written with uv 0.12.0 (lockfile revision 3); only the two SDK packages change. - **No thread hop for Prime.** `PrimeMonitor` calls `run.log_metrics` and `run.log_episodes` directly on the event loop, not through `asyncio.to_thread`. A queue put that never blocks needs no thread, and each hop had to win the GIL back from the uploader thread. - **A production-sized load test.** `tests/unit/monitors/test_monitor_load.py`: - **Episode sizes:** episodes draw their size from a lognormal fitted to a production RL run (median 6.4 MB, p90 17 MB; it also matches the mean of 8.5 MB and the p99 of 35 MB), capped at 60 MB. 1,000 episodes is ~8 GB of JSON, close to one production step of 1,024 samples. - **Memory:** episodes share 24 traces at log-spaced sizes, so the test's own memory stays small. - **Fake endpoint:** the fake Traces endpoint gzips at the SDK's level and counts lines under a lock, because the client now sends several requests at once. - **File monitor check:** it streams every trace chunk, since 8 GB spans two 5 GiB chunks. - **No expected failure:** the prime monitor at 1,000 episodes is no longer marked as one. - **A separate `log()` limit of 1 s.** A `log()` call may await the monitor's own work: the file monitor awaits its write, which encodes the episode on a worker thread, while the event loop keeps serving other tasks. The lag limits still cover the loop. - **The p99 lag limit rises from 50 to 100 ms.** At production sizes the file monitor encodes each episode in one call, which stalls the event loop with a p99 of ~60 ms on CI. Main's code shows the same stall with this test. #4366 makes the encoder work in short steps and tightens the limits to 25 ms p99 and 100 ms max. ## Verification **Headline: Prime monitor, 1,000 production-sized episodes (~8 GB of JSON), main against this branch.** Main ran this branch's test, with the fake endpoint gzipping at main's level 9. | | main | This branch | |---|---|---| | Episodes delivered | 725 of 1,000 | **1,000 of 1,000** | | Lost to a full upload queue | 27 | 0 | | Lost when the 60 s drain ran out | 248 | 0 | | Slowest `log()` | up to 5 s (blocked on the full queue) | **0.6 ms** | | Event-loop lag p99 / max | not reached (failed on delivery) | 14 / 17 ms | | Test time | 217 s | 28 s | **Scale: 10,000 episodes of 2.2 MB (~22 GB)**, from a local copy of the test with the count raised: | Setup | Result | |---|---| | No network time | 0 dropped; `log()` max 1.4 ms; event-loop lag p99 14 ms, max 21 ms; 570 MB peak RSS | | 100 ms + 100 MB/s per request, 1 request in flight | 2,595 delivered within the 60 s drain | | Same, 4 in flight | 9,822 delivered | | Same, 8 in flight (the default) | **10,000 delivered**; event-loop lag p99 12 ms; 860 MB peak RSS | The upload rate is ~180 episodes/s (~400 MB/s of JSON). Encoding the records is the limit. **Full load-test grid at production sizes, this branch (dev box):** 9 passed. | 1,000 episodes | Event-loop lag p99 / max | `log()` max | |---|---|---| | prime | 14 / 17 ms | 0.6 ms | | file | 47 / 93 ms | 177 ms | | wandb | 0.9 / 0.9 ms | 0.6 ms | **CI:** on the CI runner, the file monitor at 1,000 episodes measured event-loop lag p99 60 ms and max 157 ms, within this PR's limits (100 / 250 ms), and `log()` up to 299 ms (limit 1 s). With #4366 the same case measures p99 6 ms and max 12 ms. Other checks: - `pytest tests/unit/train/sft/test_prime_monitor.py tests/unit/utils/test_links.py tests/unit/orchestrator/test_episode_records.py tests/unit/train/rl -m "not gpu"`: 29 passed. - `uvx uv@0.12.0 lock --check` passes. ## Merge order 1. ~~PrimeIntellect-ai/prime#1136~~ and ~~PrimeIntellect-ai/prime#1140~~: merged and released as `prime-runs` 0.1.6 and `prime-traces` 0.0.12. 2. **#4363 (this PR)**: ready to merge. It does not need PrimeIntellect-ai/verifiers#3042. 3. PrimeIntellect-ai/verifiers#3042 4. #4366 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes the hot path for RL training telemetry (episode/metrics enqueue) and bumps external SDK versions; misbehavior could drop or delay uploads under load, though load tests now enforce delivery and loop responsiveness. > > **Overview** > **Prime logging no longer blocks the orchestrator event loop** by pinning `prime-runs[train]>=0.1.6` (and `prime-traces` 0.0.12 in the lockfile), which uses a non-blocking upload queue, concurrent uploads, and lighter gzip. > > **`PrimeMonitor`** now calls `run.log_metrics` and `run.log_episodes` directly instead of `asyncio.to_thread`, matching SDK behavior where enqueue never blocks. > > The **monitor load test** is reworked to use **production-like episode sizes** (lognormal ~6.4 MB median), shared trace templates for memory, a thread-safe fake Traces API, and multi-chunk file delivery checks. The **Prime monitor at 1,000 episodes is no longer expected to fail**; responsiveness asserts split event-loop lag (p99 100 ms) from per-`log()` duration (max 1 s). > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 11ae629. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
…ime (#4366) ## Summary Companion to PrimeIntellect-ai/verifiers#3042 and PrimeIntellect-ai/prime#1140 (released in `prime-runs` 0.1.6). The file monitor and the Prime monitor encode each episode in short steps, so a large episode no longer stalls the orchestrator's event loop for as long as encoding it takes. - **Pins.** - The verifiers submodule moves to verifiers main, which includes verifiers#3042: `to_json` encodes one node at a time and releases the GIL every millisecond, and `record_run` records the run on every trace. - `prime-runs` 0.1.6, already on main, contains prime#1140: the SDK uploads `episode.to_json()` as it is when the episode and its traces name the upload's run. The dispatcher's `record_run` makes that true for RL episodes. - Verifiers main also includes PrimeIntellect-ai/verifiers#3046, which removes `vf-eval`'s platform push and verifiers' `prime-runs` dependency. In `uv.lock` (uv 0.12.0), only the verifiers metadata changes. - **The file monitor writes `episode.to_json()`.** - The stream line is byte for byte what `orjson.dumps(episode.to_record())` gave before. - The index row comes from `index_record(episode)`. It dumps everything but the nodes' payloads, one trace at a time, and gives the same row. - `ChunkedJsonl.append` takes a line in parts, so the newline is written without copying a long line. - **The load test's lag limits drop** from p99 100 ms and max 250 ms to 25 ms and 100 ms, so a regression shows. Across runs at production sizes, the highest values were a p99 of 12.5 ms and a max of 22 ms. - **A hosted orchestrator stamps the platform's run id.** A hosted launch sets only `$RUN_ID`, not `$PRL_RUN_ID`, so the orchestrator stamped episodes with a random id, and the SDK re-stamped every upload. It now falls back to `$RUN_ID`, the Prime run id. ## Verification **Headline: event-loop lag at 1,000 production-sized episodes (~8 GB of JSON)**, monitor load test, dev box: | | Main (before) | This PR | |---|---|---| | prime: lag p99 / max | 14 / 17 ms | **3.1 / 3.5 ms** | | file: lag p99 / max | 47 / 93 ms | **4.8 / 15 ms** | | file: `log()` max | 177 ms | 195 ms | - **A second run:** inside the full unit suite, with more load on the box, it measured prime 11 / 18 ms, file 5.2 / 22 ms, and file `log()` max 224 ms. - **Why the GIL release matters:** encoding node by node without it left the file monitor at 15–18 / 31–66 ms, varying from run to run. A busy encoder thread keeps the GIL until CPython's 5 ms switch interval forces a handoff, and the event loop waits that out several times per wakeup. **The stall no longer grows with the episode.** One thread encodes episodes of 10 to 10,000 turns (1,000 tokens per turn) while a 5 ms ticker on the event loop records the worst stall. Medians of 5 runs, under the orchestrator's GC policy: | Turns | Episode | Encode CPU, before / after | Worst stall, before | Worst stall, after | |---|---|---|---|---| | 10 | 0.2 MB | 0.4 / 0.5 ms | 5.5 ms | 5.2 ms | | 100 | 2.2 MB | 4.0 / 4.5 ms | 8.0 ms | 1.6 ms | | 1,000 | 22 MB | 50 / 51 ms | 40 ms | 3.1 ms | | 3,000 | 65 MB | 152 / 167 ms | 100 ms | 5.2 ms | | 10,000 | 217 MB | 508 / 571 ms | 326 ms | 15 ms | - **Before:** the stall grows linearly with the episode. - **After:** it stays around 2–5 ms up to 3,000 turns, about the GIL switch interval. The 10-turn case is that interval alone, since each episode encodes in under a millisecond. - **The rest at 10,000 turns:** the trace header's pydantic dump, which holds every model call (one per turn). Alone it stalls 4.4 / 8.4 / 15 ms at 1,000 / 3,000 / 10,000 turns. Agentic episodes in production have a few hundred turns, so this part stays small. Encoding the calls one at a time, like the nodes, would remove it. - **CPU:** encoding costs up to ~12% more, from the extra calls and GIL releases. **Encoding:** `test_received_episode_encodes_as_its_record_did` sends a two-trace episode over the real delta wire into a `WireEpisode`. Then: - the file monitor's stream line equals `orjson.dumps(to_record())`, byte for byte; - its index row equals the row from the full record; - after `record_run`, the SDK takes the episode's own `to_json()`, which carries the same record as the old stamped line. **Tests:** `pytest tests/unit -m "not gpu"`: 1,077 passed, 7 failed. The same 7 fail on main on this dev box: an attention kernel without sm_120 support, and no ninja for the CPU Adam extension. **CI:** on the CI runner, the file monitor at 1,000 production-sized episodes measured event-loop lag p99 6.0 ms and max 12 ms, against 63 / 152 ms on main. Its slowest `log()` was 382 ms, under the 1 s limit. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes affect how every episode is persisted and indexed and how run identity is chosen on hosted launches; behavior is contract-tested but touches high-volume I/O on the training path. > > **Overview** > Large episodes no longer block the orchestrator event loop during trace logging. The **file monitor** writes stream lines with **`episode.to_json()`** (incremental, GIL-friendly encoding from verifiers) instead of a single **`orjson.dumps(to_record())`**, and builds index rows via new **`index_record()`** so summarization skips heavy node payloads without materializing the full record first. **`ChunkedJsonl.append`** accepts **byte parts** so the newline is appended without copying megabyte-scale lines. > > The orchestrator **`run_id`** now falls back to **`$RUN_ID`** when **`$PRL_RUN_ID`** is unset, so hosted runs align with the platform id and Prime uploads are not re-stamped with a random id. > > Monitor load tests use **stricter event-loop lag thresholds** to catch regressions. A new unit test asserts on-disk bytes and SDK upload JSON still match the legacy **`to_record()`** encoding. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit f83e849. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Summary
Episode.to_jsonandTrace.to_jsonencode one node at a time.time.sleep(0)). Without that release, a busy thread keeps the GIL until CPython's 5 ms switch interval forces a handoff, and an event loop waits that long several times per wakeup. Releasing after every node instead made encoding ~3x slower.orjson.dumps(to_record(), option=OPT_SERIALIZE_NUMPY). Each trace and node is dumped as its declared field type, as the parent's dump does. The parts are joined once at the end, so the record is copied once.Trace.run.Tracegets arunfield, andEpisode.record_run(run)sets it on the episode and on each of its traces.trace.run.id) and derives an episode's run from its members. Soto_record()andto_json()now carry what the Prime SDK used to stamp, and the SDK can uploadto_json()as it is (feat: upload an episode's own to_json when it names the run prime#1140).PolicySpan,TrainWorkInfo,EvalWorkInfo,WorkInfo,TrainRunInfo,EvalRunInfo,RunInfo) move fromepisode.pytotrace.py, whichTraceneeds.verifiers.v1exports them as before.runwith the trace header. Env servers sendNone; the consumer stamps the run when the episode arrives.Verification
pytest tests/unit: 245 passed, 10 skipped (no harbor, noPRIME_API_KEY).to_jsonequalsorjson.dumps(to_record())byte for byte: for a trace and an episode, with and without float rounding, with nestedexcludespecs (__all__plus a per-trace index), and with aTaskDatasubclass.record_runputs the run on the episode and every trace,to_jsoncarries it, and the record validates back to the same run.Worst event-loop stall while another thread encodes, and CPU per episode. The episodes are synthetic agentic episodes at the sizes of a production RL run:
End-to-end numbers are in feat: encode episodes for the file and prime monitors one node at a time prime-rl#4366, which writes and uploads episodes with this encoder under prime-rl's monitor load test.
Merge order
feat: non-blocking, concurrent run uploads in prime-runs and prime-traces prime#1136andfeat: upload an episode's own to_json when it names the run prime#1140: merged and released asprime-runs0.1.6 andprime-traces0.0.12.🤖 Generated with Claude Code
Note
Medium Risk
Serialization logic is substantially rewritten (though contract-tested); wire JSON gains optional
trace.run, and code that imported run types fromverifiers.v1.episodemust import fromtraceorvfinstead.Overview
Trace.to_jsonandEpisode.to_jsonno longer dump the full record in one pydantic/orjson call. They build JSON in pieces (splice_json,trace_json_parts), encoding each message node separately and yielding the GIL about every millisecond so large rollouts don’t block asyncio and other threads. Output stays byte-identical toorjson.dumps(to_record(), OPT_SERIALIZE_NUMPY), including nestedexcludespecs and custom task types.Run metadata moves onto each trace.
Tracegains optionalrun, andEpisode.record_runnow stamps the sameRunInfoon the episode and every trace intraces. The run-related models (PolicySpan,TrainRunInfo,EvalRunInfo, etc.) live intrace.pyinstead ofepisode.py;import verifiers.v1 as vfstill exposes them. The live trace delta stream sendsrunin the trace header (oftenNoneuntil the consumer stamps it).Tests assert
to_json↔to_recordparity and that stamped runs round-trip through JSON validation.Reviewed by Cursor Bugbot for commit dabe65d. Bugbot is set up for automated code reviews on this repo. Configure here.