feat: non-blocking, concurrent run uploads in prime-runs and prime-traces by mikasenghaas · Pull Request #1136 · PrimeIntellect-ai/prime

mikasenghaas · 2026-11-02T18:25:09Z

Summary

Uploads from prime-runs keep up with a fast producer, and logging never blocks it. Part of a stack: merge order below.

gzip level 1. prime_traces.core.client.GZIP_LEVEL = 1, in the sync and async clients. The default level 9 took ~70 ms of the ~85 ms the uploader spent per 2.2 MB episode; level 1 is ~9x faster and sends ~13% more bytes. This takes over fix: gzip trace uploads at level 1 #1129.
Concurrent requests. TracesClient.upload_lines, upload_records and upload_file keep up to max_in_flight batches in flight (default 8, the most the service contract allows per producer). gzip and hashing release the GIL, so encoding the records is the only limit.
- on_batch still runs on the calling thread, and receipts come back in batch order.
- A failed batch stops new requests. The error is raised once the requests in flight have finished.
UploadWorker.submit never blocks. The queue is bounded by records (max_queued_records, default 10,000), not by submit calls. A batch that does not fit is dropped and counted at once; a batch always fits an empty queue. Before, a full queue blocked the producer for up to 5 s, then dropped.
TracesSink prepares records lazily. It builds each record's wire mapping as the client encodes it, so a backlog of large episodes is not held twice in memory.

Breaking

UploadWorker(max_queue_size=..., put_timeout=...) becomes UploadWorker(max_queued_records=...). DEFAULT_QUEUE_SIZE and DEFAULT_PUT_TIMEOUT are removed. Run and init() are unchanged, so only direct users of UploadWorker are affected.

Verification

pytest packages/prime-runs/tests: 297 passed, 5 skipped. pytest packages/prime-traces/tests: 186 passed.
- Three requests are in flight at once: a barrier that only opens with three concurrent requests. Receipts come back in batch order, and on_batch runs on the calling thread.
- A rejected batch stops new requests, and the batch that committed before it is still reported.
- A full queue drops at once instead of blocking, and written records free their room. A batch larger than the bound fits an empty queue.
End-to-end numbers are in feat: non-blocking, concurrent prime uploads prime-rl#4363, which runs prime-rl's monitor load test against this branch.

Merge order

🤖 Generated with Claude Code

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>

chatgpt-codex-connector · 2026-11-02T18:25:21Z

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-02T18:27:40.421073Z	`d6ab9ce`	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d6ab9cefaa

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

kcoopermiller · 2026-11-02T19:35:28Z

I ran this at d6ab9cefaa against a real service and alongside #1137, which I'm closing in favor of this PR. It stops the blocking and drops, and it's about 4x faster than sequential uploads. Three things are worth fixing before it merges.

Real-service burst. 1000 episodes of about 2.3 MiB, each with its own numpy arrays, submitted as fast as possible through the real UploadWorker, TracesSink and TracesClient against a dev Prime Traces deployment, from a laptop on about 320 Mbit/s of uplink. Stored counts were read back from the service.

	Stored	Dropped	Total time	Throughput	Longest `submit()`	Peak memory
This PR	1000	0	42 s	24 episodes/s	0.1 ms	1.8 GiB
#1137 (sequential)	1000	0	178 s	5.6 episodes/s	0.2 ms	1.9 GiB

No warnings or errors on this PR's run.

Worth fixing

One bad record loses most of the backlog. A write still takes everything queued, and a failed batch stops the rest of that write, so the loss can reach the 10,000-record bound. I logged 500 of these episodes, one of them rejected with a 400, through the real worker, sink and client over a mock transport. This PR stored 130 and lost 370. fix(prime-runs): stop blocking and dropping on upload bursts #1137, which caps each write, stored 487 and lost 13. A cap of about max_in_flight × 30 MiB per write would limit the loss and keep concurrency useful. fix(prime-runs): stop blocking and dropping on upload bursts #1137's 30 MiB cap is too small here, because it would leave one request per write.
Failures are over-reported. The same run reports failed_records == {"traces": 500} although 130 were stored, because _fail_sink counts the whole write. A failure late in a large write can now overstate the loss by much more.
The Codex finding on shutdown looks right to me. Python joins ThreadPoolExecutor workers at interpreter exit, so if finish() times out with requests or retry sleeps still running, the process can stay alive well past its drain budget. I reasoned this from how the executor works and didn't reproduce it.

