Skip to content

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

Open
macandro96 wants to merge 23 commits into
mainfrom
amahishi/partial-rollout-sibling-v3
Open

feat(sc): sibling level rollout checkpointing at train boundary#3923
macandro96 wants to merge 23 commits into
mainfrom
amahishi/partial-rollout-sibling-v3

Conversation

@macandro96

@macandro96 macandro96 commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

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

  • ...

@macandro96
macandro96 requested review from a team as code owners August 31, 2026 04:54
@copy-pr-bot

copy-pr-bot Bot commented Aug 31, 2026

Copy link
Copy Markdown

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 macandro96 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread nemo_rl/experience/rollout_manager.py Outdated
Comment thread tests/unit/single_controller/test_finalizer_lifecycle.py Outdated
Comment thread tests/unit/experience/test_rollout_generation_failures.py
Comment thread nemo_rl/algorithms/single_controller.py Outdated
Comment thread nemo_rl/experience/rollout_manager.py Outdated
Comment thread nemo_rl/algorithms/single_controller.py
Comment thread nemo_rl/experience/rollout_recovery.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/config.py
Comment thread tests/unit/single_controller/test_setup.py Outdated
Comment thread nemo_rl/algorithms/single_controller.py
@macandro96
macandro96 requested a review from a team as a code owner September 1, 2026 03:25
@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Sep 1, 2026

@terrykong terrykong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread tests/unit/experience/test_rollouts.py
Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py Outdated
Comment thread nemo_rl/experience/rollout_manager.py
Comment thread nemo_rl/experience/rollout_recovery.py Outdated
Comment thread tests/unit/single_controller/test_setup.py
Comment thread docs/guides/single-controller.md
Comment thread examples/configs/ppo_math_1B_megatron_single_controller.yaml
Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py Outdated
Comment thread nemo_rl/algorithms/single_controller_utils/config.py Outdated
Comment thread nemo_rl/experience/rollout_recovery.py
@macandro96 macandro96 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Sep 3, 2026
@macandro96
macandro96 force-pushed the amahishi/partial-rollout-sibling-v3 branch 2 times, most recently from 1a68c75 to b3c90e7 Compare September 4, 2026 02:20

@terrykong terrykong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py
Comment thread nemo_rl/algorithms/single_controller.py Outdated
Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
(cherry picked from commit 52e2dbc)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
(cherry picked from commit c1f1225)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 force-pushed the amahishi/partial-rollout-sibling-v3 branch from 27e2c01 to f5044d8 Compare September 6, 2026 00:11
@macandro96
macandro96 requested review from a team as code owners September 6, 2026 00:11
@github-actions github-actions Bot added the CI Relating to CI label Sep 6, 2026
@macandro96
macandro96 changed the base branch from amahishi/partial-rollout-base-v3 to main September 6, 2026 00:15

@macandro96 macandro96 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@dcoapp recheck

@macandro96 macandro96 left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +2073 to +2077
mask_sample = bool(
(((completion.env_extras or {}).get("instance_config") or {}).get(
MASK_SAMPLE, False
))
)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. ruff-format fails at head, so the pre-commit / lint CI job is red and this PR cannot merge as-is.

uv run ruff format --diff reports 3 files would be reformatted:

file line
nemo_rl/experience/rollout_manager.py 2073 (this hunk)
nemo_rl/experience/rollout_recovery.py 1063
tests/unit/single_controller/test_finalizer_lifecycle.py 274

This is PR-introduced, and specifically a regression: thread 8 flagged the identical ruff-format failure in round 1 and it was fixed then, so something re-introduced it after the rebase. Worth checking that pre-commit is actually installed locally (uv run pre-commit install), since this is the second time it has slipped through.

The fix for this hunk (verified ruff-format-stable — the outer paren pair was redundant):

Suggested change
mask_sample = bool(
(((completion.env_extras or {}).get("instance_config") or {}).get(
MASK_SAMPLE, False
))
)
mask_sample = bool(
((completion.env_extras or {}).get("instance_config") or {}).get(
MASK_SAMPLE, False
)
)

