Repository navigation
feat: upload an episode's own to_json when it names the run - #1140
Conversation
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>
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>
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>
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>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
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>
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.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit f936a4c. Configure here.
The base branch was changed.
Main has #1136 as a squash of this branch's base, so sinks/traces.py conflicted only on that duplicate change. Main's copy equals #1136's head, so this keeps the branch's version of the file. The tree differs from main by this PR's own change alone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
## 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>
…races (#3042) ## Summary - **`Episode.to_json` and `Trace.to_json` encode one node at a time.** - Pydantic's dump and orjson each hold the GIL for a whole call. Before, each method made one call over the whole record, so encoding a large episode stalled every other thread of the process, such as an asyncio event loop, for as long as it took. - Each call now covers one node, and the encoder releases the GIL after each millisecond of encoding (`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. - The output is byte for byte `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`.** `Trace` gets a `run` field, and `Episode.record_run(run)` sets it on the episode and on each of its traces. - Prime Traces reads the run from each trace (`trace.run.id`) and derives an episode's run from its members. So `to_record()` and `to_json()` now carry what the Prime SDK used to stamp, and the SDK can upload `to_json()` as it is (PrimeIntellect-ai/prime#1140). - The run types (`PolicySpan`, `TrainWorkInfo`, `EvalWorkInfo`, `WorkInfo`, `TrainRunInfo`, `EvalRunInfo`, `RunInfo`) move from `episode.py` to `trace.py`, which `Trace` needs. `verifiers.v1` exports them as before. - The delta stream sends `run` with the trace header. Env servers send `None`; the consumer stamps the run when the episode arrives. ## Verification - `pytest tests/unit`: 245 passed, 10 skipped (no harbor, no `PRIME_API_KEY`). - `to_json` equals `orjson.dumps(to_record())` byte for byte: for a trace and an episode, with and without float rounding, with nested `exclude` specs (`__all__` plus a per-trace index), and with a `TaskData` subclass. - `record_run` puts the run on the episode and every trace, `to_json` carries 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: | Episode | Before: CPU / worst stall | After: CPU / worst stall | |---|---|---| | 8.4 MB (production mean) | 20 ms / 15 ms | 21 ms / 2.3 ms | | 35 MB (production p99) | 81 ms / 55 ms | 91 ms / 5.2 ms | | 133 MB (production max) | 319 ms / 191 ms | 359 ms / 12 ms | - End-to-end numbers are in PrimeIntellect-ai/prime-rl#4366, which writes and uploads episodes with this encoder under prime-rl's monitor load test. ## 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. **#3042 (this PR)** and PrimeIntellect-ai/prime-rl#4363, in either order. 3. PrimeIntellect-ai/prime-rl#4366 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Serialization logic is substantially rewritten (though contract-tested); wire JSON gains optional `trace.run`, and code that imported run types from `verifiers.v1.episode` must import from `trace` or `vf` instead. > > **Overview** > **`Trace.to_json` and `Episode.to_json`** no 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** to `orjson.dumps(to_record(), OPT_SERIALIZE_NUMPY)`, including nested `exclude` specs and custom task types. > > **Run metadata moves onto each trace.** `Trace` gains optional `run`, and **`Episode.record_run`** now stamps the same `RunInfo` on the episode and every trace in `traces`. The run-related models (`PolicySpan`, `TrainRunInfo`, `EvalRunInfo`, etc.) live in **`trace.py`** instead of `episode.py`; **`import verifiers.v1 as vf`** still exposes them. The live trace delta stream sends **`run` in the trace header** (often `None` until the consumer stamps it). > > Tests assert `to_json` ↔ `to_record` parity and that stamped runs round-trip through JSON validation. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit dabe65d. 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
Stacked on #1136. Needs PrimeIntellect-ai/verifiers#3042 to take effect.
TracesSinkuploads a record's ownto_json()as it is, when the record and every member trace already name this run (same id and type).Episode.record_run), and encodesto_json()in short steps. Before, the sink built each record withto_record(), stamped the run on it, and encoded it with orjson. Both calls hold the GIL for the whole episode, so a large episode stalled the producer's other threads.runsits at its field position in each trace), so content-addressed upload IDs change.upload_recordsaccepts an encoded line. Abytesrecord is used as is, with a newline added if it has none, in the sync and async clients.Verification
pytest packages/prime-runs/tests: 300 passed, 5 skipped.pytest packages/prime-traces/tests: 187 passed.episodes_writtenonly when it has traces.Merge order
🤖 Generated with Claude Code
Note
Medium Risk
Changes the trace ingestion encoding path and can alter content-addressed upload hashes when key order differs from the stamped mapping path, but fallbacks preserve behavior for records that don’t pre-name the run.
Overview
TracesSink can skip the old
to_record()→ stamp → orjson path when an episode (and every member trace) already carries the active run’s id and type and exposesto_json(). In that case it forwards the producer’s pre-encoded bytes, so large verifiers episodes can serialize incrementally instead of blocking other threads on a single GIL-heavy prepare step. Anything that doesn’t fully name this run—wrong id, traces without run metadata, or noto_json()—still goes throughstamp_runas before, so older producers stay safe.prime-traces
upload_records(sync and async) now treatsbytesas an already-serialized JSONL line: pass through unchanged and append a newline when missing, mixed with mappings andto_record()objects on the same upload.Tests cover the self-encoding fast path, fallback stamping, and byte passthrough on the wire.
Reviewed by Cursor Bugbot for commit 36ec2c6. Bugbot is set up for automated code reviews on this repo. Configure here.