feat(checkpoint): Support turn-level checkpoint

[macandro96]

Summary

Adds opt-in, turn-level checkpoint and recovery for NeMo Gym rollouts in Single Controller training.

The checkpoint is a coordinated snapshot across NeMo-RL, Gym participants, the Transfer Queue (TQ), and rollout bookkeeping. After a process or Slurm restart, unfinished rollouts are redispatched from their latest durable agent boundary instead of restarting the entire batch.

This PR also hardens checkpoint failure handling: completed-result acknowledgements must be drained before publication, failed prepares are aborted with the same checkpoint ID, and a published snapshot is not reported as failed merely because participant release is still pending.

Why

Trainer checkpoints alone do not contain enough information to resume in-flight Gym executions. A correct rollout checkpoint must keep these pieces mutually consistent:

• Gym agent continuation state and resource-session revisions;
• policy-model lineage and external TQ references;
• canonical TQ data and replay ownership;
• rollout recovery state, dataloader position, and completed-result ACK obligations;
• the compatible trainer or bootstrap anchor.

Publishing only a subset can create duplicate executions, refer to missing token rows, or restore agent state against the wrong model/resource state. This PR defines one coordinated publication contract and validates it before a snapshot becomes visible.

End-to-end flow

flowchart TD
    A[Durable trainer anchor or bootstrap state] --> B[Close Single Controller rollout admission]
    B --> C[Drain completed-result ACK outbox]
    C --> D[Gym prepare]
    D --> D1[Policy closes admission to new model calls]
    D1 --> D1A[Already accepted policy calls drain to zero]
    D --> D2[Agents park or freeze at durable boundaries]
    D --> D3[Resources freeze session revisions]
    D1A --> D4[All participants ready before the deadline]
    D2 --> D4
    D3 --> D4
    D4 --> E[Gym commit]
    E --> E1[Write agent continuation index]
    E1 --> E2[Write policy lineage and TQ reference index]
    E2 --> E3[Write resource state]
    E3 --> F[Snapshot TQ, replay ledger, recovery state, and dataloader]
    F --> G[Validate manifests, digests, continuations, and TQ keys]
    G --> H[Atomically publish rollout snapshot]
    H --> I[Resume resources]
    I --> J[Resume policy models]
    J --> K[Resume agents and reopen rollout admission]

    H --> L[Process or Slurm restart]
    L --> M[Select compatible trainer anchor and rollout snapshot]
    M --> N[Restore TQ, replay ledger, recovery state, and dataloader]
    N --> O[Validate saved Gym topology and artifacts]
    O --> P[Restore policy, agents, and resources while paused]
    P --> Q[Map source attempt N to replacement attempt N plus 1]
    Q --> R[Restart only executions that depend on restart-only resources]
    R --> I

Loading

Participant ordering is deliberate:

• Prepare: policy models → agents → resources.
• Commit: agents → policy models → resources.
• Restore: policy models → agents → resources.
• Resume: resources → policy models → agents.

Dependencies are therefore available before parked agent requests are released.

Policy-model drain behavior

Turn-level recovery does not checkpoint a policy generation in the middle of its token stream. During prepare, the policy pause closes admission to new model calls immediately, but the pause request is intentionally non-blocking. Calls that were already accepted are allowed to finish, and commit cannot begin until the policy status reports inflight_total == 0 with every configured worker acknowledged.

Agent and resource preparation can make progress while those accepted calls drain. The shared prepare_timeout_s bounds the entire operation. If a long-running policy call does not finish before that deadline, the checkpoint attempt is aborted, participants are resumed, and the previous published snapshot remains authoritative. In-generation token-prefix cuts, which avoid waiting for cuttable generations to finish, are a separate layer and are not part of this turn-level PR.

Checkpoint contract

1. The controller closes new rollout dispatch and strictly drains the Gym completed-result ACK outbox.
2. Gym prepares every participating service. Policy admission closes to new work, already accepted policy calls drain to zero, agent executions become checkpoint-safe, and resource revisions stop changing.
3. Gym commits participant artifacts and returns their manifests and external-storage references.
4. NeMo-RL snapshots TQ, replay/recovery state, and dataloader state under the data-plane checkpoint barrier.
5. Cross-artifact validation proves that continuation records, model lineage, manifests, digests, and referenced TQ keys agree.
6. The snapshot directory is atomically published only after all invariants pass.
7. Participants resume in dependency order. If release is still pending, admission remains closed and the same checkpoint release is retried.

