feat(sc): periodic rollout checkpointing

docs/guides/single-controller.md

@@ -103,6 +103,97 @@ On resume, Single-Controller validates the TQ snapshot against the trainer check
 
 Replay recovery is supported by all built-in samplers: `in_order`, `weight_fifo`, `ready_first`, and `windowed`. Custom samplers must explicitly declare `supports_buffer_checkpoint = True`. Otherwise, setup emits a warning and completed buffered groups are not restored.
 
+### Periodic rollout snapshots
+
+Normal trainer checkpoints are written at step boundaries. Periodic rollout
+snapshots preserve newer rollout progress between those trainer checkpoints,
+including while the train pump is accumulating a streamed step:
+
+```yaml
+checkpointing:
+  enabled: true
+  checkpoint_dir: /shared/checkpoints/my-run
+  save_data_plane: true
+  save_period: 1
+
+rollout_checkpointing:
+  snapshot_attempt_interval_s: 120
+  keep_latest_k: 2
+  restore_mode: latest
+  extra_fingerprint_excluded_paths: []
+
+token_capture:
+  enabled: true
+```
+
+`snapshot_attempt_interval_s` is the cadence at which Single-Controller attempts
+a rollout snapshot. It is not a guarantee that a snapshot is written at every
+interval. An attempt after step N succeeds only when the immutable trainer
+checkpoint `step_N` is already durable. Consequently, `save_period: 1` is
+recommended for continuous post-step coverage; with a larger value, attempts
+are skipped until the matching trainer checkpoint exists. Before the first
+training step, snapshots are anchored to the initial model and a fingerprint of
+the rollout-semantic configuration.
+
+The bootstrap fingerprint is fail-closed: every configuration value affects
+compatibility unless NeMo-RL's built-in denylist identifies it as operational,
+such as logging, cluster placement, checkpoint location, runtime ports, or
+credentials. This means configuration added by an external algorithm is safe
+by default—a change prevents bootstrap recovery instead of silently mixing
+incompatible rollout state.
+The bootstrap manifest also stores this credential-redacted compatibility
+identity so a rejected restart can report the exact changed dotpaths instead of
+showing only two opaque digests.
+
+An integration may use `extra_fingerprint_excluded_paths` for additional
+runtime-only values that are not part of NeMo-RL's built-in configuration:
+
+```yaml
+rollout_checkpointing:
+  extra_fingerprint_excluded_paths:
+    - custom_algo.observability
+    - env.private_agent.runtime_endpoint
+    - env.private_agent.workers.*.log_dir
+```
+
+Each dotpath removes that value and its children from the compatibility
+identity. `*` matches one mapping or list level and `**` matches any number of
+levels.
+Only exclude values that cannot affect prompts, generation, rewards, lineage,
+or the interpretation of persisted rollout data. These exclusions must be set
+on the original run as well as its restart.
+
+Periodic snapshots currently require all of the following:
+
+- `checkpointing.enabled: true` and `checkpointing.save_data_plane: true`.
+- `data_plane.backend: simple`, because native TQ save/load is required.
+- `token_capture.enabled: true`.
+- A replay-recoverable sampler with training-claim ownership. All built-in
+  samplers qualify. A custom sampler must explicitly declare both
+  `supports_buffer_checkpoint = True` and `supports_training_claims = True`.
+
+Each trainer or bootstrap anchor has a `rollout_snapshots/` directory. A
+published `snapshot_NNNNNN/` contains the native TQ snapshot and matching
+replay, dataloader, controller, replacement-reserve, and unfinished-rollout
+metadata. `keep_latest_k` retains recent committed snapshots as fallbacks;
+temporary or interrupted directories are never selected for recovery.
+
+With `restore_mode: latest`, startup selects the newest compatible committed
+snapshot under the latest trainer anchor. With `trainer_checkpoint`, it ignores
+newer periodic rollout progress and resumes from the trainer checkpoint bundle.
+Checkpoint selection is read-only: neither mode removes snapshots. If no trainer
+checkpoint exists, `trainer_checkpoint` cannot safely reuse an existing bootstrap
+namespace, so startup fails without modifying it. Recover that state with `latest`
+or choose a new `checkpoint_dir` to start a fresh bootstrap lineage. Obsolete
+bootstrap snapshots are removed only by retention after a durable trainer
+checkpoint exists.
+
+> **Bootstrap-only `trainer_checkpoint` behavior:** A bootstrap rollout snapshot
+> has no corresponding model or optimizer checkpoint. Therefore
+> `restore_mode: trainer_checkpoint` deliberately fails when bootstrap state exists
+> but no trainer checkpoint does. It does not ignore or delete that state. Use
+> `latest` to recover it, or select a new `checkpoint_dir` to start from scratch.
+
 :::{note}
 Completed groups are restored directly from the TQ snapshot. For unfinished
 token-capture groups, `rollout_recovery.default_granularity` controls both live

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -141,10 +141,28 @@ checkpointing:
   enabled: false
   checkpoint_dir: results/grpo-single-controller
   metric_name: null
