feat(sc): recover replay buffer from native TQ checkpoints

[macandro96]

What does this PR do ?

Summary

Adds crash-consistent SingleController recovery for two rollout states:

1. Completed, unconsumed groups

   • Tensor payloads remain authoritative in the native TransferQueue checkpoint.
   • SingleController stores only a metadata replay index.
   • Recovery rebuilds the local replay buffer without copying tensors back into TQ.

2. Unfinished prompt groups

   • A lightweight ownership ledger records prompts that have left the dataloader but have not reached canonical TQ commit.
   • Recovery rehydrates those prompts from the dataset and redispatches them for whole-group regeneration.
   • This prevents dataloader advancement and sampler dispatch state from silently skipping prompts after restart.

This builds on the merged TQ v0.19 bump (#3423) and SingleController checkpoint lifecycle (#3429).

Motivation

Previously, SC checkpoint recovery either serialized rollout tensors separately in replay_buffer.pt or restarted with an empty rollout buffer.

Serializing tensors in both the replay-buffer checkpoint and TQ:

• duplicates potentially large payloads;
• creates two possible sources of truth;
• requires re-putting tensors into TQ during recovery;
• cannot guarantee that the replay index and TQ snapshot describe the same rows.

This change makes the native TQ snapshot authoritative for tensors and keeps only the controller metadata required to resume consuming those rows.

Checkpoint contents

step_N/
├── policy/
│   ├── weights/
│   └── optimizer/
├── train_dataloader.pt
├── training_info.json
├── data_plane/                   # Native TQ checkpoint; owns tensors
├── replay_buffer_metadata.pt     # Completed-group metadata only
├── rollout_recovery.pt           # Unfinished prompt-group ownership
└── replacement_reserve.pt        # Optional spare prompt pool

replay_buffer_metadata.pt contains no rollout tensors or fields_data.

rollout_recovery.pt stores:

1. stable group and admission IDs;
2. dataset prompt reference and task name;
3. reserved/admitted phase;
4. expected generation count;
5. start policy weight version;
6. optional target training step;
7. batch-shortfall and sampler scheduling state.

Consistency model

TQ mutations, sampler admission commits, rollout-ledger mutations, and checkpoint capture participate in the same data-plane checkpoint barrier.

The barrier ensures a checkpoint observes transitions atomically, including:

1. dataloader advancement together with prompt ownership;
2. sampler dispatch-index advancement together with ledger admission;
3. canonical TQ commit versus unfinished ledger ownership;
4. intentional stale-rollout abort and ledger removal;
5. replay-buffer consumption versus native TQ save.

Mutation methods require a live barrier capability, preventing future call sites from accidentally changing recovery state outside the checkpoint boundary.

If a group exists in both canonical TQ metadata and the unfinished ledger, canonical TQ ownership wins and the stale ledger entry is discarded.

Save workflow

flowchart LR
    Dataloader["Dataloader fetch"] --> Ledger["Reserve prompt group in ledger"]
    Ledger --> Admission["Sampler admission"]
    Admission --> Generation["Generate rollout"]
    Generation --> Commit["Canonical TQ commit"]
    Commit --> Replay["Completed replay index"]
    Commit --> Remove["Remove unfinished ledger entry"]

    Checkpoint["SC checkpoint"] --> Barrier["Exclusive checkpoint cut"]
    Barrier --> TQ["Save native TQ state"]
    Barrier --> ReplayState["Save replay metadata"]
    Barrier --> LedgerState["Save unfinished ownership"]
    TQ --> Bundle["Finalize step_N"]
    ReplayState --> Bundle
    LedgerState --> Bundle

Loading

Generation may continue while a checkpoint is written, but commits and destructive mutations wait at the barrier.

A failed TQ save, inventory mismatch, sidecar serialization failure, or incomplete cleanup prevents the new checkpoint from becoming the latest finalized checkpoint.

Restore workflow

flowchart LR
    Bundle["Latest finalized checkpoint"] --> TQ["Restore native TQ"]
    Bundle --> Replay["Restore completed replay index"]
    Bundle --> Ledger["Restore unfinished ledger"]
    Replay --> Reconcile["Canonical TQ wins"]
    Ledger --> Reconcile
    Reconcile --> Rehydrate["Rehydrate prompts from dataset"]
    Rehydrate --> Redispatch["Redispatch unfinished groups"]
    Redispatch --> Pumps["Run rollout and train pumps"]

Loading

Recovery:

1. Restores trainer, optimizer, dataloader, controller, and sampler dispatch state.
2. Loads native TQ and validates its exact inventory against replay metadata.
3. Restores the unfinished ownership ledger.
4. Removes ledger entries already canonical in TQ.
5. Rehydrates prompts from the restored dataset.
6. Launches already-admitted groups first.
7. Re-admits reserved batches exactly once using their admission IDs.
8. Starts ordinary rollout work after prioritized recovery dispatch.

Sampler support

Replay and unfinished-group recovery are supported by all built-in samplers:

1. windowed
2. ready_first
3. weight_fifo
4. in_order

Gated samplers persist and restore their exact dispatch index so target-step scheduling does not rewind after restart.

Custom samplers must explicitly declare:

supports_buffer_checkpoint = True

NOTE: Custom gated samplers participating in unfinished recovery should implement the transactional admission interface so gate waiting happens outside the checkpoint mutation cut.
Unsupported custom samplers use shadow-mode data-plane checkpoints and do not restore buffered rollout ownership.

Configuration

checkpointing:
  enabled: true
  checkpoint_dir: /shared/checkpoints/my-run
  save_data_plane: true

data_plane:
  enabled: true
  backend: simple  # tqv0.19 checkpointing only supports simple backend

async_rl:
  sampler:
    name: in_order  # windowed, ready_first and weight_fifo also supported

Current native restore support requires the TQ simple backend. Unsupported configurations fail during setup rather than silently producing incomplete checkpoints.

Failure handling

1. Failed rollout cleanup aborts the attempt for every caller.
2. A leaked unready slot would corrupt replay-capacity accounting.
3. A post-write cleanup failure may leave orphaned TQ rows that would invalidate the next checkpoint inventory.
4. Stale in-flight aborts remove ledger ownership before task cancellation.
5. Missing, corrupt, or schema-incompatible sidecars fail loudly.
6. The previously finalized checkpoint remains the fallback when a new save fails.

Scope and limitations

This PR recovers unfinished work through whole prompt-group regeneration.

It does not preserve:

1. partially decoded token prefixes;
2. completed siblings within an otherwise unfinished GRPO group;
3. vLLM KV-cache state;
4. NeMo-Gym episode/environment state;
5. sandbox filesystem or memory state.

Prompt rehydration currently assumes:

1. a deterministic map-style dataset;
2. integer positional sample IDs;
3. identical dataset ordering across checkpoint and restart;
4. a one-row-compatible collation function.

Only completed groups already committed to TQ preserve their exact generated tokens and tensor payloads.

Partially resolves #3594 to an extent where we do not lose out on any prompt during checkpoint / restore.
The remaining ones are:

1. Restore completed sibling generations
2. Ability to periodically checkpoint
3. Checkpoint / restore at token prefix
4. Checkpoint / restore env state

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

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

[macandro96]

/ok to test e31f42c

[macandro96]

Team review — 16 inline findings across bug/guideline/test-coverage/docstring categories, plus a leader design pass.

Summary: no critical bugs. Two BUG-level actionable items: (a) is True guard silently skips CustomSamplerConfig, letting a bad config waste cluster time; (b) _validate_replay_inventory runs after the native TQ snapshot is written, leaving a stale tmp bundle on validation failure. The rest are guideline/typing polish, docstring alignment, and 3 concrete test-coverage gaps (verifier tool, back-to-back checkpoint serialization, print-monkeypatch brittleness).

Devil's advocate: 12 confirmed, 5 disputed, 2 downgraded. Notable disputes: barrier mutation-counter leak on cancellation (no await between +=1 and try:); sampler_name=None handling (call site uses config, not save_state).

Design (leader pass — the automated design-reviewer agent could not get file-read access under its permission mode):

• DataPlaneCheckpointBarrier is hand-rolled around asyncio.Condition but this is the idiomatic build (no stdlib async rwlock); exception paths on both mutation() and checkpoint() correctly notify waiters. LGTM.
• Data-plane adapters (noop.py, transfer_queue.py) are symmetric and match the interfaces.py Protocol. Clean seam.
• Only real design smell is the ClassVar capability on config classes — flagged as a GUIDELINE finding.

Pre-commit: FAIL. Details in the inline comment anchored at single_controller.py:47. The untracked local file examples/configs/recipes/llm/grpo-qwen3-1.7b-1n8g-super-swe1-sc-tq-recovery.yaml also fails configs-minimize-check, but it's not part of this PR and is excluded from findings.

Filtered findings (below confidence threshold or premise disputed):

• barrier mutation-counter leak on cancellation (no await point between the += 1 and the try: — cancellation cannot land there)
• sampler_name=None handling in _maybe_restore_native_data_plane_checkpoint (misreading — call site passes master_config.async_rl.sampler.name, not save_state.sampler_name)
• _validate_replay_inventory running after load_state_dict in the restore path (exception is terminal; partial state is discarded)
• PERF-EVIDENCE ask for barrier hold time (marginal — barrier only holds during a checkpoint, which is already blocking)
• docs/ update ask (pre-existing SC docs don't cover this area at length; reasonable follow-up)

Generated by Claude Code

[macandro96]

/ok to test c781d74

[macandro96]

/ok to test 83eafa3

[macandro96]

/ok to test da8b0a9

[macandro96]

/ok to test 0b4cecc

[macandro96]

/ok to test 608e9bd

[macandro96]

/ok to test c3b8624

[copy-pr-bot]

/ok to test bca9a7210e9f32495b1923686ac2af3fa2199385

@macandro96, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[macandro96]

/ok to test bca9a72

[terrykong]

Re-review of the rollout-recovery ledger work (commits since the last round, at bca9a721) by the same six-agent team, every finding adversarially challenged before staging.

The headline is positive. Two agents independently tried to break the single-durable-owner invariant across every dispatch, failure, substitute, abort, and checkpoint-cut interleaving — and could not. The snapshot-skew bug from the last round is properly fixed (every restart-facing read now happens in one exclusive barrier cut, and admission was split so the cursor advance takes a mutation slot while the gate wait stays outside — the exact two-part shape we asked for). The round-trip tests are real, not mocked: real buffer, real samplers, real barrier, genuine disk restart. We re-ran all 287 in-scope unit tests locally on CPU: 287 passed, 0 failed. The functional test recovers a deliberately parked rollout and asserts on its specific group id.

8 comments below: one deliberate behavior change whose justifying comment needs broadening, one stale doc, two missing test assertions (selectability and stamp preservation - both verified passing against this branch before asking), and four low-severity notes including an endorsed future-proofing pattern for the barrier. 11 team findings were raised; 1 was disputed and dropped, 3 downgraded, 2 later removed as too minor, and the raise-scoping ask was inverted after the reviewer challenged it. Plus replies on the three prior-round threads.

Explainers for reviewers joining fresh: the ledger · the atomic cut · the test evidence

Generated by Claude Code

[macandro96]

/ok to test 2dd8212

[macandro96]

/ok to test 54d9e82

[macandro96]

/ok to test b8fa75b

[macandro96]

/ok to test 2017b60

[macandro96]

/ok to test a4a2a18

[macandro96]

/ok to test ee210a0