feat(checkpoint): orchestrate Gym turn-level recovery

[macandro96]

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.

Stack	PR	Scope
1/3	#4264	Gym participant protocol
2/3	#4265	Durable rollout state and validation
3/3	This PR	Single Controller orchestration

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]

Loading

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]

Loading

Main changes

• Add the Single Controller rollout-checkpoint pump and periodic scheduling.
• Coordinate snapshotting under the data-plane checkpoint barrier.
• Persist trainer anchors before attempting the Gym-aware rollout snapshot at trainer boundaries.
• Strictly flush completed-result ACKs before snapshot publication.
• Prepare and commit Gym participants, then save TQ/replay/recovery/dataloader state into a temporary bundle.
• Validate the complete bundle before atomic publication.
• Keep rollout admission closed while abort or post-publication release remains pending.
• Distinguish published_release_pending from pre-publication I/O failure.
• Restore compatible Gym-aware snapshots and redispatch only unfinished work.
• Skip placeholder TQ partition warm-up when restoring an authoritative data-plane checkpoint.
• Enforce setup-time guards for single-Gym-actor topology, stable row-attempt identity, and required capabilities.
• Add configuration, metrics documentation, and crash/restart functional coverage.

Failure behavior

• A failed periodic attempt leaves the previous published rollout snapshot authoritative.
• At a trainer boundary, the trainer checkpoint is durable before the Gym-aware snapshot is attempted. A retryable Gym failure therefore does not discard the trainer anchor.
• Failed prepares/commits abort with the same checkpoint ID and reopen admission only after Gym confirms participant release.
• A snapshot that is already durable but still awaiting Gym release is retained and retried as published_release_pending.
• Pending completed-result ACK obligations prevent Gym-aware snapshot publication.
• Restore rejects corrupt or topology-incompatible snapshots before admitting new work.

Configuration

The feature is opt-in. A representative configuration is:

checkpointing:
  enabled: true
  save_data_plane: true
  save_period: 1

token_capture:
  enabled: true

async_rl:
  rollout_failure:
    nemo_gym:
      max_row_attempts: 1

rollout_checkpointing:
  snapshot_attempt_interval_s: 120
  restore_mode: latest
  keep_latest_k: 2
  gym:
    capability_discovery_enabled: true
    participant_checkpointing_enabled: true
    prepare_timeout_s: 300

Scope and limitations

• Single Controller and the vLLM token-capture path are covered.
• Exactly one NeMo Gym actor is currently supported; true sharded Gym checkpoint coordination is rejected at setup.
• This PR restores durable turn/external-wait boundaries. In-generation prefix restoration is not part of this stack.
• The checkpoint format is experimental and intentionally rejects earlier private schema revisions.

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:

uv run pytest -q \
  tests/unit/environments/test_gym_turn_recovery_snapshot.py \
  tests/unit/single_controller/test_checkpointing.py \
  tests/unit/single_controller/test_checkpoint_dispatch_races.py \
  tests/unit/single_controller/test_finalizer_lifecycle.py \
  tests/unit/single_controller/test_resiliency_config.py \
  tests/unit/single_controller/test_rollout_checkpoint.py \
  tests/unit/single_controller/test_rollout_pump.py \
  tests/unit/single_controller/test_setup.py \
  tests/unit/single_controller/test_single_controller_actor.py

End-to-end crash/restart coverage:

bash tests/functional/grpo_async_gym_single_controller_turn_recovery.sh
bash tests/functional/grpo_async_gym_single_controller_workplace_turn_recovery.sh
bash tests/functional/grpo_async_gym_single_controller_genrm_turn_recovery.sh

Before review

• Contributor conventions followed.
• Unit and functional coverage added for checkpoint and recovery flows.
• User-facing configuration and metrics documentation updated.
• Full cluster functional validation recorded after the latest fixture amendment.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[terrykong]

