feat(checkpoint): snapshot rollouts at the checkpoint deadline

docs/guides/single-controller.md

@@ -220,6 +220,7 @@ checkpointing:
 
 rollout_checkpointing:
   snapshot_attempt_interval_s: 120
+  snapshot_on_trainer_checkpoint: true
   telemetry_interval_s: null
   max_consecutive_failures: 3
   keep_latest_k: 2
@@ -239,6 +240,27 @@ 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.
 
+With `rollout_recovery.target_level: turn`, rollout snapshots are on even with
+`snapshot_attempt_interval_s: null`, because only a snapshot holds the
+environment's state. Two more saves then apply:
+
+- At the `checkpointing.checkpoint_must_save_by` deadline, the trainer
+  checkpoint is followed by a rollout snapshot before the run stops. The
+  deadline is moved earlier by how long the last snapshot took, so both fit.
+- `snapshot_on_trainer_checkpoint` (default `true`) takes a snapshot after
+  every other trainer checkpoint too. Set it to `false` with no interval to
+  save environment state only at the deadline, for example to carry rollouts
+  across a job's walltime without pausing the environment every step.
+
+| Saves environment state | `snapshot_on_trainer_checkpoint` | `snapshot_attempt_interval_s` |
+|---|---|---|
+| After every trainer checkpoint and at the deadline | `true` | `null` |
+| Only at the deadline | `false` | `null` |
+| After every trainer checkpoint, at the deadline, and every 300 s | `true` | `300` |
+
+Setup rejects turn-level recovery that would never save: `false`, no interval,
+and no `checkpoint_must_save_by`.
+
 `telemetry_interval_s` controls an independent wall-clock sampler for rollout
 throughput and checkpoint pressure. It is `null` (disabled) by default; set it
 to a positive number such as `30` to emit one sample every 30 seconds. This does
@@ -476,7 +498,7 @@ Field definitions:
 - `sampler.warmup_lookahead_versions` (PPO, `in_order` only) — lookahead used while `ppo.policy_training_start_step` critic warmup is in progress, shrinking back to `max_lookahead_versions` afterwards. The SC equivalent of `ppo.async_ppo.warmup_generation_lead_steps`.
 - `num_advantage_workers` — CPU Ray actors that run the advantage stage. `0` (the default) keeps it in the controller process, which is correct but both holds a whole cohort's advantage inputs in the controller's heap and blocks the controller's event loop for the duration of the computation — long enough at Ultra scale to miss Ray's actor liveness ping. A positive value moves both costs onto dedicated actors, placed off the dedicated Ray head so their host-memory peak does not land beside the controller. Only `grpo` and `opd` are sharded across the pool: every other `adv_estimator` also reduces over the whole batch, which a shard is not, so those run as one call on one actor and a larger pool buys them nothing. Each stage logs `advantage stage: N row(s) over M shard(s), pool=P` with the reason whenever `M` is 1. Under `data_plane.backend: mooncake_cpu` each actor is its own TransferQueue client and mounts a full `global_segment_size + local_buffer_size`, like each token-capture finalizer, so budget it on top of `gpus_per_node × (segment + buffer)`.
 - `log_full_train_data` — writes every consumed sample to `train_data_step<N>.jsonl` in the log dir. The advantage stage is what holds the untruncated tensors, so the stage writes the dump wherever it runs: with a pool each shard writes its own `train_data_step<N>.jsonl.part-<shard>` and the controller merges them in shard order on a completed optimizer step, assigning `idx` across the merged file. That requires the pool to share a filesystem with the controller: the log dir was controller-only before this, so a node-local one worked and now would not. The controller checks the merged row count against what the shards reported and fails the step rather than publishing a dump that is short a shard. The JSON serialization is the expensive part and scales with batch size × sequence length, so a pool also moves that cost off the controller; it is reported back per call and still shows up under the `train_data_dump` timer. Debug-only — leave it `false` for production runs.