For the other two, uv run ruff format handles them; at rollout_recovery.py:1063 it wraps the message:

                    if not isinstance(mask_sample, bool):
                        raise ValueError(
                            "sealed attempts require a boolean mask_sample"
                        )

and at test_finalizer_lifecycle.py:274 it adds the two blank lines before class _SyncDataPlaneClient.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in b890c2e

Comment on lines +197 to +200
# Leaving either map above non-empty requires token capture, which is a separate
# top-level section (not in this file, and off by default):
# token_capture:
# enabled: true

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. This footer tells the reader token_capture is not in this file, but the PR added token_capture: 18 lines above it — following the footer produces a duplicate YAML key.

The block at token_capture: (line 179 in this file) and this footer are both added by this PR, and they contradict each other. A user who reads the footer and pastes the commented-out snippet ends up with two token_capture: keys at top level; depending on the loader that is either a hard parse error or a silent last-one-wins overwrite of staging_partition.

Since the section is in this file, the footer just needs to point at it:

Suggested change
# Leaving either map above non-empty requires token capture, which is a separate
# top-level section (not in this file, and off by default):
# token_capture:
# enabled: true
# Leaving either map above non-empty requires token capture, which is the
# top-level `token_capture` section above (off by default -- set
# `token_capture.enabled: true` to use either override map).

Same defect, same wording, in examples/configs/ppo_math_1B_megatron_single_controller.yaml:235-238 — that one also has its own token_capture: block earlier in the file.

(Defect is PR-introduced.)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in b890c2e

Comment on lines +1926 to +1927
assert request.prompt_idx == 9
assert restored._impl.seen_generation_indices == [0, 1]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. This test passes identically whether the restored granularity is PROMPT_GROUP or SIBLING, so it does not test what its name claims.

The test never seals a sibling — it calls mark_group_dispatched and nothing else — so at restore time every sibling is unfinished under both granularities, and seen_generation_indices == [0, 1] holds either way. I verified this by parametrizing the persisted default_granularity and running both variants: both pass. The comment on line 1901, "The saved group policy wins over the new process configuration", therefore asserts a discrimination the test does not make — the assertion would be satisfied even if the saved policy were ignored entirely.

The fake already records exactly the value you need: _CaptureImpl.seen_recovery_granularity is set at line 1621, but grep finds only 2 references in the whole tree, both inside the fake itself. Nothing asserts it.

Cheapest fix — assert the recorded policy:

Suggested change
assert request.prompt_idx == 9
assert restored._impl.seen_generation_indices == [0, 1]
assert request.prompt_idx == 9
assert restored._impl.seen_generation_indices == [0, 1]
assert (
restored._impl.seen_recovery_granularity
is RecoveryGranularity.PROMPT_GROUP
)

Stronger fix (catches more): seal sibling 0 before checkpointing, so SIBLING would redispatch only [1] while PROMPT_GROUP redispatches [0, 1]. That makes seen_generation_indices genuinely discriminating and tests the "reuses already sealed generations" contract from the config docs. Given prompt_group regeneration is a headline behaviour of this PR, that path deserves the real assertion.

(Defect is PR-introduced — the test is new in this PR.)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in b890c2e

Comment thread nemo_rl/experience/rollout_recovery.py Outdated
Comment on lines +1046 to +1064
try:
attempt_status = RolloutAttemptStatus(raw_attempt_status)
except ValueError as error:
raise ValueError(
f"invalid rollout attempt status={raw_attempt_status!r}"
) from error
receipt = attempt_state.get("receipt")
reward = attempt_state.get("reward")
mask_sample = attempt_state.get("mask_sample")
staging_keys = attempt_state.get("staging_keys")
if not isinstance(staging_keys, list) or not all(
isinstance(key, str) for key in staging_keys
):
raise ValueError("staging_keys must be a list of strings")
if attempt_status == RolloutAttemptStatus.SEALED:
if not isinstance(reward, (int, float)):
raise ValueError("sealed attempts require a reward")
if not isinstance(mask_sample, bool):
raise ValueError("sealed attempts require a boolean mask_sample")

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. The nine attempt-level restore validators here have zero test coverage, while the five group-level ones next to them are each tested.

