feat(rollout): add partial recovery benchmark telemetry

examples/nemo_gym/grpo_qwen3_30ba3b_instruct.yaml

@@ -53,6 +53,10 @@ policy:
   generation:
     vllm_cfg:
       tensor_parallel_size: 4
+      # Required for rollout-throughput checkpoint A/B telemetry. Keep this
+      # identical in baseline and checkpoint-enabled runs.
+      enable_vllm_metrics_logger: true
+      vllm_metrics_logger_interval: 0.5
       # This is a very low GPU mem utilization. We GPU OOM in two places:
       # Refit after train, refit before validation.
       gpu_memory_utilization: 0.7
@@ -172,3 +176,8 @@ checkpointing:
   # 4. The checkpoint time for this model is around 10 mins.
   checkpoint_must_save_by: "00:03:30:00"
   save_period: 1
+
+# Optional wall-clock benchmark sampling. This is independent of the rollout
+# checkpoint save interval and may be enabled for baseline runs as well.
+rollout_checkpointing:
+  telemetry_interval_s: null

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -18,10 +18,12 @@
 import math
 import statistics
 import threading as _threading
+import time
 import uuid
-from collections import Counter
+from collections import Counter, deque
 from collections.abc import AsyncIterator, Mapping
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from numbers import Integral, Real
 from typing import Any, Iterable, Literal, Optional, TypedDict
 
@@ -42,6 +44,38 @@
 REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1
 REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint"
 
+CheckpointMutationKind = Literal[
+    "group_commits",
+    "group_removals",
+    "other",
+    "prompt_reservations",
+    "recovery_retries",
+    "sample_clears",
+    "sibling_seals",
+]
+CHECKPOINT_MUTATION_KINDS: tuple[CheckpointMutationKind, ...] = (
+    "group_commits",
+    "group_removals",
+    "prompt_reservations",
+    "recovery_retries",
+    "sample_clears",
+    "sibling_seals",
+    "other",
+)
+
+
+@dataclass(frozen=True)
+class DataPlaneCheckpointBarrierTelemetry:
+    """Bounded interval telemetry for checkpoint-induced mutation waits."""
+
+    blocked_by_kind: dict[CheckpointMutationKind, int]
+    wait_durations_s: tuple[float, ...]
+    active_mutations: int
+    waiting_mutations: int
+    max_waiting_mutations: int
+    checkpoint_active: bool
+
+
 # These TypedDicts describe the versioned, plain-mapping checkpoint wire
 # format. They are intentionally not dataclass instances: persisting a
 # dataclass would couple recovery to its Python import path and class layout.
@@ -160,14 +194,20 @@ def __init__(self) -> None:
         self._active_mutations = 0
         self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {}
         self._mutation_version = 0
+        self._waiting_mutations = 0
+        self._max_waiting_mutations = 0
+        self._blocked_by_kind: Counter[CheckpointMutationKind] = Counter()
+        self._wait_durations_s: deque[float] = deque(maxlen=10_000)
 
     @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]:
+    async def mutation(
+        self, kind: CheckpointMutationKind = "other"
+    ) -> AsyncIterator[None]:
         """Enter a commit/clear section, waiting only for an active checkpoint."""
         task = asyncio.current_task()
         if task is None:
@@ -184,7 +224,20 @@ async def mutation(self) -> AsyncIterator[None]:
                 self._mutation_depth_by_task[task] -= 1
             return
         async with self._condition:
-            await self._condition.wait_for(lambda: not self._checkpoint_active)
+            wait_started: Optional[float] = None
+            if self._checkpoint_active:
+                wait_started = time.monotonic()
+                self._waiting_mutations += 1
+                self._max_waiting_mutations = max(
+                    self._max_waiting_mutations, self._waiting_mutations
+                )
+            try:
+                await self._condition.wait_for(lambda: not self._checkpoint_active)
+            finally:
+                if wait_started is not None:
+                    self._waiting_mutations -= 1
+                    self._blocked_by_kind[kind] += 1
+                    self._wait_durations_s.append(time.monotonic() - wait_started)
             self._active_mutations += 1
             self._mutation_depth_by_task[task] = 1
         try:
@@ -200,6 +253,25 @@ async def mutation(self) -> AsyncIterator[None]:
                 if self._active_mutations == 0:
                     self._condition.notify_all()
 