At a trainer boundary, the trainer checkpoint is published first. The coordinated Gym-aware rollout snapshot is a second durability tier anchored to it. A retryable Gym snapshot failure therefore does not destroy the durable trainer checkpoint; recovery can conservatively restart unfinished rollouts from that trainer boundary.

Restore contract

Recovery:

1. selects a rollout snapshot compatible with the requested trainer checkpoint or bootstrap fingerprint;
2. restores canonical TQ data and NeMo-RL replay/recovery state;
3. validates the saved Gym topology, participant manifests, continuation index, lineage, and external TQ references;
4. restores Gym participants while admission remains closed;
5. redispatches unfinished source attempts as replacement attempts with stable recovery mappings;
6. resumes resources, policy models, and agents in dependency order.

Export/restore participants continue from durable saved state. Executions that depend on a restart_only resource are restarted selectively; unrelated continuable executions retain their saved progress.

Correctness and failure handling

• A Gym-aware snapshot is published only with zero pending completed-result ACK obligations.
• ACK retries are idempotent within the live Gym process; unresolved obligations cannot cross a published checkpoint boundary.
• Stable Gym execution identities require max_row_attempts: 1 when checkpoint/recovery and token capture are enabled.
• Prepare/commit failures fail closed and abort using the same checkpoint ID. Rollout admission reopens only after Gym confirms release.
• A durable snapshot awaiting participant release is reported as published_release_pending, not as an I/O failure.
• The previously published snapshot remains authoritative if a newer attempt fails before publication.
• Restore rejects incompatible topology, corrupt manifests, missing TQ references, and invalid continuation relationships before admitting work.
• Checkpoints from older experimental schemas are intentionally rejected with an actionable compatibility error.

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

Capability discovery validates the configured Gym topology and participant contracts during setup so unsupported recovery modes fail before training starts.

Scope and limitations

• Single Controller and the vLLM token-capture path are covered.
• Exactly one NeMo Gym actor is currently supported. An unsharded actor or one configured shard with one replica is accepted; true multi-actor/sharded Gym checkpoint coordination is intentionally rejected during setup.
• This PR restores at durable turn/external-wait boundaries. In-generation token-prefix restoration is layered separately.
• The checkpoint format is experimental and intentionally does not preserve compatibility with earlier private schema revisions.

Test coverage

Unit coverage includes checkpoint schema and wire validation, capability/topology checks, ACK outbox handling, continuation and TQ-reference validation, abort/release races, trainer-boundary fallback, and restore/redispatch behavior.

End-to-end crash/restart coverage includes:

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

