Offload finalized training publication calculations from the event loop by hubert-marek · Pull Request #3811 · PrimeIntellect-ai/prime-rl

hubert-marek · 2026-10-03T19:41:02Z

Stacked on #3806, including the pruning, encoding and payload-release changes merged from #3821. Completed arrival metrics are scalarized by the base; this change offloads shipped-cohort calculations, annotations, pruning, serialization and release while preserving one combined all/effective scalar record per shipped batch.

Effective episodes are materialized once. Mutable source counters are captured on the event loop; workers read the owned finalized cohort. Episode fields, reward formulas, transport bytes and shipment membership are unchanged.

Review fixes (2026-10-13):

Give monitors a separate episode list before background release. The asynchronous Prime uploader can retain its list without cleanup emptying a pending sample upload. Episode objects are not copied or rewritten.
Snapshot and consume policy-wait time after threaded metric computation, before the next await. Waits completed during that computation are reported; waits completed during monitor I/O remain for the next report. The warning uses the same snapshot as the metric.
Encode microbatches incrementally through BytesIO, rather than retaining every encoded fragment before joining. Empty and single-microbatch payloads use the native encoder directly. The returned type remains bytes, and MessagePack framing is unchanged.
Merge the current perf(orchestrator): finalize discarded rollouts with bounded accounting #3806 base and resolve its combined-publication conflicts while retaining asynchronous batch processing and offloaded effective metrics.

Validation of these review fixes:

Sixteen encoder regression cases pass with exact wire-byte equality and typed decoding, covering empty/single/multiple payloads, routed-expert data and fixarray/array16/array32 boundaries.
A controlled CPU replay of the actual publication method reproduces both races on the reviewed commit: deferred uploads become empty, and completed policy waits are lost. The fixed method preserves both uploaded episodes, reports the worker-period wait and retains the later monitor-period wait. Shipment payloads match, combined all/effective publication remains one scalar record, and empty batches do not advance the step. Monitor delivery and policy timing are controlled test boundaries, not live SDK/GPU measurements.
Ruff lint/format, syntax and diff checks pass. The full repository/GPU suite was not rerun for this revision; the isolated encoder runner avoids optional deployment/config imports.

Encoder allocation check, three trials per case with approximately 64 MiB of synthetic binary payload:

Microbatch composition	Previous peak traced allocation	Fixed peak traced allocation
One 64 MiB payload	160.0 MiB	96.0 MiB
Sixteen 4 MiB payloads	160.0 MiB	73.5 MiB
63 MiB + 1 MiB payloads	160.0 MiB	157.5 MiB

These are maximum tracemalloc allocation peaks excluding preallocated inputs, not process RSS or production OOM rates. Large individual fragments still require temporary encoding allocations. Every measured payload was byte-identical between variants.

The earlier threading study measured publication heartbeat p95 falling from 2.99 s to 5.3 ms for 3,360 finalized episodes, with a 1.59 s worst pause; publication wall time increased from 5.27 s to 6.70 s. Those numbers concern the earlier prototype placement, not this revision or production latency.

Note

Medium Risk
Touches the critical ship path (metrics, monitors, transport encoding) and async lifecycle around memory release; behavior is intended to be identical but threading and deferred release add concurrency edge cases during shutdown.

Overview
Moves CPU-heavy train-batch finalization off the asyncio event loop so dispatch and heartbeats stay responsive during publication.

In finalize_train_batch, trainability checks, effective-episode materialization, annotation stamping, shipped-cohort metrics, and the success log run via asyncio.to_thread. Monitor logging uses a copy of effective episodes so async upload can retain the list. Policy-wait time is snapshotted on the loop before workers read the cohort.

Memory release is deferred: after each ship, release_and_trim runs in a thread pool (shielded from loop cancel) to incrementally clear shipped sample/cohort/micro-batch containers and call malloc_trim—replacing synchronous trim_process_memory on the hot path. Teardown awaits the pending release future.

TrainSink.take_batch / process_batch are async; zero-advantage pruning when batch size is variable runs in a worker thread.

Batch transport adds BatchSender.encode to msgpack-encode micro-batches one at a time (wire-identical to a single bulk encode); ZMQ sends encode per rank in threads. trim_process_memory no longer calls gc.collect().

Unit tests updated for async take_batch; new tests assert encode byte parity across list sizes.

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

chatgpt-codex-connector · 2026-10-04T18:31:43Z

Codex Review Summary

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

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-13T11:35:40.259574Z	`c7c71dd`	New commits

ℹ️ 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.

…ccounting

mikasenghaas

this seems fine

