feat(checkpoint): trace rollout events to check what a restore cost

docs/design-docs/single-controller-rollout-checkpointing.md

@@ -173,6 +173,26 @@ original names (`gym_instance_id`, `gym_attempt`), so older checkpoints restore.
 `EnvironmentEpisode`'s row key format (`<rollout_id>-a<attempt>`) is Gym's, and a new
 backend must use it.
 
+## Tracing What a Restore Cost
+
+The metrics above are averages and counts. To check a stop and restore rollout by
+rollout, set `rollout_checkpointing.trace: true`. The controller then appends one
+JSON line per event to `rollout_trace.jsonl` in the log directory; a run that is
+killed loses only the last few milliseconds of events. Each line has the
+wall-clock time `t`, a `run` ID that is new for each process, the `event` name
+and that event's fields. The events are listed in
+`nemo_rl/algorithms/rollout_checkpointing/trace.py`. Per-rollout events need
+ledger dispatch, so they appear only with rollout recovery.
+
+Rows match across runs by `prompt_id` (the dataset index) and `generation_index`;
+group and attempt IDs are new in every run. A stop's thrown-away time is, for
+each attempt abandoned with `reason: restart`, the time from its last
+`attempt_dispatched` to the restored snapshot's `cut_time`; its lost progress is
+how much later the groups running at the stop reached `group_sealed` than the
+same prompts in runs that never stopped. A model call redone after a turn-level
+resume is not visible here: that needs the time the environment parked each
+rollout, which it does not report yet.
+
 ## Rules for New Code
 
 - Keep backend-specific code out of `single_controller.py` and `rollout_recovery.py`. Code

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -187,6 +187,9 @@ rollout_checkpointing:
   # 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: []
+  # Append one JSON line per rollout event to rollout_trace.jsonl in the log dir.
+  # The events are listed in nemo_rl/algorithms/rollout_checkpointing/trace.py.
+  trace: false
 
 policy:
   automodel_cfg:

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -198,6 +198,9 @@ rollout_checkpointing:
   # 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: []
+  # Append one JSON line per rollout event to rollout_trace.jsonl in the log dir.
+  # The events are listed in nemo_rl/algorithms/rollout_checkpointing/trace.py.
+  trace: false
 
 policy:
   tokenizer:

nemo_rl/algorithms/rollout_checkpointing/restore.py

@@ -201,6 +201,11 @@ async def restore(self) -> bool:
             Whether restore queued groups only ledger-recorded dispatch can
             launch, so the rollout pump must record prompt groups.
         """
+        self._trace.emit(
+            "restore_started",
+            path=self._resume.path,
+            from_rollout_snapshot=self._resume.from_rollout_snapshot,
+        )
         replay_restore_started = time.monotonic()
         restored_replay_groups = await self._restore_replay_buffer()
         replay_restore_seconds = time.monotonic() - replay_restore_started
@@ -216,6 +221,14 @@ async def restore(self) -> bool:
             recovery_prepare_seconds=recovery_prepare_seconds,
             restored_replay_groups=restored_replay_groups,
         )
+        if self._trace.enabled:
+            self._trace.emit(
+                "restore_finished",
+                replay_groups=restored_replay_groups,
+                unfinished_groups=[
+                    group.group_id for group in self._rollout_recovery_ledger.groups()
+                ],
+            )
         return self._requires_recovery_dispatch
 
     def _log_restore_metrics(

nemo_rl/algorithms/rollout_checkpointing/snapshots.py

@@ -239,13 +239,14 @@ async def trainer_checkpoint(self) -> AsyncIterator[None]:
         async with self._save_lock:
             yield
 
-    def trainer_checkpoint_published(self) -> None:
-        """Wake the pump once a trainer checkpoint is published.
+    def trainer_checkpoint_published(self, step: int) -> None:
+        """Wake the pump once the trainer checkpoint for ``step`` is published.
 
         A trainer checkpoint holds no environment state, so the pump waits for
         it to be durable, off the training loop, and then captures the
         environment anchored to it.
         """
