feat(checkpoint): orchestrate Gym turn-level recovery - #4266
macandro96 wants to merge 2 commits into
Conversation
github-actions
Bot
added
the
Documentation
There was a problem hiding this comment.
Withdrawn: this was posted early by mistake. The current review is #4266 (review).
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
4db8f66 to
7ebf4aa
Compare
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Thanks for splitting this out of #4117, and for adding turn_recovery to the new rollout_checkpointing.gym.mode. Every fix from that review carried over, and the publish-or-abort fence, the same-ID retry, the separate admission gate per owner, and the strict fallback-marker model are all clean.
- Blocker from #4264: the Gym submodule pin is not on Gym main. This head pins
d1abebd(71 commits ahead of Gym main, 57 behind); the earlier pinf4fcf8cfcame from open Gym PRs #3349, #3412, #3563. The stack can't merge until the pinned Gym changes land. - Please fix before merge:
- While a Gym snapshot has rollouts parked, RL's
rollout_timeout_sdeadline keeps running, so a rollout that times out while parked is never ACKed and every later Gym prepare fails. The fix is to pause the request deadlines during a Gym checkpoint, counting holders, because the colocated train switch pauses the same deadlines (comment onsingle_controller.pyL1958). - The turn-recovery functional test and the docs example still set the old Gym keys that
modereplaced. The config model rejects unknown keys, so all three new functional tests fail at startup (comment on the test script L163, docs L168). - An existing unit test fails on this PR:
test_borrow_and_repayment_survive_controller_checkpoint_restorebuilds the controller without__init__and misses three attributes this PR adds (comment onsingle_controller.pyL668).
- While a Gym snapshot has rollouts parked, RL's
- CI: only the DCO, copyright, title, and submodule checks have run on this head. Please post one green L1 SingleController run. The three new functional tests run only in full mode, each up to ~80 min worst case, and the PR checklist item is still unchecked.
- Perf numbers (not blocking, since the feature is opt-in): please compare participant checkpointing on vs off at
save_period: 1andsnapshot_attempt_interval_s: 120. Report per-step checkpoint time, rollout tokens/s, and how long admission stays closed per snapshot. Every trainer checkpoint now waits for the async weight write (finalize_pending) and a full Gym pause, and the no-mutation skip is off in this mode (a one-line comment saying why would help).total_save_secondsnow includes Gym time that no per-part metric covers; onegym_participant_secondsmetric would make the breakdown add up. - Upgrade note for the PR body:
token_capture.enabled=truenow requiresmax_row_attempts=1(default 3), so token-capture configs outside this repo fail at startup until they are updated. - Still open from #4117: the resources-session leak (thread on
nemo_gym.py, fix is on the Gym side). This PR is the first one that can hit it, because it is what runs prepare/commit over and over. - Unused code in the stack:
GymAgentStatusResponsefrom #4264 is never used (comment onsingle_controller.pyL161). The other unused code this pass found, the leftover attempt UUID and the test-onlycompleted_execution_acknowledgements(), is already removed in the latest #4265. - Checkpoint levels (design change, before release): The new
rollout_checkpointing.gym.modeis a good step, butmode: turn_recoveryand the per-environment granularity overrides are still two settings for one decision. The target-level comment onconfig.pyL1396 proposes onerollout_recovery.target_level: turn | sibling | prompt_group, lowered per environment from what Gym reports, with discovery running on every token-capture run. The related guards are onsetup.pyL1876-L1878 andsingle_controller.pyL675, and the docs section onsingle-controller.mdL164. - Follow-up, not for this PR: checkpoint and rollout-recovery code is now about a third of
SingleControllerActor. #4325 tracks moving it into aRolloutCheckpointManager(composition instead of more private methods on the actor), after this stack merges.
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4266-gym-turn-recovery.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4266-parked-deadline-stall.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4266-step-during-gym-pause.html
Generated by Claude Code
| # Independent admission owner for Gym checkpoint preparation. Weight | ||
| # synchronization and checkpointing may overlap; separate events avoid | ||
| # one operation reopening admission while the other still needs it shut. | ||
| self._gym_checkpoint_rollout_permitted: asyncio.Event = asyncio.Event() |
There was a problem hiding this comment.
1 action item — please fix in this PR.
TL;DR — this PR breaks test_borrow_and_repayment_survive_controller_checkpoint_restore. The test builds the controller with object.__new__, so it never gets three attributes this PR sets in __init__ and then reads. It fails with AttributeError: 'SingleControllerActor' object has no attribute '_gym_checkpoint_rollout_permitted'.
The test is not in the PR's test plan and has no marker, so it runs in the L0 Other lane. No unit CI has run on this head yet, which is why the failure has not shown up.
- The test builds the controller without calling
__init__(test_checkpoint_borrow_restore.py#L199). - This PR sets
_gym_checkpoint_rollout_permittedhere, and sets_gym_participant_checkpointing_enabledand_gym_restart_unfinishedin__init__. - The test path reads all three: dispatch reads the first in
_wait_for_rollout_admission, the snapshot cut reads the second (L4339), and restore reads the third (L1155).
Run on CPU: the test fails on this head. Each of the three is required, and the test passes once all three are set.
AI-1
Add the three attributes to the hand-built controller in test_checkpoint_borrow_restore.py, after L205 (asyncio is already imported there):
controller._gym_checkpoint_rollout_permitted = asyncio.Event()
controller._gym_checkpoint_rollout_permitted.set()
controller._gym_participant_checkpointing_enabled = False
controller._gym_restart_unfinished = FalseThis can't be a suggestion block because the test file is outside this PR's diff.
| from nemo_rl.distributed.batched_data_dict import BatchedDataDict | ||
| from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_context_lost | ||
| from nemo_rl.environments.nemo_gym import should_use_nemo_gym | ||
| from nemo_rl.environments.gym_checkpoint import ( |
There was a problem hiding this comment.
1 action item (low priority).
TL;DR — GymAgentStatusResponse (added in #4264) is never used. RL never calls Gym's agent status route, and no code or test refers to the class. Delete it.
The only Gym status route RL polls is the model admission /status in _wait_for_policy_model_pause. It parses GymCoordinatorModelStatusResponse or GymSingleWorkerModelStatusResponse. The agent's prepare state reaches RL only through GymAgentPrepareResponse.
AI-1
Delete GymAgentStatusResponse. This can't be a suggestion block because gym_checkpoint.py is outside this PR's diff.
| # the response, so every attempted prepare needs an idempotent abort. | ||
| prepare_attempted = True | ||
| prepare = GymCheckpointPrepareResult.model_validate( | ||
| await gym_actor.prepare_checkpoint.remote( |
There was a problem hiding this comment.
1 action item.
TL;DR — a Gym failure during a periodic snapshot reaches the pump as a RuntimeError, and the pump retries only OSError/TimeoutError. So one HTTP 5xx or control-call timeout on /commit ends the run instead of being retried.
This PR sends Gym prepare/commit failures into the pump's existing retry check. How it shows up:
- The Gym actor raises
GymControlRequestError, which subclassesRuntimeError, for any non-200 reply, and a plainRuntimeErrorwhen a control call times out.commit_checkpointconverts neither. Only prepare-past-deadline becomesTimeoutError(nemo_gym.py#L1305-L1309). - Ray re-raises it as the cause's class (object_ref.pxi#L33-L34 →
class cls(RayTaskError, cause_cls)), so SC sees aRuntimeError. _prepare_and_commit_gym_checkpointaborts Gym correctly, then re-raises the error.- The pump logs
reason=invariant_errorand re-raises anything that is notOSError/TimeoutError. The run ends.
Example: a resources server returns HTTP 500 on /commit because its artifact write hit ENOSPC. The run dies. The same ENOSPC on SC's own TQ write is an OSError and is retried. The PR body says a failed periodic attempt "leaves the previous published rollout snapshot authoritative", which reads as "retried".
The unit tests miss this because the fake actor raises TimeoutError from prepare and OSError from commit. The real actor never raises OSError from commit.
AI-1
Wrap only the two Gym RPCs (the prepare_checkpoint call here and the commit_checkpoint call) so their failures become a retryable error. For example, add class _GymCheckpointAttemptFailedError(OSError) and raise it from the original. Keep SC's own checks below them (wrong checkpoint ID, validate_checkpoint_participants, gym_checkpoint_continuations) raising as today, so real bugs still stop the run. Then make _FakeGymCheckpointActor raise GymControlRequestError from commit, like the real actor. This can't be a suggestion block because it touches two call sites, a new class, and the test file.
Why not wrap every error in the except block
Wrapping the whole except BaseException path would also turn SC's own checks in the same try (the checkpoint-ID mismatch and participant validation) into retried io_errors. A real bug would then be retried max_consecutive_failures times and logged as a storage problem.
| if ( | ||
| token_capture_config.enabled | ||
| and async_config.rollout_failure.nemo_gym.max_row_attempts != 1 | ||
| ): | ||
| raise ValueError( | ||
| "token_capture.enabled=true requires " | ||
| "async_rl.rollout_failure.nemo_gym.max_row_attempts=1. " | ||
| "Token-captured rows use stable (rollout_id, attempt_index) " | ||
| "identities, so an immediate row redispatch could overlap the old " | ||
| "Gym execution. Higher-level rollout recovery creates a new tracked " | ||
| "attempt instead." | ||
| ) |
There was a problem hiding this comment.
1 action item (design change; part of it belongs in Gym).
TL;DR — Fold rollout_checkpointing.gym.mode: turn_recovery and the per-environment recovery overrides into one target level, rollout_recovery.target_level: turn | sibling | prompt_group, and let RL lower it per environment from what Gym reports; the per-environment overrides stay only to force a coarser level. Rejecting prompt_group together with checkpointing, as this comment first asked, would leave group-scored environments such as GenRM compare with no valid setting.
Why the current shape does not work:
- The prompt-group stall is real, but it comes from when the ACK is sent, not from the combination. The feat(checkpoint): persist recoverable Gym rollout state #4265 review asks to record each sibling's ACK as soon as it arrives (#4265 comment on
rollout_manager.pyL2434). With that fix, prompt-group plus turn-level checkpointing commits normally. - A group-scored environment scores all N siblings together: each sibling waits in
/verifyuntil the whole group is there (Gymgenrm_compare/app.py#L361-L376). A sibling retried alone lands in a new, empty group and is never scored. So such an environment must always retry its whole group. That is a property of the environment, and one run can mix it with environments that score each sibling alone. - Today the user has to know this and add a per-environment override under the exact Gym name. A key that matches nothing silently falls back to
sibling(resolve_for_prompt).
AI-1
Add one target level and derive each environment's effective level from it:
rollout_recovery.target_level, finest to coarsest:turnkeeps sealed siblings and resumes unfinished ones from their saved turn;siblingkeeps sealed siblings and reruns unfinished ones from scratch;prompt_groupreruns the whole group from scratch and discards saved turns. It takes over theturn_recoveryvalue ofrollout_checkpointing.gym.mode;discoverstays as it is (earlier revision:config.py#L760). The overrides (config.py#L684-L694) stay only to force a coarser level for one environment.- Gym change first: group-scored resources servers report "group-scored, group size N" in
control_capabilities(). Discovery then runs on every token-capture run, so RL knows every environment's limits at any target. No user-supplied list of group-scored environments: the capability is the source of truth. - At setup, RL computes each environment's effective level: the finest level at or below the target that the environment supports. A group-scored environment always retries its whole group on a live failure. At target
turnits siblings still resume from their saved turns on restore, because they reach/verifytogether. Log one warning for every environment whose level was lowered. - W&B metrics, next to the existing retry counters (
RolloutManager.stats): groups retried whole because the environment is group-scored, sealed siblings discarded by those retries, and siblings resumed from a saved turn vs rerun from scratch. - At target
prompt_group, restore discards every saved turn, so the group starts over on the current weights. The discard already exists forrestart_onlyresources (_discard_restart_only_gym_continuations); extend its candidate list.
GenRM today (belongs in Gym): GenRM does not resume at turn yet. Production genrm_compare is restart_only: GenRMCompareConfig keeps Gym's conservative default (base_resources_server.py#L95-L98), and every simple_agent saved turn depends on its resources server (simple_agent/app.py#L222). So on restore #4266 discards every GenRM saved turn and the whole group restarts (single_controller.py#L1810-L1834). GenRM keeps no state that must survive a restart: its only state is the in-memory scoring groups, and a resumed group forms again and is scored once. So the fix is one line in Gym: set CHECKPOINT_RECOVERY_MODE = "stateless" on GenRMCompareConfig, as its test copy already does (checkpoint_test_app.py#L24-L27). Gym's own GenRM agent config already says "Cohort verification is stateless and may be replayed after recovery" (genrm_compare.yaml#L125-L126).
Not a suggestion block: it changes the config schema, setup, the restore path and a Gym capability.
What a GenRM group gets at each target
| Target | Live failure of one sibling | Restore |
|---|---|---|
turn |
whole group retried | all siblings resume from their saved turns together; the group forms again at /verify (after the Gym stateless fix; today the group restarts from scratch) |
sibling |
whole group retried | whole group rerun from scratch |
prompt_group |
whole group retried | whole group rerun from scratch; saved turns discarded |
A GenRM sibling cannot finish before its whole group is scored, so at a checkpoint the group is either all unfinished (parked, or waiting in /verify) or all finished (and sealed before the snapshot commits). Resuming at turn therefore never splits a group. Resumed turns come from the weights at save time, so a resumed group no longer comes entirely from one set of weights; the docs should say so. This needs checkpoint_replayable_verify on the GenRM agent (see the comment on setup.py L1877).
| if isinstance( | ||
| gym_snapshot_error, | ||
| (_GymCheckpointAbortPendingError, _GymCheckpointReleasePendingError), | ||
| ): | ||
| # The checkpoint is recoverable, but this live Gym process may | ||
| # still be paused. Fail closed instead of resuming training. | ||
| raise gym_snapshot_error |
There was a problem hiding this comment.
1 action item.
TL;DR — if Gym's resume (or abort) RPC fails once after the trainer-checkpoint snapshot, this re-raise ends the run. Rollout admission already stays closed, and the periodic pump already retries the same checkpoint ID, so the crash adds nothing.
Nothing catches this error: it goes up through _train_pump and the actor fails. Gym stays safe without the crash:
- A failed release/abort leaves
_gym_checkpoint_rollout_permittedcleared (it is set only after a successful RPC) and records_pending_gym_checkpoint_release/_pending_gym_checkpoint_abort. - The periodic pump always runs in this mode, because discovery requires
snapshot_attempt_interval_s. Its next attempt retries the same ID first. A Gym that never recovers still ends the run aftermax_consecutive_failures.
Cost today: one short-lived 503 on /resume restarts the whole job, and unfinished Gym work restarts from scratch.
AI-1
Drop the re-raise and let the periodic pump retry the release. Add a unit test where resume fails once during _save_checkpoint and the next _save_rollout_checkpoint releases and reopens admission.
| if isinstance( | |
| gym_snapshot_error, | |
| (_GymCheckpointAbortPendingError, _GymCheckpointReleasePendingError), | |
| ): | |
| # The checkpoint is recoverable, but this live Gym process may | |
| # still be paused. Fail closed instead of resuming training. | |
| raise gym_snapshot_error | |
| # A pending Gym release/abort keeps rollout admission closed; the | |
| # periodic checkpoint pump retries the same checkpoint ID first. |
| timeout_s = self._master_config.rollout_checkpointing.gym.prepare_timeout_s | ||
| gym_actor = self._nemo_gym_checkpoint_actor() | ||
| prepare_attempted = False | ||
| self._gym_checkpoint_rollout_permitted.clear() |
There was a problem hiding this comment.
1 action item — please fix in this PR.
TL;DR — A Gym snapshot parks in-flight rollouts, but RL's rollout_timeout_s clock keeps running. A rollout whose budget runs out during a slow snapshot is timed out while parked; Gym's /run keeps going, finishes after resume, nobody ACKs it, and every later prepare fails. The snapshot itself causes the lost-reply stall.
How it shows up:
- This line closes rollout admission and Gym parks each in-flight rollout at its turn boundary until RL resumes Gym (
_release_prepared_gym_checkpoint). That can take up toprepare_timeout_s(300 s by default,config.py#L761) plus the save. - RL's request deadlines are suspended only for the colocated train switch (L3202), never during a Gym snapshot, so the rollout's
_Deadlinekeeps counting. The functional test sets bothrollout_timeout_sandprepare_timeout_sto 180 (turn_recovery.sh#L166-L172). - When it fires, RL stops reading the Ray streaming call (
rollout_manager.py#L1212-L1214) but never cancels it, so Gym's/runkeeps running, finishes after resume, and stays unACKed.
A run against Gym's real participant and RL's real _Deadline (0.05 s budget):
== TODAY (a): Gym snapshot parks a rollout; RL's rollout deadline keeps running; /run is cancelled
prepare: ready_to_commit=True parked=1 (RL deadline registry suspended=False)
Gym commit after RL gave up still saves this turn: records_for_commit -> [('r7', 0, 1)]
RL outcome: {'timeout': 'NeMo-Gym prompt group exceeded 0.05s', 'state_when_rl_gave_up': 'parked'}
agent: bodies started=1 finished=0
== TODAY (b): same, but the /run keeps going inside Gym (Ray streaming call is not cancelled)
prepare: ready_to_commit=True parked=1 (RL deadline registry suspended=False)
Gym commit after RL gave up still saves this turn: records_for_commit -> [('r7', 0, 1)]
RL outcome: {'timeout': 'NeMo-Gym prompt group exceeded 0.05s', 'state_when_rl_gave_up': 'parked'}
agent: bodies started=1 finished=1
next snapshot's prepare: ready_to_commit=False completed_unacknowledged=1
== FIX: suspend RL request deadlines while the Gym checkpoint holds admission
prepare: ready_to_commit=True parked=1 (RL deadline registry suspended=True)
RL outcome: {'reply': {'reward': 1.0}}
agent: bodies started=1 finished=1
AI-1
Suspend RL request deadlines while a Gym checkpoint holds admission: call self._rollout_manager.suspend_request_deadlines() here, where admission closes, and resume_request_deadlines() wherever it reopens, both after Gym resumes (L2041) and on the path where prepare was never attempted (L2022). This reuses the existing RequestDeadlineRegistry; it is RL only, with no Gym change. The registry is a single on/off flag that the colocated train switch also uses (L3202), so make suspend()/resume() count holders: then a snapshot that ends during a train switch does not restart clocks the switch still holds. The script below is the unit test to add. Not a suggestion block: several call sites plus the registry change.
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4266-parked-deadline-stall.html
Ready-to-adapt test (real Gym participant, runs on CPU)
import asyncio
import sys
from nemo_gym._checkpoint.agent import (
AgentBoundaryRecord,
AgentCheckpointParticipant,
AgentExecutionState,
)
from nemo_rl.experience.rollout_manager import RequestDeadlineRegistry, RolloutTimeout, _Deadline
DEADLINE_S = 0.05
async def wait_state(p, execution, state):
async with p._changed:
await p._changed.wait_for(lambda: execution.state == state)
async def scenario(with_fix: bool, gym_keeps_running: bool = False) -> None:
p = AgentCheckpointParticipant()
registry = RequestDeadlineRegistry()
runs = {"agent_bodies_started": 0, "finished": 0}
reached_turn = asyncio.Event()
execution_box = {}
async def agent():
# Mimics base_responses_api_agent.run_with_rollout_context: begin, work, boundary, finish.
execution = await p.begin("r7", 0, task=asyncio.current_task())
execution_box["e"] = execution
runs["agent_bodies_started"] += 1
reached_turn.set()
# The turn ends while a Gym checkpoint is preparing -> commit_boundary parks here.
await snapshot_started.wait()
try:
await p.commit_boundary(execution, AgentBoundaryRecord(rollout_id="r7", attempt_index=0, boundary_index=1, output_items=[]))
except asyncio.CancelledError:
# base_responses_api_agent.py L236-L238: a cancelled /run reports "cancelled".
await p.finish(execution, outcome="cancelled")
raise
await p.finish(execution, outcome="completed", result={"reward": 1.0})
runs["finished"] += 1
return {"reward": 1.0}
snapshot_started = asyncio.Event()
rl_outcome = {}
async def rl_prompt_group():
# rollout_manager.py L1331: the whole prompt-group stream runs under this deadline.
try:
async with _Deadline(DEADLINE_S, "NeMo-Gym prompt group", registry=registry):
# gym_keeps_running: RL reads rows from a Ray streaming call to the NemoGym actor
# (rollout_manager.py L1212-L1214) and never ray.cancel()s it, so giving up on the
# stream does not stop the /run inside Gym.
rl_outcome["reply"] = await (asyncio.shield(agent_task) if gym_keeps_running else agent_task)
except RolloutTimeout as error:
rl_outcome["timeout"] = str(error)
rl_outcome["state_when_rl_gave_up"] = execution_box["e"].state.value
agent_task = asyncio.create_task(agent())
rl_task = asyncio.create_task(rl_prompt_group())
await reached_turn.wait()
# --- Gym snapshot starts (what #4266 _prepare_and_commit_gym_checkpoint drives) ---
if with_fix:
registry.suspend() # proposed: suspend rollout deadlines while Gym holds admission
prepare = asyncio.create_task(p.prepare(deadline_ts=10**10))
snapshot_started.set()
await wait_state(p, execution_box["e"], AgentExecutionState.PARKED)
report = await prepare
print(f" prepare: ready_to_commit={report['ready_to_commit']} parked={report['parked']} "
f"(RL deadline registry suspended={registry.suspended})")
if with_fix:
# The pause lasts longer than the whole rollout budget.
await asyncio.sleep(DEADLINE_S * 4)
else:
# Hold the checkpoint until RL's deadline fires (it will, because nothing suspends it).
await asyncio.wait({rl_task}, timeout=5)
if not gym_keeps_running:
await asyncio.wait({agent_task}, timeout=5)
saved = [(r.rollout_id, r.attempt_index, r.boundary_index) for r in p.records_for_commit()]
print(f" Gym commit after RL gave up still saves this turn: records_for_commit -> {saved}")
# --- snapshot published; RL resumes Gym (and, with the fix, the deadline clocks) ---
await p.resume()
if with_fix:
registry.resume()
await asyncio.wait({rl_task}, timeout=5)
await asyncio.wait({agent_task}, timeout=5)
print(f" RL outcome: {rl_outcome}")
print(f" agent: bodies started={runs['agent_bodies_started']} finished={runs['finished']}")
if gym_keeps_running:
# RL already gave up on try 0, so nothing seals it and nothing ever ACKs it.
nxt = await p.prepare(deadline_ts=0)
print(f" next snapshot's prepare: ready_to_commit={nxt['ready_to_commit']} "
f"completed_unacknowledged={nxt['completed_unacknowledged']}")
await p.resume()| discovered = ray.get(gym_actor.discover_checkpoint_capabilities.remote()) | ||
| GymCheckpointTopology.model_validate(discovered) | ||
| gym_checkpoint_topology = GymCheckpointTopology.model_validate(discovered) | ||
| if rollout_checkpoint_cfg.gym.participant_checkpointing_enabled: |
There was a problem hiding this comment.
1 action item (low priority).
TL;DR — In the Gym actor, per-row registration and the completion-receipt fetch turn on when Gym discovery succeeds, not when turn-level checkpointing is on. So a run with discovery on and turn-level checkpointing off still registers every row and does one extra GET per row for a receipt nothing reads.
This line runs the turn-recovery check only when rollout_checkpointing.gym.mode is turn_recovery, so discovery on with the switch off is a valid setup, and discovery says it runs "without enabling saves" (nemo_gym.py#L923). The Gym actor gates on self._gym_checkpoint_participants instead, which discovery sets: registration at nemo_gym.py#L1843 and the receipt GET at #L1904. Results stay correct; the cost is wasted work. Under the target-level comment, discovery runs on every token-capture run, so this would add the work to every run below target turn.
AI-1
Gate the registry and the receipt fetch on turn-level checkpointing: pass the switch into the Gym actor from setup, here, and check it at nemo_gym.py#L1843 and #L1904 instead of self._gym_checkpoint_participants. Today the switch is mode == 'turn_recovery' (config.py#L755; participant_checkpointing_enabled is now a property of it at #L762-L764; earlier revision: config.py#L760); the target-level comment replaces it with rollout_recovery.target_level: turn. Not a suggestion block: it spans setup and nemo_gym.py.
| timeout_s = self._master_config.rollout_checkpointing.gym.prepare_timeout_s | ||
| gym_actor = self._nemo_gym_checkpoint_actor() | ||
| prepare_attempted = False | ||
| self._gym_checkpoint_rollout_permitted.clear() |
There was a problem hiding this comment.
1 action item (low priority).
TL;DR — nothing warns when a Gym snapshot keeps new rollouts stopped for longer than snapshot_attempt_interval_s. When that happens, snapshots pause generation for a large share of the run, and the user gets no signal to raise the interval.
This PR adds the Gym pause, so the gap is new here. New rollouts stop on this line and start again when Gym is released (L2041, or L2022 when prepare never started). The periodic loop then sleeps for the interval before the next attempt (L4844-L4845). Timing how long admission stays closed takes one timestamp here and one check where admission reopens. No new settings are needed.
AI-1
Record the time here, and reopen admission through one helper that checks it. Replace self._gym_checkpoint_rollout_permitted.set() at L2022 and L2041 with self._reopen_gym_rollout_admission():
# here, after self._gym_checkpoint_rollout_permitted.clear()
self._gym_admission_closed_at = time.monotonic()
def _reopen_gym_rollout_admission(self) -> None:
"""Reopen rollout admission; warn if Gym kept it closed too long."""
closed_s = time.monotonic() - self._gym_admission_closed_at
interval_s = (
self._master_config.rollout_checkpointing.snapshot_attempt_interval_s
)
if interval_s is not None and closed_s > interval_s:
# Fixed text: Python shows it once per process from this line.
warnings.warn(
"A Gym snapshot kept new rollouts stopped for longer than "
"snapshot_attempt_interval_s, so snapshots pause generation for a "
"large share of the run. Raise snapshot_attempt_interval_s."
)
self._gym_checkpoint_rollout_permitted.set()This can't be a suggestion block because it touches three places (here, L2022, L2041) and adds a new method.
| ++rollout_checkpointing.gym.capability_discovery_enabled=true | ||
| ++rollout_checkpointing.gym.participant_checkpointing_enabled=true |
There was a problem hiding this comment.
1 action item — please fix in this PR.
TL;DR — this script still sets rollout_checkpointing.gym.capability_discovery_enabled and participant_checkpointing_enabled, but at this head the config only accepts rollout_checkpointing.gym.mode. The config model rejects unknown keys, so all three new turn-recovery functional tests fail at startup.
This PR adds these two lines. GymRolloutCheckpointConfig uses extra="forbid" and has only two fields, mode and prepare_timeout_s. The two old names are now read-only properties. Validating these overrides gives two Extra inputs are not permitted errors (checked with pydantic 2.13.5, the version in uv.lock). The workplace and GenRM scripts call this one, so all three fail the same way.
AI-1
| ++rollout_checkpointing.gym.capability_discovery_enabled=true | |
| ++rollout_checkpointing.gym.participant_checkpointing_enabled=true | |
| ++rollout_checkpointing.gym.mode=turn_recovery |
| capability_discovery_enabled: true | ||
| participant_checkpointing_enabled: true |
There was a problem hiding this comment.
1 action item.
TL;DR — this example still uses capability_discovery_enabled and participant_checkpointing_enabled, but at this head the config only accepts rollout_checkpointing.gym.mode. A user who copies it gets two Extra inputs are not permitted errors at startup.
The config model is GymRolloutCheckpointConfig (extra="forbid"). The exemplar YAMLs already use mode.
AI-1
| capability_discovery_enabled: true | |
| participant_checkpointing_enabled: true | |
| mode: turn_recovery |
This branch has not been deployed
2 participants
Summary
Integrates Gym-aware rollout checkpointing into Single Controller: periodic and trainer-boundary snapshot scheduling, data-plane coordination, atomic publication, participant release handling, crash recovery, and redispatch of unfinished Gym rollouts.
This is 3/3 in the stacked decomposition of #4117.
Review and merge this PR after its two dependencies.
Why
A trainer checkpoint is a durable model/training anchor, but it does not describe in-flight Gym executions. Single Controller must coordinate Gym participant state, the Transfer Queue, replay ownership, the dataloader, and rollout recovery metadata into one atomic rollout snapshot.
This PR turns the contracts from #4264 and #4265 into an opt-in end-to-end checkpoint/recovery path.
Checkpoint flow
flowchart TD A[Training reaches periodic or trainer-boundary trigger] --> B[Ensure durable trainer or bootstrap anchor] B --> C[Close new rollout dispatch] C --> D[Strictly drain completed-result ACK obligations] D --> E[Gym prepare] E --> E1[Policy closes admission and drains accepted calls] E --> E2[Agents park or freeze at durable boundaries] E --> E3[Resources freeze revisions] E1 --> F[All participants ready] E2 --> F E3 --> F F --> G[Gym commit participant artifacts] G --> H[Snapshot TQ, replay ledger, recovery state, and dataloader] H --> I[Validate manifests, continuations, digests, and TQ references] I --> J[Atomically publish rollout snapshot] J --> K[Resume resources] K --> L[Resume policy models] L --> M[Resume agents and reopen rollout dispatch] E -. retryable failure .-> N[Abort same checkpoint ID] G -. failure before publish .-> N H -. failure before publish .-> N N --> O[Keep previous rollout snapshot and durable trainer anchor]Turn-level recovery does not cut a generation in the middle of its token stream. Prepare closes admission to new policy calls, while calls already accepted drain to zero within
prepare_timeout_s. Token-prefix recovery remains a separate stacked layer.Recovery flow
flowchart TD A[Process or Slurm restart] --> B[Select compatible trainer anchor and rollout snapshot] B --> C[Restore TQ and replay ownership] C --> D[Restore dataloader and rollout recovery ledger] D --> E[Validate saved Gym topology and artifacts] E --> F[Restore Gym participants while paused] F --> G[Map source attempt N to replacement attempt N plus 1] G --> H[Restart only continuations requiring restart-only resources] H --> I[Resume resources] I --> J[Resume policy admission] J --> K[Resume agents] K --> L[Redispatch unfinished groups without readmitting completed work] L --> M[Continue training]Main changes
published_release_pendingfrom pre-publication I/O failure.Failure behavior
published_release_pending.Configuration
The feature is opt-in. A representative configuration is:
Scope and limitations
Test plan
Unit coverage includes snapshot publication, dispatch/checkpoint races, ACK flushing, abort/release handling, trainer-anchor fallback, setup validation, and restore/redispatch behavior.
Suggested unit command:
End-to-end crash/restart coverage:
Before review