fix(orchestrator): enforce target lag strictly at ship time

[mikasenghaas]

Summary

Strict off-policy lag control across all three places a rollout can go stale, measured with one consistent quantity: batch N trains on policy v{N-1}, so a rollout generated by v{k} shipping in batch N is (N-1) - k versions off-policy — queue time included. Previously staleness was only counted as weight updates during generation, so a rollout that finished fast but sat buffered across steps trained arbitrarily stale while reporting 0.

• Ship: batch N only ships once the trainer has published v{N-1-TARGET_LAG} (hold in finalize_train_batch). Fixes RES-1077: on tiny fast envs the orchestrator shipped buffer-fed batches past the trainer, finished, and tore down the watcher while the trainer still had pending NCCL broadcasts — which then hung forever in the STABLE → NCCL_READY handshake only the watcher can serve. Never waits on an unpublishable version (under NCCL the trainer's last broadcast is v{max_steps-2}, exactly what the final batch needs); bench runs skip the hold. Under NCCL the hold waits for the applied version (the trainer blocks in the receive handshake until the watcher applies, so this costs nothing and releasing earlier could strand the final handshake at teardown); under filesystem broadcast it releases on the published version, keeping inference weight reloads out of the ship path.
• Dispatch: demand-driven scheduling replaces the lead-based dispatch gate. The orchestrator injects a train_needed predicate — true while the batch being collected isn't covered by buffered + partial-group + in-transit rollouts (incl. the out_q backlog, so ship-holds don't over-provision) — re-evaluated per scheduled group. The coverage target is max(batch_size, max_inflight_rollouts) plus up to TARGET_LAG batches of lookahead (only as far as the staleness budget allows), so steady-state dispatch stays continuous — one-batch-exact demand synchronized dispatch into per-batch waves that stampeded the env server and idled inference through env round-trips (root cause of the CI alphabet_sort timeouts). Configured over-provisioning (oversampling_factor > 1 / raised max_inflight_rollouts) still races batch stragglers at full concurrency. Only fresh groups are demand-gated: already-open groups always finish scheduling (their buffered members count as coverage, so an unfinishable partial group would deadlock the batch a few rollouts short). This deletes the fix(orchestrator): allow final nccl batch through dispatch gate #2990 final-batch NCCL exemption along with the gate that needed it.
• Drop (max_off_policy_steps, now consumption lag): rollouts already past the bound are discarded wherever they are — in-flight groups cancelled pre-pause via cancel_stale_train_groups (aborts must land while the engine is stepping), buffered sink survivors swept whenever the cutoff moves (ship, weight update). Groups share one generation version, so stale groups drop whole rather than shrinking below group_size.
• Defensive fix surfaced by the pacing: the multi-run trainer crashed touching a token-export STABLE marker inside a run dir the integration test had already deleted (FileNotFoundError → trainer death → the resumed orchestrator's ship-hold waited forever). mark_stable now skips deleted runs like the broadcast paths do. On main the same crash risk existed but nothing waited on the trainer, so it never surfaced.
• Integration-test recalibrations for honest pacing: orchestrators now ride the trainer's real speed instead of finishing early on stale buffers, so wall time and reward-vs-step-index shift by the (bounded, visible) 1-2 version lag. Budgets: multi-run 300→900s outer + 600s inner waits, alphabet_sort 1200s, rl_sft 900s. Floors: reverse_text mismatch-KL 0.02, multi-run step-7 reward 0.1 and final reward 0.6 (siblings clear 0.69-0.75; a broken run reads ~0.1-0.3).
• max_off_policy_steps now has a field floor of 1 — the final NCCL batch necessarily trains one version off-policy (the trainer never broadcasts the last versions), so a zero budget could never be met and would hang the run collecting the final batch.

Performance

Default configs (oversampling_factor unset → max_inflight = batch_size) keep the same concurrency budget as main; the difference is that completed-but-buffered rollouts now count against it, so in-flight tapers toward the end of each batch instead of continuously over-generating rollouts that would train stale. For rollout-bound (slow) envs where inference saturation matters more than freshness, set oversampling_factor (e.g. 1.5–2) — the demand target scales with it, extra rollouts race stragglers at full concurrency, and the surplus spills exactly one batch (+1 lag, bounded by max_off_policy_steps).

Verification

uv run rl @ examples/reverse_text/rl.toml --max-steps 10 --no-wandb (2 GPUs, NCCL, ~1-5s buffer-fed steps), before and after the main merge (v0 bridge and v1 taskset paths):

• Before (main): orchestrator ships all 10 batches in ~40s and exits while the trainer is at step 6; trainer wedges silently in the step-7 broadcast handshake; launcher hangs until killed (confirmed 120s+ after orchestrator exit).
• After: batches ship paced to the trainer (Holding batch N until the trainer publishes policy vK, 8 holds incl. the final batch waiting for v8 = max_steps-2); Training finished!, launcher exit 0, 10/10 batch dirs, Max Off-Policy flat at 4 (one version from the dispatch lookahead), zero staleness drops at the default cap. pytest tests/unit -m "not gpu": 487 passed.

Latest main is merged in (the post-merge e2e runs used --trainer.model.attn flash_attention_2 on this Blackwell dev box — since made unnecessary by the attn=auto fallbacks in #3076/#3077).

🤖 Generated with Claude Code

---

Note

High Risk
Changes core orchestrator pacing, dispatch, and staleness semantics that coordinate trainer, NCCL weight broadcast, and inference lifecycle—incorrect timing could deadlock training or strand broadcasts.

Overview
Strict off-policy control now uses one metric everywhere: batch N trains on policy v{N-1}, so a rollout from v{k} that ships in batch N is (N-1)-k versions behind—including time spent buffered. Staleness is no longer tied only to weight updates during generation.

Ship pacing: finalize_train_batch blocks sending batch N until the trainer has produced v{N-1-TARGET_LAG} (skipped in bench). Under NCCL it waits on the applied policy version so the watcher stays alive for broadcast handshakes; filesystem mode can release on publish. This fixes orchestrators finishing ahead of the trainer and hanging NCCL teardown.

Dispatch: The dispatch_allowed lead gate is removed. The orchestrator injects train_needed, which tracks buffered, partial-group, and in-transit rollouts (including out_q during ship holds) plus bounded lookahead so inference stays busy without over-running max_off_policy_steps. Fresh groups only open when demand exists; in-progress groups always finish.

Drop: cancel_stale_train_groups (pre weight pause) and TrainSink.drop_stale / group finalize drop rollouts whose generation policy is below the cutoff. off_policy_steps is stamped at ship time on each rollout. max_off_policy_steps minimum is 1 (final NCCL batch is necessarily one version off).

Other: GPU CI uploads **/tmp/outputs/**/*.log on failure; pytest keeps CI output dirs for artifacts; integration timeouts/thresholds adjusted for paced multi-run training; token export mark_stable tolerates deleted run dirs.

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

[samsja]

lgtm

[cursor]

Ship-hold undercounts batch coverage

High Severity

During the new ship hold, train_needed still treats the batch being shipped as unfilled. process_batch already removed that cohort from pending_batch while progress.step has not advanced yet, so demand only sees overflow, out_q, and in-flight work. Dispatch then schedules a full replacement set on every hold, undoing the lag control this change is meant to enforce—especially on fast envs where holds are long.

Additional Locations (1)

• src/prime_rl/orchestrator/train_sink.py#L180-L213

 

^{Reviewed by Cursor Bugbot for commit a0c857f. Configure here.}

[mikasenghaas]

The mechanics are right but this is deliberate, not a lag-control leak: dispatch during a hold is the freshness-bounded pipelining that keeps inference busy through trainer-bound waits (it's what fixed the alphabet_sort engine-idle waves). Lag control is enforced at ship (hold on the trainer's version) and on data (staleness cutoff drops); dispatch-ahead is bounded by lookahead = min(TARGET_LAG, max_off_policy_steps - lag) plus the permit cap, and measured consumption lag plateaus at 4 vs the default cap of 8 with zero drops. Counting the held batch as coverage would re-idle the engine for the length of every hold. Documented at the site in 81e0e6e.

[cursor]

NCCL ship-hold deadlock on cancel

High Severity

Under NCCL, the ship hold waits on the applied policy.version, but on_version_pending only publishes then awaits cancel_stale_train_groups, which emits Cancelled markers into the bounded out_q. During a hold the main loop is not draining that queue, so it fills and the cancel put blocks; apply never runs, the hold never sees an applied version, and the trainer stays stuck in the NCCL handshake.

Additional Locations (2)

• src/prime_rl/orchestrator/orchestrator.py#L927-L932
• src/prime_rl/orchestrator/dispatcher.py#L557-L558

 

^{Reviewed by Cursor Bugbot for commit 81e0e6e. Configure here.}

[cursor]

Bench lag gates ignore frozen policy

Medium Severity

Bench skips the ship hold because no trainer versions are published, but train_needed and drop_stale still apply consumption-lag against a frozen policy.version. After a few ships, min_policy_version advances while new rollouts are still born at v0, so they are dropped as stale (or dispatch pauses when lag exceeds max_off_policy_steps), and the run cannot finish. This bites orchestrator-only --bench when max_off_policy_steps is near the new floor of 1.

Additional Locations (2)

• src/prime_rl/orchestrator/orchestrator.py#L899-L902
• src/prime_rl/orchestrator/train_sink.py#L152-L158

 

^{Reviewed by Cursor Bugbot for commit 67ec049. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 4 total unresolved issues (including 3 from previous reviews).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 9897ad7. Configure here.}

[cursor]

Final-step NCCL drain deadlock

High Severity

After shipping the final batch under NCCL with max_off_policy_steps=1, lag is permanently 2 because the trainer never broadcasts past v{max_steps-2}. train_needed then stays false and drop_stale discards every reachable rollout, so the overflow batch that alone sets draining never forms and the orchestrator hangs. The removed final-batch NCCL dispatch exemption previously kept that path alive.

Additional Locations (2)

• src/prime_rl/orchestrator/orchestrator.py#L539-L548
• src/prime_rl/orchestrator/orchestrator.py#L623-L629

 

^{Reviewed by Cursor Bugbot for commit 9897ad7. Configure here.}