[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.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 2bc70a9 (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 2561d29 (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 8768bea (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 1237e33 (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: de5b664 (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[macandro96]

Team review — PR #4117

Reviewed by a coordinated agent team (RL expert, Gym-submodule expert, bug-finder, design-reviewer, devil's-advocate) plus leader-level checks. 17 inline comments below; 4 items that don't anchor to a single line are summarized here.

🚫 BLOCKER — Gym submodule pin is not on Gym main

3rdparty/Gym-workspace/Gym is pinned to f4fcf8c, which exists only on origin/amahishi/gym-turn-level-recovery, origin/amahishi/gym-prefix-cuts-lineage, and origin/amahishi/policy-model-multiworker — not on Gym's main (which RL main currently pins at 267305e). The repo's "Submodule Fast-Forward Check" only verifies ancestry relative to the previous pin, so it structurally can't catch a pin that isn't reachable from the submodule's own default branch. Action: land the corresponding Gym PR(s) to Gym main and re-pin before merge.

CI — pre-commit run --all-files fails on PR files

• ruff: 7 import-sort errors (auto-fixable).
• ruff-format: 1 file reformatted (tests/functional/_gym_turn_recovery_snapshot.py ~L761).
• pyrefly check: 5 errors — single_controller.py:1832 (dict literal vs GymCompletionReceipt annotation), gym_checkpoint.py:966 (TypeVar bound too narrow), rollout_recovery.py:1341-1343 (Unknown | None vs non-Optional params). We traced all 5 and they are annotation-hygiene issues, not runtime bugs (_StrictWireModel lacks strict=True so dict-coercion works at runtime; PendingCompletedExecutionAcknowledgement.__post_init__ already validates before these sites are reached, so they fail closed). Still needs fixing to pass CI.

Perf-evidence request

snapshot_attempt_interval_s: 120 is more expensive than "periodic snapshot" implies: the cheap-exit is disabled when Gym checkpointing is on, the window also gates the trainer-checkpoint critical path, acknowledge_completed_executions issues serial HTTP POSTs, and _prepare_agent_checkpoint can poll at up to 10Hz for the full prepare_timeout_s (~3000 requests at the 300s default). The feature defaults off, so this isn't a blocker — but please post steady-state throughput with the flag on vs. off, and p50/p95 wall-clock of the admission-closed sub-window, so users can judge the trade-off.

Adjacent — non-Gym restore behavior changed by the same PR

setup.py's should_warm_partitions logic changes TQ-restore behavior even for users with no Gym checkpointing enabled — previously partitions were registered unconditionally on this path. Plausibly correct, but it's an unrelated semantic change riding in a 9k-line opt-in-feature PR. Please call it out explicitly in the PR description and confirm the non-Gym restore path has coverage.

[macandro96]

1 action item.

TL;DR — a resources session that never gets a successful terminal response is never retired by either side, so it leaks and accumulates across every future checkpoint cycle.

Gym's ResourcesSessionMiddleware retires a session only on a successful terminal response (terminal = successful and (kind == "terminal" or terminal_intent)). RL never calls the resources /retire route from any of prepare/commit/restore/resume (call sites at nemo_gym.py:1141, 1397, 1499, 1679).

Trigger: a rollout's /verify returns a transient 500 → successful=False → session stays live → the harness still catches it and returns a scored COMPLETED result → RL acks it, popping it from _executions → but the resources side still exports the now-orphaned session's revision on the next prepare(), forever, since nothing ever retires it.

Note: this does not reliably make the checkpoint permanently unrestorable — the orphaned rollout is COMPLETED-side and excluded from the continuation cross-checks that would otherwise raise on restore. The concrete, verified consequence is unbounded leaked resources-session state, not a guaranteed restore failure.

Action: either (Gym) make _reconcile_agent_resource_revisions tolerate a snapshot whose (rollout_id, attempt_index) has no matching agent boundary by treating it as retired; or (RL) call the resources /retire route for every execution that reaches a terminal state without a terminal resources response, inside acknowledge_completed_executions.

[macandro96]

this would need fix on the gym side - will update

[macandro96]

1 action item.

TL;DR — max_row_attempts is silently forced to 1 for every token-capture user, independent of any Gym checkpointing flag, making the configured retry count a dead setting.

stable_execution_attempts = any("_ng_attempt_index" in row for row in inputs)
max_row_attempts = (
    1
    if recovery_granularity is RecoveryGranularity.PROMPT_GROUP
    or stable_execution_attempts
    else self._max_gym_row_attempts
)

_ng_attempt_index is stamped on every row whenever rollout_ids is not None, and the token-capture dispatch path always supplies attempt_indices, so stable_execution_attempts is True on essentially every token-capture dispatch — not just when Gym participant checkpointing is enabled. The configured async_rl.rollout_failure.nemo_gym.max_row_attempts (default 3) becomes dead for any user of token_capture.enabled: true, even with Gym checkpointing off.

This is plausibly intentional — a row re-dispatch under the stable identity would reuse the same _ng_attempt_index and hit GymActorExecutionRegistry.register's "already live" error — but there's no comment and no setup-time guard, so a transient row failure that used to be absorbed by 2 in-band retries now surfaces immediately with no explanation of why the configured value was ignored.

Action: add the rationale as a comment here, and raise at setup when token_capture.enabled and max_row_attempts > 1 so the now-dead knob fails loud instead of silently degrading.

[macandro96]

Partially fixed in f70b1d3.

Done: setup-time validation now raises when token_capture.enabled and max_row_attempts != 1 (validate_single_controller_config), with coverage in test_resiliency_config.py.

Still open: the rationale comment at the confusing site. rollout_manager.py L1322-1329 still derives max_row_attempts from stable_execution_attempts with no pointer to the new setup-time constraint.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: f70b1d3 (PR #4117 from amahishi/turn-level-checkpoint)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