Thanks for splitting this out of #4117. 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 f4fcf8cf is not on Gym main (69 commits ahead, 18 behind; it comes from open Gym PRs #3349, #3412, #3563). The stack can't merge until that lands.
• Please fix before merge: while a Gym snapshot has rollouts parked, RL's rollout_timeout_s deadline 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 on single_controller.py L1958).
• 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: 1 and snapshot_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_seconds now includes Gym time that no per-part metric covers; one gym_participant_seconds metric would make the breakdown add up.
• Upgrade note for the PR body: token_capture.enabled=true now requires max_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.
• Checkpoint levels (design change, before release): capability_discovery_enabled (from #4264), participant_checkpointing_enabled, and the per-environment granularity overrides are really one decision. The target-level comment on config.py L1396 proposes one rollout_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 on setup.py L1876-L1878 and single_controller.py L675, and the docs section on single-controller.md L164.

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

Generated by Claude Code

[terrykong]

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:

1. The Gym actor raises GymControlRequestError, which subclasses RuntimeError, for any non-200 reply, and a plain RuntimeError when a control call times out. commit_checkpoint converts neither. Only prepare-past-deadline becomes TimeoutError (nemo_gym.py#L1305-L1309).
2. Ray re-raises it as the cause's class (object_ref.pxi#L33-L34 → class cls(RayTaskError, cause_cls)), so SC sees a RuntimeError.
3. _prepare_and_commit_gym_checkpoint aborts Gym correctly, then re-raises the error.
4. The pump logs reason=invariant_error and re-raises anything that is not OSError/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.

[terrykong]

1 action item (design change; part of it belongs in Gym).

TL;DR — Replace participant_checkpointing_enabled with 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.py L2434). With that fix, prompt-group plus turn-level checkpointing commits normally.
• A group-scored environment scores all N siblings together: each sibling waits in /verify until the whole group is there (Gym genrm_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:

1. rollout_recovery.target_level, finest to coarsest: turn keeps sealed siblings and resumes unfinished ones from their saved turn; sibling keeps sealed siblings and reruns unfinished ones from scratch; prompt_group reruns the whole group from scratch and discards saved turns. It replaces participant_checkpointing_enabled (config.py#L760). The overrides (config.py#L684-L694) stay only to force a coarser level for one environment.
2. 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.
3. 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 turn its siblings still resume from their saved turns on restore, because they reach /verify together. Log one warning for every environment whose level was lowered.
4. 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.
5. At target prompt_group, restore discards every saved turn, so the group starts over on the current weights. The discard already exists for restart_only resources (_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).

[terrykong]

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_permitted cleared (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 after max_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.

Suggested change

	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.

[terrykong]

1 action item.

TL;DR — a failed trainer-checkpoint Gym snapshot is reported only through warnings.warn, which Python shows once per call site, and it never logs rollout/checkpoint_outcome. A snapshot that fails at every checkpoint is invisible on dashboards.

The fallback itself (restart unfinished work from the marker) is safe. The problem is that nobody can tell it is happening. _log_rollout_checkpoint_outcome is called only from the periodic pump.

Example: a regression makes _prepare_and_commit_gym_checkpoint fail on every call. Every step_N keeps gym_restart_fallback.json, every restart throws away all in-flight turns, and the only trace is one warning line for the whole run.

AI-1

In this except, log rollout/checkpoint_outcome with the same outcome/reason mapping the pump uses. Move that mapping at single_controller.py#L4850-L4863 into a small helper that both call. Also log completed on success, so seconds_since_* stays right. Then update the rollout/checkpoint_outcome row in metrics.md ("every scheduled checkpoint attempt") to include trainer-checkpoint snapshots. This can't be a suggestion block because the helper is shared with the pump in another part of the file.

[terrykong]

1 action item.

TL;DR — the guide says trainer checkpoints hold no Gym state and startup fails without a periodic snapshot. But every trainer checkpoint now writes a restart-fallback marker and a coordinated Gym snapshot, and restore falls back to "keep completed, restart unfinished" instead of failing.

The code: marker write at single_controller.py#L5313-L5322, coordinated snapshot at #L5334-L5361, fallback restore at setup.py#L1439-L1475. The config docstring already says the snapshot is "published with every trainer checkpoint" (config.py#L751-L753).

AI-1

Replace the paragraph with this:

Suggested change

	Use `restore_mode: latest`: full trainer checkpoints do not yet contain Gym
	participant state, so startup fails safely if no compatible periodic rollout
	snapshot exists for the selected trainer anchor.
	Use `restore_mode: latest`. Every trainer checkpoint first writes a
	`gym_restart_fallback.json` marker and then tries to publish a coordinated Gym
	rollout snapshot. On restart, NeMo-RL restores the newest compatible snapshot.
	If the trainer checkpoint has no snapshot but still has the marker, NeMo-RL
	keeps completed rollouts and restarts unfinished Gym executions from their
	original task. Startup fails only if neither exists.

[terrykong]

1 action item.

TL;DR — Gym checkpointing needs about 15 settings to agree across RL, each environment and Gym, and the guide lists only three of them. Several wrong combinations fail silently (a rollout that waits forever, a snapshot that always times out, Gym memory that grows all run), so users need one place that says what to set for each checkpoint level.

This block shows capability_discovery_enabled, participant_checkpointing_enabled and prepare_timeout_s. The rest are spread across the setup checks and the Gym configs: the recovery granularity and its overrides (resolve_for_prompt), the GenRM group settings (app.py#L145-L146), and checkpoint_replayable_verify (base_responses_api_agent.py#L138). The draft below is written for the target-level design in the target-level comment, so it also shows the author what that design looks like to a user.

AI-1

Add a "Choosing a checkpoint level" section after this block: the target levels, what to set for each, how RL lowers group-scored environments, and the combinations that still fail. A draft is below, written against the target-level design. Not a suggestion block: it is a new section that depends on that design.

Draft section, ready to paste

Choosing a checkpoint level

Set one target level, rollout_recovery.target_level. It says how much work RL tries to keep when a rollout fails
or the run restarts. From finest to coarsest:

Target	Sealed siblings	Unfinished siblings
turn	kept	resume from their last saved turn
sibling	kept	rerun from scratch
prompt_group	rerun	the whole group reruns from scratch; saved turns are discarded

Each environment gets the finest level at or below the target that it supports. RL learns the limits from Gym at
setup and logs a warning for every environment it lowers.

What to set

Setting	prompt_group	sibling	turn
token_capture.enabled	true	true	true
async_rl.rollout_failure.nemo_gym.max_row_attempts	1	1	1
checkpointing.enabled, checkpointing.save_data_plane	true	true	true
rollout_checkpointing.snapshot_attempt_interval_s	set	set	set
rollout_checkpointing.restore_mode	any	any	latest
async_rl.sampler	must support buffer checkpoints	same	same
env.nemo_gym.shards	—	—	exactly one Gym actor
Gym agent class	—	—	supports continuations (simple_agent does)

rollout_recovery overrides are optional. Use one only to force a coarser level for a single environment, keyed by
its exact Gym agent name; RL rejects a name it did not discover.

Group-scored environments (GenRM compare)

GenRM compare scores the siblings of a prompt against each other: each sibling waits in /verify until all N are
there. A sibling retried alone would never be scored, so RL always retries the whole group for these environments.

Target	Live failure of one sibling	Restore
turn	whole group retried	all siblings resume from their saved turns together (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; saved turns discarded

For these environments also:

• Set the Gym server's num_rollouts_per_prompt equal to grpo.num_generations_per_prompt.
• Set checkpoint_replayable_verify: true on the agent at target turn (Gym's shipped GenRM agents already do).
• Set async_rl.rollout_failure.nemo_gym.rollout_timeout_s or the server's cohort_collection_timeout_s, so a
  group that cannot form fails instead of waiting forever.

GenRM compare resumes at turn only after Gym declares genrm_compare stateless. Until then Gym reports it as
restart_only, so on restore its saved turns are discarded and the whole group restarts from scratch, even at target
turn.

A group resumed at turn keeps turns generated with the weights from save time, so it no longer comes entirely from
one set of weights. Use target prompt_group if every group must come from the current weights.

Metrics

RL reports, per step: 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.

Combinations that fail

Combination	What happens
num_rollouts_per_prompt ≠ num_generations_per_prompt	The group never fills; the rollout waits forever (setup rejects it).
checkpoint_replayable_verify: false on a GenRM agent at target turn	Every snapshot taken while a group waits in /verify times out (setup rejects it).
No rollout_timeout_s and no cohort_collection_timeout_s for a group-scored environment	Nothing ends a group that cannot form (setup rejects it).
An override naming an unknown agent, or finer than the target	Setup rejects it.

[terrykong]

1 action item.

TL;DR — This GenRM profile runs Gym's stateless test copy of genrm_compare, but production genrm_compare is restart_only. So the test passes on the resume path, which production GenRM never takes: in production its saved turns are discarded and the whole group restarts.

L124 loads checkpoint_test_agent/configs/genrm_compare.yaml, whose resources server is checkpoint_test_app.py (Gym config L1-L4), the copy that declares stateless (checkpoint_test_app.py#L24-L27); the comment at L109 also says "the stateless cohort verifier". Production GenRMCompareConfig keeps Gym's default restart_only (base_resources_server.py#L95-L98), so #4266 discards its saved turns on restore (single_controller.py#L1810-L1834).

AI-1

Once Gym declares production genrm_compare stateless (see the target-level comment), point this profile at the production server class, with its judge replaced by a deterministic stub model, so the test covers the recovery mode production actually uses. Not a suggestion block: it needs a Gym config plus a change to this test profile.

[terrykong]

1 action item.

TL;DR — Restore correctness depends on what Gym saves and replays, and every restore test in this file mocks Gym. Nothing checks that real Gym environments restore the way RL assumes.

This adds to the continued-rollout correctness comment. The CPU harness planned there mocks the Gym servers; these tests run Gym's real participant and resources-server code, so they check Gym instead of trusting it.

AI-1

Add @pytest.mark.nemo_gym restore tests that save a checkpoint with Gym's real participant and resources-server code (only the model server is a stub), restore it into fresh Gym objects, and check the result:

1. A GenRM group at the cut, with one sibling frozen in /verify and one parked, using the production genrm_compare config: the group must form again and be scored once. Today this shows the restart_only discard; after the Gym stateless fix in the target-level comment it shows the resume.
2. The same saved turn against a stateless, a restart_only and an export_restore resources server: resumed, discarded, restored.
3. A sibling that fails after its group was scored: a whole-group retry, with no hang.
4. Under the target-level design, target prompt_group discards saved turns, and target turn keeps sealed siblings and resumes the rest.
5. The environments the functional test already uses (its three profiles: counter / example_session_state_mgmt, workplace_assistant, genrm_compare) restore to the state they had at the cut.

Not a suggestion block: it is a new test file.

[terrykong]

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:

1. 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 to prepare_timeout_s (300 s by default, config.py#L761) plus the save.
2. RL's request deadlines are suspended only for the colocated train switch (L3202), never during a Gym snapshot, so the rollout's _Deadline keeps counting. The functional test sets both rollout_timeout_s and prepare_timeout_s to 180 (turn_recovery.sh#L166-L172).
3. When it fires, RL stops reading the Ray streaming call (rollout_manager.py#L1212-L1214) but never cancels it, so Gym's /run keeps 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()

[terrykong]

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 participant_checkpointing_enabled is on, 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 participant_checkpointing_enabled (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.