…3821) Stacked on #3811, which is stacked on #3806. Move pruning, serialization and shipped-payload release off the orchestrator event loop. - Zero-advantage pruning runs in a worker after selection detaches samples; queue/cohort mutations stay on the loop. `process_batch` and its callers are async. - `BatchSender.encode` encodes one microbatch at a time under the same MessagePack array header. Output bytes and ZMQ frames are unchanged; ZMQ encoding runs in a worker. - Shipped sample, episode and microbatch lists are released incrementally by a worker. At most one release remains pending, and teardown awaits it. Explicit `gc.collect()` is removed from `trim_process_memory`; this PR does not disable automatic GC. The combined private smoke also uses Verifiers #2811, which disables automatic GC. That integration retains explicit collection in a worker to avoid having neither collection mechanism. Its live A/B is therefore a comparison of the combined private runtime, not proof of the standalone GC-removal effect. Validation of the adapted layer: 266 CPU tests passed, two skipped, across orchestrator, transport, utilities and config tests; Ruff lint/format and diff checks passed. The integrated candidate replay matched deterministic metrics, sample hashes and cohort membership. Four-node, ten-update live comparison is scheduled; results are pending. Earlier synthetic benchmark (before the #3806 API adaptation; CPU-only, 4,096 samples × 32k tokens, eight ZMQ ranks, one run per variant): | Phase | Original wall / max pause | Earlier PR wall / max pause | |---|---|---| | Pruning | 10.9 s / 11.0 s | 11.9 s / 1.1 s | | Packing, already threaded | 23.4 s / — | 23.6 s / 1.1 s | | Sending | 6.6 s / — | 7.0 s / 0.09 s | | Background release/trim | — | 6.1 s / 0.02 s | | Ship-path total | 52.9 s | 43.2 s | | Worst loop pause | 21.6 s | 1.1 s | Those historical numbers are not a fresh benchmark of this adapted diff. Automatic cyclic GC and GIL-held operations can still pause the loop. #3822 adds queue draining while these ship phases run.  --- > [!NOTE] > **Medium Risk** > Touches the hot train-ship path and memory-release timing; wire format is preserved but concurrency and teardown ordering change. > > **Overview** > Moves more of the **train ship path** off the asyncio event loop so long GIL-held work no longer stalls dispatch and rollout handling. > > **TrainSink** makes `take_batch` / `process_batch` async and runs zero-advantage pruning in a worker thread after selection; queue/cohort mutations stay on the loop. Call sites (orchestrator loop and drain handling) now `await` `take_batch`. > > **Batch transport** adds `BatchSender.encode` to MessagePack-encode micro batches one at a time (same array header and bytes as a single bulk encode). Filesystem sends use it directly; **ZMQ** encodes each rank’s grid in `asyncio.to_thread`. > > After each shipped batch, the orchestrator no longer calls `trim_process_memory` synchronously on the loop. It **awaits a shielded prior `release_future`**, then schedules `release_and_trim` in the default executor to incrementally clear shipped sample/episode/microbatch lists and call `malloc_trim`. Teardown awaits any pending release. Shutdown still trims via `asyncio.to_thread`. **`gc.collect()` is removed** from `trim_process_memory`. > > Unit tests were updated to `await` `take_batch`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 06dd796. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>  --------- Co-authored-by: Hubert Marek <163992334+hubert-marek@users.noreply.github.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2680aa4e26

ℹ️ 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".

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

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

cursor

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 2680aa4. Configure here.}

mikasenghaas · 2026-10-13T12:02:23Z

+
 def trim_process_memory() -> None:
    """Return freed heap pages to the OS on glibc systems."""
-    gc.collect()


why remove this?

resolved offline

mikasenghaas

lgtm

