feat(sc): periodic rollout checkpointing

[macandro96]

What does this PR do?

Adds periodic rollout-state checkpointing to Single Controller.

Unlike trainer checkpoints, which are written at completed training-step boundaries, these snapshots can capture newer rollout progress while generation and streamed training are active. On restart, Single Controller restores the latest compatible rollout snapshot anchored to durable trainer state.

This PR builds on #3923, which provides sibling-level rollout recovery.

Why?

Long-running asynchronous rollouts may make significant progress between trainer checkpoints. If the job is interrupted, previously completed generations and buffered rollout data would otherwise need to be regenerated.

Periodic rollout snapshots reduce the amount of rollout work lost without changing the model, optimizer, or trainer checkpoint frequency.

Design

A periodic snapshot contains matching:

• Native TQ checkpoint state.
• Metadata-only replay-buffer state.
• Rollout recovery ledger.
• Dataloader state.
• Sampler dispatch cursor.
• Replacement prompt reserve.
• Controller recovery metadata.

Snapshots are anchored to immutable trainer state:

• Before the first training step, snapshots use a bootstrap anchor tied to the initial model and a fail-closed fingerprint of rollout semantics.
• After training begins, a snapshot is only created when the trainer checkpoint for the current completed step is durable.
• The trainer checkpoint itself is not modified when newer rollout snapshots are written.

Snapshots are built under a temporary directory, flushed, and published using an atomic directory rename. Retention first renames stale snapshots out of the live namespace before deleting them.

Streaming training recovery

The trainer may already own rollout groups when a periodic checkpoint occurs during a streamed training step. Those groups remain in TQ until the optimizer step succeeds and are represented in checkpoint replay metadata. If the job crashes before the optimizer step completes, they become ordinary replayable groups after restoration.

Checkpointing is fenced against data-plane mutations and the final optimizer commit/cleanup transition so the TQ snapshot, replay metadata, and controller state describe one logical point.

Restore behavior

rollout_checkpointing.restore_mode controls snapshot selection:

• latest restores the newest compatible published periodic snapshot.
• trainer_checkpoint ignores newer periodic snapshots and restores rollout state bundled with the trainer checkpoint.

Restore selection is read-only. If bootstrap rollout state exists but no trainer checkpoint does, trainer_checkpoint fails without deleting or ignoring that state because there is no matching model/optimizer checkpoint. Use latest to recover it or choose a new checkpoint directory to start fresh.

When restoring the latest periodic snapshot:

• Completed siblings remain reusable.
• Buffered groups are restored into the replay buffer.
• Groups claimed by an unfinished streamed step are made replayable again.
• Missing siblings are redispatched according to the configured recovery granularity.
• The dataloader and sampler resume from the matching checkpoint position.

This PR does not implement mid-token prefix recovery. A generation without a sealed recovery record must still be redispatched.

Configuration

checkpointing:
  enabled: true
  checkpoint_dir: results/grpo-single-controller
  save_period: 1
  metric_name: null
  save_data_plane: true

token_capture:
  enabled: true

rollout_checkpointing:
  snapshot_attempt_interval_s: 120
  keep_latest_k: 2
  restore_mode: latest
  extra_fingerprint_excluded_paths: []

snapshot_attempt_interval_s is an attempt cadence. A periodic write still requires a durable trainer anchor for the current completed step, so checkpointing.save_period: 1 provides continuous post-step coverage.

Before the first trainer checkpoint, bootstrap compatibility is fail-closed: the full configuration participates unless a path is explicitly classified as operational. External integrations can add operational-only paths to extra_fingerprint_excluded_paths.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[terrykong]

Thanks @macandro96 — and thanks for the two pushes while this was in flight. Re-reviewed against 962a670a; several things we had written up you have already fixed, so they are not below.

The crash-consistency protocol is the part we tried hardest to break and could not. The fsync ordering (payload → COMMITTED → fsync → rename → fsync parent), _SNAPSHOT_RE anchored with .fullmatch() at every call site, max-by-integer sequence selection, and a fingerprint that is stable across runs all hold up. We checked each of those specifically and found nothing wrong with any of them.

Fixed since we started, noted so it is clear they were looked at rather than missed. The barrier nesting is gone and claim_for_training now binds as cut. _cleanup_consumed_metas_unlocked now gets its cut at both call sites — that one was worth catching, because git reported no conflict on single_controller.py during the rebase and the missed call site failed 23 tests in test_checkpointing.py on its own. rollout_recovery.py is out of the diff, so base's _reject_unknown_fields hardening survives. The duplicated bootstrap expected dict is now one helper. The PPO exemplar has the rollout_checkpointing block. And test_setup.py now exercises the restore path.

The partition test is the right call. assert declared | ignored == fields over all six schemas means a new config field cannot silently drop out of the fingerprint — someone has to decide. The one inline comment below is about a field that landed on the wrong side of it, not about the mechanism.

Other things worth saying out loud:

• Putting training ownership in the component that owns the rows (TQReplayBuffer._training_claims) rather than standing up a second lifecycle model is the right shape.
• RolloutCheckpointConfig(extra="forbid") with the reasoning in the docstring is the right call, and reads as a direct answer to the extra="allow" thread on #3923.
• The would_save() / check_save() split in timer.py is a clean seam.
• Five setup-time guards, each with a unit test — every one turns a would-be silent no-op into a startup error.

Lint. Two failures, both from this PR, one command to fix. Against ruff 0.9.9 (pinned at .pre-commit-config.yaml:11): I001 in tests/unit/single_controller/test_rollout_checkpoint.py, and a ruff-format diff in tests/unit/single_controller/test_tq_replay_buffer.py. uvx ruff@0.9.9 check --select I --fix . && uvx ruff@0.9.9 format . clears both. This came from running the pinned hooks directly, because pre-commit cannot build its venv in this checkout.

No CI:* label. Every PR needs exactly one. This looks like CI:L1 — it does add an L1 functional test.

PR description. Still the stock template. Your lead commit already has a real body covering the what, and #3923's PR body is a good model, so this reads as an oversight. The three things most worth adding: whether this is on or off by default, the on-disk format, and the failure it is meant to survive.

One caveat on all of the above. No test CI has run on this PR, we had no GPU, and nobody has executed the new functional test grpo_async_gym_single_controller_streaming_recovery.sh. Its runner registration is correct (tests/functional/L1_Functional_Tests_SingleController.sh:189) and the file mode is right — we checked both — but the test itself is unrun. Getting /ok to test issued still matters more than any single comment here.

Generated by Claude Code

[terrykong]

Re-reviewed at b1307f2a4. Everything from the last round landed, and two of the behaviour changes hold up under a real run: the fingerprint now moves on bad_words while still ignoring port_range_low, and the interrupted-delete sequence that used to strand a snapshot directory forever is swept clean. All nine earlier threads are marked fixed and resolved.

Two comments below. The first is the one that matters: 36 unit tests fail because the test doubles have not kept up with the new call sites — no assertion fails and no behaviour looks wrong. The second asks for a better message on an error that should not be reachable.

The rebase worked too — the token-capture code from #3837 that was showing up here is now upstream and out of this PR's diff.

Generated by Claude Code

[macandro96]

/ok to test 73205a6

[macandro96]

/ok to test 1557d03

[macandro96]

/ok to test ba4413f

[macandro96]

/ok to test 17b86fb

[macandro96]

/ok to test 2835680