+        self._trace.emit("trainer_checkpoint_published", step=step)
         if self._environment_checkpoint.holds_rollout_state:
             self._wakeup.set()
 
@@ -652,6 +653,7 @@ async def save_snapshot(
                     barrier_requested = time.monotonic()
                     async with self._data_plane_checkpoint_barrier.checkpoint() as cut:
                         barrier_acquired = time.monotonic()
+                        cut_wall_time = time.time()
                         if (
                             self._optimizer_commit_in_progress
                             or self._read_progress().train_steps != expected_train_step
@@ -742,6 +744,14 @@ async def save_snapshot(
                 step=expected_train_step,
                 prefix="timing/rollout_checkpoint",
             )
+            self._trace.emit(
+                "snapshot_published",
+                path=str(final_path),
+                step=expected_train_step,
+                trainer_version=expected_trainer_version,
+                # When the saved rollout state was captured.
+                cut_time=cut_wall_time,
+            )
             print(
                 "rollout checkpoint save completed: "
                 f"{final_path} (step={expected_train_step}, "

nemo_rl/algorithms/rollout_checkpointing/state.py

@@ -48,6 +48,7 @@
 if TYPE_CHECKING:
     from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer
     from nemo_rl.algorithms.async_utils.staleness_sampler import PromptGroupSampler
+    from nemo_rl.algorithms.rollout_checkpointing.trace import RolloutTrace
     from nemo_rl.experience.rollout_manager import RolloutManager
 
 
@@ -160,6 +161,8 @@ class RolloutParts:
     read_progress: Callable[[], TrainingProgress]
     # The controller's telemetry logger.
     log_metrics: Callable[..., None]
+    # The per-rollout event log; a no-op unless rollout_checkpointing.trace is set.
+    trace: RolloutTrace
 
 
 class _UsesRolloutParts:
@@ -180,3 +183,4 @@ def __init__(self, parts: RolloutParts) -> None:
         self._rollout_runtime = parts.runtime_state
         self._read_progress = parts.read_progress
         self._log_metrics = parts.log_metrics
+        self._trace = parts.trace

nemo_rl/algorithms/rollout_checkpointing/trace.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A per-rollout event log for checking what a stop and restore cost.
+
+With ``rollout_checkpointing.trace``, the controller appends one JSON line per
+event to ``rollout_trace.jsonl`` in the log directory: every rollout attempt's
+dispatch, seal, refusal, abandonment (with its reason) and resume, each prompt
+group's reservation, admission, seal and discard, each training step and
+weight sync, and each checkpoint published or restored. Every line carries the
+wall-clock time ``t`` and a ``run`` ID that is new for each process, so runs
+that resume into the same log directory stay apart.
+
+Events are queued as they happen and a background thread appends them, so
+writing never blocks the controller's event loop; a run that is killed loses
+only the events not yet written, typically the last few milliseconds' worth.
+Rows match across runs by
+``prompt_id`` (the dataset index) and ``generation_index``; group and attempt
+IDs are new in every run.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import queue
+import threading
+import time
+import uuid
+import warnings
+from pathlib import Path
+from typing import IO, Any, Optional
+
+ROLLOUT_TRACE_FILENAME = "rollout_trace.jsonl"
+
+
+class RolloutTrace:
+    """Append rollout events to a JSONL file off the caller's thread, or drop them."""
+
+    def __init__(self, path: Optional[Path]) -> None:
+        """Open ``path`` for appending; ``None`` makes every call a no-op.
+
+        Args:
+            path: The JSONL file to append to, or ``None`` to disable tracing.
+        """
+        self._run_id = uuid.uuid4().hex
+        self._file: Optional[IO[str]] = None
+        # Created only when enabled: a disabled trace holds no lock, so the
+        # controller class that defaults to one still pickles for Ray.
+        self._lines: Optional[queue.Queue[Optional[str]]] = None
+        self._writer: Optional[threading.Thread] = None
+        if path is not None:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            self._file = open(path, "a", encoding="utf-8")
+            self._lines = queue.Queue()
+            self._writer = threading.Thread(
+                target=self._write_lines, name="rollout-trace-writer", daemon=True
+            )
+            self._writer.start()
+            self.emit("trace_opened", pid=os.getpid())
+
+    @property
+    def enabled(self) -> bool:
+        return self._file is not None
+
+    def emit(self, event: str, **fields: Any) -> None:
+        """Queue one event, stamped with the current wall-clock time."""
+        if self._lines is None:
+            return
+        record = {"t": time.time(), "run": self._run_id, "event": event, **fields}
+        self._lines.put(json.dumps(record, default=str) + "\n")
+
+    def flush(self) -> None:
+        """Wait until every queued event has been written."""
+        if self._lines is not None:
+            self._lines.join()
+
+    def close(self) -> None:
+        """Write every queued event, then stop the writer and close the file."""
+        if self._file is None:
+            return
+        assert self._lines is not None and self._writer is not None
+        self._lines.put(None)
+        self._writer.join()
+        self._file.close()
+        self._file = None
+        self._lines = None
+
+    def _write_lines(self) -> None:
+        """Append queued events in batches until ``close`` queues ``None``."""
+        assert self._file is not None and self._lines is not None
+        trace_file, lines = self._file, self._lines
+        failed = False
+        while True:
+            batch = [lines.get()]
+            while True:
+                try:
+                    batch.append(lines.get_nowait())
+                except queue.Empty:
+                    break
+            done = batch[-1] is None
+            if not failed:
+                try:
+                    trace_file.write("".join(line for line in batch if line))
+                    trace_file.flush()
+                except OSError as error:
+                    # Tracing must never stop the run it observes.
+                    failed = True
+                    warnings.warn(
+                        f"rollout trace stopped writing: {error}", stacklevel=1
+                    )
+            for _ in batch:
+                lines.task_done()
+            if done:
+                return

nemo_rl/algorithms/single_controller.py

@@ -54,6 +54,7 @@
 from collections.abc import AsyncIterator, Iterator
 from dataclasses import replace
 from functools import partial
+from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -109,6 +110,10 @@
     RolloutRuntimeState,
     TrainingProgress,
 )
+from nemo_rl.algorithms.rollout_checkpointing.trace import (
+    ROLLOUT_TRACE_FILENAME,
+    RolloutTrace,
+)
 from nemo_rl.algorithms.single_controller_utils.advantage_stage import (
     AdvantageComputer,
     AdvantageOutcome,
@@ -327,6 +332,9 @@ class SingleControllerActor:
     # policy save whose async write has not been finalized and published yet.
     _defer_checkpoint_finalize: bool = False
     _pending_checkpoint: Optional[tuple[PathLike, int]] = None
+    # rollout_checkpointing.trace's event log. Disabled by default; a disabled
+    # trace holds no state, so one instance serves every controller.
+    _rollout_trace: RolloutTrace = RolloutTrace(None)
 
     # True from the moment a failed refit pulls every shard out of service until the
     # retry puts them back; the exhaustion check stands down for that window. See
@@ -716,6 +724,15 @@ def __init__(
             self._async_cfg.max_buffered_rollouts
         )
 
+        self._rollout_trace = RolloutTrace(
+            Path(self._logger.base_log_dir) / ROLLOUT_TRACE_FILENAME
+            if master_config.rollout_checkpointing.trace
+            else None
+        )
+        if self._rollout_trace.enabled:
+            self._rollout_manager.recovery_ledger.set_event_sink(
+                self._rollout_trace.emit
+            )
         # Built after every part they are given.
         self._rollout_snapshots, self._rollout_restorer = (
             self._build_rollout_checkpointing(
@@ -907,6 +924,7 @@ async def _run_pumps(self) -> dict[str, Any]:
                 except Exception as error:
                     print(f"finalizer actor termination failed: {error}", flush=True)
             await stop_storage_units(self._storage_units)
+            self._rollout_trace.close()
             if self._train_data_dump is not None:
                 try:
                     await asyncio.to_thread(self._train_data_dump.close)
@@ -968,6 +986,7 @@ def _build_rollout_checkpointing(
             runtime_state=self._rollout_runtime,
             read_progress=self._training_progress,
             log_metrics=self._log_telemetry_metrics,
+            trace=self._rollout_trace,
         )
         # Rollout snapshots, the TQ part of trainer checkpoints, the lock between
         # them and the optimizer-commit window.
@@ -3008,6 +3027,12 @@ async def _train_pump(self) -> None:
                 if not early_refit_done:
                     self._trainer_version += 1
                 self._train_steps += 1
+                self._rollout_trace.emit(
+                    "step_trained",
+                    step=self._train_steps,
+                    trainer_version=self._trainer_version,
+                    group_ids=sorted(step_finalizer_group_ids),
+                )
                 self._rollout_snapshots.end_optimizer_commit()
                 dropped_prompt_groups = self._rollout_runtime.batch_shortfall.get(
                     version_during_step, 0
@@ -4212,7 +4237,7 @@ async def _publish_when_policy_write_done(self, checkpoint_path: PathLike) -> No
             self._checkpointer,
             last_checkpoint_step=self._train_steps,
         )
-        self._rollout_snapshots.trainer_checkpoint_published()
+        self._rollout_snapshots.trainer_checkpoint_published(self._train_steps)
 
     async def _finish_pending_save(self, *, blocking: bool) -> None:
         """Finalize and publish a save deferred by defer_async_checkpoint_finalize.
@@ -4242,7 +4267,7 @@ async def _finish_pending_save(self, *, blocking: bool) -> None:
                 self._checkpointer,
                 last_checkpoint_step=step,
             )
-        self._rollout_snapshots.trainer_checkpoint_published()
+        self._rollout_snapshots.trainer_checkpoint_published(step)
 
     # Grace on top of the workers' own refit deadline, so their attributable error wins
     # the race against this bound. They need to hit the deadline, abort, unwind, raise,
@@ -4506,6 +4531,7 @@ async def _sync_weights(
             await asyncio.to_thread(
                 self._gen.set_rollout_weight_version, self._trainer_version
             )
+        self._rollout_trace.emit("weights_synced", weight_version=self._trainer_version)
         self._rollout_permitted.set()
         self._rollout_manager.resume_request_deadlines()
         return aborted_stale_inflight_groups

nemo_rl/algorithms/single_controller_utils/config.py

@@ -903,6 +903,12 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"):
     periodic-checkpoint failures are tolerated before the controller aborts the
     run. A successful or skipped attempt resets the counter; checkpoint
     invariant failures still fail immediately.
+
+    ``trace=True`` appends one JSON line per rollout event to
+    ``rollout_trace.jsonl`` in the log directory, to check what a stop and
+    restore cost. Per-rollout events need ``rollout_recovery`` dispatch. It does
+    not enable checkpointing. The events are listed in
+    ``nemo_rl/algorithms/rollout_checkpointing/trace.py``.
     """
 
     snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None
@@ -911,6 +917,7 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"):
     keep_latest_k: Annotated[int, Field(ge=1)] = 2
     restore_mode: Literal["latest", "trainer_checkpoint"] = "latest"
     extra_fingerprint_excluded_paths: list[str] = Field(default_factory=list)
+    trace: bool = False
 
     @model_validator(mode="after")
     def validate_extra_fingerprint_excluded_paths(self) -> "RolloutCheckpointConfig":