fix(orchestrator): hold batch ship on the trainer's published version

[mikasenghaas]

Summary

Minimal fix for RES-1077: the pipeline's pacing is two counters (batch N, policy version v) with three rules — START (main's existing dispatch gate), ADVANCE (this PR: batch N ships only once the trainer has published v{N-1-TARGET_LAG}), and DIE (main's existing in-flight drop aging).

• The bug: on fast envs the sink fills batches from buffered rollouts faster than the trainer consumes them; ship was ungated, so the orchestrator finished all its steps, tore down the weight watcher, and stranded the trainer inside its next in-memory (NCCL/NIXL) broadcast handshake — the watcher serves the receive side. The launcher then hung forever.
• The fix (~30 lines): hold each ship until the required version is published. The requirement is always satisfiable (the trainer skips only the final TARGET_LAG+1 in-memory broadcasts), and the hold doubles as the shutdown-liveness guarantee: the orchestrator cannot exit before the last watcher-dependent handshake has completed.
• Supporting hardening: policy.version now advances at publish confirmation (inference applies immediately and pauses during the update, so published ≡ applied for every consumer — no weight-reload time serialized into the ship path), and WeightWatcher.stop() drains an in-flight apply before cancelling so teardown can never strand the trainer mid-handshake.
• No throughput change: dispatch, drops, and batching are untouched. In the generation-bound regime the hold never binds; in the trainer-bound regime the trainer is the bottleneck regardless — the hold only changes where a finished batch waits (orchestrator instead of the trainer's queue), never when the trainer consumes it.
• Honest staleness counting: shipped rollouts are stamped with their true consumption lag — batch N trains on v{N-1}, so a rollout generated from v{k} is (N-1)-k versions off-policy, queue time included. The dispatcher's in-flight counter (unchanged, still feeds the max_off_policy_steps drop rule) only sees weight updates during generation and undercounts; Max Off-Policy in logs and rollout records now reports the real number (expect ~1-3 where runs previously showed 0 — the data was equally off-policy before, it just wasn't measured).
• Also fixes two adjacent races the longer-lived orchestrator exposes: the multi-run trainer crashing on a run dir deleted while its token-export tail flushes (mark_stable now skips deleted runs, mirroring the broadcast paths), and the multi-run test's rmtree racing that same flush (retry helper). CI uploads run logs as artifacts on integration failures.

Verification

uv run rl @ configs/basic/reverse-text/rl.toml --max-steps 10 --no-wandb (2 GPUs, NCCL, ~1-5s buffer-fed steps):

• Before (main): orchestrator ships all batches in ~40s and exits while the trainer is steps behind; trainer wedges silently in the broadcast handshake; launcher hangs until killed.
• After: the hold binds exactly twice at the buffer-fed tail (Holding batch N until the trainer publishes policy vK), Training finished!, clean exit, 10/10 batch dirs. pytest tests/unit -m "not gpu": 469 passed.

🤖 Generated with Claude Code

---

Note

High Risk
Changes orchestrator–trainer synchronization and weight-watcher lifecycle on the in-memory broadcast path; incorrect hold/version timing could deadlock shipping or leave the trainer wedged, though the change is narrowly scoped to pacing and teardown.

Overview
Fixes a hang where fast envs could ship all train batches and tear down the orchestrator before the trainer finished in-memory weight broadcasts, stranding NCCL/NIXL handshakes that need the live WeightWatcher.

Orchestrator pacing: Before sending batch N, non-bench runs now hold until policy.version reaches v{N-1-TARGET_LAG}, waking on version_advanced when the watcher publishes. Shipped rollouts get off_policy_steps stamped from batch step vs rollout policy_version (queue time included) for honest Max Off-Policy metrics.

WeightWatcher: policy.version advances at publish confirmation (before inference update_weights) so held ships release without waiting on reload; stop() waits for an in-flight apply under update_lock so teardown does not cancel mid-handshake.

Adjacent hardening: TokenExporter.mark_stable tolerates deleted multi-run dirs; multi-run integration uses retrying remove_run_dir, --clean-output-dir on RL tests, and CI uploads *.log artifacts from /tmp/outputs on GPU integration failure.

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

[samsja]

lgtm

[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.}

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