+  # Periodic rollout snapshots need an immutable trainer anchor for every
+  # completed step. Keep this at 1 when
+  # rollout_checkpointing.snapshot_attempt_interval_s is set.
+  save_period: 1
   # Include native TQ state and the metadata-only replay index. A save failure
   # aborts checkpoint finalization.
   save_data_plane: true
 
+# Frequent rollout-only snapshots are disabled unless
+# snapshot_attempt_interval_s is set. When
+# disabled, existing periodic snapshots are ignored. Enabling them requires a
+# durable trainer checkpoint for the current completed step and native TQ
+# checkpoint support. save_period=1 provides an anchor after every train step.
+rollout_checkpointing:
+  snapshot_attempt_interval_s: null
+  keep_latest_k: 2
+  restore_mode: latest
+  # Advanced escape hatch for runtime-only fields from external integrations.
+  # Every config path not excluded here or by NeMo-RL's built-in denylist must
+  # match before a bootstrap rollout snapshot can be restored.
+  extra_fingerprint_excluded_paths: []
+
 policy:
   dtensor_cfg:
     enabled: false

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -146,6 +146,27 @@ checkpointing:
   enabled: false
   checkpoint_dir: results/ppo-single-controller
   metric_name: null
+  # Periodic rollout snapshots need an immutable trainer anchor for every
+  # completed step. Keep this at 1 when
+  # rollout_checkpointing.snapshot_attempt_interval_s is set.
+  save_period: 1
+  # Include native TQ state and the metadata-only replay index. A save failure
+  # aborts checkpoint finalization.
+  save_data_plane: true
+
+# Frequent rollout-only snapshots are disabled unless
+# snapshot_attempt_interval_s is set. When
+# disabled, existing periodic snapshots are ignored. Enabling them requires a
+# durable trainer checkpoint for the current completed step and native TQ
+# checkpoint support. save_period=1 provides an anchor after every train step.
+rollout_checkpointing:
+  snapshot_attempt_interval_s: null
+  keep_latest_k: 2
+  restore_mode: latest
+  # Advanced escape hatch for runtime-only fields from external integrations.
+  # Every config path not excluded here or by NeMo-RL's built-in denylist must
+  # match before a bootstrap rollout snapshot can be restored.
+  extra_fingerprint_excluded_paths: []
 
 policy:
   tokenizer:

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import copy
 import gc
 import hashlib
 import json
@@ -231,6 +232,7 @@ def __init__(self) -> None:
         self._checkpoint_active = False
         self._active_mutations = 0
         self._section_holders: set[asyncio.Task[Any]] = set()
+        self._mutation_version = 0
 
     def _current_task(self) -> asyncio.Task[Any]:
         """Return the task entering a barrier section and reject reentrancy."""
@@ -244,6 +246,11 @@ def _current_task(self) -> asyncio.Task[Any]:
             )
         return task
 
+    @property
+    def mutation_version(self) -> int:
+        """Return a monotonic marker for completed outer mutation sections."""
+        return self._mutation_version
+
     @asynccontextmanager
     async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]:
         """Yield a live cut after any active checkpoint exits."""