-- `defer_async_checkpoint_finalize` — with `policy.megatron_cfg.checkpoint.async_save: true`, finalizes an async policy checkpoint at a later step boundary, once every rank's background write is done, instead of from a background thread whose blocking finalize holds the policy workers until the write lands. The train pump polls at the top of each step, before the step's first policy-worker call; the next save and training exit still finalize blocking, and so does teardown after a failure. `tmp_step_N` is renamed to `step_N` only after that finalize, and resume reads only `step_N`, so a crash in that window — typically one step — resumes from the previous checkpoint. It cannot be combined with rollout snapshots (`rollout_checkpointing.snapshot_attempt_interval_s`), and setup rejects that combination: a snapshot anchors on `step_N`, which a deferred finalize publishes only after training has moved past step N. Only non-colocated generation gains: a colocated refit, and PPO's per-chunk critic handoff, block on the finalize anyway. Time spent finalizing at a step boundary is reported as `timing/train/checkpoint_finalize`. Leave it `false` (the default) when `async_save` is off; setup rejects that combination.
+- `defer_async_checkpoint_finalize` — with `policy.megatron_cfg.checkpoint.async_save: true`, finalizes an async policy checkpoint at a later step boundary, once every rank's background write is done, instead of from a background thread whose blocking finalize holds the policy workers until the write lands. The train pump polls at the top of each step, before the step's first policy-worker call; the next save and training exit still finalize blocking, and so does teardown after a failure. `tmp_step_N` is renamed to `step_N` only after that finalize, and resume reads only `step_N`, so a crash in that window — typically one step — resumes from the previous checkpoint. It cannot be combined with rollout snapshots taken while training continues (`rollout_checkpointing.snapshot_attempt_interval_s`, or `snapshot_on_trainer_checkpoint` with turn-level recovery), and setup rejects that combination: a snapshot anchors on `step_N`, which a deferred finalize publishes only after training has moved past step N. The snapshot at the `checkpoint_must_save_by` deadline is allowed, since the run publishes `step_N` first and then stops. Only non-colocated generation gains: a colocated refit, and PPO's per-chunk critic handoff, block on the finalize anyway. Time spent finalizing at a step boundary is reported as `timing/train/checkpoint_finalize`. Leave it `false` (the default) when `async_save` is off; setup rejects that combination.
 
 ## Implementation Structure
 

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -172,13 +172,19 @@ checkpointing:
   # 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-only snapshots are on when snapshot_attempt_interval_s is set or
+# rollout_recovery.target_level is "turn"; otherwise existing snapshots are
+# ignored. They require 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:
+  # Attempt a snapshot every this many seconds; null for none on a timer.
   snapshot_attempt_interval_s: null
+  # With target_level "turn": also snapshot the environment after every trainer
+  # checkpoint. A snapshot always follows the trainer checkpoint that
+  # checkpointing.checkpoint_must_save_by triggers, so false with no interval
+  # saves environment state only at that deadline.
+  snapshot_on_trainer_checkpoint: true
   telemetry_interval_s: null
   max_consecutive_failures: 3
   keep_latest_k: 2
@@ -242,9 +248,8 @@ token_capture:
 # coordinated Gym participant snapshots; "sibling" preserves sealed siblings;
 # "prompt_group" regenerates every sibling in an unfinished group.
 # "turn" also requires env.should_use_nemo_gym, the vllm generation backend,
-# token_capture.enabled: true, rollout_checkpointing.snapshot_attempt_interval_s,
-# rollout_checkpointing.restore_mode: latest, both override maps left empty,
-# only legacy_agent Gym environment servers, and
+# token_capture.enabled: true, rollout_checkpointing.restore_mode: latest,
+# both override maps left empty, only legacy_agent Gym environment servers, and
 # async_rl.stall_watchdog.stall_timeout_s above
 # gym_prepare_timeout_s + 2 * token_capture.control_timeout_s.
 rollout_recovery:

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -183,13 +183,19 @@ checkpointing:
   # 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-only snapshots are on when snapshot_attempt_interval_s is set or
+# rollout_recovery.target_level is "turn"; otherwise existing snapshots are
+# ignored. They require 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:
+  # Attempt a snapshot every this many seconds; null for none on a timer.
   snapshot_attempt_interval_s: null
+  # With target_level "turn": also snapshot the environment after every trainer
+  # checkpoint. A snapshot always follows the trainer checkpoint that
+  # checkpointing.checkpoint_must_save_by triggers, so false with no interval
+  # saves environment state only at that deadline.
+  snapshot_on_trainer_checkpoint: true
   telemetry_interval_s: null
   max_consecutive_failures: 3
   keep_latest_k: 2