+    async def drain_telemetry(self) -> DataPlaneCheckpointBarrierTelemetry:
+        """Return and reset interval waits while preserving current state."""
+        async with self._condition:
+            telemetry = DataPlaneCheckpointBarrierTelemetry(
+                blocked_by_kind={
+                    kind: self._blocked_by_kind.get(kind, 0)
+                    for kind in CHECKPOINT_MUTATION_KINDS
+                },
+                wait_durations_s=tuple(self._wait_durations_s),
+                active_mutations=self._active_mutations,
+                waiting_mutations=self._waiting_mutations,
+                max_waiting_mutations=self._max_waiting_mutations,
+                checkpoint_active=self._checkpoint_active,
+            )
+            self._blocked_by_kind.clear()
+            self._wait_durations_s.clear()
+            self._max_waiting_mutations = self._waiting_mutations
+            return telemetry
+
     @asynccontextmanager
     async def checkpoint(self) -> AsyncIterator[None]:
         """Block new mutations and wait for active ones before snapshotting."""
@@ -970,7 +1042,7 @@ async def commit(
             weight_version=start_weight_version,
         )
         trace_rollout_payload(keys=sample_ids, data=train_batch)
-        async with self._data_plane_checkpoint_barrier.mutation():
+        async with self._data_plane_checkpoint_barrier.mutation("group_commits"):
             try:
                 await call_data_plane(
                     self._dp_client,
@@ -1037,7 +1109,7 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in
                 "TQReplayBuffer must be bound to the controller data-plane "
                 "checkpoint barrier before removing a group"
             )
-        async with self._data_plane_checkpoint_barrier.mutation():
+        async with self._data_plane_checkpoint_barrier.mutation("group_removals"):
             try:
                 idx = self._group_ids.index(group_id)
             except ValueError as error:
@@ -1058,7 +1130,7 @@ async def clear_staging_keys(self, staging_keys: list[str]) -> None:
                 "checkpoint barrier before clearing staging samples"
             )
         unique_keys = list(dict.fromkeys(staging_keys))