Memory bound

The queue has no byte limit. At the production sizes in prime-rl#4363 (mean 8.5 MB), 10,000 records is about 85 GB of JSON. The 570 and 860 MB peaks in this PR's description come from a load test where episodes share traces. With distinct arrays, my run reached 1.8 GiB for 1000 episodes of about 2.3 MiB, roughly 1.7 MiB of memory per episode. I'd add a byte budget on top of the record bound. #1137 has one that needs no encoding on the producer: the uploader keeps a running average of encoded size, and submit() multiplies it by the queued count. Its 4 GiB default is too low for an 8 GB step, so it would need to be well above that, and that PR measures every record, which costs extra encoding on the uploader thread. Sampling would be cheaper at this PR's upload rate.

Not tested yet

Production-sized episodes against the real service. My run used 2.3 MiB episodes, not the 8.5 MB mean.
Several runs uploading at once. Eight requests from each producer could saturate the service's writers, and repeated refusals count as strikes towards retiring the sink.
Draining at the end of a run against the real service. At the rate above, 2.3 MiB × 24 per second is about 55 MB/s, so an 8 GB step takes about 2.5 minutes, longer than prime-rl's 60 s FINISH_TIMEOUT. Whatever is left at that point is lost.

The #1137 branch stays up if the write cap or byte budget code is useful: fix/prime-runs-nonblocking-upload-queue.

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>

Ships gzip level 1 (#1129), non-blocking concurrent uploads (#1136) and self-encoded episode uploads (#1140). prime-runs now requires prime-traces>=0.0.12: its traces sink can hand TracesClient encoded lines, which 0.0.11 cannot take. 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)  --- > [!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>  --------- 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)  --- > [!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>  --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

mikasenghaas and others added 3 commits November 2, 2026 18:15

mikasenghaas requested review from DamianB-BitFlipper, JannikSt, JohannesHa, akirillo, burnpiro, cabrinha, kcoopermiller, kennethnym, mrmoxon and willccbb as code owners November 2, 2026 18:25

mikasenghaas requested a review from eexwhyzee as a code owner November 2, 2026 18:25

chatgpt-codex-connector Bot reviewed Nov 2, 2026

View reviewed changes

Comment thread packages/prime-traces/src/prime_traces/traces.py Resolved

mikasenghaas mentioned this pull request Nov 2, 2026

feat: non-blocking, concurrent prime uploads PrimeIntellect-ai/prime-rl#4363

Merged

kcoopermiller mentioned this pull request Nov 2, 2026

fix(prime-runs): stop blocking and dropping on upload bursts #1137

Closed

kcoopermiller approved these changes Nov 2, 2026

View reviewed changes

kcoopermiller merged commit d73f9ea into main Nov 2, 2026
27 checks passed

kcoopermiller deleted the feat/upload-perf branch November 2, 2026 22:59

kcoopermiller mentioned this pull request Nov 2, 2026

release: prime-traces 0.0.12 and prime-runs 0.1.6 #1146

Merged

Navigation Menu

Repository navigation

feat: non-blocking, concurrent run uploads in prime-runs and prime-traces - #1136

feat: non-blocking, concurrent run uploads in prime-runs and prime-traces#1136
kcoopermiller merged 3 commits into
mainfrom
feat/upload-perf

mikasenghaas commented Nov 2, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot commented Nov 2, 2026 •

edited

Loading

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

kcoopermiller commented Nov 2, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

mikasenghaas commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Breaking

Verification

Merge order

Uh oh!

chatgpt-codex-connector Bot commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codex Review Summary

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

kcoopermiller commented Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mikasenghaas commented Nov 2, 2026 •

edited

Loading

chatgpt-codex-connector Bot commented Nov 2, 2026 •

edited

Loading

kcoopermiller commented Nov 2, 2026 •

edited

Loading