@@ -292,9 +298,8 @@ token_capture:
 # coordinated Gym participant snapshots; "sibling" preserves sealed siblings;
 # "prompt_group" regenerates every sibling in an unfinished group.
 # "turn" also requires env.should_use_nemo_gym, the vllm generation backend,
-# token_capture.enabled: true, rollout_checkpointing.snapshot_attempt_interval_s,
-# rollout_checkpointing.restore_mode: latest, both override maps left empty,
-# only legacy_agent Gym environment servers, and
+# token_capture.enabled: true, rollout_checkpointing.restore_mode: latest,
+# both override maps left empty, only legacy_agent Gym environment servers, and
 # async_rl.stall_watchdog.stall_timeout_s above
 # gym_prepare_timeout_s + 2 * token_capture.control_timeout_s.
 rollout_recovery:

nemo_rl/algorithms/rollout_checkpointing/snapshots.py

@@ -244,10 +244,14 @@ def trainer_checkpoint_published(self, step: int) -> None:
 
         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.
+        environment anchored to it, unless
+        ``rollout_checkpointing.snapshot_on_trainer_checkpoint`` is off.
         """
         self._trace.emit("trainer_checkpoint_published", step=step)
-        if self._environment_checkpoint.holds_rollout_state:
+        if (
+            self._environment_checkpoint.holds_rollout_state
+            and self._master_config.rollout_checkpointing.snapshot_on_trainer_checkpoint
+        ):
             self._wakeup.set()
 
     async def _save_data_plane_checkpoint(
@@ -708,6 +712,10 @@ async def save_snapshot(
             self._last_snapshot_mutation_version = snapshot_cut.mutation_version
             self._last_missing_anchor = None
             save_completed = time.monotonic()
+            if self._environment_checkpoint.holds_rollout_state:
+                # The deadline's trainer checkpoint is followed by a snapshot
+                # like this one, so the deadline must leave time for both.
+                self._timeout.save_reserve_seconds = save_completed - save_started
             checkpoint_metrics = {
                 "snapshot_sequence": float(snapshot_sequence),
                 "total_save_seconds": save_completed - save_started,
@@ -823,40 +831,72 @@ def _log_outcome(
             prefix=ROLLOUT_CHECKPOINT_OUTCOME_PREFIX,
         )
 
-    async def run(self) -> None:
-        """Persist rollout state periodically, including during streamed train."""
-        snapshot_attempt_interval_s = (
+    def _failure_reason(self, error: Exception) -> RolloutCheckpointAttemptReason:
+        """Classify a failed snapshot; only ``invariant_error`` is not retryable."""
+        if isinstance(error, TimeoutError):
+            return "timeout"
+        if isinstance(error, OSError):
+            return "io_error"
+        if self._environment_checkpoint.is_unavailable(error):
+            # Nothing RL holds changed, so the next attempt can retry.
+            return "gym_unavailable"
+        return "invariant_error"
+
+    def _seconds_until_next_attempt(self) -> Optional[float]:
+        """How long the pump waits for a wakeup before attempting anyway.
+
+        ``None`` waits only for a trainer checkpoint's wakeup. Before the first
+        train step the pump also wakes at the checkpoint deadline, since no
+        train step will reach it.
+        """
+        waits = []
+        interval_s = (
             self._master_config.rollout_checkpointing.snapshot_attempt_interval_s
         )
-        if snapshot_attempt_interval_s is None:
-            raise RuntimeError("rollout checkpoint pump started while disabled")
+        if interval_s is not None:
+            waits.append(interval_s)
+        if self._read_progress().train_steps == 0:
+            until_deadline = self._timeout.seconds_until_save()
+            if until_deadline is not None:
+                # A deadline snapshot that was skipped retries once a second.
+                waits.append(max(until_deadline, 1.0))
+        return min(waits, default=None)
+
+    async def run(self) -> None:
+        """Save rollout snapshots on schedule, including during streamed train.
+
+        Attempts one every ``snapshot_attempt_interval_s`` when set, soon after
+        each trainer checkpoint that asks for one, and at the checkpoint
+        deadline when it arrives before the first train step.
+        """
         consecutive_failures = 0
         consecutive_discards = 0
         while True:
             try:
                 await asyncio.wait_for(
                     self._wakeup.wait(),
-                    timeout=snapshot_attempt_interval_s,
+                    timeout=self._seconds_until_next_attempt(),
                 )
             except TimeoutError:
                 pass
+            woken = self._wakeup.is_set()
             self._wakeup.clear()
             attempt_started = time.monotonic()
             deadline_due = (
                 self._read_progress().train_steps == 0 and self._timeout.would_save()
             )
+            if (
+                not woken
+                and not deadline_due
+                and self._master_config.rollout_checkpointing.snapshot_attempt_interval_s
+                is None
+            ):
+                # A wait timed for a deadline that the first train step now handles.
+                continue
             try:
                 result = await self.save_snapshot(force=deadline_due)
             except Exception as error:
-                if isinstance(error, TimeoutError):
-                    failure_reason: RolloutCheckpointAttemptReason = "timeout"
-                elif isinstance(error, OSError):
-                    failure_reason = "io_error"
-                elif self._environment_checkpoint.is_unavailable(error):
-                    # Nothing RL holds changed, so the next attempt can retry.
-                    failure_reason = "gym_unavailable"
-                else:
-                    failure_reason = "invariant_error"
+                failure_reason = self._failure_reason(error)
                 self._log_outcome(
                     outcome="failed",
                     reason=failure_reason,
@@ -905,3 +945,46 @@ async def run(self) -> None:
                 )
                 self._stop_event.set()
                 return
+
+    async def save_at_deadline(self) -> None:
+        """Save the environment state the deadline's trainer checkpoint lacks.
+
+        The run stops after this, so nothing later would snapshot it. A failure
+        that leaves checkpoint state consistent only warns: the trainer
+        checkpoint is durable, and a restore from it redoes the unfinished turns.
+        """
+        if not self._environment_checkpoint.holds_rollout_state:
+            return
+        # The pump need not also answer the trainer checkpoint's wakeup.
+        self._wakeup.clear()
+        attempt_started = time.monotonic()
+        try:
+            result = await self.save_snapshot(force=True)
+        except Exception as error:
+            failure_reason = self._failure_reason(error)
+            self._log_outcome(
+                outcome="failed",
+                reason=failure_reason,
+                attempt_duration_seconds=time.monotonic() - attempt_started,
+            )
+            if failure_reason == "invariant_error":
+                raise
+            warnings.warn(
+                "rollout snapshot at the checkpoint deadline failed; a restore "
+                "redoes the unfinished turns from the trainer checkpoint: "
+                f"{type(error).__name__}: {error}",
+                stacklevel=2,
+            )
+            return
+        self._log_outcome(
+            outcome="completed" if result.saved else "skipped",
+            reason=result.reason,
+            attempt_duration_seconds=time.monotonic() - attempt_started,
+        )
+        if not result.saved:
+            warnings.warn(
+                "rollout snapshot at the checkpoint deadline was skipped "
+                f"({result.reason}); a restore redoes the unfinished turns from "
+                "the trainer checkpoint",
+                stacklevel=2,
+            )

nemo_rl/algorithms/single_controller.py

@@ -127,6 +127,7 @@
     MasterConfig,
     algo_config,
     is_ppo_run,
+    rollout_snapshots_enabled,
     validate_forest_objective,
 )
 from nemo_rl.algorithms.single_controller_utils.debug_writer import (
@@ -837,8 +838,7 @@ async def _run_pumps(self) -> dict[str, Any]:
             await self._log_rollout_throughput_metrics(emit=False)
         rollout_checkpoint_task = (
             asyncio.create_task(self._rollout_snapshots.run())
-            if self._master_config.rollout_checkpointing.snapshot_attempt_interval_s
-            is not None
+            if rollout_snapshots_enabled(self._master_config)
             else None
         )
         if rollout_checkpoint_task is not None:
@@ -3270,6 +3270,11 @@ async def _train_pump(self) -> None:
             )
 
             if should_save_by_timeout:
+                if self._master_config.checkpointing["enabled"]:
+                    # The trainer checkpoint holds no environment state, and no
+                    # later step will wake the pump to snapshot it.
+                    await self._finish_pending_save(blocking=True)
+                    await self._rollout_snapshots.save_at_deadline()
                 print("Timeout has been reached, stopping training early", flush=True)
                 break
         await self._finish_pending_save(blocking=True)