-        async with self._data_plane_checkpoint_barrier.mutation():
+        async with self._data_plane_checkpoint_barrier.mutation("sample_clears"):
             await call_data_plane(
                 self._dp_client,
                 "clear_samples",
@@ -1095,45 +1167,56 @@ async def commit_finalized(
         Raises:
             ValueError: group_id has no live slot (removed or never reserved).
         """
-        try:
-            idx = self._group_ids.index(group_id)
-        except ValueError:
-            raise ValueError(
-                f"TQReplayBuffer.commit_finalized: group {group_id} has no "
-                "live slot (evicted or never reserved)"
-            ) from None
-        tagged_plans = [
-            tag[ROUTE_PLAN_TAG] for tag in (meta.tags or []) if ROUTE_PLAN_TAG in tag
-        ]
-        if tagged_plans:
-            if len(tagged_plans) != len(meta.sample_ids):
-                raise ValueError(
-                    "commit_finalized received mixed deferred/direct route plans"
-                )
-            from nemo_rl.experience.route_plan import decode_route_plan
+        if self._data_plane_checkpoint_barrier is None:
+            raise RuntimeError(
+                "TQReplayBuffer must be bound to the controller data-plane "
+                "checkpoint barrier before committing finalized samples"
+            )
 
-            plan_cleanup_keys = {
-                key
-                for encoded in tagged_plans
-                for key in decode_route_plan(encoded).cleanup_staging_keys
-            }
-            provided_staging_keys = list(staging_keys or [])
-            if len(provided_staging_keys) != len(set(provided_staging_keys)):
-                raise ValueError("commit_finalized staging_keys contains duplicates")
-            if set(provided_staging_keys) != plan_cleanup_keys:
+        async with self._data_plane_checkpoint_barrier.mutation("group_commits"):
+            try:
+                idx = self._group_ids.index(group_id)
+            except ValueError:
                 raise ValueError(
-                    "commit_finalized staging ownership does not match route plans: "
-                    f"provided={sorted(provided_staging_keys)!r}, "
-                    f"planned={sorted(plan_cleanup_keys)!r}"
-                )
-        self.meta_list[idx] = meta
-        self.start_weight_list[idx] = group_min_wv
-        self.end_weight_list[idx] = group_max_wv
-        self.ready_list[idx] = True
-        self._staging_keys_list[idx] = (
-            list(staging_keys) if staging_keys is not None else None
-        )
-        return meta
+                    f"TQReplayBuffer.commit_finalized: group {group_id} has no "
+                    "live slot (evicted or never reserved)"
+                ) from None
+            tagged_plans = [
+                tag[ROUTE_PLAN_TAG]
+                for tag in (meta.tags or [])
+                if ROUTE_PLAN_TAG in tag
+            ]
+            if tagged_plans:
+                if len(tagged_plans) != len(meta.sample_ids):
+                    raise ValueError(
+                        "commit_finalized received mixed deferred/direct route plans"
+                    )
+                from nemo_rl.experience.route_plan import decode_route_plan
+
+                plan_cleanup_keys = {
+                    key
+                    for encoded in tagged_plans
+                    for key in decode_route_plan(encoded).cleanup_staging_keys
+                }
+                provided_staging_keys = list(staging_keys or [])
+                if len(provided_staging_keys) != len(set(provided_staging_keys)):
+                    raise ValueError(
+                        "commit_finalized staging_keys contains duplicates"
+                    )
+                if set(provided_staging_keys) != plan_cleanup_keys:
+                    raise ValueError(
+                        "commit_finalized staging ownership does not match route plans: "
+                        f"provided={sorted(provided_staging_keys)!r}, "
+                        f"planned={sorted(plan_cleanup_keys)!r}"
+                    )
+            self.meta_list[idx] = meta
+            self.start_weight_list[idx] = group_min_wv
+            self.end_weight_list[idx] = group_max_wv
+            self.ready_list[idx] = True
+            self._staging_keys_list[idx] = (
+                list(staging_keys) if staging_keys is not None else None
+            )
+            return meta
 
     def abort(self, group_id: str) -> bool:
         """Drop an unready slot whose dispatch failed or was cancelled.
@@ -1188,7 +1271,7 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int:
                 "TQReplayBuffer must be bound to the controller data-plane "
                 "checkpoint barrier before removing groups"
             )
-        async with self._data_plane_checkpoint_barrier.mutation():
+        async with self._data_plane_checkpoint_barrier.mutation("group_removals"):
             drop_idxs = sorted(idxs, reverse=True)
             if drop_idxs[0] >= len(self.meta_list):
                 raise IndexError(
@@ -1478,7 +1561,7 @@ async def _clear_samples(self, *, sample_ids: list[str]) -> None:
                 "TQReplayBuffer must be bound to the controller data-plane "
                 "checkpoint barrier before clearing samples"
             )
-        async with self._data_plane_checkpoint_barrier.mutation():
+        async with self._data_plane_checkpoint_barrier.mutation("sample_clears"):
             await self._clear_samples_unlocked(sample_ids=sample_ids)
 
     async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None:

nemo_rl/algorithms/distillation.py

@@ -1145,7 +1145,12 @@ def distillation_train(
                 metrics["global_valid_toks"] / total_time / total_num_gpus
             )
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
-            logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train")
+            logger.log_metrics(
+                timing_metrics,
+                total_steps + 1,
+                prefix="timing/train",
+                step_finished=True,
+            )
 
             timer.reset()
             current_step += 1

nemo_rl/algorithms/dpo.py

@@ -792,7 +792,12 @@ def dpo_train(
                 metrics["global_valid_toks"] / total_time / total_num_gpus
             )
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
-            logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train")
+            logger.log_metrics(
+                timing_metrics,
+                total_steps + 1,
+                prefix="timing/train",
+                step_finished=True,
+            )
 
             timer.reset()
             current_step += 1

nemo_rl/algorithms/rm.py

@@ -715,7 +715,12 @@ def rm_train(
                 metrics["global_valid_toks"] / total_time / total_num_gpus
             )
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
-            logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train")
+            logger.log_metrics(
+                timing_metrics,
+                total_steps + 1,
+                prefix="timing/train",
+                step_finished=True,
+            )
 
             timer.reset()
             current_step += 1

nemo_rl/algorithms/sft.py

@@ -651,7 +651,12 @@ def sft_train(
             else:
                 timing_metrics["valid_tokens_per_sec_per_gpu"] = 0.0
             logger.log_metrics(metrics, total_steps + 1, prefix="train")
-            logger.log_metrics(timing_metrics, total_steps + 1, prefix="timing/train")
+            logger.log_metrics(
+                timing_metrics,
+                total_steps + 1,
+                prefix="timing/train",
+                step_finished=True,
+            )
 
             timer.reset()
             current_step += 1