feat(sc): sibling level rollout checkpointing at train boundary - #3923
feat(sc): sibling level rollout checkpointing at train boundary#3923macandro96 wants to merge 23 commits into
Conversation
macandro96
left a comment
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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/andtests/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
1a68c75 to
b3c90e7
Compare
terrykong
left a comment
There was a problem hiding this comment.
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
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>
27e2c01 to
f5044d8
Compare
macandro96
left a comment
There was a problem hiding this comment.
@dcoapp recheck
macandro96
left a comment
There was a problem hiding this comment.
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
RolloutReassemblerand onto the controller is the right seam — the finalizer genuinely cannot see the group-level picture, and the test comment attest_finalizer_lifecycle.py:206-211documents 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.
| mask_sample = bool( | ||
| (((completion.env_extras or {}).get("instance_config") or {}).get( | ||
| MASK_SAMPLE, False | ||
| )) | ||
| ) |
There was a problem hiding this comment.
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):
| 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.
| # 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 |
There was a problem hiding this comment.
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:
| # 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.)
| assert request.prompt_idx == 9 | ||
| assert restored._impl.seen_generation_indices == [0, 1] |
There was a problem hiding this comment.
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:
| 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.)
| 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") |
There was a problem hiding this comment.
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.
| @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, | ||
| ) |
There was a problem hiding this comment.
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:
| @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.)
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
|
/ok to test ae533ea |
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
ae533ea to
b890c2e
Compare
|
/ok to test b890c2e |
terrykong
left a comment
There was a problem hiding this comment.
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 " |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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 advantage write is at
_train_pump; _save_checkpointis called from the same function, same task, later in the same
loop iteration;- and there is exactly one
checkpoint()call site in the tree, inside_save_checkpoint.
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.
There was a problem hiding this comment.
@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
| @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 |
There was a problem hiding this comment.
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.
| 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")) |
There was a problem hiding this comment.
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_grouponly —Group.donedecides whether a group goes in the ledger at
all, but every attempt staysRESERVED. 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_PARTIALand
S_TRAINED_OUT_OF_ORDERboth leave group 12 attarget=5while the run is on step 5.
S_STALE_ONLYis 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
left a comment
There was a problem hiding this comment.
Re-anchoring one thread the rebase left outdated; the point is unchanged.
Generated by Claude Code
| @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 |
There was a problem hiding this comment.
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.
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:
Without this control-plane lineage, a restart either loses completed siblings or risks consuming stale and duplicate attempts.
What changed
Configuration
This feature requires token capture and a checkpoint containing native TQ state:
Recovery semantics
For
siblingrecovery:For
prompt_grouprecovery:Scope and non-goals
This PR recovers completed sibling-level work at a full trainer checkpoint boundary. It does not:
Those capabilities are layered in later work.
Tests
Coverage includes:
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information