feat(rollout): persist and recover partial rollout groups

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -158,16 +158,33 @@ def __init__(self) -> None:
         self._condition = asyncio.Condition()
         self._checkpoint_active = False
         self._active_mutations = 0
+        self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {}
 
     @asynccontextmanager
     async def mutation(self) -> AsyncIterator[None]:
         """Enter a commit/clear section, waiting only for an active checkpoint."""
+        task = asyncio.current_task()
+        if task is None:
+            raise RuntimeError("data-plane mutation must run inside an asyncio task")
+        depth = self._mutation_depth_by_task.get(task, 0)
+        if depth:
+            # Replay-buffer helpers may join a controller-owned mutation. Do not
+            # wait behind a checkpoint that is already waiting for this outer
+            # mutation, or the two tasks deadlock.
+            self._mutation_depth_by_task[task] = depth + 1
+            try:
+                yield
+            finally:
+                self._mutation_depth_by_task[task] -= 1
+            return
         async with self._condition:
             await self._condition.wait_for(lambda: not self._checkpoint_active)
             self._active_mutations += 1
+            self._mutation_depth_by_task[task] = 1
         try:
             yield
         finally:
+            del self._mutation_depth_by_task[task]
             async with self._condition:
                 self._active_mutations -= 1
                 if self._active_mutations == 0:

nemo_rl/algorithms/async_utils/staleness_sampler.py

@@ -41,6 +41,7 @@
 import abc
 import asyncio
 import importlib
+from collections import Counter
 from typing import (
     Annotated,
     Callable,
@@ -122,6 +123,19 @@ def set_dispatch_index(self, resume_from_trainer_version: int) -> None:
         ...
 
 
+@runtime_checkable
+class CheckpointDispatchSampler(Protocol):
+    """Sampler that can reconstruct dispatch state from restored groups."""
+
+    def restore_dispatch_state(
+        self,
+        *,
+        current_train_weight: int,
+        restored_target_steps: list[Optional[int]],
+        groups_per_step: int,
+    ) -> None: ...
+
+
 class BaseSampler(abc.ABC):
     """Shared machinery for the built-in policies.
 
@@ -205,6 +219,19 @@ def is_on_policy(self) -> bool:
     def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
         return None
 
+    def restore_dispatch_state(
+        self,
+        *,
+        current_train_weight: int,
+        restored_target_steps: list[Optional[int]],
+        groups_per_step: int,
+    ) -> None:
+        """Validate the default unstamped checkpoint representation."""
+        if any(target is not None for target in restored_target_steps):
+            raise ValueError(
+                f"{type(self).__name__} cannot restore stamped target steps"
+            )
+
     # ── shared helpers ───────────────────────────────────────────────────
     def _eviction_window(self) -> int:
         """Weight-version span kept selectable; drives the default ``evict``."""
@@ -415,6 +442,62 @@ def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> No
     def _stamp(self) -> Optional[int]:
         return self._dispatch_index
 
+    supports_buffer_checkpoint: ClassVar[bool] = True
+
+    def restore_dispatch_state(
+        self,
+        *,
+        current_train_weight: int,
+        restored_target_steps: list[Optional[int]],
+        groups_per_step: int,
+    ) -> None:
+        """Resume admission after the highest target restored from the cut."""
+        if groups_per_step < 1:
+            raise ValueError(f"groups_per_step must be positive, got {groups_per_step}")
+        if not restored_target_steps:
+            return
+        if any(target is None for target in restored_target_steps):
+            raise ValueError("restored InOrder groups must have target_step values")
+
+        target_steps = [
+            target for target in restored_target_steps if target is not None
+        ]
+        counts: Counter[int] = Counter(target_steps)
+        oldest_target = min(counts)
+        newest_target = max(counts)
+        if oldest_target < current_train_weight:
+            raise ValueError(
+                "restored InOrder target is older than the trainer: "
+                f"oldest_target={oldest_target}, "
+                f"trainer_version={current_train_weight}"
+            )
+        max_allowed_target = current_train_weight + self.max_lookahead_versions
+        if newest_target > max_allowed_target:
+            raise ValueError(
+                "restored InOrder target exceeds the configured lookahead: "
+                f"newest_target={newest_target}, "
+                f"max_allowed_target={max_allowed_target}"
+            )
+        expected_targets = set(range(current_train_weight, newest_target + 1))
+        if set(counts) != expected_targets:
+            raise ValueError(
+                "restored InOrder targets are not contiguous from the current "
+                f"trainer version: restored={sorted(counts)}, "
+                f"expected={sorted(expected_targets)}"
+            )
+        invalid_counts = {
+            target: count
+            for target, count in counts.items()
+            if count != groups_per_step
+        }
+        if invalid_counts:
+            raise ValueError(
+                "restored InOrder target batches do not contain exactly "
+                f"grpo.num_prompts_per_step={groups_per_step} groups: "
+                f"counts={dict(sorted(counts.items()))}"
+            )
+        self._dispatch_index = newest_target
+
     async def select(
         self,
         *,