diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 74f27b8189e..06ad444ebaf 100644 --- a/docs/guides/single-controller.md +++ b/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.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.jsonl.part-` 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 diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index d78c51a4aed..4901bb9be35 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/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: diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 9d91dc47fad..73ed7e44165 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/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: diff --git a/nemo_rl/algorithms/rollout_checkpointing/snapshots.py b/nemo_rl/algorithms/rollout_checkpointing/snapshots.py index 36eca824fb6..cc236891ddb 100644 --- a/nemo_rl/algorithms/rollout_checkpointing/snapshots.py +++ b/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, + ) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 3563bc93126..e792f2648d7 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/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) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index a58cdded3ac..fbcd9edabc9 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -872,8 +872,16 @@ def resolve_for_prompt( class RolloutCheckpointConfig(BaseModel, extra="forbid"): """Frequent rollout-state snapshots anchored to durable trainer state. - ``snapshot_attempt_interval_s=None`` disables saving and restoring periodic - snapshots. A snapshot taken before the first trainer checkpoint is anchored + Rollout snapshots are saved and restored when ``snapshot_attempt_interval_s`` + is set or when ``rollout_recovery.target_level='turn'``, whose environment + state no trainer checkpoint holds. ``snapshot_attempt_interval_s`` attempts + one every that many seconds. With turn-level recovery, a snapshot is also + saved after the trainer checkpoint ``checkpointing.checkpoint_must_save_by`` + triggers, before the run stops, and ``snapshot_on_trainer_checkpoint`` saves + one after every other trainer checkpoint too. Set it to ``false`` with no + interval to save environment state only at that deadline. + + A snapshot taken before the first trainer checkpoint is anchored to the initial model and a rollout-semantic configuration fingerprint. Later snapshots require the durable trainer checkpoint for the controller's current completed step; attempts are skipped until that exact anchor exists. @@ -914,6 +922,7 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): """ snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None + snapshot_on_trainer_checkpoint: bool = True telemetry_interval_s: Annotated[Optional[float], Field(gt=0)] = None max_consecutive_failures: Annotated[int, Field(ge=1)] = 3 keep_latest_k: Annotated[int, Field(ge=1)] = 2 @@ -1069,6 +1078,18 @@ def is_ppo_run(master_config: MasterConfig) -> bool: return getattr(master_config, "ppo", None) is not None +def rollout_snapshots_enabled(master_config: MasterConfig) -> bool: + """Whether the run saves and restores rollout snapshots. + + Periodic snapshots turn them on, and so does turn-level recovery, whose + environment state only a rollout snapshot holds. + """ + return ( + master_config.rollout_checkpointing.snapshot_attempt_interval_s is not None + or master_config.rollout_recovery.turn_checkpointing_enabled + ) + + def algo_config(master_config: MasterConfig) -> GRPOConfig | PPOConfig: """The active algorithm block: ``ppo`` on a PPO run, else ``grpo``. @@ -1610,12 +1631,19 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: ) # A rollout snapshot needs step_N published while training is still on # step N; a deferred finalize publishes it at a later step boundary. - if master_config.rollout_checkpointing.snapshot_attempt_interval_s is not None: + # The deadline snapshot is exempt: the run publishes it first, then stops. + rollout_checkpoint_cfg = master_config.rollout_checkpointing + if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None or ( + master_config.rollout_recovery.turn_checkpointing_enabled + and rollout_checkpoint_cfg.snapshot_on_trainer_checkpoint + ): raise ValueError( "async_rl.defer_async_checkpoint_finalize cannot be combined with " - "rollout_checkpointing.snapshot_attempt_interval_s: rollout snapshots " - "need each step's policy checkpoint published before the next " - "optimizer step, and a deferred finalize publishes it later." + "rollout_checkpointing.snapshot_attempt_interval_s or, with " + "rollout_recovery.target_level='turn', " + "rollout_checkpointing.snapshot_on_trainer_checkpoint: rollout " + "snapshots need each step's policy checkpoint published before the " + "next optimizer step, and a deferred finalize publishes it later." ) reward_penalties_enabled = any( diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 4cbcb62b39b..d9a94420d07 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -83,6 +83,7 @@ is_ppo_run, logprob_passes_required, resolve_ppo_chunk_alignment, + rollout_snapshots_enabled, validate_single_controller_config, ) from nemo_rl.algorithms.utils import set_seed @@ -1200,6 +1201,7 @@ def setup_single_controller( ) data_plane_checkpointing_supported = data_plane_supports_checkpointing(dp_config) rollout_checkpoint_cfg = master_config.rollout_checkpointing + snapshots_enabled = rollout_snapshots_enabled(master_config) if master_config.rollout_recovery.turn_checkpointing_enabled: if generation_config["backend"] != "vllm": raise NotImplementedError( @@ -1212,10 +1214,16 @@ def setup_single_controller( "rollout_recovery.target_level='turn' requires the NeMo-Gym " "rollout path (env.should_use_nemo_gym=true)" ) - if rollout_checkpoint_cfg.snapshot_attempt_interval_s is None: + if ( + rollout_checkpoint_cfg.snapshot_attempt_interval_s is None + and not rollout_checkpoint_cfg.snapshot_on_trainer_checkpoint + and master_config.checkpointing.get("checkpoint_must_save_by") is None + ): raise ValueError( - "rollout_recovery.target_level='turn' requires " - "rollout_checkpointing.snapshot_attempt_interval_s" + "rollout_recovery.target_level='turn' would never save environment " + "state: set rollout_checkpointing.snapshot_on_trainer_checkpoint=true, " + "rollout_checkpointing.snapshot_attempt_interval_s, or " + "checkpointing.checkpoint_must_save_by" ) if not master_config.token_capture.enabled: raise ValueError( @@ -1228,8 +1236,7 @@ def setup_single_controller( "rollout_checkpointing.restore_mode='latest'" ) if ( - master_config.checkpointing.get("save_data_plane") - or rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None + master_config.checkpointing.get("save_data_plane") or snapshots_enabled ) and not data_plane_checkpointing_supported: raise NotImplementedError( "SingleController data-plane checkpointing is not supported for " @@ -1367,7 +1374,7 @@ def setup_single_controller( ) if token_capture_cfg.defer_routed_experts_to_policy: raise ValueError("VLM token capture requires direct router replay assembly") - if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None: + if snapshots_enabled: if not master_config.checkpointing["enabled"]: raise ValueError( "rollout checkpointing requires checkpointing.enabled=true" @@ -1390,7 +1397,10 @@ def setup_single_controller( "rollout checkpointing requires a sampler that explicitly " "supports training-claim ownership" ) - if master_config.checkpointing["save_period"] != 1: + if ( + rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None + and master_config.checkpointing["save_period"] != 1 + ): warnings.warn( "rollout checkpointing is enabled with " f"checkpointing.save_period={master_config.checkpointing['save_period']}; " @@ -1524,10 +1534,7 @@ def setup_single_controller( restore_mode = rollout_checkpoint_cfg.restore_mode recovery_checkpoint_path = trainer_checkpoint_path bootstrap_anchor = checkpointer.checkpoint_dir / BOOTSTRAP_DIRNAME - needs_bootstrap_identity = ( - trainer_checkpoint_path is None - and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None - ) + needs_bootstrap_identity = trainer_checkpoint_path is None and snapshots_enabled bootstrap_identity = ( bootstrap_compatibility_identity(master_config) if needs_bootstrap_identity @@ -1545,7 +1552,7 @@ def setup_single_controller( snapshot_resolution_started = time.monotonic() if ( trainer_checkpoint_path is not None - and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None + and snapshots_enabled and restore_mode == "latest" ): resolved_snapshot = resolve_latest_snapshot( @@ -1555,7 +1562,7 @@ def setup_single_controller( expected_bootstrap_fingerprint=None, ) elif trainer_checkpoint_path is None: - if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None: + if snapshots_enabled: assert bootstrap_digest is not None assert bootstrap_identity is not None if bootstrap_anchor.is_dir(): @@ -1571,11 +1578,7 @@ def setup_single_controller( "restore_mode='latest' to recover it or choose a new " "checkpoint_dir. Existing checkpoint state was not modified." ) - if ( - rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None - and restore_mode == "latest" - and bootstrap_anchor.is_dir() - ): + if snapshots_enabled and restore_mode == "latest" and bootstrap_anchor.is_dir(): resolved_snapshot = resolve_latest_snapshot( bootstrap_anchor, expected_train_step=0, diff --git a/nemo_rl/utils/timer.py b/nemo_rl/utils/timer.py index 4ee295f3574..bafd8cb7e51 100644 --- a/nemo_rl/utils/timer.py +++ b/nemo_rl/utils/timer.py @@ -463,26 +463,24 @@ def __init__( self.iteration_times = [] self.previous_iteration_time: Optional[float] = None self.fit_last_save_time = fit_last_save_time + # Time the save at the deadline needs on top of a step, such as a + # rollout snapshot that follows the trainer checkpoint. + self.save_reserve_seconds = 0.0 + + def seconds_until_save(self) -> Optional[float]: + """Seconds until the deadline is due, or ``None`` once it can no longer be.""" + if self.last_saved or self.last_save_time == float("inf"): + return None + lead_time = self.save_reserve_seconds + if self.fit_last_save_time and self.iteration_times: + lead_time += sum(self.iteration_times) / len(self.iteration_times) + elapsed_time = time.time() - self.start_time + return self.last_save_time - lead_time - elapsed_time def would_save(self) -> bool: """Return whether the deadline is due without consuming the signal.""" - if self.last_saved: - return False - - current_time = time.time() - elapsed_time = current_time - self.start_time - - if self.fit_last_save_time and self.iteration_times: - average_iteration_time = sum(self.iteration_times) / len( - self.iteration_times - ) - if elapsed_time + average_iteration_time >= self.last_save_time: - return True - - if elapsed_time >= self.last_save_time: - return True - - return False + remaining = self.seconds_until_save() + return remaining is not None and remaining <= 0 def check_save(self): # Flush diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 79b2eb00490..021f404e9b6 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -124,6 +124,7 @@ from nemo_rl.experience.rollout_recovery import ( ROLLOUT_RECOVERY_SCHEMA_VERSION, ROLLOUT_RECOVERY_STATE_FILENAME, + RecoveryTargetLevel, RolloutRecoveryLedger, ) from nemo_rl.experience.route_plan import ( @@ -134,6 +135,7 @@ from nemo_rl.telemetry.vocabulary import TELEMETRY_WALL_TIME_METRIC from nemo_rl.utils.checkpoint import CheckpointManager from nemo_rl.utils.logger import LoggerConfig +from nemo_rl.utils.timer import TimeoutChecker # Reuse the factory patches from the setup tests (same cross-module fixture # import pattern as test_rollout_pump.py). @@ -1675,6 +1677,28 @@ def test_timeout_saves_and_stops_training_early(self, tmp_path): assert trainer.save_calls[0]["is_final_checkpoint"] is False assert _step_dir_names(tmp_path / "checkpoints") == {"step_1"} + def test_timeout_snapshots_rollouts_before_stopping(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=4, + save_period=100, + checkpoint_must_save_by="00:00:00:00", + ) + trainer = _FakeTrainer() + events: list[str] = [] + + def _seed(actor) -> None: + async def _save_at_deadline() -> None: + events.append(f"snapshot after step_{actor._train_steps}") + + actor._rollout_snapshots.save_at_deadline = _save_at_deadline + + actor = _run_train_pump(mc, _make_actor_args(trainer=trainer), seed=_seed) + + assert actor._train_steps == 1 + assert len(trainer.save_calls) == 1 + assert events == ["snapshot after step_1"] + def test_rollout_exhaustion_saves_final_checkpoint(self, tmp_path): # A resumed run can exhaust its data before max_num_steps: # _clamp_max_num_steps budgets against the full per-epoch batch count, @@ -1895,6 +1919,25 @@ def test_deferred_finalize_rejects_rollout_snapshots(self, tmp_path): with pytest.raises(ValueError, match="rollout_checkpointing"): validate_single_controller_config(mc) + @pytest.mark.parametrize("on_trainer_checkpoint", [True, False]) + def test_deferred_finalize_allows_only_deadline_turn_snapshots( + self, tmp_path, on_trainer_checkpoint + ): + """Only a snapshot after each trainer checkpoint needs step_N published early.""" + mc = _actor_master_config(tmp_path, defer_finalize=True) + mc.policy["megatron_cfg"] = { + "enabled": True, + "checkpoint": {"async_save": True}, + } + mc.rollout_recovery.target_level = RecoveryTargetLevel.TURN + mc.rollout_checkpointing.snapshot_on_trainer_checkpoint = on_trainer_checkpoint + + if on_trainer_checkpoint: + with pytest.raises(ValueError, match="snapshot_on_trainer_checkpoint"): + validate_single_controller_config(mc) + else: + validate_single_controller_config(mc) + class TestFinalizerStateRelease: def test_evicted_group_releases_pending_finalizer_state(self, tmp_path: Path): @@ -2212,6 +2255,8 @@ async def snapshot(_path): < logged["environment_paused_seconds"] < logged["total_save_seconds"] ) + # The deadline leaves room for a snapshot like this one. + assert actor._timeout.save_reserve_seconds == logged["total_save_seconds"] parked_warnings = [ str(warning.message) for warning in caught @@ -2809,6 +2854,176 @@ async def _save(*, force: bool = False) -> Any: finally: actor._checkpointer.shutdown() + def test_pump_without_an_interval_waits_for_a_trainer_checkpoint( + self, tmp_path: Path + ): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = None + actor._environment_checkpoint = _FakeEnvironmentCheckpoint() + attach_rollout_checkpointing(actor) + actor._train_steps = 1 + + async def _main() -> None: + attempts: list[bool] = [] + attempted = asyncio.Event() + + async def _save(*, force: bool = False) -> RolloutCheckpointSaveResult: + attempts.append(force) + attempted.set() + return RolloutCheckpointSaveResult( + saved=False, reason="no_data_plane_mutations" + ) + + actor._rollout_snapshots.save_snapshot = _save + pump = asyncio.create_task(actor._rollout_snapshots.run()) + try: + await asyncio.sleep(0.05) + assert attempts == [] + actor._rollout_snapshots.trainer_checkpoint_published(1) + await asyncio.wait_for(attempted.wait(), timeout=1.0) + finally: + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + assert attempts == [False] + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + @pytest.mark.parametrize("on_trainer_checkpoint", [True, False]) + def test_trainer_checkpoint_wakes_the_pump_only_when_configured( + self, tmp_path: Path, on_trainer_checkpoint: bool + ): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.snapshot_on_trainer_checkpoint = ( + on_trainer_checkpoint + ) + actor._environment_checkpoint = _FakeEnvironmentCheckpoint() + attach_rollout_checkpointing(actor) + try: + actor._rollout_snapshots.trainer_checkpoint_published(1) + finally: + actor._checkpointer.shutdown() + + assert actor._rollout_snapshots._wakeup.is_set() is on_trainer_checkpoint + + def test_pump_without_an_interval_snapshots_at_a_pre_step_deadline( + self, tmp_path: Path + ): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = None + actor._environment_checkpoint = _FakeEnvironmentCheckpoint() + actor._timeout = TimeoutChecker(timeout="00:00:00:00") + attach_rollout_checkpointing(actor) + forced: list[bool] = [] + + async def _save(*, force: bool = False) -> RolloutCheckpointSaveResult: + forced.append(force) + return RolloutCheckpointSaveResult(saved=True, reason="completed") + + actor._rollout_snapshots.save_snapshot = _save + try: + asyncio.run(asyncio.wait_for(actor._rollout_snapshots.run(), timeout=5.0)) + finally: + actor._checkpointer.shutdown() + + assert forced == [True] + assert actor._rollout_snapshots.stop_requested + + def test_pump_without_an_interval_ignores_a_deadline_wait_after_step_one( + self, tmp_path: Path + ): + """Past step 0 the train pump owns the deadline, so a stale wait saves nothing.""" + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = None + actor._environment_checkpoint = _FakeEnvironmentCheckpoint() + attach_rollout_checkpointing(actor) + attempts: list[bool] = [] + + async def _main() -> None: + async def _save(*, force: bool = False) -> RolloutCheckpointSaveResult: + attempts.append(force) + return RolloutCheckpointSaveResult(saved=True, reason="completed") + + actor._rollout_snapshots.save_snapshot = _save + # Timed for a deadline 1s away, then step 1 finishes first. + with patch.object( + actor._rollout_snapshots, + "_seconds_until_next_attempt", + return_value=0.01, + ): + actor._train_steps = 1 + pump = asyncio.create_task(actor._rollout_snapshots.run()) + await asyncio.sleep(0.1) + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + assert attempts == [] + + @pytest.mark.parametrize( + ("holds_rollout_state", "outcome", "expected"), + [ + (False, "completed", "no attempt"), + (True, "completed", "saved"), + (True, "skipped", "warns"), + (True, OSError("storage unavailable"), "warns"), + (True, RuntimeError("broken checkpoint invariant"), "raises"), + ], + ids=["no-environment-state", "saved", "skipped", "io-error", "invariant"], + ) + def test_deadline_snapshot( + self, tmp_path: Path, holds_rollout_state, outcome, expected + ): + actor = self._actor(tmp_path) + actor._logger = MagicMock() + actor._environment_checkpoint = ( + _FakeEnvironmentCheckpoint() + if holds_rollout_state + else NoEnvironmentCheckpoint() + ) + attach_rollout_checkpointing(actor) + actor._train_steps = 1 + actor._rollout_snapshots._wakeup.set() + forced: list[bool] = [] + + async def _save(*, force: bool = False) -> RolloutCheckpointSaveResult: + forced.append(force) + if isinstance(outcome, Exception): + raise outcome + if outcome == "skipped": + return RolloutCheckpointSaveResult( + saved=False, reason="optimizer_commit_in_progress" + ) + return RolloutCheckpointSaveResult(saved=True, reason="completed") + + actor._rollout_snapshots.save_snapshot = _save + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + if expected == "raises": + with pytest.raises(RuntimeError, match="broken checkpoint"): + asyncio.run(actor._rollout_snapshots.save_at_deadline()) + else: + asyncio.run(actor._rollout_snapshots.save_at_deadline()) + finally: + actor._checkpointer.shutdown() + + deadline_warnings = [ + w for w in caught if "checkpoint deadline" in str(w.message) + ] + if expected == "no attempt": + assert forced == [] + return + assert forced == [True] + assert not actor._rollout_snapshots._wakeup.is_set() + assert bool(deadline_warnings) is (expected == "warns") + def test_periodic_pump_reports_each_consecutive_failure( self, tmp_path: Path, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index b5723c7e463..e9dc96f5de0 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -1224,7 +1224,7 @@ def test_periodic_checkpointing_does_not_enable_gym_participants(self): assert mock_spinup.call_args.kwargs["turn_checkpointing_enabled"] is False - def test_turn_recovery_requires_periodic_snapshots(self): + def test_turn_recovery_rejects_a_config_that_never_saves(self): mc = _make_master_config(env={"should_use_nemo_gym": True}) mc.policy["generation"]["vllm_cfg"] = { "async_engine": True, @@ -1232,10 +1232,13 @@ def test_turn_recovery_requires_periodic_snapshots(self): } mc.token_capture.enabled = True mc.rollout_recovery.target_level = RecoveryTargetLevel.TURN + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_on_trainer_checkpoint=False + ) with pytest.raises( ValueError, - match="target_level='turn' requires.*snapshot_attempt_interval_s", + match="target_level='turn' would never save environment state", ): setup_single_controller(mc, MagicMock(pad_token_id=0)) diff --git a/tests/unit/utils/test_timer.py b/tests/unit/utils/test_timer.py index ac32d16c8fa..22d9945a726 100644 --- a/tests/unit/utils/test_timer.py +++ b/tests/unit/utils/test_timer.py @@ -648,6 +648,20 @@ def test_would_save_does_not_consume_deadline(self): assert checker.check_save() is True assert checker.would_save() is False + @patch("time.time", return_value=110.0) + def test_seconds_until_save_counts_down_to_the_deadline(self, _mock_time): + with patch.dict("os.environ", {"NRL_JOB_START_EPOCH": "100"}, clear=True): + checker = TimeoutChecker(timeout="00:00:00:30") + assert checker.seconds_until_save() == 20.0 + + checker.save_reserve_seconds = 25.0 + assert checker.seconds_until_save() == -5.0 + assert checker.check_save() is True + assert checker.seconds_until_save() is None + + def test_seconds_until_save_without_a_deadline(self): + assert TimeoutChecker(timeout=None).seconds_until_save() is None + def test_fit_last_save_time_enabled(self): # Create a TimeoutChecker with a 3-second timeout and enable fit_last_save_time logic checker = TimeoutChecker(timeout="00:00:00:03", fit_last_save_time=True)