grep -rl across tests/ for each of the messages raised in this block returns no files — none of these paths is exercised. Compare test_rollout_recovery.py:833-901, which has five test_restore_rejects_* cases for the group-level fields. The asymmetry looks unintentional rather than deliberate.

This matters more than usual because _SUPPORTED_ROLLOUT_RECOVERY_SCHEMA_VERSIONS = {5} with no v4 back-compat, so these validators are the only thing standing between a malformed or hand-edited checkpoint and a corrupt in-memory ledger. The newest one — "sealed attempts require a boolean mask_sample" — landed in 6f61d61bb with no accompanying test, which is how the gap grew.

Suggest mirroring the existing group-level style, one parametrized case per validator:

@pytest.mark.parametrize(
    "mutate, expected",
    [
        (lambda a: a.update(reward=None), "sealed attempts require a reward"),
        (lambda a: a.update(mask_sample="yes"),
         "sealed attempts require a boolean mask_sample"),
        # ... one row per raise in this block
    ],
)
def test_restore_rejects_malformed_attempt(mutate, expected):
    state = _valid_state_dict()
    mutate(state["groups"][0]["siblings"][0]["attempts"][0])
    with pytest.raises(ValueError, match=expected):
        _with_cut(buf, lambda cut: ledger.load_state_dict(cut, state))

I wrote this out against the real ledger while reviewing and it runs green (10 passed), so the validators are correct — they are just unpinned.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in b890c2e

Comment on lines +1163 to +1180
@staticmethod
def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord:
"""Copy mutable lineage metadata without duplicating the prompt payload."""
return PromptGroupRecoveryRecord(
group_id=record.group_id,
admission_id=record.admission_id,
prompt_id=record.prompt_id,
prompt_ref=record.prompt_ref,
agent_name=record.agent_name,
recovery_granularity=record.recovery_granularity,
runtime_prompt_payload=record.runtime_prompt_payload,
expected_generations=record.expected_generations,
target_step=record.target_step,
start_weight_version=record.start_weight_version,
siblings=copy.deepcopy(record.siblings),
phase=record.phase,
status=record.status,
)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

2 action items (both small; the second is the one I would insist on).

_copy_group hand-enumerates all 13 fields, so a field added to PromptGroupRecoveryRecord later is silently dropped on every groups() read.

PromptGroupRecoveryRecord (line 203) is a plain non-frozen @dataclass and all 13 fields are init=True, so this whole body is exactly equivalent to a one-liner — with the difference that the one-liner cannot go stale:

Suggested change
@staticmethod
def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord:
"""Copy mutable lineage metadata without duplicating the prompt payload."""
return PromptGroupRecoveryRecord(
group_id=record.group_id,
admission_id=record.admission_id,
prompt_id=record.prompt_id,
prompt_ref=record.prompt_ref,
agent_name=record.agent_name,
recovery_granularity=record.recovery_granularity,
runtime_prompt_payload=record.runtime_prompt_payload,
expected_generations=record.expected_generations,
target_step=record.target_step,
start_weight_version=record.start_weight_version,
siblings=copy.deepcopy(record.siblings),
phase=record.phase,
status=record.status,
)
@staticmethod
def _copy_group(record: PromptGroupRecoveryRecord) -> PromptGroupRecoveryRecord:
"""Copy mutable lineage metadata without duplicating the prompt payload."""
return dataclasses.replace(
record, siblings=copy.deepcopy(record.siblings)
)

(needs import dataclasses; copy is already imported.) The failure mode this removes is quiet: a new field would simply read back as its default through groups(), with no error anywhere.

