feat(rollout): add periodic partial-rollout recovery snapshots

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -71,6 +71,13 @@ checkpointing:
   checkpoint_dir: results/grpo-single-controller
   metric_name: null
 
+# Frequent rollout-only snapshots are disabled unless interval_s is set. They
+# reuse the latest durable trainer checkpoint and native TQ checkpoint support.
+rollout_checkpointing:
+  interval_s: null
+  keep_latest_k: 2
+  restore_mode: latest
+
 policy:
   dtensor_cfg:
     enabled: false

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -159,6 +159,12 @@ def __init__(self) -> None:
         self._checkpoint_active = False
         self._active_mutations = 0
         self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {}
+        self._mutation_version = 0
+
+    @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[None]:
@@ -187,6 +193,10 @@ async def mutation(self) -> AsyncIterator[None]:
             del self._mutation_depth_by_task[task]
             async with self._condition:
                 self._active_mutations -= 1
+                # Count completed mutation sections even when their body raised.
+                # A redundant periodic snapshot is safe; missing a mutation that
+                # partially changed TQ or controller metadata is not.
+                self._mutation_version += 1
                 if self._active_mutations == 0:
                     self._condition.notify_all()
 
@@ -1225,7 +1235,12 @@ async def _remove_unlocked(
 
         return len(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
@@ -1237,7 +1252,10 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState:
         across the complete publish/index or clear/remove transition. This
         includes future finalizer paths; canonical writes are not required to
         originate specifically from :meth:`commit`.
-        In-flight reservations are intentionally omitted.
+        In-flight reservations are intentionally omitted. ``additional_groups``
+        is used only by periodic snapshots to re-index canonical groups already
+        claimed by the current, uncheckpointed optimizer step. Their TQ rows are
+        still present, but normal sampler selection has removed their live slots.
         """
         groups: list[TQReplayGroupMetadata] = []
         for i, ready in enumerate(self.ready_list):
@@ -1254,6 +1272,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(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,

nemo_rl/algorithms/async_utils/staleness_sampler.py

@@ -136,6 +136,17 @@ def restore_dispatch_state(
     ) -> None: ...
 
 
+@runtime_checkable
+class TwoPhaseAdmissionSampler(Protocol):
+    """Sampler that separates a blocking gate from its admission commit."""
+
+    async def wait_until_admissible(
+        self, *, trainer_version_fn: Callable[[], int]
+    ) -> None: ...
+
+    def commit_admission(self, *, trainer_version: int) -> Optional[int]: ...
+
+
 class BaseSampler(abc.ABC):
     """Shared machinery for the built-in policies.
 
@@ -170,10 +181,19 @@ def set_dispatch_index(self, resume_from_trainer_version: int) -> None:
         self._dispatch_index = resume_from_trainer_version - 1
 
     # ── rollout-pump side ────────────────────────────────────────────────
-    @abc.abstractmethod
-    async def admit(
+    async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]:
+        """Wait for and commit one prompt-batch admission."""
+        await self.wait_until_admissible(trainer_version_fn=trainer_version_fn)
+        return self.commit_admission(trainer_version=trainer_version_fn())
+
+    async def wait_until_admissible(
         self, *, trainer_version_fn: Callable[[], int]
-    ) -> Optional[int]: ...
+    ) -> None:
+        """Wait until one batch may be admitted; ungated by default."""
+
+    def commit_admission(self, *, trainer_version: int) -> Optional[int]:
+        """Commit one already-admissible batch; unstamped by default."""
+        return None
 
     # ── train-pump side ──────────────────────────────────────────────────
     @abc.abstractmethod
@@ -377,9 +397,20 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
             gate_window=self._gate_window,
         )
 
-    async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]:
+    async def wait_until_admissible(
+        self, *, trainer_version_fn: Callable[[], int]
+    ) -> None:
         while self._dispatch_index >= trainer_version_fn() + self._gate_window:
             await asyncio.sleep(_GATE_POLL_SECONDS)
+
+    def commit_admission(self, *, trainer_version: int) -> Optional[int]:
+        if self._dispatch_index >= trainer_version + self._gate_window:
+            raise RuntimeError(
+                "sampler admission was committed before its gate opened: "
+                f"dispatch_index={self._dispatch_index}, "
+                f"trainer_version={trainer_version}, "
+                f"gate_window={self._gate_window}"
+            )
         self._dispatch_index += 1
         return self._stamp()