feat(sc): sibling level rollout checkpointing at train boundary

[macandro96]

What does this PR do ?

Summary

Adds controller-owned recovery lineage for unfinished NeMo-Gym prompt groups, allowing completed siblings to survive a full SingleController checkpoint and restart.

Recovery behavior is configurable as either:

• sibling: reuse sealed siblings and redispatch only unfinished siblings.
• prompt_group: discard and regenerate the complete group when any sibling fails.

The selected policy is applied consistently to both live failures and checkpoint recovery.

Builds on top of #3480 and #3837. Note token capture must be enabled for this feature to work.

Why

A native TQ checkpoint preserves token payloads, but TQ alone does not describe:

• which prompt group owns each staged row;
• which generations in the group completed;
• which physical attempt produced a result;
• whether finalization completed;
• which siblings must be redispatched after restart.

Without this control-plane lineage, a restart either loses completed siblings or risks consuming stale and duplicate attempts.

What changed

• Adds a versioned, metadata-only rollout recovery ledger.
• Tracks stable prompt-group and sibling identities across physical retries.
• Persists sibling attempts, sealed receipts, rewards, recovery granularity, sampler admission and finalization ownership.
• Restores unfinished prompt groups from the dataset using stable prompt references.
• Validates restored lineage against the staged rows in TQ.
• Clears unreferenced staging rows while failing closed if referenced rows are missing.
• Supports sealed placeholder siblings with a nullable receipt.
• Applies prompt-group recovery atomically:
  • live failures redispatch the complete cohort;
  • late results from superseded attempts are rejected;
  • the group is sealed only after every sibling in the same attempt completes.
• Places reservation, admission, sibling sealing, finalization and cleanup transitions behind the data-plane checkpoint mutation barrier.
• Removes replay groups by stable group ID after asynchronous data-plane operations, avoiding index-shift races.
• Propagates stable prompt indices through finalization.

Configuration

rollout_recovery:
  default_granularity: sibling

  # Matches extra_env_info.agent_ref.name.
  # Agent overrides take precedence over task overrides.
  agent_granularity_overrides:
    some_agent: prompt_group

  # Matches the prompt's task_name.
  task_granularity_overrides:
    some_task: prompt_group

This feature requires token capture and a checkpoint containing native TQ state:

checkpointing:
  enabled: true
  save_data_plane: true

token_capture:
  enabled: true

Recovery semantics

For sibling recovery:

1. Sealed siblings retain their original logical result and staged token rows.
2. Unfinished sibling attempts are abandoned.
3. Only missing generations receive a new physical attempt.
4. The group is finalized once all siblings are sealed.

For prompt_group recovery:

1. If any sibling fails, every result from the current group attempt is discarded.
2. The complete group receives a new attempt.
3. Results from older attempts are rejected.
4. All siblings are sealed atomically.

Scope and non-goals

This PR recovers completed sibling-level work at a full trainer checkpoint boundary. It does not:

• checkpoint an unfinished vLLM token prefix;
• persist Gym environment or sandbox state;
• create periodic snapshots during an active streamed optimizer step.

Those capabilities are layered in later work.

Tests

Coverage includes:

• sibling reuse across process restart;
• prompt-group recovery for both live and restart failures;
• atomic prompt-group sealing;
• nullable receipt placeholders;
• malformed recovery configuration and state;
• missing and orphaned TQ staging rows;
• finalization ownership;
• reentrant mutation barriers;
• stable-ID replay removal under concurrent mutation.

Issues

List issues that this PR closes (syntax):

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]

PR #3923 review — sibling-level rollout checkpointing

Reviewed by a 5-agent team (RL codebase, NeMo-Gym API, tests, bug-finder, design) plus an adversarial devil's-advocate pass that ran the repo's own test suite and pre-commit hooks at both the diff base (9469fd593dff) and PR head (a03f311c46225da2242b904cbe5d88305c72a05f), and independently verified every claim below against the source (not against the other reviewers' summaries).

This is a substantial, largely well-built feature — the recovery ledger itself is genuinely pure (no Ray/torch/TQ client dependency) with a strong, mock-light unit suite, and the design choice to persist resolved granularity per-group (rather than re-reading it from YAML at restore) is good judgment that avoids a real footgun.

4 blocking issues (a reward-correctness regression and 3 tests that fail on HEAD) and 8 should-fix items are inline below. A larger set of FYI/nit findings and design affirmations were filtered out as non-actionable or already-handled by the code.

Every inline finding states whether it's PR-introduced or pre-existing, and links the exact lines it depends on.

Generated by Claude Code

[terrykong]

Re-review at 0a93f0ffd, after the 13 threads from the previous round were answered. Six reviewers plus an adversarial pass; every finding below was run, not reasoned about, and each names what was run.

Nine of the thirteen prior threads are genuinely fixed, each re-checked against the code rather than taken at face value. Worth calling out two:

• The effort-shaping reward bug is properly fixed, and the fix quietly repairs a second latent bug: the old call paired the compacted results list against the full inputs list, so rows were mismatched whenever one was missing. The new per-row pairing fixes both, and there is a direct regression test.
• Deleting the train-step lifecycle was clean — no dangling references anywhere, every remaining enum member still used, and a test that asserts the removed key is gone.

Four are partly fixed, two of which are inline below (lint, and the test_setup assertion). The other two: the stats stub fix reached the stub it named but a sibling stub in the same file hits the same line, and the finalizer-barrier thread got a good explanation but no before/after checkpoint-start numbers. Neither is worth a new comment — the first is masked by a failure that predates this PR, and the second is answered on the design half.

Design is in good shape and two choices deserve credit. Resolving the recovery policy once and persisting it per group — rather than re-reading YAML on the restore path — avoids a real footgun, and most implementations of this would get it wrong. The stable-ID rework in the replay buffer closes a whole class of index-shift bug rather than patching one instance, and it fixes a live misalignment on the restore path.

Two things not worth an inline comment:

• DCO is failing. Two mid-stack commits have no Signed-off-by, so it needs an amend or rebase rather than a new commit on top.
• 48 unit tests under tests/unit/single_controller/ and tests/unit/experience/ already fail at the diff base, so they belong to the base branch. I mention it only because it is the control that isolates the one regression flagged below.

How this was run, so you can weight it. Everything was checked on CPU: no GPU, no NeMo-Gym, no Megatron. The ledger round trip, the barrier's cancellation behaviour, the lint hooks at both base and head, and the config typo case were all executed. The functional recovery test and anything needing a real TransferQueue were not — those are read from source only, and the one finding that rests purely on a code trace says so.

Generated by Claude Code

[terrykong]

Follow-up on the cut rollout in 6ceafb405. The two threads it answered
(the escaping sweep and the nested cut) are both properly fixed, and
two conversions landed that were never asked for in a comment — commit_finalized and
_remove_groups_unlocked both take a cut now.

The removal of the task-keyed join is safe. That was the risk worth checking: a nested
mutation() no longer joins, it waits, so if a checkpoint were pending it would wait for a
checkpoint waiting on this task's own cut. Instrumented the barrier to record same-task nested
acquisition and ran 1039 tests across tests/unit/single_controller and
tests/unit/experience — no task ever opens a section while already holding one. 991 passed,
48 failed, and the 48 are the same base-branch failures called out last round, all in files this
round did not touch. ruff format, ruff check and ruff check --select I are clean repo-wide.

What is below is the tail of the same rule: two helpers that still state the requirement in prose
instead of taking the cut, and one guard that is now cheap to add. None of it is blocking.

Generated by Claude Code

[macandro96]

@dcoapp recheck

[macandro96]

Round 4 review

Focused on the three commits not previously reviewed (5412881f2..HEAD), with a lighter pass over the rest after the rebase.

Verification performed: ran the unit suite at head — 606 passed, plus 9 vLLM-mode and 16 NeMo-Gym-mode tests, no failures. Separately re-derived the four open threads (T27, T29, T30, T31) against the working tree: all four are correctly fixed, and I found no regressions among the 27 already-resolved threads. Replies going on those threads separately.

5 comments, 1 of them blocking (ruff-format is red at head — everything else is a test-coverage or docs/robustness item).

Two things I want to call out as good, since they were non-obvious calls:

• Moving the low-valid-row-fraction drop decision out of RolloutReassembler and onto the controller is the right seam — the finalizer genuinely cannot see the group-level picture, and the test comment at test_finalizer_lifecycle.py:206-211 documents the distinction well.
• Refusing v4 back-compat (_SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {5}) rather than writing a migration shim is the correct trade for a checkpoint format nobody has in production yet.

I also chased a possible double-application of the reward penalty through the streamed-vs-batched completion paths and satisfied myself it is not reachable: the non-receipt path raises at rollout_manager.py:2043 before the second conversion at :1205 can run. Noting it only so the next reviewer does not have to re-derive it.

[macandro96]

/ok to test e7447f7