Second item — the same drift risk applies to the serializer's field sets, and there the round-trip test cannot catch it. _GROUP_STATE_FIELDS / _ATTEMPT_STATE_FIELDS (lines 49-76) are six hand-maintained frozensets. A field omitted from both the serializer and the frozenset round-trips perfectly, so test_state_dict_round_trip stays green while the field is silently not persisted across a checkpoint — the exact class of bug this PR exists to prevent. Four lines pin it:

def test_group_state_fields_match_dataclass():
    assert _GROUP_STATE_FIELDS == {f.name for f in dataclasses.fields(PromptGroupRecoveryRecord)} - _GROUP_DERIVED_FIELDS

(Both are PR-introduced. Flagging per the repo's "silent-on-misconfiguration / drift-prone duplicated state" guidance, not as an architecture opinion — each has a concrete one-to-four-line fix.)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addressed in b890c2e

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@github-actions github-actions Bot removed the CI Relating to CI label Sep 6, 2026
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test ae533ea

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 force-pushed the amahishi/partial-rollout-sibling-v3 branch from ae533ea to b890c2e Compare September 6, 2026 01:37
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test b890c2e

@terrykong terrykong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-review at b890c2e0c, after the barrier-ownership work and the rebase onto merged #3837.

The three barrier threads from last round are properly fixed and are resolved with the commit
that fixed them. _clear_samples_unlocked and _cleanup_consumed_metas_unlocked both take a cut and
check it, and the nesting guard rejects a second section on both mutation() and checkpoint(). Ran
tests/unit/single_controller and tests/unit/experience at this head: 1144 passed, 0 failed, and
an instrumented barrier shows the guard never fires, so no existing path opens a nested section.
ruff format, ruff check and ruff check --select I are clean across nemo_rl and tests.

Two open threads have replies rather than new comments. The naming one is down to a preference and
can be closed either way. The advantage-stage one turned out not to be a defect — that write and the
only checkpoint() call site are the same coroutine, so there is nothing for a cut to serialize;
the reply explains what would break if it ever moved off the train pump.

Three things below are new. One is a follow-up guard on the mask_sample fix in this round, one
is a test-coverage gap in the checkpoint recovery matrix, and one is the original merge-order note now
that the rebase has landed. None of them block.

Generated by Claude Code

raise RuntimeError("data-plane barrier sections require an asyncio task")
if task in self._section_holders:
raise RuntimeError(
"this task already holds a data-plane barrier section; pass the "

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. Not a defect at this head — a merge-order hazard that lands silently on rebase.

TL;DR — after you rebase onto merged #3837 you inherit a contract that says "No writer is exempt", and _advantage_stage is the one canonical writer this PR does not give a cut.

This PR does the barrier-ownership work thoroughly — 16 functions gain a cut parameter, 17 call cut.require_live(), commit_finalized takes the caller's cut, and the guard on this line makes nesting an error instead of something you can write by accident. That covers the finalizer path, which is the one that actually needed it.

One canonical writer is left out: _advantage_stage (:3589) writes with await self._call_dp("put_samples", ...) at :3741, and there is no mutation(), cut or require_live anywhere between those two lines.

That is fine at this head, and that is exactly why it is easy to miss. Your metadata_state_dict docstring (:1491-1507) still carries the original disjunction, which sanctions it by name — "must either participate in that barrier ... or run in the same asyncio task as the checkpoint save. The advantage stage relies on the latter". Code and contract agree.

#3837 rewrote that paragraph. At its head the same docstring reads "No writer is exempt", with both the same-task arm and the "future finalizer paths" sentence removed. Since this PR does not edit that paragraph, the rebase takes #3837's version cleanly, with no conflict — so nothing will prompt anyone to re-check _advantage_stage against the stricter text it just inherited.

AI-1

Pick one when you rebase. Either give _advantage_stage a cut like the other writers, or keep the disjunction and let the docstring say why that writer does not need one. If you keep it, the durable reason is better than the same-task one that #3837 removed: those rows have already left the replay index by then, because BaseSampler._finalize_selection removes them at select via remove(selected_idxs, remove_in_dp=False) — which stays true even if the advantage stage ever moves to another task. No suggestion block: the edit depends on which option you pick and lands after a rebase that has not happened yet.

Context — no action. Raised here rather than on #3837 because this PR owns the barrier-ownership work and is where the fix belongs. The matching comment on #3837 has been removed so this is not asked twice.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Traced it at b890c2e0c: no gap today, and the reason is exactly what you said — the
write and the only checkpoint are the same coroutine, so a save cannot be in progress while
advantages are being computed.

The barrier orders concurrent tasks. There is one task here, so Python already orders it, and a cut
would be decoration.

On restore: the rows being written are already out of the replay index —
select() drops them with remove_in_dp=False before advantages run, so they stay in TQ
with nothing referencing them until the post-step clear. A crash mid-step therefore falls
back to the previous checkpoint, which predates the selection: the group is still in the restored
index, its rows are still in TQ, and the train pump re-selects it and recomputes. Nothing skipped,
nothing counted twice.

The future-proofing point is bigger than this one write, though, and it is worth keeping for that
reason. The dependency is not "the advantage stage must not race a checkpoint" — it is the entire
window from select() to the post-step clear must not overlap a save
. Throughout that window TQ
holds rows the saved replay index does not list, because the index is rebuilt from
meta_list and select() already removed them.

So if this moved into its own pump or coroutine, the failure would not be a silent gap — it would be
_validate_replay_inventory raising on unexpected sample ids and the checkpoint failing
outright. Loud, but it would take checkpointing down, and the same is true for anything else lifted
out of that window, not just advantages.

Suggested close: nothing to change in the code, but one line near
_advantage_stage saying it may write with no cut because it only ever runs inline in the
train pump, inside the select-to-clear window
would keep the next person from either re-deriving
this or moving it somewhere that breaks saves. Happy for this to be closed with or without that.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@macandro96 maybe just add a comment about this on why that particular TQ advantage write doesn't need to be in a mutation cut just in case since it doesn't follow the other situations that required a mutationcut

Comment thread nemo_rl/algorithms/single_controller_utils/config.py Outdated
Comment on lines +275 to +284
@dataclass(frozen=True)
class SiblingSealResult:
"""One terminal sibling result waiting for an atomic prompt-group seal."""

gate_rollout_id: str
# None is an explicit terminal capture failure. The finalizer turns it
# into a masked placeholder, matching the base token-capture contract.
receipt: Optional[dict[str, Any]]
reward: float
mask_sample: bool = False

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. Low severity — a guard on the bug this commit just fixed.

TL;DR — mask_sample is the per-sibling "produced, but do not train on this one" flag. Persisting
it in 6f61d61bb fixes a real restore bug; the two in-memory constructors still default it to
False, which is the same silent unmask one call site away.

What the flag is, since it is easy to read as something else. The environment sets it per rollout
in extra_env_info.instance_config.mask_sample. It does not control what gets saved and
it does not drop the rollout: the row is still built, still stored, still counted in the group. It
rides through as a bool column on the training row
(MASK_SAMPLE in the batch), and the only place it acts is the advantage stage, which
zeroes that sample's weight with sample_mask * (~mask_sample) and
counts it as num_mask_sample_filtered. So it is one bit per sibling, saved alongside
that sibling either way, and read at training time.

What was broken. Before this commit the bit was never persisted — it lived in a
mask_sample_by_index dict built while rollouts streamed in, and finalization read it back with
.get(index, False). Restart between sealing a sibling and finalizing its group and that dict is
empty, so every restored sibling finalized as False: the environment said do not train on this, the
flag was silently dropped, and the sample got a full-weight training row. Persisting it per attempt
and returning it from finalization_inputs closes that.

What is left. The parser demands a bool for a sealed attempt — that half is
strict. The Python side is not: this dataclass field and mark_sibling_sealed both
default to False. Both current callers pass it explicitly, so nothing is wrong today, but the
default reproduces exactly the behaviour that was just removed — a sibling sealed without the flag
trains at full weight, with nothing to notice it.

Action: drop both defaults so mask_sample is required. A caller that omits it then fails at the
call instead of producing an unmasked training row. Not a suggestion block: it edits this dataclass
and a signature in another region.

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py
Comment thread nemo_rl/algorithms/single_controller.py Outdated
Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py
Comment on lines +675 to +690
def test_restart_preserves_sealed_sibling_and_retries_only_interrupted_one() -> None:
ledger = RolloutRecoveryLedger()
group = _reserve(
ledger,
group_id="g7",
admission_id="batch-7",
prompt_id="7",
prompt_payload=_prompt(),
expected_generations=2,
target_step=7,
start_weight_version=6,
agent_name=None,
recovery_granularity=RecoveryGranularity.SIBLING,
admitted=True,
)
_mutate(lambda cut: ledger.mark_group_dispatched(cut, "g7"))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. Test coverage — the matrix is missing the shape this PR exists for.

TL;DR — the recovery matrix reserves ledger groups but never seals a sibling, and no scenario
strands an unfinished group behind the current step. So "some siblings finished at step 5, checkpoint
at step 10, restore and reuse exactly those" is not covered anywhere across samplers.

This test is the unit-level version of it and it is a good one. What is missing is the same shape in
the checkpoint recovery matrix, which is where it would run against all four
samplers and both lag settings.

Two separate holes there:

  • No sibling is ever sealed. The matrix builds its ledger with
    reserve_group onlyGroup.done decides whether a group goes in the ledger at
    all, but every attempt stays RESERVED. So the matrix proves an unfinished group is owned across
    a restart; it never proves the finished siblings inside it are reused rather than regenerated,
    which is the behaviour this PR adds.
  • An unfinished group is never far behind. In every in-flight scenario the partly-generated group
    carries the newest target — S_PARTIAL and
    S_TRAINED_OUT_OF_ORDER both leave group 12 at target=5 while the run is on step 5.
    S_STALE_ONLY is the one case with a wide target gap, and there both groups are fully
    generated. Nothing combines "still generating" with "stamped several steps back".

That combination is exactly what ready_first produces: it trains whatever is ready and leaves an
earlier partly-generated group behind, so that group can sit at target=5 while training reaches
step 10 — and its sealed siblings have to come back attached to the right target step, not be
regenerated or re-stamped.

Action: add a scenario with a partly-generated group whose target is several steps behind the
newest one, and seal a subset of its siblings before the snapshot, then assert after restore that the
sealed generation indices come back unchanged and only the missing ones redispatch. Running it over
the existing sampler list would cover the ready_first path specifically. Not a suggestion block:
it adds a scenario and a sealing step to files outside this PR's diff.

@terrykong terrykong left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-anchoring one thread the rebase left outdated; the point is unchanged.

Generated by Claude Code

Comment on lines +625 to +634
@dataclass(frozen=True)
class RecoveryGranularityResolution:
"""Recovery granularity selected for a prompt-group reservation.

``agent_name`` is copied from the prompt when present. ``granularity`` is
selected from an agent override, task override, or the global default.
"""

agent_name: Optional[str]
granularity: RecoveryGranularity

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

1 action item. Naming only, no behavior change. Re-anchored from
the original thread, which the rebase left pointing at a line that no longer exists.

TL;DR — the docstring half of that thread is fully fixed; the name still says something is resolved
when only one of the two fields is.

Fixed and worth saying so: Policy is gone, and the docstring now spells out what each field is
rather than calling them "policy coordinates".

What is left is small. agent_name is not resolved — it is read off the prompt and
returned unchanged on every path, including the one where no override matched.
Only granularity is selected, by resolve_for_prompt walking agent override, then
task override, then the default. So Resolution carries the same overreach Resolved did, one word
further along.

AgentRecoveryGranularity names the two fields it holds and claims nothing about how they got there.

Action: rename, or close this — it is cosmetic and the substantive half is already done. Flagging
it only because the rebase hid the earlier reply from the Files changed view. Not a suggestion
block: it renames a symbol used in four other places in this file.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants