feat: RolloutGenerationMixin + per-prompt prefetch

[shuangwu]

Branch: feat/rollout-generation-mixin (currently at c9b2183)
Base: upstream/main (c35f517), but stacked on PR #689 and the UCXX PR (see below)
Relationship:

• Stacked on PR feat: gym example trainer #689 (feat/gym-example-trainer) — the gym example is the first in-tree adopter of the mixin. Two commits from feat: gym example trainer #689 (7706326, 95b35db) appear in this branch's upstream/main..HEAD until feat: gym example trainer #689 merges; they drop naturally on rebase.
• Stacked on the UCXX PR (feat/ucxx-single-chunk-rotation) — the UCXX single-chunk commit (c9b2183 here, acec19a on the standalone branch — same content, different SHA due to rebase) is included so this branch is self-sufficient for bench runs. Drops on rebase once the UCXX PR merges.
• Final standalone contribution: one commit (96941ac) once both prerequisite PRs land.

Suggested merge order: PR #689 → UCXX PR → this PR.

---

Summary

Adds RolloutGenerationMixin for RolloutBase subclasses that structures rollout_generation via five hooks and overlaps per-prompt preprocessing with in-flight generation when config.rollout.prefetch_rollout is set. Bundles in eight independent controller-rollout robustness fixes debugged alongside the mixin during the May-2026 bench cycle.

The two concerns are combined because they share hunks in cosmos_rl/rollout/worker/rollout_control.py (prefetch-loop instrumentation: Trace A/B/C, branch counters, _prompt_queue peeks) and were bench-validated together. Splitting them on a hunk basis produced an awkward "subset" commit that wasn't a clean cherry-pick; combining them keeps the provenance honest.

Background

Why a mixin

Today every RolloutBase subclass implements its own rollout_generation loop end-to-end, duplicating the prompt-pull → preflight → generate → postprocess scaffolding. This makes it hard to add cross-cutting features (prefetch, instrumentation, retry policy) without N-way drift. RolloutGenerationMixin reifies the loop in five hooks (_preflight_check, _prepare_sample, _collate_batch, _generate, _postprocess) so subclasses contribute the workload-specific pieces and the mixin owns the orchestration.

Why prefetch

On the disaggregated single-prompt path the rollout worker spends N ms per prompt on preprocessing (tokenisation, image decode, prompt rewrites) before reaching generate(). When generate() itself is also CPU-bound between kernel launches, the prep cost serialises the loop and inflates per-iteration latency. Prefetch overlaps the next prompt's preprocessing with the current generation, on the same worker thread, using a small bounded queue. Opt-in via config.rollout.prefetch_rollout so subclasses that don't want it pay no cost.

Why eight robustness fixes

The May-2026 bench cycle surfaced eight independent controller-rollout correctness / stability issues, each reproducible and small enough to land alongside the mixin work:

• samples_on_the_fly over-counting drove the controller's soft throttle to engage spuriously.
• Wedged controllers left workers hot-spinning the prompt_consume_end branch indefinitely.
• A wedged controller's unregister could block worker exit.
• Heartbeat subprocesses outlived their parent on SIGSEGV and lied to the controller.
• The end-of-run mesh-rebuild could try to include dead replicas in the new mesh.
• Etc.

Each is described in the commit body and is self-contained (could be reverted independently).

Changes

Commit 96941ac — RolloutGenerationMixin + prefetch + robustness

Prefetch mixin (cosmos_rl/rollout/worker/rollout_generation_mixin.py + plumbing):

• Five hooks: _preflight_check, _prepare_sample, _collate_batch, _generate, _postprocess.
• Opt-in via config.rollout.prefetch_rollout; reads the flag live (not snapshotted in __init__) so subclasses can flip it from post_init_hook.
• Single-producer-mode property + bounded _prompt_queue for overlap.
• Trace A (per-iteration branch counters), Trace B (queue depth), Trace C (per-generation entry/exit) at DEBUG.

Controller-rollout robustness fixes:

• Controller samples_on_the_fly accounting — decrement on filter_outdated_rollouts discards; decrement train_ack by the actual per-step dispatch count via dispatched_rollouts_by_step (was sized by train_batch_per_replica * arrived_replicas, equal on normal steps but zero on is_fake_last_cmd); include in-flight samples in recompute_total_steps.
• Rollout main_loop — set shutdown_signal on prompt_consume_end() (mirrors the async path; previously a crashed controller left workers hot-spinning the consume_end branch indefinitely); 50 ms backoff after weight-version rejection (was busy-spinning until the next P→R broadcast).
• Rollout teardown — bound unregister / post_heartbeat / daemon joins with timeouts so a wedged controller can't prevent worker exit; skip BuildMeshCommand when this replica is draining (defence-in-depth for the controller-side filter below).
• Controller end-of-run mesh management — filter RolloutStatusManager.unregister's BuildMesh recipient set to live survivors only (not status.ended) and require >= 2 survivors before triggering rebuild.
• Process lifecycle — PR_SET_PDEATHSIG on the heartbeat subprocess so a parent SIGSEGV doesn't leave it orphaned and lying to the controller; opt-in COSMOS_SHUTDOWN_ON_NO_POLICY_REPLICAS to exit the controller when all policy replicas are dead (off by default to preserve dynamic-replica deployments).
• Log hygiene — soft-throttle engagement / heartbeat / release events + paired Trace D dispatch-state snapshot for the silent throttle path; Trace E for the controller-side filter decisions.

Commit c9b2183 — UCXX single-chunk + client rotation

Same content as feat/ucxx-single-chunk-rotation's acec19a; see that PR. Included here so this branch is self-sufficient for bench runs. Drops on rebase once that PR merges.

Validation

Local CPU test suite (driver: pytest -q --tb=short on branch tip c9b2183)

File	Result
tests/test_rollout_generation_mixin.py (16 tests incl. test_prefetch_throughput_synthetic, test_subclass_preflight_check_propagates)	passed
tests/test_rollout_prefetch_loop_integration.py	passed
tests/test_put_rollouts.py	passed
tests/test_gym_example.py::TestGymBackendStructured + ::TestGymBackendPrefetch	passed
tests/test_ucxx_* (UCXX commit)	55 passed / 1 skipped
Stack total	171 passed / 1 skipped / 180 subtests

Bench validation

rl-gym companion branch feat/rollout-generation-mixin carries the SimpleRolloutWorker / ModularRolloutWorker adopters plus the bench matrix. May-2026 cycle's headline numbers are paste-ready in rl-gym/docs/BENCHMARKS.md. Five disaggregated reruns confirmed each robustness fix eliminates its corresponding wedge mode.

Caveats

• The robustness portion has thin direct unit-test coverage — test_put_rollouts.py is the only CPU-runnable test that imports any of the eight changed controller/dispatcher modules. The dispatch-accounting / soft-throttle / RolloutStatusManager.unregister filter changes are validated by the May-2026 bench cycle (see commit body). Multi-rank integration tests (test_policy_overfit, test_policy_to_rollout) gate this on the CI side.
• Local-only Failed to import VLLM Rollout log line — benign, pre-existing, caused by the local venv not having vllm installed. Not introduced by this PR.

Known follow-up (not blocking)

Bench data revealed a separate NCCL P2R recv backlog stalls the inference stream on the rollout side when the trainer is fast. This is a downstream effect that limits prefetch's measured upside in the May-2026 matrix to 5–7% (vs. the historical 17%), not a correctness issue. Investigation is complete; fix lives outside this stack and will ship as a separate PR. Cross-references:

• rl-gym/docs/PREFETCH.md carries a note that prefetch's observed tail-latency cost is downstream of this bug, not intrinsic to prefetch.
• Smoking-gun bench artefact: rollout_0 generate_ms=50026.17 on job 28243035 (vs. typical ~2,500 ms).
• Three candidate fixes documented in cosmos_rl_in_flight.md.

Test plan

• full CI matrix green — especially test_policy_overfit, test_policy_to_rollout, test_context_parallel (multi-GPU NCCL paths that exercise the controller-rollout protocol).
• reviewer sanity-check on the eight robustness fixes — each is self-contained and could be reverted independently if needed.
• reviewer sanity-check on the live prefetch_rollout flag read (deliberately not snapshotted in __init__ so subclasses can flip from post_init_hook — verify this is the contract you want).
• rebase once PR feat: gym example trainer #689 lands (drops the two gym-example commits).
• rebase once the UCXX PR lands (drops c9b2183).
• confirm bench numbers reproduce on a clean cluster via rl-gym feat/rollout-generation-mixin companion branch.

[shuangwu]

The job did not fail on a Python exception in the test code; it timed out after 2 hours