@@ -260,6 +267,9 @@ async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]:
             async with self._condition:
                 self._section_holders.discard(task)
                 self._active_mutations -= 1
+                # Count the section even when its body raised. A redundant
+                # snapshot is safe; skipping a partially applied mutation is not.
+                self._mutation_version += 1
                 if self._active_mutations == 0:
                     self._condition.notify_all()
 
@@ -1043,6 +1053,11 @@ def __init__(
         self._post_write_enricher: Optional[
             Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]]
         ] = None
+        # Sampler selection removes ready slots from the live replay index but
+        # deliberately leaves their rows in TQ until optimizer completion.
+        # Retain their metadata here so a periodic checkpoint can make an open
+        # streamed step replayable without depending on the sibling lineage.
+        self._training_claims: dict[str, TQReplayGroupMetadata] = {}
 
     def set_data_plane_checkpoint_barrier(
         self, barrier: DataPlaneCheckpointBarrier
@@ -1425,15 +1440,75 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int:
                 cut, drop_group_ids, clear_data_plane=remove_in_dp
             )
 
+    async def claim_for_training(self, idxs: list[int]) -> int:
+        """Transfer ready groups from sampler ownership to an open train step.
+
+        The canonical rows remain in TQ. Their metadata stays checkpoint-visible
+        until :meth:`release_training_claims` runs after optimizer success and
+        data-plane cleanup.
+        """
+        if len(idxs) == 0:
+            return 0
+        if len(idxs) != len(set(idxs)):
+            raise ValueError("training claim contains duplicate replay indices")
+        if min(idxs) < 0:
+            raise IndexError("training claim indices must be non-negative")
+        if self._data_plane_checkpoint_barrier is None:
+            raise RuntimeError(
+                "TQReplayBuffer must be bound to the controller data-plane "
+                "checkpoint barrier before claiming groups for training"
+            )
+        claim_idxs = sorted(idxs, reverse=True)
+        if claim_idxs[0] >= len(self.meta_list):
+            raise IndexError(
+                "TQReplayBuffer.claim_for_training: indices out of range: "
+                f"{claim_idxs[0]}; size={len(self.meta_list)}"
+            )
+        claim_group_ids = [self._group_ids[i] for i in claim_idxs]
+        async with self._data_plane_checkpoint_barrier.mutation() as cut:
+            return await self._remove_groups_unlocked(
+                cut,
+                claim_group_ids,
+                clear_data_plane=False,
+                retain_training_claims=True,
+            )
+
+    def training_owned_replay_groups(self) -> list[TQReplayGroupMetadata]:
+        """Return metadata for canonical rows owned by the open train step."""
+        return copy.deepcopy(list(self._training_claims.values()))
+
+    def training_owned_group_ids(self) -> set[str]:
+        """Return stable IDs currently owned by the open train step."""
+        return set(self._training_claims)
+
+    def release_training_claims(self, group_ids: list[str]) -> None:
+        """Release checkpoint ownership after consumed TQ rows are cleared."""
+        if len(group_ids) != len(set(group_ids)):
+            raise ValueError("training claim release contains duplicate group IDs")
+        claimed_group_ids = set(self._training_claims)
+        released_group_ids = set(group_ids)
+        unknown = sorted(released_group_ids - claimed_group_ids)
+        unreleased = sorted(claimed_group_ids - released_group_ids)
+        if unknown or unreleased:
+            raise ValueError(
+                "training claim release does not match current ownership: "
+                f"unknown={unknown!r}, unreleased={unreleased!r}"
+            )
+        for group_id in group_ids:
+            del self._training_claims[group_id]
+
     async def _remove_groups_unlocked(
         self,
         cut: DataPlaneMutationCut,
         group_ids: list[str],
         *,
         clear_data_plane: bool,
+        retain_training_claims: bool = False,
     ) -> int:
         """Remove stable groups while the caller owns a live mutation cut."""
         cut.require_live()
+        if clear_data_plane and retain_training_claims:
+            raise ValueError("cleared rows cannot be retained as training claims")
         if len(group_ids) != len(set(group_ids)):
             raise ValueError("replay removal contains duplicate group IDs")
         index_by_group_id = {group_id: i for i, group_id in enumerate(self._group_ids)}