…ng (#3806) Completed rollout groups with no training signal could keep large episode graphs alive until another batch shipped. This PR records their scalar statistics when the group finishes and releases the discarded graph immediately. Queued, stale, pruned, and cancelled traces release their references at the appropriate boundary. Training metrics remain batch-aligned: each shipped batch emits one record with both all-rollout (`train/*/all`, including rejected/error/cancelled groups) and shipped-trace (`train/*/effective`) statistics. Exact percentile observations and rate denominators are retained. Training charts use custom `step`; `rollout/window` existed in an intermediate PR revision but was published by neither pre-PR main nor this revision. Shutdown synchronizes with pending admissions before cancelling groups, including attempts that were reserved but never dispatched. Archive turn counts now exclude copied assistant history; existing saved indexes need rebuilding to correct older counts. The current trainer-requested filesystem/ZMQ batch handshakes are preserved. ### Validation - [CPU CI after merging main](https://github.com/PrimeIntellect-ai/prime-rl/actions/runs/38550476797): **803 passed, 1 skipped**; affected-area checks: **130 passed, 1 skipped**. A paired replay of 50 scenarios matched 412,085 metric values, shipped samples, and step counters against the previous implementation. - Nebius Qwen 3.5 smokes resumed the same step-110 checkpoint. The table covers completed steps **111–119**. A monitor sampled the orchestrator process's proportional set size (PSS) on its compute node and probed event-loop lag every 50 ms. Max stall is the longest observed lag. | Smoke | Orchestrator PSS mean / peak | Mean event-loop lag | Max stall | | --- | ---: | ---: | ---: | | [4190: with #3806 stack](https://wandb.ai/primeintellect/perplexity-qwen-travel/runs/3cba96dda403413eae860e28bc8d432d) | 65.2 / 210.7 GiB | 6.44 ms | 13.97 s | | [4192: without #3806 stack](https://wandb.ai/primeintellect/perplexity-qwen-travel/runs/c1716a73a0924f86b77e4c96d6f2f29c) | 68.9 / 181.5 GiB | 5.83 ms | 60.00 s | An earlier three-update smoke checked the integrated runtime separately: | Smoke | Stack | Completed steps | Max full-run event-loop lag | | --- | --- | ---: | ---: | | [3761](https://wandb.ai/primeintellect/perplexity-qwen-travel/runs/e4ceddff45d84e45a9342075dddd3918) | #3806; no #3811/#3821/#3822 | 96–98 | 78.48 s | | [3762](https://wandb.ai/primeintellect/perplexity-qwen-travel/runs/2a464e25d2a0417ca1e9b19f61f2f1fc) | #3806 + #3811/#3821/#3822 | 96–98 | 4.84 s | In the newer pair, 4190 used #3811/#3821 through this PR while 4192 did not; neither used #3822. Both kept the same other model/runtime PRs, fractional rewards, and 2 trainer + 2 inference nodes. The with-stack run had lower mean PSS and shorter worst stalls, but higher peak PSS and slightly higher mean lag. The older full-run lag includes startup; the newer table is step-scoped. These are independent rollouts with bundled changes, so neither pair isolates #3806 or proves a memory or speed improvement. --------- Co-authored-by: hallerite <git@hallerite.com> Co-authored-by: faresobeid <111092724+faresobeid@users.noreply.github.com> Co-authored-by: Mika Senghaas <mail@mikasenghaas.de> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

hubert-marek marked this pull request as ready for review October 4, 2026 18:31

hubert-marek marked this pull request as draft October 8, 2026 22:49

perf(orchestrator): offload shipped metrics above finalized arrival a…

3794ce4

…ccounting

hubert-marek force-pushed the fix/offload-train-publication branch from 3f7c38b to 3794ce4 Compare October 10, 2026 16:12

hubert-marek changed the base branch from main to fix/release-discarded-episode-payloads October 10, 2026 16:13

hubert-marek mentioned this pull request Oct 10, 2026

perf(orchestrator): build and ship train batches in the background while draining out_q #3822

Closed

mikasenghaas approved these changes Oct 12, 2026

View reviewed changes

hubert-marek marked this pull request as ready for review October 12, 2026 18:15

chatgpt-codex-connector Bot reviewed Oct 12, 2026

View reviewed changes

Comment thread src/prime_rl/orchestrator/orchestrator.py Resolved

Comment thread src/prime_rl/transports/batch/base.py Resolved

Comment thread src/prime_rl/orchestrator/orchestrator.py Resolved

cursor Bot reviewed Oct 12, 2026

View reviewed changes

Comment thread src/prime_rl/orchestrator/orchestrator.py Resolved

fix(orchestrator): preserve publication uploads and policy waits

c7c71dd

mikasenghaas reviewed Oct 13, 2026

View reviewed changes

mikasenghaas approved these changes Oct 13, 2026

View reviewed changes

mikasenghaas merged commit 8b07484 into fix/release-discarded-episode-payloads Oct 13, 2026
15 checks passed

mikasenghaas deleted the fix/offload-train-publication branch October 13, 2026 12:10

hubert-marek mentioned this pull request Oct 14, 2026

perf(orchestrator): finalize discarded rollouts with bounded accounting #3806

Merged

Navigation Menu

Repository navigation

Offload finalized training publication calculations from the event loop - #3811

Offload finalized training publication calculations from the event loop#3811
mikasenghaas merged 3 commits into
fix/release-discarded-episode-payloadsfrom
fix/offload-train-publication

hubert-marek commented Oct 3, 2026 •

edited by cursor Bot

Loading

Uh oh!

chatgpt-codex-connector Bot commented Oct 4, 2026 •

edited

Loading

Uh oh!

mikasenghaas left a comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Uh oh!

Uh oh!

mikasenghaas Oct 13, 2026

Uh oh!

mikasenghaas Oct 13, 2026

Uh oh!

mikasenghaas Oct 13, 2026

Uh oh!

mikasenghaas left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

hubert-marek commented Oct 3, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

Uh oh!

Codex Review Summary

Uh oh!

mikasenghaas left a comment

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mikasenghaas Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

mikasenghaas Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

mikasenghaas Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

mikasenghaas left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hubert-marek commented Oct 3, 2026 •

edited by cursor Bot

Loading

chatgpt-codex-connector Bot commented Oct 4, 2026 •

edited

Loading