[mikasenghaas]

Summary

• Training arrays never stream. Nodes' per-token training arrays (TRAINING_ARRAY_FIELDS: routing rows, sampling masks, top-k heads) are left out of every delta, so they never cross the env server → client hop. The routing-row repairs, which only patched rows already streamed, are gone on both sides.
• set_episode_hook(hook) (verifiers.v1.serve) is the only way the arrays leave a worker. It installs an EpisodeHook with two methods:
  • keep(episode, sampling) runs off the event loop on each finished episode, before its final flush. A host keeps the arrays itself and leaves references in trace.info; what it writes to the traces streams with the final flush.
  • discard(episode) frees what keep kept for an episode that was not delivered.
• Strict contract: every episode passed to keep is either delivered in its run's successful reply or handed to discard, exactly once and never while keep still runs. This covers a keep that raises, a cancel that lands during or after keep (shutdown included), a reply that fails to encode or send, and a client that cancels the run after the reply left, including while it assembles or validates the reply: the client always sends a cancel, the pool broker forwards a cancel to the worker that replied within LATE_CANCEL_SECONDS, and that worker discards the kept episode. The host can then free what it keeps without any sweep.

Used by PrimeIntellect-ai/prime-rl#4196: its train env-server workers keep each trainable branch's arrays and serve them to the trainer, so the orchestrator never holds them.

Breaking

• Deltas carry no training arrays. A client that read routed_experts, sampling_mask or top_logprobs from streamed nodes now gets None. Migration: keep the arrays on the server with set_episode_hook and fetch them from there.
• routing_repairs deltas are gone. EpisodeAssembly no longer applies them, and no server sends them.

Tests

tests/unit/serve/test_delta.py (an episode streams and reassembles without its training arrays) and tests/unit/serve/test_episode_hook.py (a failed, abandoned or finished keep, a cancel before or after the reply leaves, and a cancel during the client's final decode each discard exactly once; an expired delivery is never discarded) and tests/unit/serve/test_pool_cancel.py (the broker routes a late cancel to the worker that replied); full tests/unit: 239 passed. prime-rl#4196 covers keep / discard end to end.

🤖 Generated with Claude Code

---

Note

High Risk
Breaking wire semantics for training arrays and cancel/discard lifecycle; incorrect hook or cancel handling could leak GPU memory or drop data trainers depend on.

Overview
Per-token training arrays no longer cross the env-server wire. Streamed node deltas now omit routed_experts, sampling_mask, and top_logprobs (TRAINING_ARRAY_FIELDS), and the routing_repairs mechanism is removed from both DeltaStreamer and EpisodeAssembly. Clients that assemble WireEpisode from deltas should expect those fields to be None.

set_episode_hook is the supported way for a host to retain arrays on the worker: keep(episode, sampling) runs in a background thread after rollout finishes (before the final delta flush); discard(episode) runs when delivery fails or the client abandons the run. The server tracks kept episodes through reply send and a LATE_CANCEL_SECONDS window so a late cancel can still trigger discard.

Client and pool behavior is tightened so abandonment is signaled reliably: EnvClient.run reuses a stable request_id and fires _cancel_run if decode/validation fails after the reply, or on request cancellation; the pool broker forwards cancels to the worker that still holds an in-flight run or recently replied within the late-cancel window.

Tests now assert reassembly without training arrays instead of routing-row repair/parametrized dtypes.

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

Note

Let host keep finished episodes' training arrays via episode hooks

• Adds EpisodeHook with keep(episode, sampling) and discard(episode), plus set_episode_hook to configure it. When a run finishes, the server runs keep in a worker thread before the final delta flush and tracks the kept episode for cleanup.
• Streamed node deltas now omit the routed_experts, sampling_mask, and top_logprobs training arrays (see TRAINING_ARRAY_FIELDS), and routing_repairs delta logic is removed from both streaming and assembly (delta.py).
• The client and pool support late cancellation: EnvClient can cancel a run after its reply arrived, and EnvServerPool forwards such cancels to the worker that handled the request within LATE_CANCEL_SECONDS (client.py, pool.py).
• Cancellation and failed/unsendable replies now discard hook-kept episode data.
• Risk: reassembled streamed nodes no longer contain routed_experts (the updated test asserts it is None); consumers relying on the old streaming repair path in test_delta.py behavior will need the hook to obtain training arrays.

^{Macroscope summarized 1cee4be.}

[macroscopeapp]

Approvability

Verdict: Needs Human Approval

Macroscope's review found this PR requires human approval — This PR introduces host-side retention and cleanup of finished episodes while removing training arrays from the existing streamed episode contract. The worker, broker, client, and cancellation lifecycle all change, creating a cross-component runtime impact that merits human review.

Notes:

• All code in this push has already been reviewed. Approvability was decided on eligibility alone.

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

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[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 098cbdd. Configure here.}