@@ -1484,6 +1559,25 @@ async def _remove_groups_unlocked(
                         "may already be cleared"
                     ) from error
 
+        new_training_claims: dict[str, TQReplayGroupMetadata] = {}
+        if retain_training_claims:
+            for group_id in group_ids:
+                i = index_by_group_id[group_id]
+                meta = self.meta_list[i]
+                if meta is None or not self.ready_list[i]:
+                    raise RuntimeError(
+                        "only ready replay groups may be claimed for training"
+                    )
+                if group_id in self._training_claims:
+                    raise ValueError(f"duplicate training-owned group_id={group_id!r}")
+                new_training_claims[group_id] = {
+                    "meta": copy.deepcopy(meta),
+                    "start_weight": self.start_weight_list[i],
+                    "end_weight": self.end_weight_list[i],
+                    "target_step": self.target_step_list[i],
+                    "group_id": group_id,
+                }
+
         # A different mutation may have removed a lower list slot while the
         # DataPlane calls were awaiting. Resolve the original stable IDs again;
         # never apply pre-await indices to the now-shifted parallel lists. A group
@@ -1499,12 +1593,20 @@ async def _remove_groups_unlocked(
             ),
             reverse=True,
         )
+        if retain_training_claims and len(current_drop_idxs) != len(group_ids):
+            raise RuntimeError("training claim ownership changed during mutation")
+        self._training_claims.update(new_training_claims)
         for i in current_drop_idxs:
             self._delete_slot(i)
 
         return len(current_drop_idxs)
 
-    def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState:
+    def metadata_state_dict(
+        self,
+        *,
+        saved_capacity: int,
+        additional_groups: Optional[list[TQReplayGroupMetadata]] = None,
+    ) -> TQReplayMetadataState:
         """Capture the controller index for ready groups without tensor payloads.
 
         The caller must hold the exclusive side of the shared data-plane
@@ -1516,7 +1618,11 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState:
         complete publish/index or clear/remove transition. No writer is exempt,
         including post-train cleanup in ``_train_pump``; canonical writes are
         not required to originate specifically from :meth:`commit`.
-        In-flight reservations are intentionally omitted.
+        The advantage stage also takes a mutation slot because the periodic
+        checkpoint pump runs concurrently with ``_train_pump``.
+        In-flight reservations are intentionally omitted. ``additional_groups``
+        is used by periodic snapshots to re-index rows claimed by an unfinished
+        streamed optimizer step.
         """
         groups: list[TQReplayGroupMetadata] = []
         for i, ready in enumerate(self.ready_list):
@@ -1533,6 +1639,27 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState:
                     "group_id": self._group_ids[i],
                 }
             )
+        existing_group_ids = {group["group_id"] for group in groups}
+        existing_sample_ids = {
+            sample_id for group in groups for sample_id in group["meta"].sample_ids
+        }
+        for group in additional_groups or []:
+            group_id = group["group_id"]
+            if group_id in existing_group_ids:
+                raise ValueError(
+                    f"additional replay metadata duplicates group_id={group_id!r}"
+                )
+            duplicate_sample_ids = existing_sample_ids.intersection(
+                group["meta"].sample_ids
+            )
+            if duplicate_sample_ids:
+                raise ValueError(
+                    "additional replay metadata duplicates sample IDs: "
+                    f"{sorted(duplicate_sample_ids)!r}"
+                )
+            groups.append(copy.deepcopy(group))
+            existing_group_ids.add(group_id)
+            existing_sample_ids.update(group["meta"].sample_ids)
         return {
             "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
             "storage": REPLAY_BUFFER_METADATA_STORAGE,
@@ -1584,7 +1711,7 @@ async def load_state_dict(
                 sample_ids), disagrees with the native TQ snapshot, or exceeds
                 ``max_groups``.
         """
-        if self.meta_list or self._group_ids:
+        if self.meta_list or self._group_ids or self._training_claims:
             raise RuntimeError(
                 "Replay-buffer checkpoint loading requires an empty local buffer"
             )