From c8877bc78ad0b29a22babd246187bc925819f879 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 17:03:36 -0400 Subject: [PATCH 01/14] feat(rollout): checkpoint active streamed steps Add periodic TQ, replay, lineage, and dataloader snapshots anchored to durable trainer state. Roll back groups claimed by an unfinished optimizer step into replayable state during persistence, restore them on restart, and cover partial-sibling plus 2/8 and 6/8 streamed recovery paths. (cherry picked from commit 98c196f4b0793541c4bde4e1a01ba0b5c7d21cb3) Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 358 +++++++++++- .../single_controller_utils/__init__.py | 2 + .../single_controller_utils/config.py | 17 + .../rollout_checkpoint.py | 516 ++++++++++++++++++ .../single_controller_utils/setup.py | 158 +++++- nemo_rl/utils/timer.py | 21 +- pyrefly.toml | 1 + .../single_controller/test_checkpointing.py | 69 +++ .../test_rollout_checkpoint.py | 395 ++++++++++++++ tests/unit/utils/test_timer.py | 8 + 10 files changed, 1511 insertions(+), 34 deletions(-) create mode 100644 nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py create mode 100644 tests/unit/single_controller/test_rollout_checkpoint.py diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index af55ab9b493..a1bf4de9fd1 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -43,9 +43,11 @@ import contextlib import hashlib import io +import json import logging import math import os +import shutil import statistics import threading import time @@ -53,6 +55,7 @@ import warnings from collections import deque from collections.abc import Iterator +from dataclasses import dataclass from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Union, cast @@ -95,6 +98,15 @@ validate_sampler_buffer_capacity, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + RolloutSnapshotManifest, + commit_snapshot, + ensure_bootstrap_anchor, + prepare_snapshot_paths, + prune_bootstrap_snapshots, +) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, @@ -150,6 +162,18 @@ log = logging.getLogger(__name__) +@dataclass(frozen=True) +class _RolloutCheckpointCut: + """Controller sidecars captured with one native TQ snapshot.""" + + dataloader_state: dict[str, Any] + replacement_reserve: list[DatumSpec] + replay_metadata: Optional[TQReplayMetadataState] + rollout_recovery_payload: Optional[bytes] + rollout_recovery_group_count: Optional[int] + mutation_version: int + + def _pooled_opd_metrics( stat_sum: float, stat_sumsq: float, count: int ) -> dict[str, float]: @@ -405,6 +429,17 @@ def __init__( self._rollout_manager.set_data_plane_checkpoint_barrier( self._data_plane_checkpoint_barrier ) + # Full trainer checkpoints and lightweight rollout snapshots share one + # namespace and must never publish concurrently. + self._checkpoint_save_lock = asyncio.Lock() + self._last_rollout_snapshot_mutation_version: Optional[int] = None + self._last_missing_rollout_snapshot_anchor: Optional[tuple[int, int]] = None + self._bootstrap_fingerprint = getattr(actor_args, "bootstrap_fingerprint", None) + self._rollout_checkpoint_stop_requested = asyncio.Event() + # Periodic snapshots intentionally skip an optimizer step in progress. + # The transition to/from idle joins the data-plane barrier below. + self._train_step_idle = asyncio.Event() + self._train_step_idle.set() # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() @@ -508,6 +543,16 @@ async def run(self) -> dict[str, Any]: train_task = asyncio.create_task(self._train_pump()) watchdog_task = asyncio.create_task(self._stall_watchdog_pump()) tasks = [rollout_task, train_task, watchdog_task] + rollout_checkpoint_cfg = getattr( + self._master_config, "rollout_checkpointing", None + ) + rollout_checkpoint_task = ( + asyncio.create_task(self._rollout_checkpoint_pump()) + if getattr(rollout_checkpoint_cfg, "interval_s", None) is not None + else None + ) + if rollout_checkpoint_task is not None: + tasks.append(rollout_checkpoint_task) # Only with fleet health on. Created unconditionally it would be a timer firing # every probe_interval_s for every run that does not use the feature, which is # the default. @@ -522,19 +567,31 @@ async def run(self) -> dict[str, Any]: done, _ = await asyncio.wait( set(tasks), return_when=asyncio.FIRST_COMPLETED ) + stop_after_rollout_checkpoint = False + if ( + rollout_checkpoint_task is not None + and rollout_checkpoint_task in done + ): + await rollout_checkpoint_task + if not self._rollout_checkpoint_stop_requested.is_set(): + raise RuntimeError( + "rollout checkpoint pump exited without requesting stop" + ) + stop_after_rollout_checkpoint = True if probe_task is not None and probe_task in done: # Loops forever like the watchdog, so finishing at all means it raised. await probe_task - if watchdog_task in done: + if not stop_after_rollout_checkpoint and watchdog_task in done: # The watchdog loops forever, so finishing at all means it raised -- # a stall or an unhealthy environment. Surface that ahead of the # pumps, whose own symptom would just be "waiting". await watchdog_task - if rollout_task in done: + if not stop_after_rollout_checkpoint and rollout_task in done: # Propagate rollout failures immediately. A normally exhausted # rollout pump leaves the train pump to drain committed groups. await rollout_task - await train_task + if not stop_after_rollout_checkpoint: + await train_task finally: for task in tasks: task.cancel() @@ -2248,11 +2305,18 @@ async def _train_pump(self) -> None: self._async_cfg.min_groups_for_streaming_train, max_prompt_groups, ) - train_meta, num_groups = await self._sampler.select( - current_train_weight=self._trainer_version, - min_prompt_groups=min_prompt_groups, - max_prompt_groups=max_prompt_groups, - ) + async with self._data_plane_checkpoint_barrier.mutation(): + train_meta, num_groups = await self._sampler.select( + current_train_weight=self._trainer_version, + min_prompt_groups=min_prompt_groups, + max_prompt_groups=max_prompt_groups, + ) + if train_meta is not None: + train_step_idle = getattr( + self, "_train_step_idle", None + ) + if train_step_idle is not None: + train_step_idle.clear() # If no batch is selectable, sleep and retry if train_meta is None: @@ -2716,6 +2780,11 @@ async def _train_pump(self) -> None: flush=True, ) + train_step_idle = getattr(self, "_train_step_idle", None) + if train_step_idle is not None: + async with self._data_plane_checkpoint_barrier.mutation(): + train_step_idle.set() + if should_save_by_timeout: print("Timeout has been reached, stopping training early", flush=True) break @@ -3275,11 +3344,281 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: ) return len(stale_tasks) + async def _capture_rollout_checkpoint_cut( + self, + checkpoint_path: PathLike, + ) -> _RolloutCheckpointCut: + """Save TQ and capture matching rollout sidecars under the barrier.""" + if not self._train_step_idle.is_set(): + raise RuntimeError( + "periodic rollout snapshot attempted after training consumed rows" + ) + + save_state = self._save_state + save_state.current_step = self._train_steps + save_state.total_steps = self._train_steps + save_state.trainer_version = self._trainer_version + save_state.current_epoch = self._current_epoch + save_state.consumed_samples = self._consumed_samples + save_state.total_valid_tokens = self._total_valid_tokens + save_state.sampler_name = self._async_cfg.sampler.name + save_state.sampler_dispatch_index = self._sampler.dispatch_index + + dataloader_state = self._dataloader.state_dict() + replacement_reserve = list(self._replacement_reserve) + replay_metadata = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + await self._validate_replay_inventory(replay_metadata) + + recovery_state = self._rollout_manager.recovery_ledger.state_dict() + recovery_state["batch_shortfall"] = self._batch_shortfall.copy() + recovery_state["sampler_stamps_target_steps"] = ( + self._sampler_stamps_target_steps + ) + canonical_group_ids = { + group["group_id"] for group in replay_metadata["groups"] + } + recovery_state["groups"] = [ + group + for group in recovery_state["groups"] + if group["group_id"] not in canonical_group_ids + ] + payload_buffer = io.BytesIO() + await asyncio.to_thread(torch.save, recovery_state, payload_buffer) + recovery_payload = payload_buffer.getvalue() + recovery_digest = hashlib.sha256(recovery_payload).hexdigest() + + if self._master_config.token_capture.enabled: + await self._validate_rollout_recovery_inventory( + replay_metadata=replay_metadata, + clear_unreferenced=False, + ) + await self._save_data_plane_checkpoint( + checkpoint_path, + replay_metadata=replay_metadata, + rollout_recovery_payload_sha256=recovery_digest, + rollout_recovery_group_count=len(recovery_state["groups"]), + ) + return _RolloutCheckpointCut( + dataloader_state=dataloader_state, + replacement_reserve=replacement_reserve, + replay_metadata=replay_metadata, + rollout_recovery_payload=recovery_payload, + rollout_recovery_group_count=len(recovery_state["groups"]), + mutation_version=self._data_plane_checkpoint_barrier.mutation_version, + ) + + async def _write_rollout_checkpoint_sidecars( + self, + checkpoint_path: Path, + cut: _RolloutCheckpointCut, + ) -> None: + """Write metadata-only controller state beside a native TQ snapshot.""" + await asyncio.to_thread( + torch.save, + cut.dataloader_state, + checkpoint_path / "train_dataloader.pt", + ) + if cut.replacement_reserve: + await asyncio.to_thread( + torch.save, + cut.replacement_reserve, + checkpoint_path / "replacement_reserve.pt", + ) + if cut.replay_metadata is not None: + await asyncio.to_thread( + torch.save, + cut.replay_metadata, + checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME, + ) + if cut.rollout_recovery_payload is not None: + await asyncio.to_thread( + (checkpoint_path / ROLLOUT_RECOVERY_STATE_FILENAME).write_bytes, + cut.rollout_recovery_payload, + ) + + def _write_config() -> None: + import yaml + + dumped = self._master_config.model_dump(mode="json") + with (checkpoint_path / "config.yaml").open("w") as config_file: + yaml.safe_dump(dumped, config_file) + + await asyncio.to_thread(_write_config) + + async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: + """Publish one rollout-only snapshot anchored to durable trainer state.""" + # Never hold the filesystem publication lock while waiting for the train + # step. A full checkpoint runs before the step becomes idle and needs the + # same lock, so reversing these waits would deadlock both save paths. + await self._train_step_idle.wait() + async with self._checkpoint_save_lock: + if ( + not force + and self._last_rollout_snapshot_mutation_version + == self._data_plane_checkpoint_barrier.mutation_version + ): + return False + + await asyncio.to_thread(self._checkpointer.finalize_pending) + if self._train_steps == 0: + if self._trainer_version != 0: + raise RuntimeError( + "bootstrap rollout snapshot requires trainer version zero" + ) + if self._bootstrap_fingerprint is None: + raise RuntimeError( + "rollout snapshotting requires a bootstrap fingerprint" + ) + anchor = await asyncio.to_thread( + ensure_bootstrap_anchor, + self._checkpointer.checkpoint_dir, + fingerprint=self._bootstrap_fingerprint, + ) + snapshot_fingerprint = self._bootstrap_fingerprint + else: + if self._trainer_version != self._train_steps: + raise RuntimeError( + "rollout snapshot trainer identity is ambiguous: " + f"step={self._train_steps}, " + f"trainer_version={self._trainer_version}" + ) + anchor = self._checkpointer.checkpoint_dir / f"step_{self._train_steps}" + if not anchor.is_dir(): + skip_key = (self._train_steps, self._trainer_version) + if self._last_missing_rollout_snapshot_anchor != skip_key: + print( + "rollout checkpoint skipped: matching trainer " + f"checkpoint is not durable yet: {anchor}", + flush=True, + ) + self._last_missing_rollout_snapshot_anchor = skip_key + return False + try: + await asyncio.to_thread( + prune_bootstrap_snapshots, + self._checkpointer.checkpoint_dir, + durable_trainer_checkpoint=anchor, + ) + except OSError as error: + warnings.warn( + "Failed to prune obsolete bootstrap rollout snapshots: " + f"{type(error).__name__}: {error}", + stacklevel=2, + ) + snapshot_fingerprint = None + + expected_train_step = self._train_steps + expected_trainer_version = self._trainer_version + tmp_path, final_path, _ = await asyncio.to_thread( + prepare_snapshot_paths, anchor + ) + try: + async with self._data_plane_checkpoint_barrier.checkpoint(): + if ( + not self._train_step_idle.is_set() + or self._train_steps != expected_train_step + or self._trainer_version != expected_trainer_version + ): + await asyncio.to_thread(shutil.rmtree, tmp_path) + return False + snapshot_epoch = self._current_epoch + cut = await self._capture_rollout_checkpoint_cut(tmp_path) + + await self._write_rollout_checkpoint_sidecars(tmp_path, cut) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=expected_train_step, + trainer_version=expected_trainer_version, + current_epoch=snapshot_epoch, + mutation_version=cut.mutation_version, + rolled_back_train_group_count=0, + bootstrap_fingerprint=snapshot_fingerprint, + ) + await asyncio.to_thread( + (tmp_path / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text, + json.dumps(manifest.to_dict(), sort_keys=True, indent=2) + "\n", + ) + await asyncio.to_thread( + commit_snapshot, + tmp_path, + final_path, + keep_latest_k=( + self._master_config.rollout_checkpointing.keep_latest_k + ), + ) + except BaseException: + if tmp_path.exists(): + await asyncio.to_thread(shutil.rmtree, tmp_path) + raise + + self._last_rollout_snapshot_mutation_version = cut.mutation_version + self._last_missing_rollout_snapshot_anchor = None + print( + "rollout checkpoint save completed: " + f"{final_path} (step={expected_train_step}, " + f"trainer_version={expected_trainer_version}, " + f"ledger_groups={cut.rollout_recovery_group_count or 0})", + flush=True, + ) + return True + + async def _rollout_checkpoint_pump(self) -> None: + """Persist rollout state periodically at trainer-safe boundaries.""" + interval_s = self._master_config.rollout_checkpointing.interval_s + if interval_s is None: + raise RuntimeError("rollout checkpoint pump started while disabled") + while True: + await asyncio.sleep(interval_s) + deadline_due = self._train_steps == 0 and self._timeout.would_save() + try: + saved = await self._save_rollout_checkpoint(force=deadline_due) + except Exception as error: + if deadline_due: + raise RuntimeError( + "failed to save the required pre-step rollout checkpoint" + ) from error + warnings.warn( + "Periodic rollout checkpoint failed; retaining the previous " + f"committed snapshot: {type(error).__name__}: {error}", + stacklevel=2, + ) + continue + if deadline_due and saved and self._timeout.check_save(): + print( + "Checkpoint deadline reached before the first train step; " + "stopping after a durable rollout snapshot", + flush=True, + ) + self._rollout_checkpoint_stop_requested.set() + return + async def _save_checkpoint( self, step_metrics: dict[str, Any], *, is_policy_training_step: bool, + ) -> None: + """Serialize full and rollout-only checkpoint publication.""" + lock = getattr(self, "_checkpoint_save_lock", None) + if lock is None: + await self._save_checkpoint_impl( + step_metrics, + is_policy_training_step=is_policy_training_step, + ) + return + async with lock: + await self._save_checkpoint_impl( + step_metrics, + is_policy_training_step=is_policy_training_step, + ) + + async def _save_checkpoint_impl( + self, + step_metrics: dict[str, Any], + *, + is_policy_training_step: bool, ) -> None: """Write a full checkpoint for the just-finished train step. @@ -3400,6 +3739,9 @@ async def _save_checkpoint( else None ), ) + self._last_rollout_snapshot_mutation_version = ( + self._data_plane_checkpoint_barrier.mutation_version + ) # Save value model if self._is_ppo: diff --git a/nemo_rl/algorithms/single_controller_utils/__init__.py b/nemo_rl/algorithms/single_controller_utils/__init__.py index 425c1a9f48f..d2d182cf539 100644 --- a/nemo_rl/algorithms/single_controller_utils/__init__.py +++ b/nemo_rl/algorithms/single_controller_utils/__init__.py @@ -18,6 +18,7 @@ AdvantageConfig, AsyncRLConfig, MasterConfig, + RolloutCheckpointConfig, RolloutFailureConfig, WatchdogConfig, algo_config, @@ -32,6 +33,7 @@ "AdvantageConfig", "AsyncRLConfig", "MasterConfig", + "RolloutCheckpointConfig", "RolloutFailureConfig", "SingleControllerActorArgs", "WatchdogConfig", diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 72419d0e009..628985c4896 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -714,6 +714,20 @@ def resolve_for_prompt( return TaskSourceRecoveryGranularity(task_source, self.default_granularity) +class RolloutCheckpointConfig(BaseModel, extra="forbid"): + """Frequent rollout-state snapshots anchored to durable trainer state. + + ``interval_s=None`` disables the periodic pump. ``latest`` restores the + newest compatible rollout snapshot, ``trainer_checkpoint`` ignores newer + rollout-only snapshots, and ``none`` restores trainer state without any + replay, lineage, or dataloader state. + """ + + interval_s: Optional[float] = Field(default=None, gt=0) + keep_latest_k: int = Field(default=2, ge=1) + restore_mode: Literal["latest", "trainer_checkpoint", "none"] = "latest" + + class MasterConfig(BaseModel, extra="allow"): # algo configs grpo: Optional[GRPOConfig] = None @@ -734,6 +748,9 @@ class MasterConfig(BaseModel, extra="allow"): rollout_recovery: RolloutRecoveryConfig = Field( default_factory=RolloutRecoveryConfig ) + rollout_checkpointing: RolloutCheckpointConfig = Field( + default_factory=RolloutCheckpointConfig + ) on_policy_distillation: Optional[OnPolicyDistillationConfig] = None token_capture: TokenCaptureConfig = Field(default_factory=TokenCaptureConfig) diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py new file mode 100644 index 00000000000..4f2de3f08df --- /dev/null +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -0,0 +1,516 @@ +# 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. + +"""Filesystem contract for frequent Single Controller rollout snapshots.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Mapping, Optional + +from nemo_rl.algorithms.single_controller_utils.config import MasterConfig + +ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 1 +BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 2 +BOOTSTRAP_DIRNAME = "bootstrap" +BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" +ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" +ROLLOUT_SNAPSHOT_MANIFEST_FILENAME = "manifest.json" +ROLLOUT_SNAPSHOT_COMMITTED_FILENAME = "COMMITTED" +ROLLOUT_SNAPSHOT_LATEST_FILENAME = "LATEST" + +_SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") + +# Keep this projection limited to values needed to interpret persisted rollout +# state or execute missing siblings. Trainer-only and post-rollout settings do +# not belong in bootstrap compatibility. +_BOOTSTRAP_POLICY_FIELDS = frozenset( + { + "model_name", + "pretrained_checkpoint", + "hf_config_overrides", + "max_total_sequence_length", + "tokenizer", + } +) +_BOOTSTRAP_GENERATION_FIELDS = frozenset( + { + "backend", + "max_new_tokens", + "stop_strings", + "stop_token_ids", + "temperature", + "top_k", + "top_p", + } +) +_BOOTSTRAP_VLLM_FIELDS = frozenset( + { + "http_server_serving_chat_kwargs", + "max_model_len", + "reasoning_parser_plugin", + } +) +_BOOTSTRAP_GRPO_FIELDS = frozenset( + { + "max_rollout_turns", + "num_generations_per_prompt", + "num_prompts_per_step", + "seed", + } +) +_BOOTSTRAP_DATA_FIELDS = frozenset( + { + "default", + "max_input_seq_length", + "shuffle", + "train", + } +) +_BOOTSTRAP_TOKEN_CAPTURE_FIELDS = frozenset( + { + "enabled", + "min_valid_fraction_per_group", + "mixed_weight_version_policy", + "staging_partition", + } +) +_BOOTSTRAP_ENV_RUNTIME_FIELDS = frozenset( + { + "apptainer_memory_limit_mb", + "concurrency", + "nemo_gym_log_dir", + "num_gpu_nodes", + "port_range_high", + "port_range_low", + "should_log_nemo_gym_responses", + "skip_venv_if_present", + "use_absolute_ip", + } +) + + +def _select_fields( + mapping: Mapping[str, Any] | None, + fields: frozenset[str], +) -> dict[str, Any]: + """Select explicitly rollout-semantic fields from one config section.""" + if mapping is None: + return {} + return {key: mapping[key] for key in sorted(fields) if key in mapping} + + +def _drop_runtime_fields(value: Any, runtime_fields: frozenset[str]) -> Any: + """Recursively strip known operational leaves from an environment.""" + if isinstance(value, Mapping): + return { + key: _drop_runtime_fields(child, runtime_fields) + for key, child in value.items() + if key not in runtime_fields + } + if isinstance(value, list): + return [_drop_runtime_fields(child, runtime_fields) for child in value] + return value + + +def _fsync_file(path: Path) -> None: + """Flush one completed regular file to its backing filesystem.""" + with path.open("rb") as file_obj: + os.fsync(file_obj.fileno()) + + +def _fsync_directory(path: Path) -> None: + """Flush directory-entry updates such as rename and replace.""" + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + directory_fd = os.open(path, flags) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _fsync_tree(root: Path) -> None: + """Flush every snapshot payload before publishing its commit marker.""" + for directory, _, filenames in os.walk(root, topdown=False): + directory_path = Path(directory) + for filename in filenames: + file_path = directory_path / filename + if not file_path.is_symlink() and file_path.is_file(): + _fsync_file(file_path) + _fsync_directory(directory_path) + + +def _snapshot_sequence(path: Path) -> int: + match = _SNAPSHOT_RE.fullmatch(path.name) + if match is None: + raise ValueError(f"not a rollout snapshot directory: {path}") + return int(match.group(1)) + + +@dataclass(frozen=True) +class RolloutSnapshotManifest: + """Identity binding one rollout-state cut to reconstructable trainer state.""" + + schema_version: int + base_train_step: int + trainer_version: int + current_epoch: int + mutation_version: int + rolled_back_train_group_count: int + bootstrap_fingerprint: Optional[str] + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> RolloutSnapshotManifest: + """Parse and validate a committed snapshot manifest.""" + required_ints = ( + "schema_version", + "base_train_step", + "trainer_version", + "current_epoch", + "mutation_version", + "rolled_back_train_group_count", + ) + for key in required_ints: + value = raw.get(key) + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError( + f"rollout snapshot manifest {key!r} must be an integer" + ) + fingerprint = raw.get("bootstrap_fingerprint") + if fingerprint is not None and not isinstance(fingerprint, str): + raise ValueError( + "rollout snapshot bootstrap_fingerprint must be a string or null" + ) + manifest = cls( + schema_version=raw["schema_version"], + base_train_step=raw["base_train_step"], + trainer_version=raw["trainer_version"], + current_epoch=raw["current_epoch"], + mutation_version=raw["mutation_version"], + rolled_back_train_group_count=raw["rolled_back_train_group_count"], + bootstrap_fingerprint=fingerprint, + ) + if manifest.schema_version != ROLLOUT_SNAPSHOT_SCHEMA_VERSION: + raise ValueError( + "unsupported rollout snapshot schema version: " + f"{manifest.schema_version}" + ) + if ( + min( + manifest.base_train_step, + manifest.trainer_version, + manifest.current_epoch, + manifest.mutation_version, + manifest.rolled_back_train_group_count, + ) + < 0 + ): + raise ValueError("rollout snapshot counters must be non-negative") + return manifest + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ResolvedRolloutCheckpoint: + """A committed rollout snapshot selected for startup recovery.""" + + path: Path + manifest: RolloutSnapshotManifest + + +@dataclass(frozen=True) +class BootstrapCompatibilityIdentity: + """Rollout-semantic inputs that must match a trainer-version-zero cut.""" + + schema_version: int + model: Mapping[str, Any] + generation: Mapping[str, Any] + rollout: Mapping[str, Any] + dataset: Mapping[str, Any] + environment: Mapping[str, Any] + sampler: Mapping[str, Any] + token_capture: Mapping[str, Any] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def bootstrap_compatibility_identity( + master_config: MasterConfig, +) -> BootstrapCompatibilityIdentity: + """Project a full run config onto inputs that affect recovered rollouts. + + Dataset identity is intentionally retained because a bootstrap snapshot + restores the dataloader cursor together with its unfinished prompt ledger. + Cluster shape, logging, checkpoint paths, worker counts, and other runtime + tuning are excluded so they may change across a restart. + """ + dumped = master_config.model_dump(mode="json") + policy = dumped.get("policy", {}) + generation = policy.get("generation", {}) + generation_identity = _select_fields( + generation, + _BOOTSTRAP_GENERATION_FIELDS, + ) + vllm_identity = _select_fields( + generation.get("vllm_cfg", {}), + _BOOTSTRAP_VLLM_FIELDS, + ) + if vllm_identity: + generation_identity["vllm_cfg"] = vllm_identity + + async_rl = dumped.get("async_rl", {}) + sampler = async_rl.get("sampler", {}) + if not isinstance(sampler, Mapping): + sampler = {} + + rollout = dumped.get("ppo") or dumped.get("grpo") or {} + return BootstrapCompatibilityIdentity( + schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + model=_select_fields(policy, _BOOTSTRAP_POLICY_FIELDS), + generation=generation_identity, + rollout=_select_fields(rollout, _BOOTSTRAP_GRPO_FIELDS), + dataset=_select_fields(dumped.get("data", {}), _BOOTSTRAP_DATA_FIELDS), + environment=_drop_runtime_fields( + dumped.get("env", {}), + _BOOTSTRAP_ENV_RUNTIME_FIELDS, + ), + sampler=dict(sampler), + token_capture=_select_fields( + dumped.get("token_capture", {}), + _BOOTSTRAP_TOKEN_CAPTURE_FIELDS, + ), + ) + + +def bootstrap_fingerprint(master_config: MasterConfig) -> str: + """Hash rollout-semantic inputs needed to reuse a bootstrap snapshot. + + This is a compatibility guard, not a hash of the full training recipe. + Operational settings are deliberately excluded so a restart may use a + different cluster shape, checkpoint interval, or logging destination. + """ + payload = json.dumps( + bootstrap_compatibility_identity(master_config).to_dict(), + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def prune_bootstrap_snapshots( + checkpoint_dir: Path, + *, + durable_trainer_checkpoint: Path, +) -> bool: + """Remove trainer-version-zero snapshots once a trainer checkpoint exists.""" + if not durable_trainer_checkpoint.is_dir(): + raise FileNotFoundError( + "cannot prune bootstrap snapshots without a durable trainer " + f"checkpoint: {durable_trainer_checkpoint}" + ) + snapshot_root = checkpoint_dir / BOOTSTRAP_DIRNAME / ROLLOUT_SNAPSHOTS_DIRNAME + if not snapshot_root.is_dir(): + return False + shutil.rmtree(snapshot_root) + return True + + +def ensure_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: + """Create or validate the lightweight trainer-version-zero anchor.""" + anchor = checkpoint_dir / BOOTSTRAP_DIRNAME + anchor.mkdir(parents=True, exist_ok=True) + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + expected = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "bootstrap_fingerprint": fingerprint, + } + if manifest_path.is_file(): + raw = json.loads(manifest_path.read_text()) + if raw != expected: + raise ValueError( + "existing rollout bootstrap anchor does not match the current " + f"trainer configuration: checkpoint={raw!r}, expected={expected!r}" + ) + return anchor + + tmp_path = manifest_path.with_suffix(".json.tmp") + tmp_path.write_text(json.dumps(expected, sort_keys=True, indent=2) + "\n") + _fsync_file(tmp_path) + os.replace(tmp_path, manifest_path) + _fsync_directory(anchor) + _fsync_directory(anchor.parent) + return anchor + + +def reset_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: + """Discard skipped pre-step snapshots and start a new bootstrap lineage. + + This is used only when restore mode deliberately skips partial-rollout + recovery and no trainer checkpoint exists. The user's restore choice makes + the previous state intentionally unreachable; removing it also prevents a + later periodic save from appending to an incompatible bootstrap anchor. + """ + anchor = checkpoint_dir / BOOTSTRAP_DIRNAME + snapshot_root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + if snapshot_root.exists(): + shutil.rmtree(snapshot_root) + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + if manifest_path.exists(): + manifest_path.unlink() + return ensure_bootstrap_anchor(checkpoint_dir, fingerprint=fingerprint) + + +def validate_bootstrap_anchor(anchor: Path, *, fingerprint: str) -> None: + """Fail loudly when bootstrap snapshots belong to different initial state.""" + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + if not manifest_path.is_file(): + raise FileNotFoundError( + f"rollout bootstrap manifest is missing at {manifest_path}" + ) + raw = json.loads(manifest_path.read_text()) + expected = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "bootstrap_fingerprint": fingerprint, + } + if raw != expected: + raise ValueError( + "rollout bootstrap anchor does not match the current trainer " + f"configuration: checkpoint={raw!r}, expected={expected!r}" + ) + + +def prepare_snapshot_paths(anchor: Path) -> tuple[Path, Path, int]: + """Allocate the next temporary/final snapshot directory pair.""" + root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + root.mkdir(parents=True, exist_ok=True) + sequences = [ + int(match.group(1)) + for child in root.iterdir() + if child.is_dir() and (match := _SNAPSHOT_RE.fullmatch(child.name)) + ] + sequence = max(sequences, default=0) + 1 + final_path = root / f"snapshot_{sequence:06d}" + tmp_path = root / f"tmp_snapshot_{sequence:06d}" + if tmp_path.exists(): + shutil.rmtree(tmp_path) + tmp_path.mkdir(parents=True) + return tmp_path, final_path, sequence + + +def commit_snapshot( + tmp_path: Path, + final_path: Path, + *, + keep_latest_k: int, +) -> None: + """Atomically publish one validated snapshot and retain recent fallbacks.""" + if keep_latest_k < 1: + raise ValueError("rollout snapshot retention must keep at least one snapshot") + _fsync_tree(tmp_path) + committed_path = tmp_path / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME + committed_path.write_text("committed\n") + _fsync_file(committed_path) + _fsync_directory(tmp_path) + os.rename(tmp_path, final_path) + + root = final_path.parent + _fsync_directory(root) + latest_path = root / ROLLOUT_SNAPSHOT_LATEST_FILENAME + latest_tmp = latest_path.with_suffix(".tmp") + latest_tmp.write_text(final_path.name + "\n") + _fsync_file(latest_tmp) + os.replace(latest_tmp, latest_path) + _fsync_directory(root) + + committed = sorted( + ( + child + for child in root.iterdir() + if child.is_dir() + and _SNAPSHOT_RE.fullmatch(child.name) + and (child / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() + ), + key=_snapshot_sequence, + reverse=True, + ) + stale_snapshots = committed[keep_latest_k:] + for stale in stale_snapshots: + shutil.rmtree(stale) + if stale_snapshots: + _fsync_directory(root) + + +def resolve_latest_snapshot( + anchor: Path, + *, + expected_train_step: int, + expected_trainer_version: int, + expected_bootstrap_fingerprint: Optional[str], +) -> Optional[ResolvedRolloutCheckpoint]: + """Select the newest complete snapshot compatible with its trainer anchor.""" + root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + if not root.is_dir(): + return None + + candidates = sorted( + ( + child + for child in root.iterdir() + if child.is_dir() and _SNAPSHOT_RE.fullmatch(child.name) + ), + key=_snapshot_sequence, + reverse=True, + ) + errors: list[str] = [] + for candidate in candidates: + if not (candidate / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file(): + continue + manifest_path = candidate / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME + if not manifest_path.is_file(): + errors.append(f"{candidate.name}: missing manifest") + continue + try: + raw = json.loads(manifest_path.read_text()) + manifest = RolloutSnapshotManifest.from_mapping(raw) + except (json.JSONDecodeError, OSError, ValueError) as error: + errors.append(f"{candidate.name}: {error}") + continue + if ( + manifest.base_train_step != expected_train_step + or manifest.trainer_version != expected_trainer_version + or manifest.bootstrap_fingerprint != expected_bootstrap_fingerprint + ): + errors.append(f"{candidate.name}: trainer-anchor mismatch") + continue + return ResolvedRolloutCheckpoint(candidate, manifest) + + if errors: + raise ValueError( + "no committed rollout snapshot matches the selected trainer anchor: " + + "; ".join(errors) + ) + return None diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 41f53c89d3d..47daf3d7f66 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -68,6 +68,15 @@ is_ppo_run, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_DIRNAME, + ROLLOUT_SNAPSHOTS_DIRNAME, + bootstrap_fingerprint, + ensure_bootstrap_anchor, + reset_bootstrap_anchor, + resolve_latest_snapshot, + validate_bootstrap_anchor, +) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS @@ -155,8 +164,9 @@ class SingleControllerActorArgs: finalizer_actors: list[Any] # Defaulted fields must follow the required ones above, so these stay last. data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None - # None when async_rl.generation_fleet_health is disabled; the SingleController drives the - # probe loop when it is present. + bootstrap_fingerprint: Optional[str] = None + # None when async_rl.generation_fleet_health is disabled; the SingleController + # drives the probe loop when it is present. fleet_monitor: Optional[GenerationFleetHealth] = None # None unless async_rl.generation_router is enabled. generation_router: Optional[ray.actor.ActorHandle[GenerationRouterImpl]] = None @@ -941,10 +951,11 @@ def setup_single_controller( "SingleController path is built on the TransferQueue data plane." ) data_plane_checkpointing_supported = data_plane_supports_checkpointing(dp_config) + rollout_checkpoint_cfg = master_config.rollout_checkpointing if ( master_config.checkpointing.get("save_data_plane") - and not data_plane_checkpointing_supported - ): + or rollout_checkpoint_cfg.interval_s is not None + ) and not data_plane_checkpointing_supported: raise NotImplementedError( "SingleController data-plane checkpointing is not supported for " f"data_plane.backend={dp_config['backend']!r}." @@ -953,8 +964,10 @@ def setup_single_controller( sampler_supports_replay_recovery = sampler_supports_buffer_checkpoint( master_config.async_rl.sampler ) - if sampler_supports_replay_recovery and not master_config.checkpointing.get( - "save_data_plane" + if ( + sampler_supports_replay_recovery + and rollout_checkpoint_cfg.restore_mode != "none" + and not master_config.checkpointing.get("save_data_plane") ): error_message = ( "SingleController checkpointing with a replay-checkpoint-capable " @@ -969,7 +982,10 @@ def setup_single_controller( "checkpointing.enabled=false." ) raise ValueError(error_message) - if not sampler_supports_replay_recovery: + if ( + rollout_checkpoint_cfg.restore_mode != "none" + and not sampler_supports_replay_recovery + ): warnings.warn( f"Sampler {master_config.async_rl.sampler.name!r} cannot recover " "completed buffered rollouts. On resume, the dataloader cursor " @@ -1003,6 +1019,24 @@ def setup_single_controller( # ray_actor_environment_registry.py), so nothing here needs to change the # worker's environment. token_capture_cfg = master_config.token_capture + if rollout_checkpoint_cfg.interval_s is not None: + if not master_config.checkpointing["enabled"]: + raise ValueError( + "rollout checkpointing requires checkpointing.enabled=true" + ) + if not master_config.checkpointing.get("save_data_plane"): + raise ValueError( + "rollout checkpointing requires checkpointing.save_data_plane=true" + ) + if not token_capture_cfg.enabled: + raise ValueError( + "rollout checkpointing currently requires token_capture.enabled=true" + ) + if not sampler_supports_buffer_checkpoint(master_config.async_rl.sampler): + raise ValueError( + "rollout checkpointing requires a sampler that supports " + "replay-buffer recovery" + ) if token_capture_cfg.enabled: if not should_use_nemo_gym(master_config): raise ValueError( @@ -1045,24 +1079,108 @@ def setup_single_controller( # Checkpointing # ========================== checkpointer = CheckpointManager(master_config.checkpointing) - last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + trainer_checkpoint_path = checkpointer.get_latest_checkpoint_path() loaded_state = cast( - Optional[dict[str, Any]], checkpointer.load_training_info(last_checkpoint_path) + Optional[dict[str, Any]], + checkpointer.load_training_info(trainer_checkpoint_path), ) save_state = _get_grpo_save_state(loaded_state) - weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path) + weights_path, optimizer_path = checkpointer.get_resume_paths( + trainer_checkpoint_path + ) if is_ppo_run(master_config): # Only a fresh run reads this; a resume ignores it and restores the critic # from its own checkpoint, so the key can stay in the config. warm_start = master_config.ppo.warm_start_value_checkpoint - if last_checkpoint_path is None and warm_start is not None: + if trainer_checkpoint_path is None and warm_start is not None: validate_warm_start_checkpoint(warm_start) print(f"🔥 Warm-starting the value model from {warm_start}") value_weights_path, value_optimizer_path = checkpointer.get_resume_paths( - last_checkpoint_path or warm_start, + trainer_checkpoint_path or warm_start, model_component="value", ) + restore_mode = rollout_checkpoint_cfg.restore_mode + recovery_checkpoint_path = ( + trainer_checkpoint_path if restore_mode != "none" else None + ) + bootstrap_anchor = checkpointer.checkpoint_dir / BOOTSTRAP_DIRNAME + needs_bootstrap_identity = trainer_checkpoint_path is None and ( + rollout_checkpoint_cfg.interval_s is not None + or (restore_mode == "latest" and bootstrap_anchor.is_dir()) + ) + bootstrap_digest = ( + bootstrap_fingerprint(master_config) if needs_bootstrap_identity else None + ) + resolved_snapshot = None + restored_trainer_version = ( + save_state.trainer_version + if save_state.trainer_version is not None + else save_state.current_step + ) + if trainer_checkpoint_path is not None and restore_mode == "latest": + resolved_snapshot = resolve_latest_snapshot( + Path(trainer_checkpoint_path), + expected_train_step=save_state.current_step, + expected_trainer_version=restored_trainer_version, + expected_bootstrap_fingerprint=None, + ) + elif trainer_checkpoint_path is None: + if rollout_checkpoint_cfg.interval_s is not None: + assert bootstrap_digest is not None + if restore_mode == "latest": + bootstrap_anchor = ensure_bootstrap_anchor( + checkpointer.checkpoint_dir, + fingerprint=bootstrap_digest, + ) + else: + had_bootstrap_snapshots = ( + bootstrap_anchor / ROLLOUT_SNAPSHOTS_DIRNAME + ).is_dir() + bootstrap_anchor = reset_bootstrap_anchor( + checkpointer.checkpoint_dir, + fingerprint=bootstrap_digest, + ) + if had_bootstrap_snapshots: + print( + "📦 Ignored existing bootstrap rollout snapshots and " + "started a new bootstrap lineage because " + f"rollout_checkpointing.restore_mode={restore_mode!r}.", + flush=True, + ) + elif restore_mode == "latest" and bootstrap_anchor.is_dir(): + assert bootstrap_digest is not None + validate_bootstrap_anchor( + bootstrap_anchor, + fingerprint=bootstrap_digest, + ) + if restore_mode == "latest" and bootstrap_anchor.is_dir(): + resolved_snapshot = resolve_latest_snapshot( + bootstrap_anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint=bootstrap_digest, + ) + if resolved_snapshot is not None: + recovery_checkpoint_path = str(resolved_snapshot.path) + save_state.current_epoch = resolved_snapshot.manifest.current_epoch + print( + f"📦 Selected rollout recovery snapshot: {recovery_checkpoint_path}", + flush=True, + ) + elif restore_mode == "trainer_checkpoint" and trainer_checkpoint_path: + print( + "📦 Restoring rollout state from the durable trainer checkpoint " + f"without considering newer periodic snapshots: {trainer_checkpoint_path}", + flush=True, + ) + elif restore_mode == "none" and trainer_checkpoint_path: + print( + "📦 Resuming trainer state without rollout, replay, lineage, or " + "dataloader recovery.", + flush=True, + ) + # ========================== # Setup Dataset & Environments # ========================== @@ -1103,9 +1221,12 @@ def setup_single_controller( drop_last=True, num_workers=data_config["num_workers"], ) - if last_checkpoint_path is not None: - print(f"📦 Restoring dataloader state from checkpoint: {last_checkpoint_path}") - load_dataloader_state(dataloader, last_checkpoint_path, data_config) + if recovery_checkpoint_path is not None: + print( + "📦 Restoring dataloader state from checkpoint: " + f"{recovery_checkpoint_path}" + ) + load_dataloader_state(dataloader, recovery_checkpoint_path, data_config) _clamp_max_num_steps(master_config, dataloader) _maybe_inject_megatron_train_iters(master_config) @@ -1368,7 +1489,7 @@ def _build_generation_then_trainer( # operation starts. data_plane_checkpoint_metadata = _maybe_restore_native_data_plane_checkpoint( trainer, - last_checkpoint_path=last_checkpoint_path, + last_checkpoint_path=recovery_checkpoint_path, save_state=save_state, partition_id=partition_id, sampler_name=master_config.async_rl.sampler.name, @@ -1605,9 +1726,10 @@ def _build_generation_then_trainer( tq_buffer=tq_buffer, partition_id=partition_id, save_state=save_state, - last_checkpoint_path=last_checkpoint_path, - finalizer_actors=finalizer_actors, + last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + bootstrap_fingerprint=bootstrap_digest, + finalizer_actors=finalizer_actors, fleet_monitor=fleet_monitor, generation_router=generation_router, teacher_worker_groups=teacher_worker_groups, diff --git a/nemo_rl/utils/timer.py b/nemo_rl/utils/timer.py index 8413c5db079..69a5d79989c 100644 --- a/nemo_rl/utils/timer.py +++ b/nemo_rl/utils/timer.py @@ -436,12 +436,8 @@ def __init__( self.previous_iteration_time: Optional[float] = None self.fit_last_save_time = fit_last_save_time - def check_save(self): - # Flush - sys.stdout.flush() - sys.stderr.flush() - - # Already saved after timeout + def would_save(self) -> bool: + """Return whether the deadline is due without consuming the signal.""" if self.last_saved: return False @@ -453,15 +449,24 @@ def check_save(self): self.iteration_times ) if elapsed_time + average_iteration_time >= self.last_save_time: - self.last_saved = True return True if elapsed_time >= self.last_save_time: - self.last_saved = True return True return False + def check_save(self): + # Flush + sys.stdout.flush() + sys.stderr.flush() + + if not self.would_save(): + return False + + self.last_saved = True + return True + def start_iterations(self): self.previous_iteration_time = time.time() diff --git a/pyrefly.toml b/pyrefly.toml index c9cdedd5e46..c442f2c44df 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -66,6 +66,7 @@ project-includes = [ "nemo_rl/algorithms/single_controller.py", "nemo_rl/algorithms/single_controller_utils/__init__.py", "nemo_rl/algorithms/single_controller_utils/config.py", + "nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py", "nemo_rl/algorithms/single_controller_utils/setup.py", "nemo_rl/algorithms/single_controller_utils/utils.py", "nemo_rl/algorithms/utils.py", diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index a813610185f..c354343c9a1 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -75,8 +75,14 @@ from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, MasterConfig, + RolloutCheckpointConfig, setup_single_controller, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_DIRNAME, + ROLLOUT_SNAPSHOT_COMMITTED_FILENAME, + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, +) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.data.utils import load_dataloader_state from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta @@ -525,6 +531,7 @@ def _actor_master_config( max_num_epochs: int = 1, buffer_checkpoint: bool = False, data_plane_checkpoint: bool = True, + rollout_checkpoint_interval_s: Optional[float] = None, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -584,6 +591,9 @@ def _actor_master_config( max_inflight_prompts=4, max_buffered_rollouts=4, ), + rollout_checkpointing=RolloutCheckpointConfig( + interval_s=rollout_checkpoint_interval_s + ), ) @@ -596,6 +606,7 @@ def _make_actor_args( dp_client: Optional[_FakeDPClient] = None, last_checkpoint_path: Optional[str] = None, data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None, + bootstrap_fingerprint: Optional[str] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=_FakeGeneration(), @@ -617,6 +628,7 @@ def _make_actor_args( last_checkpoint_path=last_checkpoint_path, finalizer_actors=[], data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + bootstrap_fingerprint=bootstrap_fingerprint, ) @@ -1048,6 +1060,63 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} +class TestPeriodicRolloutCheckpoint: + def _actor(self, tmp_path: Path): + config = _actor_master_config( + tmp_path, + buffer_checkpoint=True, + rollout_checkpoint_interval_s=120.0, + ) + return _ACTOR_CLS( + config, + _make_actor_args(bootstrap_fingerprint="bootstrap-digest"), + SetupTimingMetrics(), + ) + + def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): + actor = self._actor(tmp_path) + try: + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + finally: + actor._checkpointer.shutdown() + + snapshot = ( + tmp_path + / "checkpoints" + / BOOTSTRAP_DIRNAME + / "rollout_snapshots" + / "snapshot_000001" + ) + assert (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() + assert (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).is_file() + assert (snapshot / "data_plane" / "metadata.json").is_file() + assert (snapshot / "train_dataloader.pt").is_file() + assert (snapshot / REPLAY_BUFFER_METADATA_FILENAME).is_file() + assert (snapshot / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() + assert not (snapshot / "policy").exists() + + def test_snapshot_waits_until_no_training_rows_are_owned(self, tmp_path: Path): + actor = self._actor(tmp_path) + + async def exercise() -> None: + actor._train_step_idle.clear() + save_task = asyncio.create_task( + actor._save_rollout_checkpoint(force=True) + ) + await asyncio.sleep(0) + assert not save_task.done() + assert actor._dp_client.save_calls == [] + + async with actor._data_plane_checkpoint_barrier.mutation(): + actor._train_step_idle.set() + assert await asyncio.wait_for(save_task, timeout=5.0) + + try: + asyncio.run(exercise()) + finally: + actor._checkpointer.shutdown() + + class TestDataPlaneCheckpoint: def test_metadata_uses_pre_await_save_state_snapshot(self, tmp_path): mc = _actor_master_config( diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py new file mode 100644 index 00000000000..1226071e741 --- /dev/null +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -0,0 +1,395 @@ +# 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. + +import json +from pathlib import Path +from typing import Any, cast +from unittest.mock import Mock, call + +import pytest + +from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + RolloutSnapshotManifest, + bootstrap_fingerprint, + commit_snapshot, + ensure_bootstrap_anchor, + prepare_snapshot_paths, + prune_bootstrap_snapshots, + reset_bootstrap_anchor, + resolve_latest_snapshot, + validate_bootstrap_anchor, +) + + +class _DumpedConfig: + def __init__(self, dumped: dict[str, Any]): + self._dumped = dumped + + def model_dump(self, *, mode: str) -> dict[str, Any]: + assert mode == "json" + return self._dumped + + +def _commit_snapshot( + anchor, + *, + mutation_version: int, + trainer_version: int = 0, + fingerprint: str | None = "fingerprint-v1", +): + tmp_path, final_path, _ = prepare_snapshot_paths(anchor) + manifest = RolloutSnapshotManifest( + schema_version=1, + base_train_step=trainer_version, + trainer_version=trainer_version, + current_epoch=2, + mutation_version=mutation_version, + rolled_back_train_group_count=0, + bootstrap_fingerprint=fingerprint, + ) + (tmp_path / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + commit_snapshot(tmp_path, final_path, keep_latest_k=3) + return final_path + + +def test_bootstrap_anchor_rejects_different_initial_state(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v1") + + with pytest.raises(ValueError, match="does not match"): + validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v2") + + +def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: + base = { + "policy": { + "model_name": "model-a", + "optimizer": {"lr": 1.0e-6}, + "generation": { + "backend": "vllm", + "temperature": 1.0, + "colocated": {"enabled": True, "resources": {"gpus": 8}}, + "vllm_cfg": { + "kv_cache_dtype": "auto", + "precision": "bfloat16", + "skip_tokenizer_init": False, + }, + }, + }, + "data": { + "train": [{"data_path": "/datasets/train.jsonl"}], + "num_workers": 4, + }, + "grpo": { + "num_generations_per_prompt": 4, + "max_num_steps": 10, + "batch_multiplier": 1, + "use_dynamic_sampling": False, + "reward_shaping": {"enabled": False}, + "reward_scaling": {"enabled": False}, + }, + "loss_fn": { + "reference_policy_kl_penalty": 0.01, + "use_kl_in_reward": False, + }, + "reward_penalties": {"penalize_unwanted_tokens": False}, + "token_capture": { + "enabled": True, + "staging_partition": "rollout_staging", + "on_capture_failure": "continue", + }, + "cluster": {"num_nodes": 2}, + "logger": {"log_dir": "/run/one"}, + } + compatible_changed = { + **base, + "policy": { + **base["policy"], + "optimizer": {"lr": 5.0e-7}, + "generation": { + **base["policy"]["generation"], + "colocated": {"enabled": False, "resources": {"gpus": 16}}, + "vllm_cfg": { + "kv_cache_dtype": "fp8", + "precision": "float16", + "skip_tokenizer_init": True, + }, + }, + }, + "data": {**base["data"], "num_workers": 16}, + "grpo": { + **base["grpo"], + "max_num_steps": 100, + "batch_multiplier": 2, + "use_dynamic_sampling": True, + "reward_shaping": {"enabled": True}, + "reward_scaling": {"enabled": True}, + }, + "loss_fn": { + "reference_policy_kl_penalty": 0.1, + "use_kl_in_reward": True, + }, + "reward_penalties": {"penalize_unwanted_tokens": True}, + "token_capture": { + **base["token_capture"], + "on_capture_failure": "abort", + }, + "cluster": {"num_nodes": 8}, + "logger": {"log_dir": "/run/two"}, + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(compatible_changed)) + ) + + +@pytest.mark.parametrize( + ("section", "changed"), + [ + ("policy", {"model_name": "model-b"}), + ("data", {"train": [{"data_path": "/datasets/other.jsonl"}]}), + ("grpo", {"num_generations_per_prompt": 8}), + ("token_capture", {"mixed_weight_version_policy": "reject"}), + ( + "async_rl", + {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, + ), + ], +) +def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( + section: str, + changed: dict[str, Any], +) -> None: + base = { + "policy": { + "model_name": "model-a", + "tokenizer": {"name": "tokenizer-a"}, + "generation": {"backend": "vllm", "temperature": 1.0}, + }, + "data": {"train": [{"data_path": "/datasets/train.jsonl"}]}, + "grpo": {"num_generations_per_prompt": 4}, + "token_capture": { + "enabled": True, + "mixed_weight_version_policy": "allow", + }, + "async_rl": {"sampler": {"name": "windowed", "max_staleness_versions": 1}}, + } + modified = {**base, section: {**base[section], **changed}} + + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(modified))) + ) + + +def test_bootstrap_fingerprint_rejects_generation_semantic_changes() -> None: + base = { + "policy": { + "model_name": "model-a", + "generation": { + "backend": "vllm", + "temperature": 1.0, + "vllm_cfg": {"max_model_len": 4096}, + }, + } + } + sampling_changed = { + **base, + "policy": { + **base["policy"], + "generation": { + **base["policy"]["generation"], + "temperature": 0.5, + }, + }, + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(sampling_changed))) + ) + + context_changed = { + **base, + "policy": { + **base["policy"], + "generation": { + **base["policy"]["generation"], + "vllm_cfg": {"max_model_len": 8192}, + }, + }, + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(context_changed))) + ) + + +def test_bootstrap_fingerprint_ignores_nested_gym_log_directory() -> None: + base = { + "policy": {"model": "model-a"}, + "env": { + "should_use_nemo_gym": True, + "nemo_gym": { + "nemo_gym_log_dir": "/run/one/nemo_gym", + "should_log_nemo_gym_responses": True, + "policy_model": {"temperature": 1.0}, + "agent": {"concurrency": 16, "max_turns": 20}, + }, + }, + } + runtime_changed = { + **base, + "env": { + **base["env"], + "nemo_gym": { + **base["env"]["nemo_gym"], + "nemo_gym_log_dir": "/run/two/nemo_gym", + "should_log_nemo_gym_responses": False, + "agent": {"concurrency": 64, "max_turns": 20}, + }, + }, + } + semantic_changed = { + **base, + "env": { + **base["env"], + "nemo_gym": { + **base["env"]["nemo_gym"], + "policy_model": {"temperature": 0.5}, + }, + }, + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(runtime_changed)) + ) + assert fingerprint != bootstrap_fingerprint( + cast(Any, _DumpedConfig(semantic_changed)) + ) + assert base["env"]["nemo_gym"]["nemo_gym_log_dir"] == "/run/one/nemo_gym" + + +def test_prune_bootstrap_snapshots_requires_durable_trainer_checkpoint(tmp_path): + snapshot_root = tmp_path / "bootstrap" / "rollout_snapshots" + snapshot_root.mkdir(parents=True) + (snapshot_root / "snapshot_000001").mkdir() + durable_anchor = tmp_path / "step_1" + + with pytest.raises(FileNotFoundError, match="durable trainer checkpoint"): + prune_bootstrap_snapshots( + tmp_path, + durable_trainer_checkpoint=durable_anchor, + ) + + assert snapshot_root.is_dir() + durable_anchor.mkdir() + assert prune_bootstrap_snapshots( + tmp_path, + durable_trainer_checkpoint=durable_anchor, + ) + assert not snapshot_root.exists() + + +def test_reset_bootstrap_anchor_discards_skipped_snapshot_lineage( + tmp_path: Path, +) -> None: + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="old-fingerprint") + snapshot_root = anchor / "rollout_snapshots" + (snapshot_root / "snapshot_000001").mkdir(parents=True) + + reset = reset_bootstrap_anchor(tmp_path, fingerprint="new-fingerprint") + + assert reset == anchor + assert not snapshot_root.exists() + validate_bootstrap_anchor(anchor, fingerprint="new-fingerprint") + + +def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + first = _commit_snapshot(anchor, mutation_version=1) + second = _commit_snapshot(anchor, mutation_version=2) + + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + assert resolved is not None + assert resolved.path == second + assert resolved.manifest.mutation_version == 2 + assert first.is_dir() + + +def test_resolver_falls_back_from_corrupt_newest_snapshot(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + first = _commit_snapshot(anchor, mutation_version=1) + second = _commit_snapshot(anchor, mutation_version=2) + (second / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text("not-json") + + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + assert resolved is not None + assert resolved.path == first + + +def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + _commit_snapshot(anchor, mutation_version=1, fingerprint="different") + + with pytest.raises(ValueError, match="trainer-anchor mismatch"): + resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + +def test_commit_snapshot_flushes_payload_before_publication(tmp_path, monkeypatch): + anchor = tmp_path / "step_1" + anchor.mkdir() + tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(anchor) + (tmp_snapshot / "payload").write_text("payload") + fsync_tree = Mock() + fsync_file = Mock() + fsync_directory = Mock() + monkeypatch.setattr(rollout_checkpoint, "_fsync_tree", fsync_tree) + monkeypatch.setattr(rollout_checkpoint, "_fsync_file", fsync_file) + monkeypatch.setattr(rollout_checkpoint, "_fsync_directory", fsync_directory) + + commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=1) + + fsync_tree.assert_called_once_with(tmp_snapshot) + assert fsync_file.call_args_list == [ + call(tmp_snapshot / "COMMITTED"), + call(anchor / "rollout_snapshots" / "LATEST.tmp"), + ] + assert fsync_directory.call_args_list[:2] == [ + call(tmp_snapshot), + call(anchor / "rollout_snapshots"), + ] + assert (final_snapshot / "COMMITTED").is_file() + assert ( + anchor / "rollout_snapshots" / "LATEST" + ).read_text().strip() == final_snapshot.name diff --git a/tests/unit/utils/test_timer.py b/tests/unit/utils/test_timer.py index a274e5a075e..e89cde6fd70 100644 --- a/tests/unit/utils/test_timer.py +++ b/tests/unit/utils/test_timer.py @@ -625,6 +625,14 @@ def test_double_save_prevented(self): assert checker.check_save() is True assert checker.check_save() is False + def test_would_save_does_not_consume_deadline(self): + checker = TimeoutChecker(timeout="00:00:00:00") + + assert checker.would_save() is True + assert checker.would_save() is True + assert checker.check_save() is True + assert checker.would_save() is False + 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) From 56f9787bb21df4c8bd9faad2a9b59a9671d2537c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 20:52:27 -0400 Subject: [PATCH 02/14] docs(rollout): surface periodic checkpoint defaults (cherry picked from commit 4cd773d826a87c6ec8b444b9eec5e50064c06658) Signed-off-by: Anish Mahishi --- ...grpo_math_1B_megatron_single_controller.yaml | 7 +++++++ .../single_controller_utils/config.py | 17 +++++++++++++---- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 96f2d1ea69d..71df265c016 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -145,6 +145,13 @@ checkpointing: # aborts checkpoint finalization. save_data_plane: true +# Frequent rollout-only snapshots are disabled unless interval_s is set. They +# reuse the latest durable trainer checkpoint and native TQ checkpoint support. +rollout_checkpointing: + interval_s: null + keep_latest_k: 2 + restore_mode: latest + policy: dtensor_cfg: enabled: false diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 628985c4896..0453c45dffd 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -717,10 +717,19 @@ def resolve_for_prompt( class RolloutCheckpointConfig(BaseModel, extra="forbid"): """Frequent rollout-state snapshots anchored to durable trainer state. - ``interval_s=None`` disables the periodic pump. ``latest`` restores the - newest compatible rollout snapshot, ``trainer_checkpoint`` ignores newer - rollout-only snapshots, and ``none`` restores trainer state without any - replay, lineage, or dataloader state. + ``interval_s=None`` disables periodic snapshots. A snapshot taken before + the first trainer checkpoint is anchored to the initial model and a + rollout-semantic configuration fingerprint. Later snapshots are anchored + to the most recent durable trainer checkpoint. + + ``restore_mode="latest"`` selects the newest compatible periodic snapshot. + ``trainer_checkpoint`` ignores newer periodic snapshots, while ``none`` + resumes trainer state without restoring rollout, replay, lineage, or the + dataloader cursor. + + SingleController has no validation loop, so checkpoint selection must use + ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited + ``val:`` settings are rejected during setup. """ interval_s: Optional[float] = Field(default=None, gt=0) From e03d2b9475dc70139cdafeb2a072c2dfd556680c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 16:22:29 -0400 Subject: [PATCH 03/14] fix(rollout): recover active streamed checkpoints Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 3 +- .../algorithms/async_utils/replay_buffer.py | 127 +++++++++++- .../async_utils/staleness_sampler.py | 2 +- nemo_rl/algorithms/single_controller.py | 109 ++++++----- .../single_controller_utils/config.py | 5 +- .../L1_Functional_Tests_SingleController.sh | 3 + ...ym_single_controller_streaming_recovery.sh | 185 ++++++++++++++++++ .../single_controller/test_checkpointing.py | 84 ++++++-- .../test_sampler_interface.py | 3 + .../test_tq_replay_buffer.py | 68 +++++++ 10 files changed, 520 insertions(+), 69 deletions(-) create mode 100755 tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 71df265c016..ee38f058600 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -146,7 +146,8 @@ checkpointing: save_data_plane: true # Frequent rollout-only snapshots are disabled unless interval_s is set. They -# reuse the latest durable trainer checkpoint and native TQ checkpoint support. +# 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: interval_s: null keep_latest_k: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 60f94efd3eb..31aa391e8f9 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -13,6 +13,7 @@ # limitations under the License. import asyncio +import copy import gc import hashlib import json @@ -231,6 +232,7 @@ def __init__(self) -> None: self._checkpoint_active = False self._active_mutations = 0 self._section_holders: set[asyncio.Task[Any]] = set() + self._mutation_version = 0 def _current_task(self) -> asyncio.Task[Any]: """Return the task entering a barrier section and reject reentrancy.""" @@ -244,6 +246,11 @@ def _current_task(self) -> asyncio.Task[Any]: ) return task + @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[DataPlaneMutationCut]: """Yield a live cut after any active checkpoint exits.""" @@ -260,6 +267,9 @@ async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: async with self._condition: self._section_holders.discard(task) self._active_mutations -= 1 + # Count the section even when its body raised. A redundant + # snapshot is safe; skipping a partially applied mutation is not. + self._mutation_version += 1 if self._active_mutations == 0: self._condition.notify_all() @@ -1043,6 +1053,11 @@ def __init__( self._post_write_enricher: Optional[ Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]] ] = None + # Sampler selection removes ready slots from the live replay index but + # deliberately leaves their rows in TQ until optimizer completion. + # Retain their metadata here so a periodic checkpoint can make an open + # streamed step replayable without depending on the sibling lineage. + self._training_claims: dict[str, TQReplayGroupMetadata] = {} def set_data_plane_checkpoint_barrier( self, barrier: DataPlaneCheckpointBarrier @@ -1425,15 +1440,71 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: cut, drop_group_ids, clear_data_plane=remove_in_dp ) + async def claim_for_training(self, idxs: list[int]) -> int: + """Transfer ready groups from sampler ownership to an open train step. + + The canonical rows remain in TQ. Their metadata stays checkpoint-visible + until :meth:`release_training_claims` runs after optimizer success and + data-plane cleanup. + """ + if len(idxs) == 0: + return 0 + if len(idxs) != len(set(idxs)): + raise ValueError("training claim contains duplicate replay indices") + if min(idxs) < 0: + raise IndexError("training claim indices must be non-negative") + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before claiming groups for training" + ) + claim_idxs = sorted(idxs, reverse=True) + if claim_idxs[0] >= len(self.meta_list): + raise IndexError( + "TQReplayBuffer.claim_for_training: indices out of range: " + f"{claim_idxs[0]}; size={len(self.meta_list)}" + ) + claim_group_ids = [self._group_ids[i] for i in claim_idxs] + async with self._data_plane_checkpoint_barrier.mutation() as cut: + return await self._remove_groups_unlocked( + cut, + claim_group_ids, + clear_data_plane=False, + retain_training_claims=True, + ) + + def training_owned_replay_groups(self) -> list[TQReplayGroupMetadata]: + """Return metadata for canonical rows owned by the open train step.""" + return copy.deepcopy(list(self._training_claims.values())) + + def release_training_claims(self, group_ids: list[str]) -> None: + """Release checkpoint ownership after consumed TQ rows are cleared.""" + if len(group_ids) != len(set(group_ids)): + raise ValueError("training claim release contains duplicate group IDs") + claimed_group_ids = set(self._training_claims) + released_group_ids = set(group_ids) + unknown = sorted(released_group_ids - claimed_group_ids) + unreleased = sorted(claimed_group_ids - released_group_ids) + if unknown or unreleased: + raise ValueError( + "training claim release does not match current ownership: " + f"unknown={unknown!r}, unreleased={unreleased!r}" + ) + for group_id in group_ids: + del self._training_claims[group_id] + async def _remove_groups_unlocked( self, cut: DataPlaneMutationCut, group_ids: list[str], *, clear_data_plane: bool, + retain_training_claims: bool = False, ) -> int: """Remove stable groups while the caller owns a live mutation cut.""" cut.require_live() + if clear_data_plane and retain_training_claims: + raise ValueError("cleared rows cannot be retained as training claims") if len(group_ids) != len(set(group_ids)): raise ValueError("replay removal contains duplicate group IDs") index_by_group_id = {group_id: i for i, group_id in enumerate(self._group_ids)} @@ -1484,6 +1555,25 @@ async def _remove_groups_unlocked( "may already be cleared" ) from error + new_training_claims: dict[str, TQReplayGroupMetadata] = {} + if retain_training_claims: + for group_id in group_ids: + i = index_by_group_id[group_id] + meta = self.meta_list[i] + if meta is None or not self.ready_list[i]: + raise RuntimeError( + "only ready replay groups may be claimed for training" + ) + if group_id in self._training_claims: + raise ValueError(f"duplicate training-owned group_id={group_id!r}") + new_training_claims[group_id] = { + "meta": copy.deepcopy(meta), + "start_weight": self.start_weight_list[i], + "end_weight": self.end_weight_list[i], + "target_step": self.target_step_list[i], + "group_id": group_id, + } + # A different mutation may have removed a lower list slot while the # DataPlane calls were awaiting. Resolve the original stable IDs again; # never apply pre-await indices to the now-shifted parallel lists. A group @@ -1499,12 +1589,20 @@ async def _remove_groups_unlocked( ), reverse=True, ) + if retain_training_claims and len(current_drop_idxs) != len(group_ids): + raise RuntimeError("training claim ownership changed during mutation") + self._training_claims.update(new_training_claims) for i in current_drop_idxs: self._delete_slot(i) return len(current_drop_idxs) - def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: + def metadata_state_dict( + self, + *, + saved_capacity: int, + additional_groups: Optional[list[TQReplayGroupMetadata]] = None, + ) -> TQReplayMetadataState: """Capture the controller index for ready groups without tensor payloads. The caller must hold the exclusive side of the shared data-plane @@ -1516,7 +1614,9 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: complete publish/index or clear/remove transition. No writer is exempt, including post-train cleanup in ``_train_pump``; canonical writes are not required to originate specifically from :meth:`commit`. - In-flight reservations are intentionally omitted. + In-flight reservations are intentionally omitted. ``additional_groups`` + is used by periodic snapshots to re-index rows claimed by an unfinished + streamed optimizer step. """ groups: list[TQReplayGroupMetadata] = [] for i, ready in enumerate(self.ready_list): @@ -1533,6 +1633,27 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: "group_id": self._group_ids[i], } ) + existing_group_ids = {group["group_id"] for group in groups} + existing_sample_ids = { + sample_id for group in groups for sample_id in group["meta"].sample_ids + } + for group in additional_groups or []: + group_id = group["group_id"] + if group_id in existing_group_ids: + raise ValueError( + f"additional replay metadata duplicates group_id={group_id!r}" + ) + duplicate_sample_ids = existing_sample_ids.intersection( + group["meta"].sample_ids + ) + if duplicate_sample_ids: + raise ValueError( + "additional replay metadata duplicates sample IDs: " + f"{sorted(duplicate_sample_ids)!r}" + ) + groups.append(copy.deepcopy(group)) + existing_group_ids.add(group_id) + existing_sample_ids.update(group["meta"].sample_ids) return { "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, "storage": REPLAY_BUFFER_METADATA_STORAGE, @@ -1584,7 +1705,7 @@ async def load_state_dict( sample_ids), disagrees with the native TQ snapshot, or exceeds ``max_groups``. """ - if self.meta_list or self._group_ids: + if self.meta_list or self._group_ids or self._training_claims: raise RuntimeError( "Replay-buffer checkpoint loading requires an empty local buffer" ) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 6f65d2225cc..1afba47afcf 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -287,7 +287,7 @@ async def _finalize_selection( ] selected_meta = selected_metas[0].concat(*selected_metas[1:]) # type: ignore[union-attr] selected_meta.extra_info[ROLLOUT_METRICS] = selected_rollout_metrics - await self._buffer.remove(selected_idxs, remove_in_dp=False) + await self._buffer.claim_for_training(selected_idxs) return selected_meta, len(selected_idxs) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index a1bf4de9fd1..4ada4eef865 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -171,6 +171,7 @@ class _RolloutCheckpointCut: replay_metadata: Optional[TQReplayMetadataState] rollout_recovery_payload: Optional[bytes] rollout_recovery_group_count: Optional[int] + rolled_back_train_group_count: int mutation_version: int @@ -436,10 +437,11 @@ def __init__( self._last_missing_rollout_snapshot_anchor: Optional[tuple[int, int]] = None self._bootstrap_fingerprint = getattr(actor_args, "bootstrap_fingerprint", None) self._rollout_checkpoint_stop_requested = asyncio.Event() - # Periodic snapshots intentionally skip an optimizer step in progress. - # The transition to/from idle joins the data-plane barrier below. - self._train_step_idle = asyncio.Event() - self._train_step_idle.set() + # Narrow unsafe window after an optimizer mutates model state and before + # SC publishes the matching TQ cleanup and trainer counters. Gradient + # accumulation remains snapshot-safe because selected rows stay owned by + # the replay buffer and are re-indexed in periodic snapshots. + self._optimizer_commit_in_progress = False # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() @@ -568,10 +570,7 @@ async def run(self) -> dict[str, Any]: set(tasks), return_when=asyncio.FIRST_COMPLETED ) stop_after_rollout_checkpoint = False - if ( - rollout_checkpoint_task is not None - and rollout_checkpoint_task in done - ): + if rollout_checkpoint_task is not None and rollout_checkpoint_task in done: await rollout_checkpoint_task if not self._rollout_checkpoint_stop_requested.is_set(): raise RuntimeError( @@ -1520,6 +1519,22 @@ async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: async with self._data_plane_checkpoint_barrier.mutation() as cut: await self._cleanup_consumed_metas_unlocked(cut, metas) + @staticmethod + def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: + """Return stable prompt-group IDs in canonical sample order.""" + group_ids: list[str] = [] + seen_group_ids: set[str] = set() + for sample_id in meta.sample_ids: + group_id = sample_id + if "_g" in sample_id: + candidate, generation_index = sample_id.rsplit("_g", 1) + if candidate and generation_index.isdigit(): + group_id = candidate + if group_id not in seen_group_ids: + group_ids.append(group_id) + seen_group_ids.add(group_id) + return group_ids + # ── the three pumps + the inline advantage stage ─────────────────────── async def _rollout_pump(self) -> None: @@ -2262,6 +2277,8 @@ async def _train_pump(self) -> None: # Always True off the PPO path: the start step is pinned to 0 there. is_policy_training_step = self._train_steps >= policy_training_start_step consumed_metas: list[KVBatchMeta] = [] + consumed_group_ids: list[str] = [] + consumed_group_count = 0 step_finalizer_metrics: dict[str, list[float]] = {} with self._timer.time("total_step_time"): @@ -2305,6 +2322,7 @@ async def _train_pump(self) -> None: self._async_cfg.min_groups_for_streaming_train, max_prompt_groups, ) + selected_group_ids: list[str] = [] async with self._data_plane_checkpoint_barrier.mutation(): train_meta, num_groups = await self._sampler.select( current_train_weight=self._trainer_version, @@ -2312,11 +2330,9 @@ async def _train_pump(self) -> None: max_prompt_groups=max_prompt_groups, ) if train_meta is not None: - train_step_idle = getattr( - self, "_train_step_idle", None + selected_group_ids = self._group_ids_from_meta( + train_meta ) - if train_step_idle is not None: - train_step_idle.clear() # If no batch is selectable, sleep and retry if train_meta is None: @@ -2344,16 +2360,8 @@ async def _train_pump(self) -> None: continue consumed_metas.append(train_meta) - - # Release buffer capacity. The rows stay in TQ until the - # post-step clear, but every reservation uses a fresh - # uuid group id, so nothing can collide with them. - for _ in range(num_groups): - self._buffer_capacity.release() - selected_group_ids = { - sample_id.rsplit("_g", 1)[0] - for sample_id in train_meta.sample_ids - } + consumed_group_ids.extend(selected_group_ids) + consumed_group_count += num_groups for group_id in selected_group_ids: for name, value in self._finalizer_metrics_by_group.pop( group_id, {} @@ -2442,6 +2450,10 @@ async def _train_pump(self) -> None: # TODO(#2625): value_result, policy_result only record the last epoch's metrics. # That matches ppo.py for the losses; total_flops is additive and undercounted. if self._is_ppo: + # A critic optimizer update is already irreversible. Keep + # periodic snapshots out until this whole training step is + # published as consumed below. + self._optimizer_commit_in_progress = True with self._timer.time("value_training"): value_result = await self._value_train_epochs( train_meta, @@ -2569,10 +2581,7 @@ async def _train_pump(self) -> None: policy_result = await asyncio.to_thread( self._trainer.finish_train_step ) - - # Clear consumed canonical rows (and their staged capture deltas) - # now that the step's training dispatches are complete. - await self._cleanup_consumed_metas(consumed_metas) + self._optimizer_commit_in_progress = True # Aggregate step metrics step_metrics = {} @@ -2580,6 +2589,11 @@ async def _train_pump(self) -> None: step_metrics.update(aggregate_step_metrics(policy_result)) if value_result is not None: step_metrics.update(_compute_critic_metrics(value_result)) + async with self._data_plane_checkpoint_barrier.mutation() as cut: + await self._cleanup_consumed_metas_unlocked(cut, consumed_metas) + self._buffer.release_training_claims(consumed_group_ids) + for _ in range(consumed_group_count): + self._buffer_capacity.release() step_metrics.update( { name: statistics.fmean(values) @@ -2622,6 +2636,7 @@ async def _train_pump(self) -> None: self._trainer_version += 1 self._train_steps += 1 + self._optimizer_commit_in_progress = False dropped_prompt_groups = self._batch_shortfall.get( version_during_step, 0 ) @@ -2780,11 +2795,6 @@ async def _train_pump(self) -> None: flush=True, ) - train_step_idle = getattr(self, "_train_step_idle", None) - if train_step_idle is not None: - async with self._data_plane_checkpoint_barrier.mutation(): - train_step_idle.set() - if should_save_by_timeout: print("Timeout has been reached, stopping training early", flush=True) break @@ -3348,12 +3358,12 @@ async def _capture_rollout_checkpoint_cut( self, checkpoint_path: PathLike, ) -> _RolloutCheckpointCut: - """Save TQ and capture matching rollout sidecars under the barrier.""" - if not self._train_step_idle.is_set(): - raise RuntimeError( - "periodic rollout snapshot attempted after training consumed rows" - ) + """Save TQ and capture matching restart state under the barrier. + Groups selected by an unfinished streamed step are absent from the live + replay index but remain in TQ. Re-index them only in this persisted cut; + the live trainer keeps accumulating gradients without modification. + """ save_state = self._save_state save_state.current_step = self._train_steps save_state.total_steps = self._train_steps @@ -3366,8 +3376,10 @@ async def _capture_rollout_checkpoint_cut( dataloader_state = self._dataloader.state_dict() replacement_reserve = list(self._replacement_reserve) + training_owned_groups = self._buffer.training_owned_replay_groups() replay_metadata = self._buffer.metadata_state_dict( - saved_capacity=self._async_cfg.max_buffered_rollouts + saved_capacity=self._async_cfg.max_buffered_rollouts, + additional_groups=training_owned_groups, ) await self._validate_replay_inventory(replay_metadata) @@ -3376,9 +3388,7 @@ async def _capture_rollout_checkpoint_cut( recovery_state["sampler_stamps_target_steps"] = ( self._sampler_stamps_target_steps ) - canonical_group_ids = { - group["group_id"] for group in replay_metadata["groups"] - } + canonical_group_ids = {group["group_id"] for group in replay_metadata["groups"]} recovery_state["groups"] = [ group for group in recovery_state["groups"] @@ -3406,6 +3416,7 @@ async def _capture_rollout_checkpoint_cut( replay_metadata=replay_metadata, rollout_recovery_payload=recovery_payload, rollout_recovery_group_count=len(recovery_state["groups"]), + rolled_back_train_group_count=len(training_owned_groups), mutation_version=self._data_plane_checkpoint_barrier.mutation_version, ) @@ -3449,11 +3460,9 @@ def _write_config() -> None: async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: """Publish one rollout-only snapshot anchored to durable trainer state.""" - # Never hold the filesystem publication lock while waiting for the train - # step. A full checkpoint runs before the step becomes idle and needs the - # same lock, so reversing these waits would deadlock both save paths. - await self._train_step_idle.wait() async with self._checkpoint_save_lock: + if self._optimizer_commit_in_progress: + return False if ( not force and self._last_rollout_snapshot_mutation_version @@ -3517,7 +3526,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: try: async with self._data_plane_checkpoint_barrier.checkpoint(): if ( - not self._train_step_idle.is_set() + self._optimizer_commit_in_progress or self._train_steps != expected_train_step or self._trainer_version != expected_trainer_version ): @@ -3533,7 +3542,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: trainer_version=expected_trainer_version, current_epoch=snapshot_epoch, mutation_version=cut.mutation_version, - rolled_back_train_group_count=0, + rolled_back_train_group_count=(cut.rolled_back_train_group_count), bootstrap_fingerprint=snapshot_fingerprint, ) await asyncio.to_thread( @@ -3565,7 +3574,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: return True async def _rollout_checkpoint_pump(self) -> None: - """Persist rollout state periodically at trainer-safe boundaries.""" + """Persist rollout state periodically, including during streamed train.""" interval_s = self._master_config.rollout_checkpointing.interval_s if interval_s is None: raise RuntimeError("rollout checkpoint pump started while disabled") @@ -3689,6 +3698,12 @@ async def _save_checkpoint_impl( ) if self._master_config.checkpointing.get("save_data_plane"): + training_owned_groups = self._buffer.training_owned_replay_groups() + if training_owned_groups: + raise RuntimeError( + "full trainer checkpoint still owns streamed training rows: " + f"groups={[group['group_id'] for group in training_owned_groups]!r}" + ) if self._sampler.supports_buffer_checkpoint: replay_metadata = self._buffer.metadata_state_dict( saved_capacity=self._async_cfg.max_buffered_rollouts diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 0453c45dffd..cb1476669fe 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -719,8 +719,9 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): ``interval_s=None`` disables periodic snapshots. A snapshot taken before the first trainer checkpoint is anchored to the initial model and a - rollout-semantic configuration fingerprint. Later snapshots are anchored - to the most recent durable trainer checkpoint. + rollout-semantic configuration fingerprint. Later snapshots require the + durable trainer checkpoint for the controller's current completed step; + interval attempts are skipped until that exact anchor exists. ``restore_mode="latest"`` selects the newest compatible periodic snapshot. ``trainer_checkpoint`` ignores newer periodic snapshots, while ``none`` diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index ded20da8876..b991087eaf5 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -185,6 +185,9 @@ run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controll # Two-process token-capture recovery: preserve one sealed sibling in TQ and # redispatch only its unfinished peer after restoring the step checkpoint. run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +# Periodic native-TQ snapshot while a streamed step owns only part of its +# rollout batch, followed by SIGKILL and rollback to the durable trainer anchor. +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh new file mode 100755 index 00000000000..100fc041068 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -0,0 +1,185 @@ +#!/bin/bash +# Crash/restart coverage for a periodic cut taken during streamed GRPO train. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +BASE_TEST=$SCRIPT_DIR/grpo_async_gym_single_controller.sh +BASE_RUN_LOG=$SCRIPT_DIR/grpo_async_gym_single_controller/run.log +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_streaming_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +PHASE1_LOG=$TEST_DIR/phase1.log +PHASE2_LOG=$TEST_DIR/phase2.log +SELECTION_FILE=$TEST_DIR/selected_snapshot +PHASE1_PID="" + +NUM_PROMPTS=${SC_STREAMING_RECOVERY_NUM_PROMPTS:-8} +NUM_GENERATIONS=${SC_STREAMING_RECOVERY_NUM_GENERATIONS:-2} +MIN_STREAMING_GROUPS=${SC_STREAMING_RECOVERY_MIN_GROUPS:-2} +CLAIMED_GROUPS=${SC_STREAMING_RECOVERY_CLAIMED_GROUPS:-2} +MAX_STEPS=${SC_STREAMING_RECOVERY_MAX_STEPS:-3} +SNAPSHOT_INTERVAL_S=${SC_STREAMING_RECOVERY_INTERVAL_S:-0.05} +SNAPSHOT_TIMEOUT_S=${SC_STREAMING_RECOVERY_TIMEOUT_S:-2400} +TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) + +if (( CLAIMED_GROUPS < MIN_STREAMING_GROUPS || CLAIMED_GROUPS >= NUM_PROMPTS )); then + echo "CLAIMED_GROUPS must be in [MIN_STREAMING_GROUPS, NUM_PROMPTS)" + exit 2 +fi + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +stop_phase1() { + if [[ -z "$PHASE1_PID" ]]; then + return + fi + kill -KILL -- "-$PHASE1_PID" 2>/dev/null || true + wait "$PHASE1_PID" 2>/dev/null || true + PHASE1_PID="" +} + +cleanup() { + stop_phase1 + rm -rf "$CHECKPOINT_DIR" +} +trap cleanup EXIT + +COMMON_OVERRIDES=( + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + checkpointing.metric_name=null + +checkpointing.save_data_plane=true + ++token_capture.enabled=true + ++rollout_recovery.default_granularity=sibling + ++rollout_checkpointing.interval_s="$SNAPSHOT_INTERVAL_S" + ++rollout_checkpointing.keep_latest_k=256 + ++rollout_checkpointing.restore_mode=latest + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=0 + async_rl.min_groups_for_streaming_train="$MIN_STREAMING_GROUPS" + async_rl.max_inflight_prompts="$MIN_STREAMING_GROUPS" + async_rl.max_buffered_rollouts=$((NUM_PROMPTS + MIN_STREAMING_GROUPS)) + grpo.num_prompts_per_step="$NUM_PROMPTS" + grpo.num_generations_per_prompt="$NUM_GENERATIONS" + grpo.max_num_steps="$MAX_STEPS" + policy.train_global_batch_size="$TRAIN_GLOBAL_BATCH_SIZE" +) + +echo "=== Phase 1: crash with $CLAIMED_GROUPS/$NUM_PROMPTS groups claimed ===" +command -v setsid >/dev/null +setsid env RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" & +PHASE1_PID=$! + +uv run --directory "$PROJECT_ROOT" --no-sync python - \ + "$CHECKPOINT_DIR/step_1/rollout_snapshots" \ + "$SELECTION_FILE" \ + "$PHASE1_PID" \ + "$BASE_RUN_LOG" \ + "$CLAIMED_GROUPS" \ + "$SNAPSHOT_TIMEOUT_S" <<'PY' +import json +import os +import sys +import time +from pathlib import Path + +root = Path(sys.argv[1]) +selection = Path(sys.argv[2]) +phase_pid = int(sys.argv[3]) +phase_log = Path(sys.argv[4]) +expected_claimed = int(sys.argv[5]) +deadline = time.monotonic() + float(sys.argv[6]) + +while time.monotonic() < deadline: + for snapshot in sorted(root.glob("snapshot_*"), reverse=True): + manifest_path = snapshot / "manifest.json" + if not (snapshot / "COMMITTED").is_file() or not manifest_path.is_file(): + continue + manifest = json.loads(manifest_path.read_text()) + if ( + manifest["base_train_step"] == 1 + and manifest["trainer_version"] == 1 + and manifest["rolled_back_train_group_count"] == expected_claimed + ): + selection.write_text(snapshot.name + "\n") + raise SystemExit(0) + try: + os.kill(phase_pid, 0) + except ProcessLookupError as error: + tail = "" + if phase_log.is_file(): + tail = "\n".join(phase_log.read_text(errors="replace").splitlines()[-40:]) + raise RuntimeError( + "phase one exited before producing the requested streamed cut:\n" + tail + ) from error + time.sleep(0.1) +raise TimeoutError(f"no snapshot captured {expected_claimed} claimed groups") +PY + +stop_phase1 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" +SNAPSHOT_NAME=$(tr -d '\n' < "$SELECTION_FILE") +SNAPSHOT_ROOT=$CHECKPOINT_DIR/step_1/rollout_snapshots +SNAPSHOT_DIR=$SNAPSHOT_ROOT/$SNAPSHOT_NAME + +# Force restore to use the exact fault-injection cut. +for candidate in "$SNAPSHOT_ROOT"/snapshot_*; do + if [[ -d "$candidate" && "$(basename "$candidate")" != "$SNAPSHOT_NAME" ]]; then + rm -rf "$candidate" + fi +done +printf '%s\n' "$SNAPSHOT_NAME" > "$SNAPSHOT_ROOT/LATEST" + +uv run --directory "$PROJECT_ROOT" --no-sync python - \ + "$SNAPSHOT_DIR/manifest.json" \ + "$SNAPSHOT_DIR/replay_buffer_metadata.pt" \ + "$SNAPSHOT_DIR/rollout_recovery.pt" \ + "$CLAIMED_GROUPS" <<'PY' +import json +import sys + +import torch + +manifest = json.load(open(sys.argv[1])) +replay = torch.load(sys.argv[2], weights_only=False) +lineage = torch.load(sys.argv[3], weights_only=False) +expected_claimed = int(sys.argv[4]) + +assert manifest["rolled_back_train_group_count"] == expected_claimed, manifest +assert len(replay["groups"]) >= expected_claimed, replay +assert lineage["open_train_step"] is None, lineage +PY + +echo "=== Phase 2: restore claimed rows and finish without duplicate steps ===" +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "train step $MAX_STEPS/$MAX_STEPS" "$PHASE2_LOG" + +uv run --directory "$PROJECT_ROOT" --no-sync python - \ + "$PHASE2_LOG" "$CHECKPOINT_DIR/step_$MAX_STEPS/training_info.json" \ + "$MAX_STEPS" <<'PY' +import json +import re +import sys +from pathlib import Path + +log = Path(sys.argv[1]).read_text() +training_info = json.loads(Path(sys.argv[2]).read_text()) +max_steps = int(sys.argv[3]) + +assert training_info["current_step"] == max_steps, training_info +assert training_info["trainer_version"] == max_steps, training_info +assert not re.search(r"train step 1/", log), "restored run repeated anchor step 1" +for step in range(2, max_steps + 1): + matches = re.findall(rf"train step {step}/{max_steps}(?:\s|$)", log) + assert len(matches) == 1, (step, len(matches)) +PY + +echo "Streamed-step periodic recovery functional test passed." diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index c354343c9a1..33be6338951 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -449,6 +449,7 @@ def __init__( self.metadata_state_dict_calls: list[int] = [] self.load_calls: list[dict[str, Any]] = [] self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None + self.training_claims: list[dict[str, Any]] = [] @property def group_ids(self) -> tuple[str, ...]: @@ -459,9 +460,28 @@ def set_data_plane_checkpoint_barrier( ) -> None: self.checkpoint_barrier = barrier - def metadata_state_dict(self, *, saved_capacity: int) -> dict[str, Any]: + def metadata_state_dict( + self, + *, + saved_capacity: int, + additional_groups: Optional[list[dict[str, Any]]] = None, + ) -> dict[str, Any]: self.metadata_state_dict_calls.append(saved_capacity) - return dict(self._metadata_state) + state = dict(self._metadata_state) + state["groups"] = [ + *self._metadata_state["groups"], + *(additional_groups or []), + ] + return state + + def training_owned_replay_groups(self) -> list[dict[str, Any]]: + return list(self.training_claims) + + def release_training_claims(self, group_ids: list[str]) -> None: + claimed = {group["group_id"] for group in self.training_claims} + if claimed: + assert set(group_ids) == claimed + self.training_claims = [] def count_for_target_step(self, target_step: int) -> int: """Return the number of ready fake groups owned by one gated step.""" @@ -1095,24 +1115,58 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): assert (snapshot / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() assert not (snapshot / "policy").exists() - def test_snapshot_waits_until_no_training_rows_are_owned(self, tmp_path: Path): + def test_snapshot_reindexes_rows_owned_by_active_streamed_step( + self, tmp_path: Path + ): actor = self._actor(tmp_path) + claimed_meta = KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=["claimed-group_g0"], + sequence_lengths=[16], + tags=[{"weight_version": 0}], + ) + actor._buffer.training_claims = [ + { + "meta": claimed_meta, + "start_weight": 0, + "end_weight": 0, + "target_step": 0, + "group_id": "claimed-group", + } + ] + actor._dp_client.sample_ids = list(claimed_meta.sample_ids) - async def exercise() -> None: - actor._train_step_idle.clear() - save_task = asyncio.create_task( - actor._save_rollout_checkpoint(force=True) - ) - await asyncio.sleep(0) - assert not save_task.done() - assert actor._dp_client.save_calls == [] + try: + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + finally: + actor._checkpointer.shutdown() - async with actor._data_plane_checkpoint_barrier.mutation(): - actor._train_step_idle.set() - assert await asyncio.wait_for(save_task, timeout=5.0) + snapshot = ( + tmp_path + / "checkpoints" + / BOOTSTRAP_DIRNAME + / "rollout_snapshots" + / "snapshot_000001" + ) + manifest = json.loads( + (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).read_text() + ) + replay_state = torch.load( + snapshot / REPLAY_BUFFER_METADATA_FILENAME, + weights_only=False, + ) + assert manifest["rolled_back_train_group_count"] == 1 + assert [group["group_id"] for group in replay_state["groups"]] == [ + "claimed-group" + ] + def test_snapshot_skips_optimizer_commit_window(self, tmp_path: Path): + actor = self._actor(tmp_path) + actor._optimizer_commit_in_progress = True try: - asyncio.run(exercise()) + assert not asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert actor._dp_client.save_calls == [] finally: actor._checkpointer.shutdown() diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 28640eeae30..041466d9157 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -93,6 +93,9 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: del self.ready_list[i] return len(idxs) + async def claim_for_training(self, idxs: list[int]) -> int: + return await self.remove(idxs, remove_in_dp=False) + def _run(coro): return asyncio.run(coro) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 1d2f8fa6ee1..c7a907e23c8 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -270,6 +270,21 @@ def _add_group( class TestDataPlaneCheckpointBarrier: + def test_mutation_version_counts_completed_outer_sections(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + assert barrier.mutation_version == 0 + async with barrier.mutation(): + async with barrier.mutation(): + assert barrier.mutation_version == 0 + assert barrier.mutation_version == 1 + with pytest.raises(RuntimeError, match="injected"): + async with barrier.mutation(): + raise RuntimeError("injected") + assert barrier.mutation_version == 2 + + asyncio.run(exercise()) + def test_mutation_and_checkpoint_cuts_expire_on_context_exit(self): async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() @@ -1121,6 +1136,59 @@ def test_ignores_rollout_metrics_logging_sidecar(self): class TestTQReplayBufferStateDict: + def test_training_claim_is_reindexed_only_for_periodic_snapshot(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + group_id = buf.reserve( + weight_version=3, + target_step=4, + group_id="claimed-group", + ) + claimed_meta = _run( + buf.commit( + group_id, + _make_record(), + start_weight_version=3, + end_weight_version=3, + ) + ) + + assert _run(buf.claim_for_training([0])) == 1 + assert buf.size() == 0 + claims = buf.training_owned_replay_groups() + assert [group["group_id"] for group in claims] == ["claimed-group"] + assert buf.metadata_state_dict(saved_capacity=8)["groups"] == [] + + periodic_state = buf.metadata_state_dict( + saved_capacity=8, + additional_groups=claims, + ) + assert [group["meta"].sample_ids for group in periodic_state["groups"]] == [ + list(claimed_meta.sample_ids) + ] + assert dp.depth() == _N_GENS + + with pytest.raises(ValueError, match="unreleased=\\['claimed-group'\\]"): + buf.release_training_claims([]) + assert [group["group_id"] for group in buf.training_owned_replay_groups()] == [ + "claimed-group" + ] + + buf.release_training_claims([claims[0]["group_id"]]) + assert buf.training_owned_replay_groups() == [] + + def test_duplicate_training_claim_indices_do_not_change_ownership(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + _add_group(buf, weight=3) + + with pytest.raises(ValueError, match="duplicate replay indices"): + _run(buf.claim_for_training([0, 0])) + + assert buf.size() == 1 + assert buf.training_owned_replay_groups() == [] + assert dp.depth() == _N_GENS + def test_metadata_state_dict_omits_tensors_and_data_plane_reads(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) From b3f920393de4d84f388d825a8d0f595ba43bd332 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 17:40:45 -0400 Subject: [PATCH 04/14] fix(checkpoint): require consistent rollout restore Signed-off-by: Anish Mahishi --- .../single_controller_utils/config.py | 7 +++--- .../single_controller_utils/setup.py | 24 ++++--------------- .../single_controller/test_checkpointing.py | 5 ++++ 3 files changed, 13 insertions(+), 23 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index cb1476669fe..1e07e9fe490 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -724,9 +724,8 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): interval attempts are skipped until that exact anchor exists. ``restore_mode="latest"`` selects the newest compatible periodic snapshot. - ``trainer_checkpoint`` ignores newer periodic snapshots, while ``none`` - resumes trainer state without restoring rollout, replay, lineage, or the - dataloader cursor. + ``trainer_checkpoint`` ignores newer periodic snapshots and restores the + rollout state bundled with the durable trainer checkpoint. SingleController has no validation loop, so checkpoint selection must use ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited @@ -735,7 +734,7 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): interval_s: Optional[float] = Field(default=None, gt=0) keep_latest_k: int = Field(default=2, ge=1) - restore_mode: Literal["latest", "trainer_checkpoint", "none"] = "latest" + restore_mode: Literal["latest", "trainer_checkpoint"] = "latest" class MasterConfig(BaseModel, extra="allow"): diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 47daf3d7f66..b058e6066be 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -964,10 +964,8 @@ def setup_single_controller( sampler_supports_replay_recovery = sampler_supports_buffer_checkpoint( master_config.async_rl.sampler ) - if ( - sampler_supports_replay_recovery - and rollout_checkpoint_cfg.restore_mode != "none" - and not master_config.checkpointing.get("save_data_plane") + if sampler_supports_replay_recovery and not master_config.checkpointing.get( + "save_data_plane" ): error_message = ( "SingleController checkpointing with a replay-checkpoint-capable " @@ -982,10 +980,7 @@ def setup_single_controller( "checkpointing.enabled=false." ) raise ValueError(error_message) - if ( - rollout_checkpoint_cfg.restore_mode != "none" - and not sampler_supports_replay_recovery - ): + if not sampler_supports_replay_recovery: warnings.warn( f"Sampler {master_config.async_rl.sampler.name!r} cannot recover " "completed buffered rollouts. On resume, the dataloader cursor " @@ -1101,9 +1096,7 @@ def setup_single_controller( ) restore_mode = rollout_checkpoint_cfg.restore_mode - recovery_checkpoint_path = ( - trainer_checkpoint_path if restore_mode != "none" else None - ) + 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.interval_s is not None @@ -1174,12 +1167,6 @@ def setup_single_controller( f"without considering newer periodic snapshots: {trainer_checkpoint_path}", flush=True, ) - elif restore_mode == "none" and trainer_checkpoint_path: - print( - "📦 Resuming trainer state without rollout, replay, lineage, or " - "dataloader recovery.", - flush=True, - ) # ========================== # Setup Dataset & Environments @@ -1223,8 +1210,7 @@ def setup_single_controller( ) if recovery_checkpoint_path is not None: print( - "📦 Restoring dataloader state from checkpoint: " - f"{recovery_checkpoint_path}" + f"📦 Restoring dataloader state from checkpoint: {recovery_checkpoint_path}" ) load_dataloader_state(dataloader, recovery_checkpoint_path, data_config) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 33be6338951..1cfc92b9606 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -48,6 +48,7 @@ import pytest import torch import yaml +from pydantic import ValidationError from torchdata.stateful_dataloader import StatefulDataLoader from nemo_rl.algorithms.async_utils.replay_buffer import ( @@ -1081,6 +1082,10 @@ def test_ft_save_period_triggers_saves(self, tmp_path): class TestPeriodicRolloutCheckpoint: + def test_restore_mode_rejects_inconsistent_trainer_only_resume(self): + with pytest.raises(ValidationError, match="restore_mode"): + RolloutCheckpointConfig.model_validate({"restore_mode": "none"}) + def _actor(self, tmp_path: Path): config = _actor_master_config( tmp_path, From 9d72e3cc4ffb83c9468746e51ce535d8f4dd1fd4 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 19:16:05 -0400 Subject: [PATCH 05/14] fix(rollout): harden periodic checkpoint consistency Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 6 + .../async_utils/staleness_sampler.py | 29 ++++- nemo_rl/algorithms/single_controller.py | 101 +++++++++++----- .../single_controller_utils/config.py | 8 +- .../rollout_checkpoint.py | 16 +-- .../single_controller_utils/setup.py | 9 ++ ...ym_single_controller_streaming_recovery.sh | 22 ++-- .../single_controller/test_checkpointing.py | 109 ++++++++++++++++- .../test_rollout_checkpoint.py | 111 +++++++++++++++++- .../test_sampler_interface.py | 16 +++ tests/unit/single_controller/test_setup.py | 79 +++++++++++++ .../test_single_controller_actor.py | 7 ++ .../test_tq_replay_buffer.py | 46 ++++++++ 13 files changed, 503 insertions(+), 56 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 31aa391e8f9..608da389533 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1477,6 +1477,10 @@ def training_owned_replay_groups(self) -> list[TQReplayGroupMetadata]: """Return metadata for canonical rows owned by the open train step.""" return copy.deepcopy(list(self._training_claims.values())) + def training_owned_group_ids(self) -> set[str]: + """Return stable IDs currently owned by the open train step.""" + return set(self._training_claims) + def release_training_claims(self, group_ids: list[str]) -> None: """Release checkpoint ownership after consumed TQ rows are cleared.""" if len(group_ids) != len(set(group_ids)): @@ -1614,6 +1618,8 @@ def metadata_state_dict( complete publish/index or clear/remove transition. No writer is exempt, including post-train cleanup in ``_train_pump``; canonical writes are not required to originate specifically from :meth:`commit`. + The advantage stage also takes a mutation slot because the periodic + checkpoint pump runs concurrently with ``_train_pump``. In-flight reservations are intentionally omitted. ``additional_groups`` is used by periodic snapshots to re-index rows claimed by an unfinished streamed optimizer step. diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 1afba47afcf..19a406a46a9 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -93,7 +93,12 @@ async def select( min_prompt_groups: int, max_prompt_groups: int, ) -> tuple[Optional[KVBatchMeta], int]: - """Pick up to ``max_prompt_groups`` eligible groups; drop them locally.""" + """Pick up to ``max_prompt_groups`` eligible groups for training. + + Claim-aware samplers transfer the groups from ordinary replay-buffer + selection into training ownership until the controller releases them. + Legacy custom samplers may still remove selected groups immediately. + """ ... async def evict(self, *, current_train_weight: int) -> int: @@ -159,6 +164,7 @@ class BaseSampler(abc.ABC): """ supports_buffer_checkpoint: ClassVar[bool] = False + supports_training_claims: ClassVar[bool] = True def __init__(self, buffer: TQReplayBuffer) -> None: self._buffer = buffer @@ -747,6 +753,27 @@ def sampler_supports_buffer_checkpoint(cfg: SamplerConfig) -> bool: return capability +def sampler_supports_training_claims(cfg: SamplerConfig) -> bool: + """Return whether selection transfers rows into training ownership. + + Built-in samplers use :meth:`TQReplayBuffer.claim_for_training`. Custom + samplers retain the legacy local-removal contract unless they explicitly + opt in, so enabling periodic snapshots cannot silently assume ownership + metadata that the sampler never created. + """ + sampler_cls = _sampler_class_for_config(cfg) + if isinstance(cfg, CustomSamplerConfig): + capability = sampler_cls.__dict__.get("supports_training_claims", False) + else: + capability = getattr(sampler_cls, "supports_training_claims", None) + if not isinstance(capability, bool): + raise TypeError( + f"{sampler_cls.__name__}.supports_training_claims must be a " + f"boolean class attribute, got {capability!r}" + ) + return capability + + def create_sampler( buffer: TQReplayBuffer, cfg: SamplerConfig, diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 4ada4eef865..fd21124cae1 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -167,6 +167,7 @@ class _RolloutCheckpointCut: """Controller sidecars captured with one native TQ snapshot.""" dataloader_state: dict[str, Any] + sampler_dispatch_index: int replacement_reserve: list[DatumSpec] replay_metadata: Optional[TQReplayMetadataState] rollout_recovery_payload: Optional[bytes] @@ -577,7 +578,18 @@ async def run(self) -> dict[str, Any]: "rollout checkpoint pump exited without requesting stop" ) stop_after_rollout_checkpoint = True - if probe_task is not None and probe_task in done: + if stop_after_rollout_checkpoint: + # FIRST_COMPLETED may return several tasks. Do not let the + # orderly pre-step checkpoint stop hide a rollout/train failure + # that completed in the same event-loop turn. + for task in done: + if task is not rollout_checkpoint_task: + await task + if ( + not stop_after_rollout_checkpoint + and probe_task is not None + and probe_task in done + ): # Loops forever like the watchdog, so finishing at all means it raised. await probe_task if not stop_after_rollout_checkpoint and watchdog_task in done: @@ -2277,7 +2289,7 @@ async def _train_pump(self) -> None: # Always True off the PPO path: the start step is pinned to 0 there. is_policy_training_step = self._train_steps >= policy_training_start_step consumed_metas: list[KVBatchMeta] = [] - consumed_group_ids: list[str] = [] + consumed_training_claim_ids: list[str] = [] consumed_group_count = 0 step_finalizer_metrics: dict[str, list[float]] = {} @@ -2323,16 +2335,54 @@ async def _train_pump(self) -> None: max_prompt_groups, ) selected_group_ids: list[str] = [] + selected_training_claim_ids: list[str] = [] async with self._data_plane_checkpoint_barrier.mutation(): + training_claim_ids_before = ( + self._buffer.training_owned_group_ids() + ) train_meta, num_groups = await self._sampler.select( current_train_weight=self._trainer_version, min_prompt_groups=min_prompt_groups, max_prompt_groups=max_prompt_groups, ) + training_claim_ids_after = ( + self._buffer.training_owned_group_ids() + ) + removed_training_claim_ids = ( + training_claim_ids_before - training_claim_ids_after + ) + if removed_training_claim_ids: + raise RuntimeError( + "sampler selection removed existing training " + "claims: " + f"{sorted(removed_training_claim_ids)!r}" + ) + new_training_claim_ids = ( + training_claim_ids_after - training_claim_ids_before + ) if train_meta is not None: selected_group_ids = self._group_ids_from_meta( train_meta ) + if new_training_claim_ids: + if ( + set(selected_group_ids) + != new_training_claim_ids + ): + raise RuntimeError( + "sampler selection does not match its new " + "training claims: " + f"selected={selected_group_ids!r}, " + "claimed=" + f"{sorted(new_training_claim_ids)!r}" + ) + selected_training_claim_ids = selected_group_ids + elif new_training_claim_ids: + raise RuntimeError( + "sampler selection created training claims without " + "returning batch metadata: " + f"{sorted(new_training_claim_ids)!r}" + ) # If no batch is selectable, sleep and retry if train_meta is None: @@ -2360,7 +2410,7 @@ async def _train_pump(self) -> None: continue consumed_metas.append(train_meta) - consumed_group_ids.extend(selected_group_ids) + consumed_training_claim_ids.extend(selected_training_claim_ids) consumed_group_count += num_groups for group_id in selected_group_ids: for name, value in self._finalizer_metrics_by_group.pop( @@ -2591,7 +2641,7 @@ async def _train_pump(self) -> None: step_metrics.update(_compute_critic_metrics(value_result)) async with self._data_plane_checkpoint_barrier.mutation() as cut: await self._cleanup_consumed_metas_unlocked(cut, consumed_metas) - self._buffer.release_training_claims(consumed_group_ids) + self._buffer.release_training_claims(consumed_training_claim_ids) for _ in range(consumed_group_count): self._buffer_capacity.release() step_metrics.update( @@ -3412,6 +3462,7 @@ async def _capture_rollout_checkpoint_cut( ) return _RolloutCheckpointCut( dataloader_state=dataloader_state, + sampler_dispatch_index=self._sampler.dispatch_index, replacement_reserve=replacement_reserve, replay_metadata=replay_metadata, rollout_recovery_payload=recovery_payload, @@ -3530,7 +3581,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: or self._train_steps != expected_train_step or self._trainer_version != expected_trainer_version ): - await asyncio.to_thread(shutil.rmtree, tmp_path) + await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) return False snapshot_epoch = self._current_epoch cut = await self._capture_rollout_checkpoint_cut(tmp_path) @@ -3541,6 +3592,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: base_train_step=expected_train_step, trainer_version=expected_trainer_version, current_epoch=snapshot_epoch, + sampler_dispatch_index=cut.sampler_dispatch_index, mutation_version=cut.mutation_version, rolled_back_train_group_count=(cut.rolled_back_train_group_count), bootstrap_fingerprint=snapshot_fingerprint, @@ -3559,7 +3611,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: ) except BaseException: if tmp_path.exists(): - await asyncio.to_thread(shutil.rmtree, tmp_path) + await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) raise self._last_rollout_snapshot_mutation_version = cut.mutation_version @@ -3578,6 +3630,7 @@ async def _rollout_checkpoint_pump(self) -> None: interval_s = self._master_config.rollout_checkpointing.interval_s if interval_s is None: raise RuntimeError("rollout checkpoint pump started while disabled") + consecutive_failures = 0 while True: await asyncio.sleep(interval_s) deadline_due = self._train_steps == 0 and self._timeout.would_save() @@ -3588,12 +3641,16 @@ async def _rollout_checkpoint_pump(self) -> None: raise RuntimeError( "failed to save the required pre-step rollout checkpoint" ) from error - warnings.warn( + consecutive_failures += 1 + print( "Periodic rollout checkpoint failed; retaining the previous " - f"committed snapshot: {type(error).__name__}: {error}", - stacklevel=2, + "committed snapshot: " + f"consecutive_failures={consecutive_failures}, " + f"{type(error).__name__}: {error}", + flush=True, ) continue + consecutive_failures = 0 if deadline_due and saved and self._timeout.check_save(): print( "Checkpoint deadline reached before the first train step; " @@ -3610,14 +3667,7 @@ async def _save_checkpoint( is_policy_training_step: bool, ) -> None: """Serialize full and rollout-only checkpoint publication.""" - lock = getattr(self, "_checkpoint_save_lock", None) - if lock is None: - await self._save_checkpoint_impl( - step_metrics, - is_policy_training_step=is_policy_training_step, - ) - return - async with lock: + async with self._checkpoint_save_lock: await self._save_checkpoint_impl( step_metrics, is_policy_training_step=is_policy_training_step, @@ -4278,16 +4328,13 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: fields_to_put[adv_cfg.returns_field] = returns new_fields.append(adv_cfg.returns_field) - # Trainer-step checkpointing runs later in this same train-pump task, so - # this publication cannot race a checkpoint save. If advantage staging - # moves to another task, the write must participate in the data-plane - # mutation barrier. - await self._call_dp( - "put_samples", - sample_ids=meta.sample_ids, - partition_id=meta.partition_id, - fields=fields_for_put(meta, fields_to_put), - ) + async with self._data_plane_checkpoint_barrier.mutation(): + await self._call_dp( + "put_samples", + sample_ids=meta.sample_ids, + partition_id=meta.partition_id, + fields=fields_for_put(meta, fields_to_put), + ) return ( meta.with_fields(new_fields), has_valid_training_tokens, diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index 1e07e9fe490..d0e3cfaa889 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -729,11 +729,13 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): SingleController has no validation loop, so checkpoint selection must use ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited - ``val:`` settings are rejected during setup. + ``val:`` settings are rejected during setup. Unknown keys are + forbidden because a misspelled interval, retention, or restore option can + silently disable the durability behavior the operator intended. """ - interval_s: Optional[float] = Field(default=None, gt=0) - keep_latest_k: int = Field(default=2, ge=1) + interval_s: Annotated[Optional[float], Field(gt=0)] = None + keep_latest_k: Annotated[int, Field(ge=1)] = 2 restore_mode: Literal["latest", "trainer_checkpoint"] = "latest" diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index 4f2de3f08df..7e74de029a3 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -27,14 +27,13 @@ from nemo_rl.algorithms.single_controller_utils.config import MasterConfig -ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 1 +ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 2 BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 2 BOOTSTRAP_DIRNAME = "bootstrap" BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" ROLLOUT_SNAPSHOT_MANIFEST_FILENAME = "manifest.json" ROLLOUT_SNAPSHOT_COMMITTED_FILENAME = "COMMITTED" -ROLLOUT_SNAPSHOT_LATEST_FILENAME = "LATEST" _SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") @@ -172,6 +171,7 @@ class RolloutSnapshotManifest: base_train_step: int trainer_version: int current_epoch: int + sampler_dispatch_index: int mutation_version: int rolled_back_train_group_count: int bootstrap_fingerprint: Optional[str] @@ -184,6 +184,7 @@ def from_mapping(cls, raw: Mapping[str, Any]) -> RolloutSnapshotManifest: "base_train_step", "trainer_version", "current_epoch", + "sampler_dispatch_index", "mutation_version", "rolled_back_train_group_count", ) @@ -203,6 +204,7 @@ def from_mapping(cls, raw: Mapping[str, Any]) -> RolloutSnapshotManifest: base_train_step=raw["base_train_step"], trainer_version=raw["trainer_version"], current_epoch=raw["current_epoch"], + sampler_dispatch_index=raw["sampler_dispatch_index"], mutation_version=raw["mutation_version"], rolled_back_train_group_count=raw["rolled_back_train_group_count"], bootstrap_fingerprint=fingerprint, @@ -223,6 +225,10 @@ def from_mapping(cls, raw: Mapping[str, Any]) -> RolloutSnapshotManifest: < 0 ): raise ValueError("rollout snapshot counters must be non-negative") + if manifest.sampler_dispatch_index < -1: + raise ValueError( + "rollout snapshot sampler_dispatch_index must be at least -1" + ) return manifest def to_dict(self) -> dict[str, Any]: @@ -439,12 +445,6 @@ def commit_snapshot( root = final_path.parent _fsync_directory(root) - latest_path = root / ROLLOUT_SNAPSHOT_LATEST_FILENAME - latest_tmp = latest_path.with_suffix(".tmp") - latest_tmp.write_text(final_path.name + "\n") - _fsync_file(latest_tmp) - os.replace(latest_tmp, latest_path) - _fsync_directory(root) committed = sorted( ( diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index b058e6066be..c49cc47d02f 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -47,6 +47,7 @@ ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( sampler_supports_buffer_checkpoint, + sampler_supports_training_claims, ) from nemo_rl.algorithms.grpo import ( GRPOSaveState, @@ -1032,6 +1033,11 @@ def setup_single_controller( "rollout checkpointing requires a sampler that supports " "replay-buffer recovery" ) + if not sampler_supports_training_claims(master_config.async_rl.sampler): + raise ValueError( + "rollout checkpointing requires a sampler that explicitly " + "supports training-claim ownership" + ) if token_capture_cfg.enabled: if not should_use_nemo_gym(master_config): raise ValueError( @@ -1157,6 +1163,9 @@ def setup_single_controller( if resolved_snapshot is not None: recovery_checkpoint_path = str(resolved_snapshot.path) save_state.current_epoch = resolved_snapshot.manifest.current_epoch + save_state.sampler_dispatch_index = ( + resolved_snapshot.manifest.sampler_dispatch_index + ) print( f"📦 Selected rollout recovery snapshot: {recovery_checkpoint_path}", flush=True, diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index 100fc041068..b1c2b6d07da 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -19,8 +19,9 @@ NUM_GENERATIONS=${SC_STREAMING_RECOVERY_NUM_GENERATIONS:-2} MIN_STREAMING_GROUPS=${SC_STREAMING_RECOVERY_MIN_GROUPS:-2} CLAIMED_GROUPS=${SC_STREAMING_RECOVERY_CLAIMED_GROUPS:-2} MAX_STEPS=${SC_STREAMING_RECOVERY_MAX_STEPS:-3} -SNAPSHOT_INTERVAL_S=${SC_STREAMING_RECOVERY_INTERVAL_S:-0.05} +SNAPSHOT_INTERVAL_S=${SC_STREAMING_RECOVERY_INTERVAL_S:-0.2} SNAPSHOT_TIMEOUT_S=${SC_STREAMING_RECOVERY_TIMEOUT_S:-2400} +PHASE2_TIMEOUT_S=${SC_STREAMING_RECOVERY_PHASE2_TIMEOUT_S:-2400} TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) if (( CLAIMED_GROUPS < MIN_STREAMING_GROUPS || CLAIMED_GROUPS >= NUM_PROMPTS )); then @@ -55,13 +56,17 @@ COMMON_OVERRIDES=( ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling ++rollout_checkpointing.interval_s="$SNAPSHOT_INTERVAL_S" - ++rollout_checkpointing.keep_latest_k=256 + ++rollout_checkpointing.keep_latest_k=8 ++rollout_checkpointing.restore_mode=latest async_rl.sampler.name=in_order async_rl.sampler.max_lookahead_versions=0 async_rl.min_groups_for_streaming_train="$MIN_STREAMING_GROUPS" async_rl.max_inflight_prompts="$MIN_STREAMING_GROUPS" async_rl.max_buffered_rollouts=$((NUM_PROMPTS + MIN_STREAMING_GROUPS)) + ++async_rl.rollout_failure.nemo_gym.rollout_timeout_s=120 + ++async_rl.stall_watchdog.interval_s=10 + ++async_rl.stall_watchdog.stall_timeout_s=300 + ++async_rl.stall_watchdog.stall_action=abort grpo.num_prompts_per_step="$NUM_PROMPTS" grpo.num_generations_per_prompt="$NUM_GENERATIONS" grpo.max_num_steps="$MAX_STEPS" @@ -103,7 +108,7 @@ while time.monotonic() < deadline: if ( manifest["base_train_step"] == 1 and manifest["trainer_version"] == 1 - and manifest["rolled_back_train_group_count"] == expected_claimed + and manifest["rolled_back_train_group_count"] >= expected_claimed ): selection.write_text(snapshot.name + "\n") raise SystemExit(0) @@ -117,7 +122,9 @@ while time.monotonic() < deadline: "phase one exited before producing the requested streamed cut:\n" + tail ) from error time.sleep(0.1) -raise TimeoutError(f"no snapshot captured {expected_claimed} claimed groups") +raise TimeoutError( + f"no snapshot captured at least {expected_claimed} claimed groups" +) PY stop_phase1 @@ -132,8 +139,6 @@ for candidate in "$SNAPSHOT_ROOT"/snapshot_*; do rm -rf "$candidate" fi done -printf '%s\n' "$SNAPSHOT_NAME" > "$SNAPSHOT_ROOT/LATEST" - uv run --directory "$PROJECT_ROOT" --no-sync python - \ "$SNAPSHOT_DIR/manifest.json" \ "$SNAPSHOT_DIR/replay_buffer_metadata.pt" \ @@ -149,13 +154,14 @@ replay = torch.load(sys.argv[2], weights_only=False) lineage = torch.load(sys.argv[3], weights_only=False) expected_claimed = int(sys.argv[4]) -assert manifest["rolled_back_train_group_count"] == expected_claimed, manifest +assert manifest["rolled_back_train_group_count"] >= expected_claimed, manifest assert len(replay["groups"]) >= expected_claimed, replay assert lineage["open_train_step"] is None, lineage PY echo "=== Phase 2: restore claimed rows and finish without duplicate steps ===" -RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" +timeout --signal=TERM --kill-after=30s "${PHASE2_TIMEOUT_S}s" \ + env RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" cp "$BASE_RUN_LOG" "$PHASE2_LOG" grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 1cfc92b9606..e4080a0f7b0 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -83,6 +83,10 @@ BOOTSTRAP_DIRNAME, ROLLOUT_SNAPSHOT_COMMITTED_FILENAME, ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + RolloutSnapshotManifest, + commit_snapshot, + prepare_snapshot_paths, ) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.data.utils import load_dataloader_state @@ -478,11 +482,14 @@ def metadata_state_dict( def training_owned_replay_groups(self) -> list[dict[str, Any]]: return list(self.training_claims) + def training_owned_group_ids(self) -> set[str]: + return {group["group_id"] for group in self.training_claims} + def release_training_claims(self, group_ids: list[str]) -> None: claimed = {group["group_id"] for group in self.training_claims} - if claimed: - assert set(group_ids) == claimed - self.training_claims = [] + assert len(group_ids) == len(set(group_ids)) + assert set(group_ids) == claimed + self.training_claims = [] def count_for_target_step(self, target_step: int) -> int: """Return the number of ready fake groups owned by one gated step.""" @@ -1082,10 +1089,22 @@ def test_ft_save_period_triggers_saves(self, tmp_path): class TestPeriodicRolloutCheckpoint: - def test_restore_mode_rejects_inconsistent_trainer_only_resume(self): + def test_restore_mode_rejects_removed_none_value(self): with pytest.raises(ValidationError, match="restore_mode"): RolloutCheckpointConfig.model_validate({"restore_mode": "none"}) + @pytest.mark.parametrize( + "config", + [ + {"interval_s": 0}, + {"keep_latest_k": 0}, + {"unknown_option": True}, + ], + ) + def test_rejects_invalid_periodic_checkpoint_config(self, config): + with pytest.raises(ValidationError): + RolloutCheckpointConfig.model_validate(config) + def _actor(self, tmp_path: Path): config = _actor_master_config( tmp_path, @@ -1101,6 +1120,7 @@ def _actor(self, tmp_path: Path): def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): actor = self._actor(tmp_path) try: + actor._sampler.restore_dispatch_index(5) assert asyncio.run(actor._save_rollout_checkpoint(force=True)) finally: actor._checkpointer.shutdown() @@ -1114,6 +1134,10 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): ) assert (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() assert (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).is_file() + manifest = json.loads( + (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).read_text() + ) + assert manifest["sampler_dispatch_index"] == 5 assert (snapshot / "data_plane" / "metadata.json").is_file() assert (snapshot / "train_dataloader.pt").is_file() assert (snapshot / REPLAY_BUFFER_METADATA_FILENAME).is_file() @@ -1175,6 +1199,43 @@ def test_snapshot_skips_optimizer_commit_window(self, tmp_path: Path): finally: actor._checkpointer.shutdown() + def test_periodic_pump_reports_each_consecutive_failure( + self, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + ): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.interval_s = 0.001 + actor._train_steps = 1 + + async def _main() -> None: + two_failures = asyncio.Event() + calls = 0 + + async def _failing_save(*, force: bool = False) -> bool: + nonlocal calls + del force + calls += 1 + if calls == 2: + two_failures.set() + raise OSError("storage unavailable") + + actor._save_rollout_checkpoint = _failing_save + pump = asyncio.create_task(actor._rollout_checkpoint_pump()) + await asyncio.wait_for(two_failures.wait(), timeout=1.0) + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + output = capsys.readouterr().out + assert output.count("Periodic rollout checkpoint failed") == 2 + assert "consecutive_failures=1" in output + assert "consecutive_failures=2" in output + class TestDataPlaneCheckpoint: def test_metadata_uses_pre_await_save_state_snapshot(self, tmp_path): @@ -1625,6 +1686,7 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): checkpoint_path.mkdir(parents=True, exist_ok=True) actor._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + actor._checkpoint_save_lock = asyncio.Lock() actor._save_state = SimpleNamespace() actor._train_steps = 1 actor._trainer_version = 1 @@ -1923,6 +1985,45 @@ def test_setup_forwards_latest_resume_paths( assert actor_args.save_state == _get_grpo_save_state(dict(_STEP_3_SAVE_STATE)) assert actor_args.last_checkpoint_path == str(step_3) + def test_periodic_snapshot_restores_exact_dispatch_cursor( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + ckpt_dir = tmp_path / "ckpts" + step_3 = _write_checkpoint( + ckpt_dir, + 3, + _STEP_3_SAVE_STATE, + dataloader_state={"fake_position": 3}, + ) + tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(step_3) + torch.save( + {"fake_position": 7}, + tmp_snapshot / "train_dataloader.pt", + ) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=3, + trainer_version=3, + current_epoch=4, + sampler_dispatch_index=6, + mutation_version=9, + rolled_back_train_group_count=0, + bootstrap_fingerprint=None, + ) + (tmp_snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=2) + mc = _setup_master_config(str(ckpt_dir)) + + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert actor_args.save_state.current_epoch == 4 + assert actor_args.save_state.sampler_dispatch_index == 6 + assert actor_args.last_checkpoint_path == str(final_snapshot) + def test_setup_fresh_start_passes_none_paths( self, patched_factories, # noqa: F811 diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index 1226071e741..11e9f7e9d9c 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -20,8 +20,15 @@ import pytest from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint +from nemo_rl.algorithms.grpo import GRPOConfig +from nemo_rl.algorithms.single_controller_utils.config import TokenCaptureConfig +from nemo_rl.data import DataConfig +from nemo_rl.models.generation.interfaces import GenerationConfig +from nemo_rl.models.generation.vllm.config import VllmSpecificArgs +from nemo_rl.models.policy import PolicyConfig from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + ROLLOUT_SNAPSHOT_SCHEMA_VERSION, RolloutSnapshotManifest, bootstrap_fingerprint, commit_snapshot, @@ -43,6 +50,26 @@ def model_dump(self, *, mode: str) -> dict[str, Any]: return self._dumped +@pytest.mark.parametrize( + ("declared", "schema"), + [ + (rollout_checkpoint._BOOTSTRAP_POLICY_FIELDS, PolicyConfig), + (rollout_checkpoint._BOOTSTRAP_GENERATION_FIELDS, GenerationConfig), + (rollout_checkpoint._BOOTSTRAP_VLLM_FIELDS, VllmSpecificArgs), + (rollout_checkpoint._BOOTSTRAP_GRPO_FIELDS, GRPOConfig), + (rollout_checkpoint._BOOTSTRAP_DATA_FIELDS, DataConfig), + (rollout_checkpoint._BOOTSTRAP_TOKEN_CAPTURE_FIELDS, TokenCaptureConfig), + ], +) +def test_bootstrap_projection_fields_exist_in_config_schema(declared, schema): + fields = ( + set(schema.model_fields) + if hasattr(schema, "model_fields") + else set(schema.__annotations__) + ) + assert declared <= fields + + def _commit_snapshot( anchor, *, @@ -52,10 +79,11 @@ def _commit_snapshot( ): tmp_path, final_path, _ = prepare_snapshot_paths(anchor) manifest = RolloutSnapshotManifest( - schema_version=1, + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, base_train_step=trainer_version, trainer_version=trainer_version, current_epoch=2, + sampler_dispatch_index=trainer_version + 1, mutation_version=mutation_version, rolled_back_train_group_count=0, bootstrap_fingerprint=fingerprint, @@ -353,6 +381,24 @@ def test_resolver_falls_back_from_corrupt_newest_snapshot(tmp_path): assert resolved.path == first +def test_resolver_ignores_snapshot_without_commit_marker(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + committed = _commit_snapshot(anchor, mutation_version=1) + incomplete = anchor / "rollout_snapshots" / "snapshot_000002" + incomplete.mkdir() + (incomplete / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text("{}") + + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + assert resolved is not None + assert resolved.path == committed + + def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") _commit_snapshot(anchor, mutation_version=1, fingerprint="different") @@ -383,13 +429,68 @@ def test_commit_snapshot_flushes_payload_before_publication(tmp_path, monkeypatc fsync_tree.assert_called_once_with(tmp_snapshot) assert fsync_file.call_args_list == [ call(tmp_snapshot / "COMMITTED"), - call(anchor / "rollout_snapshots" / "LATEST.tmp"), ] assert fsync_directory.call_args_list[:2] == [ call(tmp_snapshot), call(anchor / "rollout_snapshots"), ] assert (final_snapshot / "COMMITTED").is_file() - assert ( - anchor / "rollout_snapshots" / "LATEST" - ).read_text().strip() == final_snapshot.name + + +def test_commit_snapshot_prunes_oldest_committed_snapshot(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + first = _commit_snapshot(anchor, mutation_version=1) + second = _commit_snapshot(anchor, mutation_version=2) + third_tmp, third, _ = prepare_snapshot_paths(anchor) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=0, + trainer_version=0, + current_epoch=2, + sampler_dispatch_index=2, + mutation_version=3, + rolled_back_train_group_count=0, + bootstrap_fingerprint="fingerprint-v1", + ) + (third_tmp / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + + commit_snapshot(third_tmp, third, keep_latest_k=2) + + assert not first.exists() + assert second.is_dir() + assert third.is_dir() + + +def test_manifest_rejects_bool_for_integer_field(): + raw = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "current_epoch": 0, + "sampler_dispatch_index": -1, + "mutation_version": 0, + "rolled_back_train_group_count": 0, + "bootstrap_fingerprint": "fingerprint-v1", + } + raw["mutation_version"] = True + + with pytest.raises(ValueError, match="mutation_version.*integer"): + RolloutSnapshotManifest.from_mapping(raw) + + +def test_manifest_rejects_dispatch_index_below_initial_state(): + raw = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "current_epoch": 0, + "sampler_dispatch_index": -2, + "mutation_version": 0, + "rolled_back_train_group_count": 0, + "bootstrap_fingerprint": "fingerprint-v1", + } + + with pytest.raises(ValueError, match="sampler_dispatch_index.*at least -1"): + RolloutSnapshotManifest.from_mapping(raw) diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 041466d9157..8143270fee2 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -44,6 +44,7 @@ create_sampler, required_buffer_capacity_for_config, sampler_supports_buffer_checkpoint, + sampler_supports_training_claims, ) from nemo_rl.data_plane import KVBatchMeta from nemo_rl.data_plane.schema import ROLLOUT_METRICS @@ -275,6 +276,20 @@ def test_custom_checkpoint_capability_is_discoverable_without_construction(self) ) assert not CheckpointingEchoSampler.constructed + @pytest.mark.parametrize( + ("config", "expected"), + [ + (InOrderSamplerConfig(), True), + (CustomSamplerConfig(target=f"{__name__}:EchoSampler"), False), + ( + CustomSamplerConfig(target=f"{__name__}:CheckpointingEchoSampler"), + True, + ), + ], + ) + def test_training_claim_capability_requires_custom_opt_in(self, config, expected): + assert sampler_supports_training_claims(config) is expected + def test_ready_first_config_builds_ready_first_sampler(self): s = create_sampler( FakeBuffer(), @@ -696,6 +711,7 @@ class CheckpointingEchoSampler(EchoSampler): """Custom sampler with a static replay-checkpoint capability.""" supports_buffer_checkpoint = True + supports_training_claims = True constructed = False def __init__(self, *args, **kwargs) -> None: diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index b370bf2b8d6..5f46db2ac0f 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -57,6 +57,8 @@ setup_single_controller, ) from nemo_rl.algorithms.single_controller_utils.config import ( + RolloutCheckpointConfig, + TokenCaptureConfig, validate_single_controller_config, ) from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS @@ -79,6 +81,7 @@ class _CheckpointingCustomSampler(WindowedSampler): """Custom sampler whose static capability must be validated during setup.""" supports_buffer_checkpoint = True + supports_training_claims = True def __init__(self, buffer: Any) -> None: super().__init__(buffer, max_staleness_versions=1) @@ -93,6 +96,15 @@ def __init__(self, buffer: Any) -> None: super().__init__(buffer, max_staleness_versions=1) +class _CheckpointingNonClaimingCustomSampler(WindowedSampler): + """Replay-capable custom sampler that retains legacy local selection.""" + + supports_buffer_checkpoint = True + + def __init__(self, buffer: Any) -> None: + super().__init__(buffer, max_staleness_versions=1) + + def _make_master_config( *, dp_enabled: bool = True, @@ -683,6 +695,73 @@ def test_rejects_mooncake_data_plane_checkpointing(self): with pytest.raises(NotImplementedError, match="backend='mooncake_cpu'"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_periodic_checkpointing_requires_trainer_checkpointing(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = False + mc.checkpointing["save_data_plane"] = True + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + + with pytest.raises(ValueError, match="requires checkpointing.enabled=true"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_periodic_checkpointing_requires_data_plane_save(self): + mc = _make_master_config( + sampler_cfg=CustomSamplerConfig( + target=f"{__name__}:_NonCheckpointingCustomSampler" + ) + ) + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = False + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + + with ( + pytest.warns(UserWarning, match="cannot recover completed buffered"), + pytest.raises( + ValueError, match="requires checkpointing.save_data_plane=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_periodic_checkpointing_requires_token_capture(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = True + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + + with pytest.raises(ValueError, match="requires token_capture.enabled=true"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_periodic_checkpointing_requires_replay_capable_sampler(self): + mc = _make_master_config( + sampler_cfg=CustomSamplerConfig( + target=f"{__name__}:_NonCheckpointingCustomSampler" + ) + ) + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = True + mc.token_capture = TokenCaptureConfig(enabled=True) + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + + with ( + pytest.warns(UserWarning, match="cannot recover completed buffered"), + pytest.raises(ValueError, match="supports replay-buffer recovery"), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_periodic_checkpointing_requires_claim_aware_custom_sampler(self): + mc = _make_master_config( + sampler_cfg=CustomSamplerConfig( + target=(f"{__name__}:_CheckpointingNonClaimingCustomSampler") + ) + ) + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = True + mc.token_capture = TokenCaptureConfig(enabled=True) + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + + with pytest.raises(ValueError, match="supports training-claim ownership"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_rejects_windowed_checkpointing_without_native_tq(self): mc = _make_master_config() mc.checkpointing["enabled"] = True diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 9416e9989c1..5d61313b21e 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -608,6 +608,7 @@ def test_advantage_stage_composes_all_filters_before_computing_advantages( ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -668,6 +669,7 @@ def test_advantage_stage_composes_all_filters_before_computing_advantages( assert metrics[0]["max_seq_mult_prob_error"] == pytest.approx(math.e) assert metrics[0]["max_seq_mult_prob_error_after_mask"] == pytest.approx(1.0) assert "advantages" in (result_meta.fields or []) + assert ctrl._data_plane_checkpoint_barrier.mutation_version == 1 @pytest.mark.parametrize( @@ -770,6 +772,7 @@ def test_advantage_stage_reports_seq_logprob_metrics_without_masking() -> None: ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -902,6 +905,7 @@ def test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk( ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -963,6 +967,7 @@ def test_advantage_stage_skips_preexisting_empty_mask_without_seq_threshold() -> ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -1043,6 +1048,7 @@ def put_samples(self, sample_ids, partition_id, fields): ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = FakeEstimator() + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = True ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = True @@ -2264,6 +2270,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index c7a907e23c8..cd485deabec 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1189,6 +1189,52 @@ def test_duplicate_training_claim_indices_do_not_change_ownership(self): assert buf.training_owned_replay_groups() == [] assert dp.depth() == _N_GENS + def test_training_claim_rejects_negative_index(self): + buf = _make_buffer(FakeDataPlaneClient()) + _add_group(buf, weight=3) + + with pytest.raises(IndexError, match="must be non-negative"): + _run(buf.claim_for_training([-1])) + + assert buf.size() == 1 + assert buf.training_owned_group_ids() == set() + + def test_training_claim_rejects_out_of_range_index(self): + buf = _make_buffer(FakeDataPlaneClient()) + _add_group(buf, weight=3) + + with pytest.raises(IndexError, match=r"out of range: 2; size=1"): + _run(buf.claim_for_training([2])) + + assert buf.size() == 1 + assert buf.training_owned_group_ids() == set() + + def test_training_claim_requires_bound_checkpoint_barrier(self): + buf = _make_buffer(FakeDataPlaneClient()) + _add_group(buf, weight=3) + buf._data_plane_checkpoint_barrier = None + + with pytest.raises(RuntimeError, match="must be bound"): + _run(buf.claim_for_training([0])) + + def test_training_claim_rejects_non_ready_group(self): + buf = _make_buffer(FakeDataPlaneClient()) + buf.reserve(weight_version=3) + + with pytest.raises(RuntimeError, match="only ready replay groups"): + _run(buf.claim_for_training([0])) + + assert buf.size() == 1 + assert buf.training_owned_group_ids() == set() + + def test_training_claim_release_rejects_unknown_and_duplicate_ids(self): + buf = _make_buffer(FakeDataPlaneClient()) + + with pytest.raises(ValueError, match=r"unknown=\['unknown'\]"): + buf.release_training_claims(["unknown"]) + with pytest.raises(ValueError, match="duplicate group IDs"): + buf.release_training_claims(["unknown", "unknown"]) + def test_metadata_state_dict_omits_tensors_and_data_plane_reads(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) From 9c5bfdd2c7cd6e87b9b31375a9ce029692f2881b Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 21:11:32 -0400 Subject: [PATCH 06/14] test(rollout): cover stable training claims Signed-off-by: Anish Mahishi --- .../test_tq_replay_buffer.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index cd485deabec..caa806f3c94 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1189,6 +1189,36 @@ def test_duplicate_training_claim_indices_do_not_change_ownership(self): assert buf.training_owned_replay_groups() == [] assert dp.depth() == _N_GENS + def test_training_claim_keeps_stable_selection_while_barrier_waits(self): + async def exercise() -> None: + checkpoint_barrier = DataPlaneCheckpointBarrier() + buf = _make_buffer( + FakeDataPlaneClient(), checkpoint_barrier=checkpoint_barrier + ) + unready_group_id = buf.reserve(weight_version=0, group_id="unready") + ready_group_ids = [ + buf.reserve(weight_version=i, group_id=f"ready-{i}") + for i in (1, 2) + ] + for i, group_id in enumerate(ready_group_ids, start=1): + await buf.commit( + group_id, + _make_record(), + start_weight_version=i, + end_weight_version=i, + ) + + async with checkpoint_barrier.checkpoint(): + claim_task = asyncio.create_task(buf.claim_for_training([2])) + await asyncio.sleep(0) + assert buf.abort(unready_group_id) is True + + assert await claim_task == 1 + assert buf.group_ids == (ready_group_ids[0],) + assert buf.training_owned_group_ids() == {ready_group_ids[1]} + + asyncio.run(exercise()) + def test_training_claim_rejects_negative_index(self): buf = _make_buffer(FakeDataPlaneClient()) _add_group(buf, weight=3) From 8e7079365ae3e98487c9d04dd5bf6f74bfc1eea6 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sat, 29 Aug 2026 23:06:33 -0400 Subject: [PATCH 07/14] fix(rollout): ignore runtime gym bootstrap fields Signed-off-by: Anish Mahishi --- .../rollout_checkpoint.py | 68 ++++++++++- .../test_rollout_checkpoint.py | 115 ++++++++++++++++++ 2 files changed, 177 insertions(+), 6 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index 7e74de029a3..625a1f7b6b1 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -28,7 +28,7 @@ from nemo_rl.algorithms.single_controller_utils.config import MasterConfig ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 2 -BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 2 +BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 3 BOOTSTRAP_DIRNAME = "bootstrap" BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" @@ -91,17 +91,53 @@ "staging_partition", } ) +# Deployment, placement, logging, and credential leaves do not change the +# meaning of persisted rollout state. The recursive filter keeps semantic +# neighbors such as model names, parser settings, timeouts, and reward config. _BOOTSTRAP_ENV_RUNTIME_FIELDS = frozenset( { + "api_key", + "api_server_count", "apptainer_memory_limit_mb", "concurrency", + "debug", + "default_host", + "global_aiohttp_connector_limit", + "global_aiohttp_connector_limit_per_host", "nemo_gym_log_dir", "num_gpu_nodes", + "num_processes", + "num_workers", + "policy_base_url", "port_range_high", "port_range_low", + "ray_head_node_address", + "ray_worker_py_executable", "should_log_nemo_gym_responses", "skip_venv_if_present", "use_absolute_ip", + "uv_cache_dir", + "uv_venv_dir", + } +) +# These remote services keep their model identity in the fingerprint, but the +# endpoint may legitimately move when a Slurm job is restarted. +_BOOTSTRAP_ENV_RUNTIME_PATHS = frozenset( + { + ( + "nemo_gym", + "genrm_model", + "responses_api_models", + "genrm_model", + "base_url", + ), + ( + "nemo_gym", + "nl2bash_judge_model", + "responses_api_models", + "local_vllm_model", + "base_url", + ), } ) @@ -116,16 +152,35 @@ def _select_fields( return {key: mapping[key] for key in sorted(fields) if key in mapping} -def _drop_runtime_fields(value: Any, runtime_fields: frozenset[str]) -> Any: - """Recursively strip known operational leaves from an environment.""" +def _drop_runtime_fields( + value: Any, + runtime_fields: frozenset[str], + *, + runtime_paths: frozenset[tuple[str, ...]], + path: tuple[str, ...] = (), +) -> Any: + """Recursively strip known operational leaves and paths from an environment.""" if isinstance(value, Mapping): return { - key: _drop_runtime_fields(child, runtime_fields) + key: _drop_runtime_fields( + child, + runtime_fields, + runtime_paths=runtime_paths, + path=(*path, key), + ) for key, child in value.items() - if key not in runtime_fields + if key not in runtime_fields and (*path, key) not in runtime_paths } if isinstance(value, list): - return [_drop_runtime_fields(child, runtime_fields) for child in value] + return [ + _drop_runtime_fields( + child, + runtime_fields, + runtime_paths=runtime_paths, + path=path, + ) + for child in value + ] return value @@ -299,6 +354,7 @@ def bootstrap_compatibility_identity( environment=_drop_runtime_fields( dumped.get("env", {}), _BOOTSTRAP_ENV_RUNTIME_FIELDS, + runtime_paths=_BOOTSTRAP_ENV_RUNTIME_PATHS, ), sampler=dict(sampler), token_capture=_select_fields( diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index 11e9f7e9d9c..c97bb88c24d 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -30,6 +30,7 @@ ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, RolloutSnapshotManifest, + bootstrap_compatibility_identity, bootstrap_fingerprint, commit_snapshot, ensure_bootstrap_anchor, @@ -311,6 +312,120 @@ def test_bootstrap_fingerprint_ignores_nested_gym_log_directory() -> None: assert base["env"]["nemo_gym"]["nemo_gym_log_dir"] == "/run/one/nemo_gym" +def test_bootstrap_fingerprint_ignores_nemo_gym_service_routing() -> None: + base = { + "env": { + "nemo_gym": { + "genrm_model": { + "responses_api_models": { + "genrm_model": { + "base_url": "http://genrm-one/v1", + "model": "genrm-model-a", + } + } + }, + "nl2bash_judge_model": { + "responses_api_models": { + "local_vllm_model": { + "base_url": "http://nl2bash-one/v1", + "model": "nl2bash-model-a", + } + } + }, + } + } + } + routing_changed = { + "env": { + "nemo_gym": { + "genrm_model": { + "responses_api_models": { + "genrm_model": { + "base_url": "http://genrm-two/v1", + "model": "genrm-model-a", + } + } + }, + "nl2bash_judge_model": { + "responses_api_models": { + "local_vllm_model": { + "base_url": "http://nl2bash-two/v1", + "model": "nl2bash-model-a", + } + } + }, + } + } + } + model_changed = { + "env": { + "nemo_gym": { + **base["env"]["nemo_gym"], + "genrm_model": { + "responses_api_models": { + "genrm_model": { + "base_url": "http://genrm-two/v1", + "model": "genrm-model-b", + } + } + }, + } + } + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(routing_changed)) + ) + assert fingerprint != bootstrap_fingerprint(cast(Any, _DumpedConfig(model_changed))) + identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) + assert ( + "base_url" + not in identity.environment["nemo_gym"]["genrm_model"]["responses_api_models"][ + "genrm_model" + ] + ) + assert ( + "base_url" + not in identity.environment["nemo_gym"]["nl2bash_judge_model"][ + "responses_api_models" + ]["local_vllm_model"] + ) + + +@pytest.mark.parametrize( + "field", + sorted(rollout_checkpoint._BOOTSTRAP_ENV_RUNTIME_FIELDS), +) +def test_bootstrap_fingerprint_ignores_declared_environment_runtime_fields( + field: str, +) -> None: + base = { + "env": { + "nemo_gym": { + "service": { + "model": "model-a", + field: "runtime-one", + } + } + } + } + runtime_changed = { + "env": { + "nemo_gym": { + "service": { + "model": "model-a", + field: "runtime-two", + } + } + } + } + + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) == ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(runtime_changed))) + ) + + def test_prune_bootstrap_snapshots_requires_durable_trainer_checkpoint(tmp_path): snapshot_root = tmp_path / "bootstrap" / "rollout_snapshots" snapshot_root.mkdir(parents=True) From e1915a7ca38a99334edeb33d3de1569ca1e864b3 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 30 Aug 2026 16:01:12 -0400 Subject: [PATCH 08/14] test(rollout): update streamed recovery ledger assertion Signed-off-by: Anish Mahishi --- .../grpo_async_gym_single_controller_streaming_recovery.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index b1c2b6d07da..53ed80f1eac 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -156,7 +156,7 @@ expected_claimed = int(sys.argv[4]) assert manifest["rolled_back_train_group_count"] >= expected_claimed, manifest assert len(replay["groups"]) >= expected_claimed, replay -assert lineage["open_train_step"] is None, lineage +assert "open_train_step" not in lineage, lineage PY echo "=== Phase 2: restore claimed rows and finish without duplicate steps ===" From 451cd0841b4c9564c0d5612d0708858848ddc8f1 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 4 Sep 2026 15:47:11 -0400 Subject: [PATCH 09/14] fix(rollout): reconcile periodic checkpoint mutation cuts Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 83 +++++++++---------- .../test_tq_replay_buffer.py | 3 +- 2 files changed, 39 insertions(+), 47 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index fd21124cae1..944187b10c3 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -2336,53 +2336,46 @@ async def _train_pump(self) -> None: ) selected_group_ids: list[str] = [] selected_training_claim_ids: list[str] = [] - async with self._data_plane_checkpoint_barrier.mutation(): - training_claim_ids_before = ( - self._buffer.training_owned_group_ids() - ) - train_meta, num_groups = await self._sampler.select( - current_train_weight=self._trainer_version, - min_prompt_groups=min_prompt_groups, - max_prompt_groups=max_prompt_groups, - ) - training_claim_ids_after = ( - self._buffer.training_owned_group_ids() - ) - removed_training_claim_ids = ( - training_claim_ids_before - training_claim_ids_after + training_claim_ids_before = ( + self._buffer.training_owned_group_ids() + ) + train_meta, num_groups = await self._sampler.select( + current_train_weight=self._trainer_version, + min_prompt_groups=min_prompt_groups, + max_prompt_groups=max_prompt_groups, + ) + training_claim_ids_after = ( + self._buffer.training_owned_group_ids() + ) + removed_training_claim_ids = ( + training_claim_ids_before - training_claim_ids_after + ) + if removed_training_claim_ids: + raise RuntimeError( + "sampler selection removed existing training claims: " + f"{sorted(removed_training_claim_ids)!r}" ) - if removed_training_claim_ids: - raise RuntimeError( - "sampler selection removed existing training " - "claims: " - f"{sorted(removed_training_claim_ids)!r}" - ) - new_training_claim_ids = ( - training_claim_ids_after - training_claim_ids_before + new_training_claim_ids = ( + training_claim_ids_after - training_claim_ids_before + ) + if train_meta is not None: + selected_group_ids = self._group_ids_from_meta(train_meta) + if new_training_claim_ids: + if set(selected_group_ids) != new_training_claim_ids: + raise RuntimeError( + "sampler selection does not match its new " + "training claims: " + f"selected={selected_group_ids!r}, " + "claimed=" + f"{sorted(new_training_claim_ids)!r}" + ) + selected_training_claim_ids = selected_group_ids + elif new_training_claim_ids: + raise RuntimeError( + "sampler selection created training claims without " + "returning batch metadata: " + f"{sorted(new_training_claim_ids)!r}" ) - if train_meta is not None: - selected_group_ids = self._group_ids_from_meta( - train_meta - ) - if new_training_claim_ids: - if ( - set(selected_group_ids) - != new_training_claim_ids - ): - raise RuntimeError( - "sampler selection does not match its new " - "training claims: " - f"selected={selected_group_ids!r}, " - "claimed=" - f"{sorted(new_training_claim_ids)!r}" - ) - selected_training_claim_ids = selected_group_ids - elif new_training_claim_ids: - raise RuntimeError( - "sampler selection created training claims without " - "returning batch metadata: " - f"{sorted(new_training_claim_ids)!r}" - ) # If no batch is selectable, sleep and retry if train_meta is None: diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index caa806f3c94..aa67fdb3ec1 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -275,8 +275,7 @@ async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() assert barrier.mutation_version == 0 async with barrier.mutation(): - async with barrier.mutation(): - assert barrier.mutation_version == 0 + assert barrier.mutation_version == 0 assert barrier.mutation_version == 1 with pytest.raises(RuntimeError, match="injected"): async with barrier.mutation(): From 6c621132072dc4574af1c09f973de4866ce30e39 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 4 Sep 2026 17:31:09 -0400 Subject: [PATCH 10/14] fix(rollout): address periodic checkpoint review Signed-off-by: Anish Mahishi --- ...po_math_1B_megatron_single_controller.yaml | 7 +- ...po_math_1B_megatron_single_controller.yaml | 12 + nemo_rl/algorithms/single_controller.py | 68 ++--- .../single_controller_utils/config.py | 11 +- .../rollout_checkpoint.py | 80 ++++-- .../single_controller_utils/setup.py | 35 ++- .../single_controller/test_checkpointing.py | 164 +++++++++-- .../test_rollout_checkpoint.py | 268 ++++++++++++++++-- tests/unit/single_controller/test_setup.py | 74 +++++ .../test_tq_replay_buffer.py | 3 +- 10 files changed, 597 insertions(+), 125 deletions(-) diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index ee38f058600..cf2af7978bb 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -145,9 +145,10 @@ checkpointing: # aborts checkpoint finalization. save_data_plane: true -# Frequent rollout-only snapshots are disabled unless interval_s is set. 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. +# Frequent rollout-only snapshots are disabled unless 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_checkpointing: interval_s: null keep_latest_k: 2 diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 0c414a80216..1a803f6579e 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -146,6 +146,18 @@ checkpointing: enabled: false checkpoint_dir: results/ppo-single-controller metric_name: null + # Include native TQ state and the metadata-only replay index. A save failure + # aborts checkpoint finalization. + save_data_plane: true + +# Frequent rollout-only snapshots are disabled unless 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_checkpointing: + interval_s: null + keep_latest_k: 2 + restore_mode: latest policy: tokenizer: diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 944187b10c3..f906c64273b 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -161,6 +161,8 @@ # Logger this module also uses as `self._logger`. log = logging.getLogger(__name__) +_MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES = 3 + @dataclass(frozen=True) class _RolloutCheckpointCut: @@ -546,12 +548,9 @@ async def run(self) -> dict[str, Any]: train_task = asyncio.create_task(self._train_pump()) watchdog_task = asyncio.create_task(self._stall_watchdog_pump()) tasks = [rollout_task, train_task, watchdog_task] - rollout_checkpoint_cfg = getattr( - self._master_config, "rollout_checkpointing", None - ) rollout_checkpoint_task = ( asyncio.create_task(self._rollout_checkpoint_pump()) - if getattr(rollout_checkpoint_cfg, "interval_s", None) is not None + if self._master_config.rollout_checkpointing.interval_s is not None else None ) if rollout_checkpoint_task is not None: @@ -1203,6 +1202,10 @@ async def _call_dp(self, method_name: str, **kwargs) -> Any: async def _save_data_plane_checkpoint( self, checkpoint_path: PathLike, + *, + train_steps: int, + trainer_version: int, + current_epoch: int, replay_metadata: Optional[TQReplayMetadataState] = None, rollout_recovery_payload_sha256: Optional[str] = None, rollout_recovery_group_count: Optional[int] = None, @@ -1219,20 +1222,13 @@ async def _save_data_plane_checkpoint( checkpoint_path, DATA_PLANE_CHECKPOINT_DIR, ) - save_state = self._save_state - checkpoint_trainer_version = save_state.trainer_version - if checkpoint_trainer_version is None: - raise RuntimeError( - "Cannot save a data-plane checkpoint before trainer_version " - "is captured in the controller save state" - ) metadata: DataPlaneCheckpointMetadata = { "data_plane_checkpoint_schema_version": ( DATA_PLANE_CHECKPOINT_SCHEMA_VERSION ), - "single_controller_train_steps": save_state.current_step, - "single_controller_trainer_version": checkpoint_trainer_version, - "single_controller_epoch": save_state.current_epoch, + "single_controller_train_steps": train_steps, + "single_controller_trainer_version": trainer_version, + "single_controller_epoch": current_epoch, "partition_id": self._partition_id, "sampler_name": self._async_cfg.sampler.name, "mode": "authoritative" if replay_metadata is not None else "shadow", @@ -3399,6 +3395,7 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: async def _capture_rollout_checkpoint_cut( self, + cut: DataPlaneMutationCut, checkpoint_path: PathLike, ) -> _RolloutCheckpointCut: """Save TQ and capture matching restart state under the barrier. @@ -3407,16 +3404,7 @@ async def _capture_rollout_checkpoint_cut( replay index but remain in TQ. Re-index them only in this persisted cut; the live trainer keeps accumulating gradients without modification. """ - save_state = self._save_state - save_state.current_step = self._train_steps - save_state.total_steps = self._train_steps - save_state.trainer_version = self._trainer_version - save_state.current_epoch = self._current_epoch - save_state.consumed_samples = self._consumed_samples - save_state.total_valid_tokens = self._total_valid_tokens - save_state.sampler_name = self._async_cfg.sampler.name - save_state.sampler_dispatch_index = self._sampler.dispatch_index - + cut.require_live() dataloader_state = self._dataloader.state_dict() replacement_reserve = list(self._replacement_reserve) training_owned_groups = self._buffer.training_owned_replay_groups() @@ -3444,11 +3432,15 @@ async def _capture_rollout_checkpoint_cut( if self._master_config.token_capture.enabled: await self._validate_rollout_recovery_inventory( + cut, replay_metadata=replay_metadata, clear_unreferenced=False, ) await self._save_data_plane_checkpoint( checkpoint_path, + train_steps=self._train_steps, + trainer_version=self._trainer_version, + current_epoch=self._current_epoch, replay_metadata=replay_metadata, rollout_recovery_payload_sha256=recovery_digest, rollout_recovery_group_count=len(recovery_state["groups"]), @@ -3568,7 +3560,7 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: prepare_snapshot_paths, anchor ) try: - async with self._data_plane_checkpoint_barrier.checkpoint(): + async with self._data_plane_checkpoint_barrier.checkpoint() as cut: if ( self._optimizer_commit_in_progress or self._train_steps != expected_train_step @@ -3577,17 +3569,21 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) return False snapshot_epoch = self._current_epoch - cut = await self._capture_rollout_checkpoint_cut(tmp_path) + snapshot_cut = await self._capture_rollout_checkpoint_cut( + cut, tmp_path + ) - await self._write_rollout_checkpoint_sidecars(tmp_path, cut) + await self._write_rollout_checkpoint_sidecars(tmp_path, snapshot_cut) manifest = RolloutSnapshotManifest( schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, base_train_step=expected_train_step, trainer_version=expected_trainer_version, current_epoch=snapshot_epoch, - sampler_dispatch_index=cut.sampler_dispatch_index, - mutation_version=cut.mutation_version, - rolled_back_train_group_count=(cut.rolled_back_train_group_count), + sampler_dispatch_index=snapshot_cut.sampler_dispatch_index, + mutation_version=snapshot_cut.mutation_version, + rolled_back_train_group_count=( + snapshot_cut.rolled_back_train_group_count + ), bootstrap_fingerprint=snapshot_fingerprint, ) await asyncio.to_thread( @@ -3607,13 +3603,13 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) raise - self._last_rollout_snapshot_mutation_version = cut.mutation_version + self._last_rollout_snapshot_mutation_version = snapshot_cut.mutation_version self._last_missing_rollout_snapshot_anchor = None print( "rollout checkpoint save completed: " f"{final_path} (step={expected_train_step}, " f"trainer_version={expected_trainer_version}, " - f"ledger_groups={cut.rollout_recovery_group_count or 0})", + f"ledger_groups={snapshot_cut.rollout_recovery_group_count or 0})", flush=True, ) return True @@ -3642,6 +3638,11 @@ async def _rollout_checkpoint_pump(self) -> None: f"{type(error).__name__}: {error}", flush=True, ) + if consecutive_failures >= _MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES: + raise RuntimeError( + "periodic rollout checkpoint failed " + f"{consecutive_failures} consecutive times" + ) from error continue consecutive_failures = 0 if deadline_due and saved and self._timeout.check_save(): @@ -3789,6 +3790,9 @@ async def _save_checkpoint_impl( await self._save_data_plane_checkpoint( checkpoint_path, + train_steps=save_state.current_step, + trainer_version=self._trainer_version, + current_epoch=save_state.current_epoch, replay_metadata=replay_metadata, rollout_recovery_payload_sha256=(rollout_recovery_payload_sha256), rollout_recovery_group_count=( diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index d0e3cfaa889..b86cdba230d 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -717,11 +717,12 @@ def resolve_for_prompt( class RolloutCheckpointConfig(BaseModel, extra="forbid"): """Frequent rollout-state snapshots anchored to durable trainer state. - ``interval_s=None`` disables periodic snapshots. 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; - interval attempts are skipped until that exact anchor exists. + ``interval_s=None`` disables saving and restoring periodic snapshots. 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; interval attempts are skipped until that exact + anchor exists. ``restore_mode="latest"`` selects the newest compatible periodic snapshot. ``trainer_checkpoint`` ignores newer periodic snapshots and restores the diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index 625a1f7b6b1..31a4b6df3b4 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -28,7 +28,7 @@ from nemo_rl.algorithms.single_controller_utils.config import MasterConfig ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 2 -BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 3 +BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 4 BOOTSTRAP_DIRNAME = "bootstrap" BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" @@ -85,9 +85,11 @@ ) _BOOTSTRAP_TOKEN_CAPTURE_FIELDS = frozenset( { + "defer_routed_experts_to_policy", "enabled", "min_valid_fraction_per_group", "mixed_weight_version_policy", + "on_capture_failure", "staging_partition", } ) @@ -96,30 +98,71 @@ # neighbors such as model names, parser settings, timeouts, and reward config. _BOOTSTRAP_ENV_RUNTIME_FIELDS = frozenset( { + "_copy", + "_inherit_from", + "allow_openai_version_skew", "api_key", "api_server_count", "apptainer_memory_limit_mb", + "cache_dir", + "component_name", "concurrency", + "config_paths", "debug", "default_host", + "disallowed_ports", + "dry_run", + "entrypoint", "global_aiohttp_connector_limit", "global_aiohttp_connector_limit_per_host", + "head_server", + "head_server_deps", + "hf_token", + "json", + "model_call_capture_dir", + "model_endpoint_readiness_timeout_seconds", "nemo_gym_log_dir", "num_gpu_nodes", "num_processes", "num_workers", + "observability_enabled", + "pip_install_verbose", "policy_base_url", "port_range_high", "port_range_low", + "python_version", + "query", "ray_head_node_address", "ray_worker_py_executable", + "results_dir", "should_log_nemo_gym_responses", "skip_venv_if_present", + "token_id_capture", "use_absolute_ip", "uv_cache_dir", + "uv_pip_set_python", "uv_venv_dir", + "verbose", } ) +# Gym is optional in NeMo-RL, so keep credential recognition local rather than +# importing its global-config helpers. Match credential-shaped leaves without +# dropping semantic neighbors such as ``max_tokens`` or ``tokenizer``. +_BOOTSTRAP_ENV_CREDENTIAL_FIELDS = frozenset( + { + "api_key", + "apikey", + "password", + "secret", + "token", + } +) +_BOOTSTRAP_ENV_CREDENTIAL_FIELD_SUFFIXES = ( + "_api_key", + "_password", + "_secret", + "_token", +) # These remote services keep their model identity in the fingerprint, but the # endpoint may legitimately move when a Slurm job is restarted. _BOOTSTRAP_ENV_RUNTIME_PATHS = frozenset( @@ -152,6 +195,14 @@ def _select_fields( return {key: mapping[key] for key in sorted(fields) if key in mapping} +def _is_environment_credential_field(field: str) -> bool: + """Whether one environment leaf is credential-shaped, not semantic config.""" + normalized = field.casefold() + return normalized in _BOOTSTRAP_ENV_CREDENTIAL_FIELDS or normalized.endswith( + _BOOTSTRAP_ENV_CREDENTIAL_FIELD_SUFFIXES + ) + + def _drop_runtime_fields( value: Any, runtime_fields: frozenset[str], @@ -159,7 +210,7 @@ def _drop_runtime_fields( runtime_paths: frozenset[tuple[str, ...]], path: tuple[str, ...] = (), ) -> Any: - """Recursively strip known operational leaves and paths from an environment.""" + """Recursively strip operational and credential leaves from an environment.""" if isinstance(value, Mapping): return { key: _drop_runtime_fields( @@ -169,7 +220,9 @@ def _drop_runtime_fields( path=(*path, key), ) for key, child in value.items() - if key not in runtime_fields and (*path, key) not in runtime_paths + if key not in runtime_fields + and not _is_environment_credential_field(key) + and (*path, key) not in runtime_paths } if isinstance(value, list): return [ @@ -444,27 +497,6 @@ def reset_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: return ensure_bootstrap_anchor(checkpoint_dir, fingerprint=fingerprint) -def validate_bootstrap_anchor(anchor: Path, *, fingerprint: str) -> None: - """Fail loudly when bootstrap snapshots belong to different initial state.""" - manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME - if not manifest_path.is_file(): - raise FileNotFoundError( - f"rollout bootstrap manifest is missing at {manifest_path}" - ) - raw = json.loads(manifest_path.read_text()) - expected = { - "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, - "base_train_step": 0, - "trainer_version": 0, - "bootstrap_fingerprint": fingerprint, - } - if raw != expected: - raise ValueError( - "rollout bootstrap anchor does not match the current trainer " - f"configuration: checkpoint={raw!r}, expected={expected!r}" - ) - - def prepare_snapshot_paths(anchor: Path) -> tuple[Path, Path, int]: """Allocate the next temporary/final snapshot directory pair.""" root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index c49cc47d02f..ebffdbc2699 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -76,7 +76,6 @@ ensure_bootstrap_anchor, reset_bootstrap_anchor, resolve_latest_snapshot, - validate_bootstrap_anchor, ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn @@ -1038,6 +1037,16 @@ def setup_single_controller( "rollout checkpointing requires a sampler that explicitly " "supports training-claim ownership" ) + if master_config.checkpointing["save_period"] != 1: + warnings.warn( + "rollout checkpointing is enabled with " + f"checkpointing.save_period={master_config.checkpointing['save_period']}; " + "periodic rollout snapshots can only be saved while a matching " + "trainer checkpoint exists. Set checkpointing.save_period=1 for " + "continuous post-step coverage.", + UserWarning, + stacklevel=2, + ) if token_capture_cfg.enabled: if not should_use_nemo_gym(master_config): raise ValueError( @@ -1104,9 +1113,9 @@ 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.interval_s is not None - or (restore_mode == "latest" and bootstrap_anchor.is_dir()) + needs_bootstrap_identity = ( + trainer_checkpoint_path is None + and rollout_checkpoint_cfg.interval_s is not None ) bootstrap_digest = ( bootstrap_fingerprint(master_config) if needs_bootstrap_identity else None @@ -1117,7 +1126,11 @@ def setup_single_controller( if save_state.trainer_version is not None else save_state.current_step ) - if trainer_checkpoint_path is not None and restore_mode == "latest": + if ( + trainer_checkpoint_path is not None + and rollout_checkpoint_cfg.interval_s is not None + and restore_mode == "latest" + ): resolved_snapshot = resolve_latest_snapshot( Path(trainer_checkpoint_path), expected_train_step=save_state.current_step, @@ -1147,13 +1160,11 @@ def setup_single_controller( f"rollout_checkpointing.restore_mode={restore_mode!r}.", flush=True, ) - elif restore_mode == "latest" and bootstrap_anchor.is_dir(): - assert bootstrap_digest is not None - validate_bootstrap_anchor( - bootstrap_anchor, - fingerprint=bootstrap_digest, - ) - if restore_mode == "latest" and bootstrap_anchor.is_dir(): + if ( + rollout_checkpoint_cfg.interval_s is not None + and restore_mode == "latest" + and bootstrap_anchor.is_dir() + ): resolved_snapshot = resolve_latest_snapshot( bootstrap_anchor, expected_train_step=0, diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index e4080a0f7b0..491be53a79e 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -79,6 +79,7 @@ RolloutCheckpointConfig, setup_single_controller, ) +from nemo_rl.algorithms.single_controller_utils.config import TokenCaptureConfig from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( BOOTSTRAP_DIRNAME, ROLLOUT_SNAPSHOT_COMMITTED_FILENAME, @@ -115,6 +116,7 @@ _ACTOR_CLS = SingleControllerActor.__ray_metadata__.modified_class _PARTITION_ID = "rollout_data" +_STAGING_PARTITION_ID = "rollout_staging" def _consumed_meta(*sample_ids: str) -> KVBatchMeta: @@ -338,8 +340,11 @@ def __init__( self.sample_ids = list(sample_ids or []) def list_sample_ids(self, partition_id: str) -> list[str]: - assert partition_id == _PARTITION_ID - return sorted(self.sample_ids) + if partition_id == _PARTITION_ID: + return sorted(self.sample_ids) + if partition_id == _STAGING_PARTITION_ID: + return [] + raise AssertionError(f"unexpected partition_id={partition_id!r}") def clear_samples(self, sample_ids: list[str], partition_id: str) -> None: self.clear_thread_ids.append(threading.get_ident()) @@ -560,6 +565,7 @@ def _actor_master_config( buffer_checkpoint: bool = False, data_plane_checkpoint: bool = True, rollout_checkpoint_interval_s: Optional[float] = None, + token_capture_enabled: bool = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -622,6 +628,7 @@ def _actor_master_config( rollout_checkpointing=RolloutCheckpointConfig( interval_s=rollout_checkpoint_interval_s ), + token_capture=TokenCaptureConfig(enabled=token_capture_enabled), ) @@ -1110,6 +1117,7 @@ def _actor(self, tmp_path: Path): tmp_path, buffer_checkpoint=True, rollout_checkpoint_interval_s=120.0, + token_capture_enabled=True, ) return _ACTOR_CLS( config, @@ -1236,9 +1244,41 @@ async def _failing_save(*, force: bool = False) -> bool: assert "consecutive_failures=1" in output assert "consecutive_failures=2" in output + def test_periodic_pump_aborts_after_repeated_failures( + self, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + ): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.interval_s = 0.001 + actor._train_steps = 1 + + async def _main() -> None: + async def _failing_save(*, force: bool = False) -> bool: + del force + raise OSError("storage unavailable") + + actor._save_rollout_checkpoint = _failing_save + with pytest.raises( + RuntimeError, + match="periodic rollout checkpoint failed 3 consecutive times", + ): + await asyncio.wait_for( + actor._rollout_checkpoint_pump(), + timeout=1.0, + ) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + output = capsys.readouterr().out + assert output.count("Periodic rollout checkpoint failed") == 3 + class TestDataPlaneCheckpoint: - def test_metadata_uses_pre_await_save_state_snapshot(self, tmp_path): + def test_metadata_uses_explicit_snapshot_identity(self, tmp_path): mc = _actor_master_config( tmp_path, max_num_steps=1, @@ -1256,12 +1296,16 @@ async def _main() -> None: _make_actor_args(save_state=save_state, dp_client=dp_client), SetupTimingMetrics(), ) - # Simulate live fields diverging after _save_checkpoint captured - # save_state. The rollout pump can advance _current_epoch while - # checkpoint I/O awaits; both fields must come from one snapshot. + # The helper receives one explicit identity instead of reading + # mutable controller fields after checkpoint I/O has started. actor._trainer_version = 11 actor._current_epoch = 5 - await actor._save_data_plane_checkpoint(str(tmp_path / "tmp_step_3")) + await actor._save_data_plane_checkpoint( + str(tmp_path / "tmp_step_3"), + train_steps=3, + trainer_version=7, + current_epoch=2, + ) actor._checkpointer.shutdown() asyncio.run(_main()) @@ -1870,6 +1914,30 @@ def _write_checkpoint( return step_dir +def _write_periodic_snapshot(step_dir: Path) -> Path: + """Write one committed rollout snapshot newer than its trainer anchor.""" + tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(step_dir) + torch.save( + {"fake_position": 7}, + tmp_snapshot / "train_dataloader.pt", + ) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=3, + trainer_version=3, + current_epoch=4, + sampler_dispatch_index=6, + mutation_version=9, + rolled_back_train_group_count=0, + bootstrap_fingerprint=None, + ) + (tmp_snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=2) + return final_snapshot + + def _setup_master_config(checkpoint_dir: str) -> MasterConfig: """Partially-populated MasterConfig for setup_single_controller tests. @@ -1997,33 +2065,75 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( _STEP_3_SAVE_STATE, dataloader_state={"fake_position": 3}, ) - tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(step_3) - torch.save( - {"fake_position": 7}, - tmp_snapshot / "train_dataloader.pt", - ) - manifest = RolloutSnapshotManifest( - schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, - base_train_step=3, - trainer_version=3, - current_epoch=4, - sampler_dispatch_index=6, - mutation_version=9, - rolled_back_train_group_count=0, - bootstrap_fingerprint=None, - ) - (tmp_snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( - json.dumps(manifest.to_dict()) - ) - commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=2) + final_snapshot = _write_periodic_snapshot(step_3) mc = _setup_master_config(str(ckpt_dir)) + mc.checkpointing["save_period"] = 1 + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=120.0) + mc.token_capture = TokenCaptureConfig(enabled=True) + mc.policy["generation"].update( + { + "model_name": "test-model", + "stop_strings": None, + "stop_token_ids": None, + "top_k": None, + "vllm_cfg": {"async_engine": True}, + } + ) + mc.logger = {"log_dir": str(tmp_path / "logs")} + patched_factories["setup_response_data"].return_value = ( + list(range(8)), + None, + ) - actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + with ( + patch( + "nemo_rl.algorithms.single_controller_utils.setup.should_use_nemo_gym", + return_value=True, + ), + patch( + "nemo_rl.algorithms.single_controller_utils.setup.spinup_nemo_gym_actor", + return_value=MagicMock(), + ), + patch( + "nemo_rl.algorithms.single_controller_utils.setup.router_replay_enabled", + return_value=False, + ), + patch( + "nemo_rl.experience.finalizer_actor.create_finalizer_actors", + return_value=[MagicMock(name="finalizer")], + ), + ): + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) assert actor_args.save_state.current_epoch == 4 assert actor_args.save_state.sampler_dispatch_index == 6 assert actor_args.last_checkpoint_path == str(final_snapshot) + def test_disabled_periodic_checkpointing_uses_trainer_anchor( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + ckpt_dir = tmp_path / "ckpts" + step_3 = _write_checkpoint( + ckpt_dir, + 3, + _STEP_3_SAVE_STATE, + dataloader_state={"fake_position": 3}, + ) + _write_periodic_snapshot(step_3) + mc = _setup_master_config(str(ckpt_dir)) + mc.rollout_checkpointing = RolloutCheckpointConfig( + interval_s=None, + restore_mode="latest", + ) + + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert actor_args.save_state.current_epoch == 1 + assert actor_args.save_state.sampler_dispatch_index is None + assert actor_args.last_checkpoint_path == str(step_3) + def test_setup_fresh_start_passes_none_paths( self, patched_factories, # noqa: F811 diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index c97bb88c24d..a8dfd05c780 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -19,13 +19,9 @@ import pytest -from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint from nemo_rl.algorithms.grpo import GRPOConfig +from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint from nemo_rl.algorithms.single_controller_utils.config import TokenCaptureConfig -from nemo_rl.data import DataConfig -from nemo_rl.models.generation.interfaces import GenerationConfig -from nemo_rl.models.generation.vllm.config import VllmSpecificArgs -from nemo_rl.models.policy import PolicyConfig from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, @@ -38,8 +34,11 @@ prune_bootstrap_snapshots, reset_bootstrap_anchor, resolve_latest_snapshot, - validate_bootstrap_anchor, ) +from nemo_rl.data import DataConfig +from nemo_rl.models.generation.interfaces import GenerationConfig +from nemo_rl.models.generation.vllm.config import VllmSpecificArgs +from nemo_rl.models.policy import PolicyConfig class _DumpedConfig: @@ -51,24 +50,179 @@ def model_dump(self, *, mode: str) -> dict[str, Any]: return self._dumped +_BOOTSTRAP_PROJECTION_CASES = ( + ( + rollout_checkpoint._BOOTSTRAP_POLICY_FIELDS, + frozenset( + { + "disable_modelopt_layer_spec", + "draft", + "dtensor_cfg", + "dynamic_batching", + "generation", + "generation_batch_size", + "is_vlm", + "logprob_batch_size", + "logprob_chunk_size", + "make_sequence_length_divisible_by", + "max_grad_norm", + "megatron_cfg", + "optimizer", + "precision", + "quant_batch_size", + "quant_calib_data", + "quant_calib_size", + "quant_cfg", + "quant_sequence_length", + "refit_buffer_size_gb", + "reward_model_cfg", + "router_replay", + "scheduler", + "sequence_packing", + "train_global_batch_size", + "train_micro_batch_size", + } + ), + PolicyConfig, + ), + ( + rollout_checkpoint._BOOTSTRAP_GENERATION_FIELDS, + frozenset( + { + "_debug_payload_metrics", + "_mtp_weights_from_refit", + "_pad_token_id", + "bad_words", + "colocated", + "model_name", + "port_range_high", + "port_range_low", + "use_async_rollouts", + "val_temperature", + "val_top_k", + "val_top_p", + } + ), + GenerationConfig, + ), + ( + rollout_checkpoint._BOOTSTRAP_VLLM_FIELDS, + frozenset( + { + "async_engine", + "cap_max_tokens_to_context", + "enable_return_routed_experts", + "enforce_eager", + "env_vars", + "expert_parallel_size", + "expose_http_server", + "gpu_memory_utilization", + "http_refit_api_key_env_var", + "http_refit_server_port", + "is_mx", + "kv_cache_dtype", + "load_format", + "logprobs_mode", + "pipeline_parallel_size", + "precision", + "quantization_ignore_patterns", + "quantization_ignored_layer_kws", + "reset_encoder_cache_after_weight_update", + "skip_tokenizer_init", + "tensor_parallel_size", + "tool_parser_plugin", + "use_tqdm", + "video", + "zmq_refit_server_port", + } + ), + VllmSpecificArgs, + ), + ( + rollout_checkpoint._BOOTSTRAP_GRPO_FIELDS, + frozenset( + { + "adv_estimator", + "advantage_clip_high", + "advantage_clip_low", + "async_grpo", + "batch_multiplier", + "calculate_advantages_on_gpu", + "debug_payload_metrics", + "deduplicate_multimodal_data", + "dynamic_sampling_max_gen_batches", + "invalid_tool_call_advantage", + "malformed_thinking_advantage", + "max_num_epochs", + "max_num_steps", + "max_val_samples", + "normalize_rewards", + "overlong_filtering", + "reward_scaling", + "reward_shaping", + "seq_logprob_error_threshold", + "skip_reference_policy_logprobs_calculation", + "stop_at_validation_metric", + "stop_at_validation_threshold", + "use_dynamic_sampling", + "use_leave_one_out_baseline", + "val_at_end", + "val_at_start", + "val_batch_size", + "val_num_generations_per_prompt", + "val_period", + "val_start_at", + } + ), + GRPOConfig, + ), + ( + rollout_checkpoint._BOOTSTRAP_DATA_FIELDS, + frozenset( + { + "add_bos", + "add_eos", + "add_generation_prompt", + "add_system_prompt", + "custom_dataloader", + "num_prompts_per_dataloader", + "num_workers", + "use_multiple_dataloader", + "validation", + } + ), + DataConfig, + ), + ( + rollout_checkpoint._BOOTSTRAP_TOKEN_CAPTURE_FIELDS, + frozenset( + { + "capture_dir", + "control_auth_token", + "control_timeout_s", + "num_finalizer_workers", + } + ), + TokenCaptureConfig, + ), +) + + @pytest.mark.parametrize( - ("declared", "schema"), + ("declared", "ignored", "schema"), [ - (rollout_checkpoint._BOOTSTRAP_POLICY_FIELDS, PolicyConfig), - (rollout_checkpoint._BOOTSTRAP_GENERATION_FIELDS, GenerationConfig), - (rollout_checkpoint._BOOTSTRAP_VLLM_FIELDS, VllmSpecificArgs), - (rollout_checkpoint._BOOTSTRAP_GRPO_FIELDS, GRPOConfig), - (rollout_checkpoint._BOOTSTRAP_DATA_FIELDS, DataConfig), - (rollout_checkpoint._BOOTSTRAP_TOKEN_CAPTURE_FIELDS, TokenCaptureConfig), + pytest.param(declared, ignored, schema, id=schema.__name__) + for declared, ignored, schema in _BOOTSTRAP_PROJECTION_CASES ], ) -def test_bootstrap_projection_fields_exist_in_config_schema(declared, schema): +def test_bootstrap_projection_covers_config_schema(declared, ignored, schema): fields = ( set(schema.model_fields) if hasattr(schema, "model_fields") else set(schema.__annotations__) ) - assert declared <= fields + assert declared.isdisjoint(ignored) + assert declared | ignored == fields def _commit_snapshot( @@ -98,10 +252,10 @@ def _commit_snapshot( def test_bootstrap_anchor_rejects_different_initial_state(tmp_path): anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") - validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v1") + assert ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") == anchor with pytest.raises(ValueError, match="does not match"): - validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v2") + ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v2") def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: @@ -138,9 +292,11 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: }, "reward_penalties": {"penalize_unwanted_tokens": False}, "token_capture": { + "control_timeout_s": 60.0, "enabled": True, - "staging_partition": "rollout_staging", + "num_finalizer_workers": 2, "on_capture_failure": "continue", + "staging_partition": "rollout_staging", }, "cluster": {"num_nodes": 2}, "logger": {"log_dir": "/run/one"}, @@ -176,7 +332,8 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: "reward_penalties": {"penalize_unwanted_tokens": True}, "token_capture": { **base["token_capture"], - "on_capture_failure": "abort", + "control_timeout_s": 15.0, + "num_finalizer_workers": 8, }, "cluster": {"num_nodes": 8}, "logger": {"log_dir": "/run/two"}, @@ -195,6 +352,8 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: ("data", {"train": [{"data_path": "/datasets/other.jsonl"}]}), ("grpo", {"num_generations_per_prompt": 8}), ("token_capture", {"mixed_weight_version_policy": "reject"}), + ("token_capture", {"defer_routed_experts_to_policy": True}), + ("token_capture", {"on_capture_failure": "abort"}), ( "async_rl", {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, @@ -426,6 +585,75 @@ def test_bootstrap_fingerprint_ignores_declared_environment_runtime_fields( ) +@pytest.mark.parametrize( + "field", + ["hf_token", "judge_api_key", "policy_api_key", "wandb_api_key"], +) +def test_bootstrap_fingerprint_ignores_nested_environment_credentials( + field: str, +) -> None: + base = { + "env": { + "nemo_gym": { + "service": { + "model": "model-a", + field: "credential-one", + } + } + } + } + credential_changed = { + "env": { + "nemo_gym": { + "service": { + "model": "model-a", + field: "credential-two", + } + } + } + } + + identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) + assert field not in identity.environment["nemo_gym"]["service"] + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) == ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(credential_changed))) + ) + + +def test_bootstrap_fingerprint_keeps_environment_token_semantics() -> None: + base = { + "env": { + "nemo_gym": { + "service": { + "model": "model-a", + "max_tokens": 1024, + "tokenizer": "tokenizer-a", + } + } + } + } + max_tokens_changed = { + "env": { + "nemo_gym": { + "service": { + **base["env"]["nemo_gym"]["service"], + "max_tokens": 2048, + } + } + } + } + + identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) + assert identity.environment["nemo_gym"]["service"] == { + "max_tokens": 1024, + "model": "model-a", + "tokenizer": "tokenizer-a", + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(max_tokens_changed))) + ) + + def test_prune_bootstrap_snapshots_requires_durable_trainer_checkpoint(tmp_path): snapshot_root = tmp_path / "bootstrap" / "rollout_snapshots" snapshot_root.mkdir(parents=True) @@ -458,7 +686,7 @@ def test_reset_bootstrap_anchor_discards_skipped_snapshot_lineage( assert reset == anchor assert not snapshot_root.exists() - validate_bootstrap_anchor(anchor, fingerprint="new-fingerprint") + assert ensure_bootstrap_anchor(tmp_path, fingerprint="new-fingerprint") == anchor def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 5f46db2ac0f..f0f90001f76 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -762,6 +762,80 @@ def test_periodic_checkpointing_requires_claim_aware_custom_sampler(self): with pytest.raises(ValueError, match="supports training-claim ownership"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( + self, + tmp_path: Path, + patched_factories, + ): + mc = _make_master_config(colocated=False, backend="vllm") + mc.checkpointing.update( + { + "checkpoint_dir": str(tmp_path / "checkpoints"), + "enabled": True, + "save_data_plane": True, + "save_period": 2, + } + ) + mc.policy["generation"].update( + { + "model_name": "test-model", + "stop_strings": None, + "stop_token_ids": None, + "top_k": None, + "vllm_cfg": {"async_engine": True}, + } + ) + mc.logger = {"log_dir": str(tmp_path / "logs")} + mc.token_capture.enabled = True + mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + fake_finalizers = [MagicMock(name="finalizer")] + patched_factories["setup_response_data"].return_value = ( + list(range(8)), + None, + ) + + with ( + pytest.warns(UserWarning, match="checkpointing.save_period=2"), + patch.object(sc_setup_mod, "should_use_nemo_gym", return_value=True), + patch.object( + sc_setup_mod, "spinup_nemo_gym_actor", return_value=MagicMock() + ), + patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), + patch( + "nemo_rl.experience.finalizer_actor.create_finalizer_actors", + return_value=fake_finalizers, + ), + ): + actor_args, _ = setup_single_controller( + mc, + MagicMock(pad_token_id=0), + ) + + assert actor_args.finalizer_actors == fake_finalizers + + def test_disabled_periodic_checkpointing_ignores_existing_snapshots( + self, + tmp_path: Path, + patched_factories, + ): + mc = _make_master_config() + checkpoint_dir = tmp_path / "checkpoints" + mc.checkpointing["checkpoint_dir"] = str(checkpoint_dir) + mc.rollout_checkpointing = RolloutCheckpointConfig( + interval_s=None, + restore_mode="latest", + ) + (checkpoint_dir / "bootstrap" / "rollout_snapshots").mkdir(parents=True) + + with patch.object(sc_setup_mod, "resolve_latest_snapshot") as resolve: + actor_args, _ = setup_single_controller( + mc, + MagicMock(pad_token_id=0), + ) + + resolve.assert_not_called() + assert actor_args.last_checkpoint_path is None + def test_rejects_windowed_checkpointing_without_native_tq(self): mc = _make_master_config() mc.checkpointing["enabled"] = True diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index aa67fdb3ec1..46504442116 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -1196,8 +1196,7 @@ async def exercise() -> None: ) unready_group_id = buf.reserve(weight_version=0, group_id="unready") ready_group_ids = [ - buf.reserve(weight_version=i, group_id=f"ready-{i}") - for i in (1, 2) + buf.reserve(weight_version=i, group_id=f"ready-{i}") for i in (1, 2) ] for i, group_id in enumerate(ready_group_ids, start=1): await buf.commit( From b466886e40976c37da6acdf793825c83ba0de528 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 6 Sep 2026 00:35:13 -0400 Subject: [PATCH 11/14] fix(rollout): harden periodic checkpoint lifecycle Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 91 +++ ...po_math_1B_megatron_single_controller.yaml | 13 +- ...po_math_1B_megatron_single_controller.yaml | 13 +- .../async_utils/staleness_sampler.py | 14 +- nemo_rl/algorithms/single_controller.py | 31 +- .../single_controller_utils/config.py | 41 +- .../rollout_checkpoint.py | 566 ++++++++-------- .../single_controller_utils/setup.py | 60 +- ...ym_single_controller_streaming_recovery.sh | 4 +- .../single_controller/test_checkpointing.py | 77 ++- .../test_finalizer_lifecycle.py | 6 +- .../unit/single_controller/test_ppo_setup.py | 3 +- .../test_rollout_checkpoint.py | 620 ++++++++++-------- tests/unit/single_controller/test_setup.py | 95 ++- .../test_single_controller_actor.py | 3 + 15 files changed, 1003 insertions(+), 634 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 30262c4abf5..8e66f2ed79b 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -103,6 +103,97 @@ On resume, Single-Controller validates the TQ snapshot against the trainer check Replay recovery is supported by all built-in samplers: `in_order`, `weight_fifo`, `ready_first`, and `windowed`. Custom samplers must explicitly declare `supports_buffer_checkpoint = True`. Otherwise, setup emits a warning and completed buffered groups are not restored. +### Periodic rollout snapshots + +Normal trainer checkpoints are written at step boundaries. Periodic rollout +snapshots preserve newer rollout progress between those trainer checkpoints, +including while the train pump is accumulating a streamed step: + +```yaml +checkpointing: + enabled: true + checkpoint_dir: /shared/checkpoints/my-run + save_data_plane: true + save_period: 1 + +rollout_checkpointing: + snapshot_attempt_interval_s: 120 + keep_latest_k: 2 + restore_mode: latest + extra_fingerprint_excluded_paths: [] + +token_capture: + enabled: true +``` + +`snapshot_attempt_interval_s` is the cadence at which Single-Controller attempts +a rollout snapshot. It is not a guarantee that a snapshot is written at every +interval. An attempt after step N succeeds only when the immutable trainer +checkpoint `step_N` is already durable. Consequently, `save_period: 1` is +recommended for continuous post-step coverage; with a larger value, attempts +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. + +The bootstrap fingerprint is fail-closed: every configuration value affects +compatibility unless NeMo-RL's built-in denylist identifies it as operational, +such as logging, cluster placement, checkpoint location, runtime ports, or +credentials. This means configuration added by an external algorithm is safe +by default—a change prevents bootstrap recovery instead of silently mixing +incompatible rollout state. +The bootstrap manifest also stores this credential-redacted compatibility +identity so a rejected restart can report the exact changed dotpaths instead of +showing only two opaque digests. + +An integration may use `extra_fingerprint_excluded_paths` for additional +runtime-only values that are not part of NeMo-RL's built-in configuration: + +```yaml +rollout_checkpointing: + extra_fingerprint_excluded_paths: + - custom_algo.observability + - env.private_agent.runtime_endpoint + - env.private_agent.workers.*.log_dir +``` + +Each dotpath removes that value and its children from the compatibility +identity. `*` matches one mapping or list level and `**` matches any number of +levels. +Only exclude values that cannot affect prompts, generation, rewards, lineage, +or the interpretation of persisted rollout data. These exclusions must be set +on the original run as well as its restart. + +Periodic snapshots currently require all of the following: + +- `checkpointing.enabled: true` and `checkpointing.save_data_plane: true`. +- `data_plane.backend: simple`, because native TQ save/load is required. +- `token_capture.enabled: true`. +- A replay-recoverable sampler with training-claim ownership. All built-in + samplers qualify. A custom sampler must explicitly declare both + `supports_buffer_checkpoint = True` and `supports_training_claims = True`. + +Each trainer or bootstrap anchor has a `rollout_snapshots/` directory. A +published `snapshot_NNNNNN/` contains the native TQ snapshot and matching +replay, dataloader, controller, replacement-reserve, and unfinished-rollout +metadata. `keep_latest_k` retains recent committed snapshots as fallbacks; +temporary or interrupted directories are never selected for recovery. + +With `restore_mode: latest`, startup selects the newest compatible committed +snapshot under the latest trainer anchor. With `trainer_checkpoint`, it ignores +newer periodic rollout progress and resumes from the trainer checkpoint bundle. +Checkpoint selection is read-only: neither mode removes snapshots. If no trainer +checkpoint exists, `trainer_checkpoint` cannot safely reuse an existing bootstrap +namespace, so startup fails without modifying it. Recover that state with `latest` +or choose a new `checkpoint_dir` to start a fresh bootstrap lineage. Obsolete +bootstrap snapshots are removed only by retention after a durable trainer +checkpoint exists. + +> **Bootstrap-only `trainer_checkpoint` behavior:** A bootstrap rollout snapshot +> has no corresponding model or optimizer checkpoint. Therefore +> `restore_mode: trainer_checkpoint` deliberately fails when bootstrap state exists +> but no trainer checkpoint does. It does not ignore or delete that state. Use +> `latest` to recover it, or select a new `checkpoint_dir` to start from scratch. + :::{note} Completed groups are restored directly from the TQ snapshot. For unfinished token-capture groups, `rollout_recovery.default_granularity` controls both live diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index cf2af7978bb..38cf8d4ece6 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -141,18 +141,27 @@ checkpointing: enabled: false checkpoint_dir: results/grpo-single-controller metric_name: null + # Periodic rollout snapshots need an immutable trainer anchor for every + # completed step. Keep this at 1 when + # rollout_checkpointing.snapshot_attempt_interval_s is set. + save_period: 1 # Include native TQ state and the metadata-only replay index. A save failure # aborts checkpoint finalization. save_data_plane: true -# Frequent rollout-only snapshots are disabled unless interval_s is set. When +# 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_checkpointing: - interval_s: null + snapshot_attempt_interval_s: null keep_latest_k: 2 restore_mode: latest + # Advanced escape hatch for runtime-only fields from external integrations. + # 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: [] policy: dtensor_cfg: diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 1a803f6579e..3a13053c7f7 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -146,18 +146,27 @@ checkpointing: enabled: false checkpoint_dir: results/ppo-single-controller metric_name: null + # Periodic rollout snapshots need an immutable trainer anchor for every + # completed step. Keep this at 1 when + # rollout_checkpointing.snapshot_attempt_interval_s is set. + save_period: 1 # Include native TQ state and the metadata-only replay index. A save failure # aborts checkpoint finalization. save_data_plane: true -# Frequent rollout-only snapshots are disabled unless interval_s is set. When +# 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_checkpointing: - interval_s: null + snapshot_attempt_interval_s: null keep_latest_k: 2 restore_mode: latest + # Advanced escape hatch for runtime-only fields from external integrations. + # 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: [] policy: tokenizer: diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 19a406a46a9..5809e57f31a 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -71,6 +71,11 @@ class PromptGroupSampler(Protocol): Implement this (or subclass ``BaseSampler``) to add a custom sampling algorithm; point ``async_rl.sampler`` at ``module:ClassName`` to load it. + A custom sampler that supports replay recovery must explicitly declare + ``supports_buffer_checkpoint = True``. It must additionally declare + ``supports_training_claims = True`` before periodic rollout snapshots may + be enabled; omitting that optional capability preserves the legacy + remove-on-selection behavior. """ async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: @@ -122,6 +127,9 @@ def is_on_policy(self) -> bool: supports_buffer_checkpoint: ClassVar[bool] """Whether completed buffered groups can be restored safely.""" + supports_training_claims: ClassVar[bool] + """Whether selected groups remain owned until the train step commits.""" + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" ... @@ -661,7 +669,8 @@ class CustomSamplerConfig(BaseModel, extra="allow"): # Extra keys are forwarded to the constructor (after ``buffer``). The # target class must declare a boolean ``supports_buffer_checkpoint`` class # attribute so setup can validate recovery requirements before allocating - # cluster resources. + # cluster resources. Periodic rollout snapshots additionally require an + # explicit boolean ``supports_training_claims = True`` declaration. target: str @@ -817,7 +826,8 @@ def create_sampler( f"interface (needs admit/select/evict/should_abort_inflight, " f"dispatch_index, set_dispatch_index, restore_dispatch_index, " f"is_on_policy, supports_buffer_checkpoint, " - f"required_buffer_capacity)" + f"required_buffer_capacity; periodic rollout snapshots also " + f"require supports_training_claims=True)" ) else: raise ValueError(f"unknown sampler config {type(cfg).__name__}") diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index f906c64273b..cf1e039b31c 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -101,6 +101,7 @@ from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + BootstrapCompatibilityIdentity, RolloutSnapshotManifest, commit_snapshot, ensure_bootstrap_anchor, @@ -438,7 +439,9 @@ def __init__( self._checkpoint_save_lock = asyncio.Lock() self._last_rollout_snapshot_mutation_version: Optional[int] = None self._last_missing_rollout_snapshot_anchor: Optional[tuple[int, int]] = None - self._bootstrap_fingerprint = getattr(actor_args, "bootstrap_fingerprint", None) + self._bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = ( + actor_args.bootstrap_identity + ) self._rollout_checkpoint_stop_requested = asyncio.Event() # Narrow unsafe window after an optimizer mutates model state and before # SC publishes the matching TQ cleanup and trainer counters. Gradient @@ -550,7 +553,8 @@ async def run(self) -> dict[str, Any]: tasks = [rollout_task, train_task, watchdog_task] rollout_checkpoint_task = ( asyncio.create_task(self._rollout_checkpoint_pump()) - if self._master_config.rollout_checkpointing.interval_s is not None + if self._master_config.rollout_checkpointing.snapshot_attempt_interval_s + is not None else None ) if rollout_checkpoint_task is not None: @@ -1522,11 +1526,6 @@ async def _cleanup_consumed_metas_unlocked( if errors: raise BaseExceptionGroup("post-train DataPlane cleanup failed", errors) - async def _cleanup_consumed_metas(self, metas: list[KVBatchMeta]) -> None: - """Clear consumed rows without racing a native TQ checkpoint.""" - async with self._data_plane_checkpoint_barrier.mutation() as cut: - await self._cleanup_consumed_metas_unlocked(cut, metas) - @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: """Return stable prompt-group IDs in canonical sample order.""" @@ -3512,16 +3511,16 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: raise RuntimeError( "bootstrap rollout snapshot requires trainer version zero" ) - if self._bootstrap_fingerprint is None: + if self._bootstrap_identity is None: raise RuntimeError( - "rollout snapshotting requires a bootstrap fingerprint" + "rollout snapshotting requires a bootstrap identity" ) anchor = await asyncio.to_thread( ensure_bootstrap_anchor, self._checkpointer.checkpoint_dir, - fingerprint=self._bootstrap_fingerprint, + identity=self._bootstrap_identity, ) - snapshot_fingerprint = self._bootstrap_fingerprint + snapshot_fingerprint = self._bootstrap_identity.fingerprint() else: if self._trainer_version != self._train_steps: raise RuntimeError( @@ -3616,16 +3615,18 @@ async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: async def _rollout_checkpoint_pump(self) -> None: """Persist rollout state periodically, including during streamed train.""" - interval_s = self._master_config.rollout_checkpointing.interval_s - if interval_s is None: + snapshot_attempt_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") consecutive_failures = 0 while True: - await asyncio.sleep(interval_s) + await asyncio.sleep(snapshot_attempt_interval_s) deadline_due = self._train_steps == 0 and self._timeout.would_save() try: saved = await self._save_rollout_checkpoint(force=deadline_due) - except Exception as error: + except (OSError, TimeoutError) as error: if deadline_due: raise RuntimeError( "failed to save the required pre-step rollout checkpoint" diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index b86cdba230d..bc381273b09 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -717,16 +717,24 @@ def resolve_for_prompt( class RolloutCheckpointConfig(BaseModel, extra="forbid"): """Frequent rollout-state snapshots anchored to durable trainer state. - ``interval_s=None`` disables saving and restoring periodic snapshots. A - snapshot taken before the first trainer checkpoint is anchored to the - initial model and a rollout-semantic configuration fingerprint. Later + ``snapshot_attempt_interval_s=None`` disables saving and restoring periodic + snapshots. 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; interval attempts are skipped until that exact - anchor exists. + current completed step; attempts are skipped until that exact anchor exists. ``restore_mode="latest"`` selects the newest compatible periodic snapshot. ``trainer_checkpoint`` ignores newer periodic snapshots and restores the - rollout state bundled with the durable trainer checkpoint. + rollout state bundled with the durable trainer checkpoint. Restore + selection never deletes checkpoint state. If no trainer checkpoint exists, + ``trainer_checkpoint`` rejects an occupied bootstrap namespace; use + ``latest`` or a new checkpoint directory instead. + + Bootstrap compatibility is fail-closed: every configuration value affects + the fingerprint unless it is on the built-in operational denylist. + ``extra_fingerprint_excluded_paths`` lets integrations exclude additional + runtime-only dotpaths. ``*`` matches one mapping or list level and ``**`` + matches any number of levels. SingleController has no validation loop, so checkpoint selection must use ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited @@ -735,9 +743,28 @@ class RolloutCheckpointConfig(BaseModel, extra="forbid"): silently disable the durability behavior the operator intended. """ - interval_s: Annotated[Optional[float], Field(gt=0)] = None + snapshot_attempt_interval_s: Annotated[Optional[float], Field(gt=0)] = None 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) + + @model_validator(mode="after") + def validate_extra_fingerprint_excluded_paths(self) -> "RolloutCheckpointConfig": + """Reject ambiguous paths that could silently fail to exclude a value.""" + invalid = [ + path + for path in self.extra_fingerprint_excluded_paths + if not path + or path != path.strip() + or any(not segment for segment in path.split(".")) + or path in {"*", "**"} + ] + if invalid: + raise ValueError( + "extra_fingerprint_excluded_paths must contain non-empty dotpaths " + f"and cannot exclude the whole config, got {invalid!r}" + ) + return self class MasterConfig(BaseModel, extra="allow"): diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index 31a4b6df3b4..c7dec0c191c 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -22,218 +22,163 @@ import re import shutil from dataclasses import asdict, dataclass +from fnmatch import fnmatchcase from pathlib import Path from typing import Any, Mapping, Optional from nemo_rl.algorithms.single_controller_utils.config import MasterConfig -ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 2 -BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 4 +ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 3 +BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 7 BOOTSTRAP_DIRNAME = "bootstrap" BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" ROLLOUT_SNAPSHOT_MANIFEST_FILENAME = "manifest.json" -ROLLOUT_SNAPSHOT_COMMITTED_FILENAME = "COMMITTED" _SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") - -# Keep this projection limited to values needed to interpret persisted rollout -# state or execute missing siblings. Trainer-only and post-rollout settings do -# not belong in bootstrap compatibility. -_BOOTSTRAP_POLICY_FIELDS = frozenset( - { - "model_name", - "pretrained_checkpoint", - "hf_config_overrides", - "max_total_sequence_length", - "tokenizer", - } -) -_BOOTSTRAP_GENERATION_FIELDS = frozenset( - { - "backend", - "max_new_tokens", - "stop_strings", - "stop_token_ids", - "temperature", - "top_k", - "top_p", - } -) -_BOOTSTRAP_VLLM_FIELDS = frozenset( - { - "http_server_serving_chat_kwargs", - "max_model_len", - "reasoning_parser_plugin", - } -) -_BOOTSTRAP_GRPO_FIELDS = frozenset( - { - "max_rollout_turns", - "num_generations_per_prompt", - "num_prompts_per_step", - "seed", - } -) -_BOOTSTRAP_DATA_FIELDS = frozenset( - { - "default", - "max_input_seq_length", - "shuffle", - "train", - } -) -_BOOTSTRAP_TOKEN_CAPTURE_FIELDS = frozenset( - { - "defer_routed_experts_to_policy", - "enabled", - "min_valid_fraction_per_group", - "mixed_weight_version_policy", - "on_capture_failure", - "staging_partition", - } -) -# Deployment, placement, logging, and credential leaves do not change the -# meaning of persisted rollout state. The recursive filter keeps semantic -# neighbors such as model names, parser settings, timeouts, and reward config. -_BOOTSTRAP_ENV_RUNTIME_FIELDS = frozenset( - { - "_copy", - "_inherit_from", - "allow_openai_version_skew", - "api_key", - "api_server_count", - "apptainer_memory_limit_mb", - "cache_dir", - "component_name", - "concurrency", - "config_paths", - "debug", - "default_host", - "disallowed_ports", - "dry_run", - "entrypoint", - "global_aiohttp_connector_limit", - "global_aiohttp_connector_limit_per_host", - "head_server", - "head_server_deps", - "hf_token", - "json", - "model_call_capture_dir", - "model_endpoint_readiness_timeout_seconds", - "nemo_gym_log_dir", - "num_gpu_nodes", - "num_processes", - "num_workers", - "observability_enabled", - "pip_install_verbose", - "policy_base_url", - "port_range_high", - "port_range_low", - "python_version", - "query", - "ray_head_node_address", - "ray_worker_py_executable", - "results_dir", - "should_log_nemo_gym_responses", - "skip_venv_if_present", - "token_id_capture", - "use_absolute_ip", - "uv_cache_dir", - "uv_pip_set_python", - "uv_venv_dir", - "verbose", - } -) -# Gym is optional in NeMo-RL, so keep credential recognition local rather than -# importing its global-config helpers. Match credential-shaped leaves without -# dropping semantic neighbors such as ``max_tokens`` or ``tokenizer``. -_BOOTSTRAP_ENV_CREDENTIAL_FIELDS = frozenset( +_TMP_SNAPSHOT_RE = re.compile(r"tmp_snapshot_(\d+)") +_TRASH_SNAPSHOT_RE = re.compile(r"trash_snapshot_(\d+)") + +# A bootstrap fingerprint is fail-closed: every config value participates unless +# this one denylist says it is operational. ``**`` matches any number of mapping +# levels, which keeps nested runtime fields and credential-shaped keys concise. +# User-defined exclusions are appended from RolloutCheckpointConfig. +_BOOTSTRAP_FINGERPRINT_EXCLUDED_PATHS = frozenset( { - "api_key", - "apikey", - "password", - "secret", - "token", - } -) -_BOOTSTRAP_ENV_CREDENTIAL_FIELD_SUFFIXES = ( - "_api_key", - "_password", - "_secret", - "_token", -) -# These remote services keep their model identity in the fingerprint, but the -# endpoint may legitimately move when a Slurm job is restarted. -_BOOTSTRAP_ENV_RUNTIME_PATHS = frozenset( - { - ( - "nemo_gym", - "genrm_model", - "responses_api_models", - "genrm_model", - "base_url", - ), - ( - "nemo_gym", - "nl2bash_judge_model", - "responses_api_models", - "local_vllm_model", - "base_url", - ), + "async_rl.diagnostics", + "async_rl.generation_fleet_health", + "async_rl.generation_router", + "async_rl.stall_watchdog", + "checkpointing", + "cluster", + "data.num_workers", + "data.validation", + "logger", + "policy.generation.colocated", + "policy.generation.port_range_high", + "policy.generation.port_range_low", + "policy.generation.val_temperature", + "policy.generation.val_top_k", + "policy.generation.val_top_p", + "policy.generation.vllm_cfg.env_vars", + "policy.generation.vllm_cfg.http_refit_api_key_env_var", + "policy.generation.vllm_cfg.http_refit_server_port", + "policy.generation.vllm_cfg.zmq_refit_server_port", + "policy.optimizer", + "policy.scheduler", + "rollout_checkpointing", + "token_capture.capture_dir", + "token_capture.control_auth_token", + "token_capture.control_timeout_s", + "token_capture.num_reassembler_workers", + "**.api_key", + "**.apikey", + "**.password", + "**.secret", + "**.token", + "**.*_api_key", + "**.*_password", + "**.*_secret", + "**.*_token", + "env.**._copy", + "env.**._inherit_from", + "env.**.allow_openai_version_skew", + "env.**.api_server_count", + "env.**.apptainer_memory_limit_mb", + "env.**.cache_dir", + "env.**.component_name", + "env.**.concurrency", + "env.**.config_paths", + "env.**.debug", + "env.**.default_host", + "env.**.disallowed_ports", + "env.**.dry_run", + "env.**.entrypoint", + "env.**.global_aiohttp_connector_limit", + "env.**.global_aiohttp_connector_limit_per_host", + "env.**.head_server", + "env.**.head_server_deps", + "env.**.json", + "env.**.model_call_capture_dir", + "env.**.model_endpoint_readiness_timeout_seconds", + "env.**.nemo_gym_log_dir", + "env.**.num_gpu_nodes", + "env.**.num_processes", + "env.**.num_workers", + "env.**.observability_enabled", + "env.**.pip_install_verbose", + "env.**.policy_base_url", + "env.**.port_range_high", + "env.**.port_range_low", + "env.**.python_version", + "env.**.query", + "env.**.ray_head_node_address", + "env.**.ray_worker_py_executable", + "env.**.results_dir", + "env.**.should_log_nemo_gym_responses", + "env.**.skip_venv_if_present", + "env.**.token_id_capture", + "env.**.use_absolute_ip", + "env.**.uv_cache_dir", + "env.**.uv_pip_set_python", + "env.**.uv_venv_dir", + "env.**.verbose", + "env.nemo_gym.genrm_model.responses_api_models.genrm_model.base_url", + "env.nemo_gym.nl2bash_judge_model.responses_api_models.local_vllm_model.base_url", } ) -def _select_fields( - mapping: Mapping[str, Any] | None, - fields: frozenset[str], -) -> dict[str, Any]: - """Select explicitly rollout-semantic fields from one config section.""" - if mapping is None: - return {} - return {key: mapping[key] for key in sorted(fields) if key in mapping} - - -def _is_environment_credential_field(field: str) -> bool: - """Whether one environment leaf is credential-shaped, not semantic config.""" - normalized = field.casefold() - return normalized in _BOOTSTRAP_ENV_CREDENTIAL_FIELDS or normalized.endswith( - _BOOTSTRAP_ENV_CREDENTIAL_FIELD_SUFFIXES +def _path_matches(pattern: tuple[str, ...], path: tuple[str, ...]) -> bool: + """Return whether one segmented dotpath pattern matches a concrete path.""" + if not pattern: + return not path + if pattern[0] == "**": + return _path_matches(pattern[1:], path) or ( + bool(path) and _path_matches(pattern, path[1:]) + ) + return ( + bool(path) + and fnmatchcase(path[0], pattern[0]) + and _path_matches(pattern[1:], path[1:]) ) -def _drop_runtime_fields( +def _drop_excluded_paths( value: Any, - runtime_fields: frozenset[str], *, - runtime_paths: frozenset[tuple[str, ...]], + excluded_paths: tuple[tuple[str, ...], ...], path: tuple[str, ...] = (), ) -> Any: - """Recursively strip operational and credential leaves from an environment.""" + """Recursively remove denylisted mapping paths from a JSON config dump.""" if isinstance(value, Mapping): - return { - key: _drop_runtime_fields( + projected: dict[str, Any] = {} + for key, child in value.items(): + if not isinstance(key, str): + raise TypeError("fingerprinted config mappings must use string keys") + child_path = (*path, key) + if any(_path_matches(pattern, child_path) for pattern in excluded_paths): + continue + projected[key] = _drop_excluded_paths( child, - runtime_fields, - runtime_paths=runtime_paths, - path=(*path, key), + excluded_paths=excluded_paths, + path=child_path, ) - for key, child in value.items() - if key not in runtime_fields - and not _is_environment_credential_field(key) - and (*path, key) not in runtime_paths - } + return projected if isinstance(value, list): - return [ - _drop_runtime_fields( - child, - runtime_fields, - runtime_paths=runtime_paths, - path=path, + projected_list: list[Any] = [] + for index, child in enumerate(value): + child_path = (*path, str(index)) + if any(_path_matches(pattern, child_path) for pattern in excluded_paths): + continue + projected_list.append( + _drop_excluded_paths( + child, + excluded_paths=excluded_paths, + path=child_path, + ) ) - for child in value - ] + return projected_list return value @@ -254,7 +199,7 @@ def _fsync_directory(path: Path) -> None: def _fsync_tree(root: Path) -> None: - """Flush every snapshot payload before publishing its commit marker.""" + """Flush every snapshot payload before publishing its directory.""" for directory, _, filenames in os.walk(root, topdown=False): directory_path = Path(directory) for filename in filenames: @@ -356,63 +301,53 @@ class BootstrapCompatibilityIdentity: """Rollout-semantic inputs that must match a trainer-version-zero cut.""" schema_version: int - model: Mapping[str, Any] - generation: Mapping[str, Any] - rollout: Mapping[str, Any] - dataset: Mapping[str, Any] - environment: Mapping[str, Any] - sampler: Mapping[str, Any] - token_capture: Mapping[str, Any] + excluded_paths: tuple[str, ...] + config: Mapping[str, Any] def to_dict(self) -> dict[str, Any]: - return asdict(self) + return { + "schema_version": self.schema_version, + "excluded_paths": list(self.excluded_paths), + "config": self.config, + } + + def fingerprint(self) -> str: + """Return the canonical digest stored in snapshot manifests.""" + payload = json.dumps( + self.to_dict(), + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() def bootstrap_compatibility_identity( master_config: MasterConfig, ) -> BootstrapCompatibilityIdentity: - """Project a full run config onto inputs that affect recovered rollouts. - - Dataset identity is intentionally retained because a bootstrap snapshot - restores the dataloader cursor together with its unfinished prompt ledger. - Cluster shape, logging, checkpoint paths, worker counts, and other runtime - tuning are excluded so they may change across a restart. - """ + """Remove explicitly operational paths from the fail-closed run identity.""" dumped = master_config.model_dump(mode="json") - policy = dumped.get("policy", {}) - generation = policy.get("generation", {}) - generation_identity = _select_fields( - generation, - _BOOTSTRAP_GENERATION_FIELDS, + rollout_checkpointing = dumped.get("rollout_checkpointing", {}) + if not isinstance(rollout_checkpointing, Mapping): + raise TypeError("rollout_checkpointing config must be a mapping") + extra_excluded_paths = rollout_checkpointing.get( + "extra_fingerprint_excluded_paths", [] ) - vllm_identity = _select_fields( - generation.get("vllm_cfg", {}), - _BOOTSTRAP_VLLM_FIELDS, + if not isinstance(extra_excluded_paths, list) or not all( + isinstance(path, str) for path in extra_excluded_paths + ): + raise TypeError("extra_fingerprint_excluded_paths must be a list of strings") + excluded_paths = tuple( + sorted(_BOOTSTRAP_FINGERPRINT_EXCLUDED_PATHS | frozenset(extra_excluded_paths)) + ) + parsed_excluded_paths = tuple( + tuple(excluded_path.split(".")) for excluded_path in excluded_paths ) - if vllm_identity: - generation_identity["vllm_cfg"] = vllm_identity - - async_rl = dumped.get("async_rl", {}) - sampler = async_rl.get("sampler", {}) - if not isinstance(sampler, Mapping): - sampler = {} - - rollout = dumped.get("ppo") or dumped.get("grpo") or {} return BootstrapCompatibilityIdentity( schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, - model=_select_fields(policy, _BOOTSTRAP_POLICY_FIELDS), - generation=generation_identity, - rollout=_select_fields(rollout, _BOOTSTRAP_GRPO_FIELDS), - dataset=_select_fields(dumped.get("data", {}), _BOOTSTRAP_DATA_FIELDS), - environment=_drop_runtime_fields( - dumped.get("env", {}), - _BOOTSTRAP_ENV_RUNTIME_FIELDS, - runtime_paths=_BOOTSTRAP_ENV_RUNTIME_PATHS, - ), - sampler=dict(sampler), - token_capture=_select_fields( - dumped.get("token_capture", {}), - _BOOTSTRAP_TOKEN_CAPTURE_FIELDS, + excluded_paths=excluded_paths, + config=_drop_excluded_paths( + dumped, + excluded_paths=parsed_excluded_paths, ), ) @@ -424,12 +359,68 @@ def bootstrap_fingerprint(master_config: MasterConfig) -> str: Operational settings are deliberately excluded so a restart may use a different cluster shape, checkpoint interval, or logging destination. """ - payload = json.dumps( - bootstrap_compatibility_identity(master_config).to_dict(), - sort_keys=True, - separators=(",", ":"), - ).encode() - return hashlib.sha256(payload).hexdigest() + return bootstrap_compatibility_identity(master_config).fingerprint() + + +_MISSING = object() + + +def _format_changed_value(value: Any) -> str: + """Format one redacted compatibility value without flooding an error.""" + if value is _MISSING: + return "" + rendered = repr(value) + if len(rendered) > 160: + return rendered[:157] + "..." + return rendered + + +def _compatibility_differences( + checkpoint: Any, + expected: Any, + *, + path: tuple[str, ...] = (), +) -> list[str]: + """Describe changed compatibility leaves using user-facing dotpaths.""" + if isinstance(checkpoint, Mapping) and isinstance(expected, Mapping): + differences: list[str] = [] + for key in sorted(set(checkpoint) | set(expected)): + if key == "bootstrap_fingerprint" and not path: + continue + differences.extend( + _compatibility_differences( + checkpoint.get(key, _MISSING), + expected.get(key, _MISSING), + path=(*path, key), + ) + ) + return differences + if checkpoint == expected: + return [] + + display_path = path + if display_path[:2] == ("bootstrap_identity", "config"): + display_path = display_path[2:] + elif display_path[:1] == ("bootstrap_identity",): + display_path = display_path[1:] + name = ".".join(display_path) or "bootstrap_identity" + return [ + f"{name}: {_format_changed_value(checkpoint)} -> " + f"{_format_changed_value(expected)}" + ] + + +def _bootstrap_anchor_manifest( + identity: BootstrapCompatibilityIdentity, +) -> dict[str, Any]: + """Build the bootstrap manifest from one self-consistent identity.""" + return { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "bootstrap_fingerprint": identity.fingerprint(), + "bootstrap_identity": identity.to_dict(), + } def prune_bootstrap_snapshots( @@ -450,24 +441,18 @@ def prune_bootstrap_snapshots( return True -def ensure_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: +def ensure_bootstrap_anchor( + checkpoint_dir: Path, + *, + identity: BootstrapCompatibilityIdentity, +) -> Path: """Create or validate the lightweight trainer-version-zero anchor.""" anchor = checkpoint_dir / BOOTSTRAP_DIRNAME anchor.mkdir(parents=True, exist_ok=True) manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME - expected = { - "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, - "base_train_step": 0, - "trainer_version": 0, - "bootstrap_fingerprint": fingerprint, - } + expected = _bootstrap_anchor_manifest(identity) if manifest_path.is_file(): - raw = json.loads(manifest_path.read_text()) - if raw != expected: - raise ValueError( - "existing rollout bootstrap anchor does not match the current " - f"trainer configuration: checkpoint={raw!r}, expected={expected!r}" - ) + validate_bootstrap_anchor(anchor, identity=identity) return anchor tmp_path = manifest_path.with_suffix(".json.tmp") @@ -479,28 +464,63 @@ def ensure_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: return anchor -def reset_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: - """Discard skipped pre-step snapshots and start a new bootstrap lineage. - - This is used only when restore mode deliberately skips partial-rollout - recovery and no trainer checkpoint exists. The user's restore choice makes - the previous state intentionally unreachable; removing it also prevents a - later periodic save from appending to an incompatible bootstrap anchor. - """ - anchor = checkpoint_dir / BOOTSTRAP_DIRNAME - snapshot_root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME - if snapshot_root.exists(): - shutil.rmtree(snapshot_root) +def validate_bootstrap_anchor( + anchor: Path, + *, + identity: BootstrapCompatibilityIdentity, +) -> None: + """Validate a bootstrap anchor without modifying checkpoint state.""" manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME - if manifest_path.exists(): - manifest_path.unlink() - return ensure_bootstrap_anchor(checkpoint_dir, fingerprint=fingerprint) + if not manifest_path.is_file(): + raise FileNotFoundError( + f"rollout bootstrap manifest is missing at {manifest_path}" + ) + raw = json.loads(manifest_path.read_text()) + if not isinstance(raw, Mapping): + raise ValueError( + f"rollout bootstrap manifest at {manifest_path} must be a mapping" + ) + expected = _bootstrap_anchor_manifest(identity) + if raw != expected: + differences = _compatibility_differences(raw, expected) + if not differences: + differences = [ + "bootstrap_fingerprint: checkpoint digest does not match its " + "persisted compatibility identity" + ] + visible = differences[:20] + if len(differences) > len(visible): + visible.append(f"... and {len(differences) - len(visible)} more change(s)") + details = "\n".join(f" {difference}" for difference in visible) + raise ValueError( + f"rollout bootstrap anchor at {manifest_path} is incompatible " + "with the current rollout-semantic configuration. Changed " + f"compatibility fields:\n{details}\n" + "If a changed field is operational only, list its dotpath in " + "rollout_checkpointing.extra_fingerprint_excluded_paths in both " + "the original and restarted configurations. Otherwise reuse the " + "original configuration or choose a new checkpoint_dir. Existing " + "checkpoint state was not modified." + ) def prepare_snapshot_paths(anchor: Path) -> tuple[Path, Path, int]: """Allocate the next temporary/final snapshot directory pair.""" root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME root.mkdir(parents=True, exist_ok=True) + garbage = [ + child + for child in root.iterdir() + if child.is_dir() + and ( + _TMP_SNAPSHOT_RE.fullmatch(child.name) + or _TRASH_SNAPSHOT_RE.fullmatch(child.name) + ) + ] + for child in garbage: + shutil.rmtree(child) + if garbage: + _fsync_directory(root) sequences = [ int(match.group(1)) for child in root.iterdir() @@ -525,10 +545,6 @@ def commit_snapshot( if keep_latest_k < 1: raise ValueError("rollout snapshot retention must keep at least one snapshot") _fsync_tree(tmp_path) - committed_path = tmp_path / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME - committed_path.write_text("committed\n") - _fsync_file(committed_path) - _fsync_directory(tmp_path) os.rename(tmp_path, final_path) root = final_path.parent @@ -538,17 +554,19 @@ def commit_snapshot( ( child for child in root.iterdir() - if child.is_dir() - and _SNAPSHOT_RE.fullmatch(child.name) - and (child / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() + if child.is_dir() and _SNAPSHOT_RE.fullmatch(child.name) ), key=_snapshot_sequence, reverse=True, ) stale_snapshots = committed[keep_latest_k:] for stale in stale_snapshots: - shutil.rmtree(stale) - if stale_snapshots: + trash = root / f"trash_{stale.name}" + if trash.exists(): + shutil.rmtree(trash) + os.rename(stale, trash) + _fsync_directory(root) + shutil.rmtree(trash) _fsync_directory(root) @@ -575,8 +593,6 @@ def resolve_latest_snapshot( ) errors: list[str] = [] for candidate in candidates: - if not (candidate / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file(): - continue manifest_path = candidate / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME if not manifest_path.is_file(): errors.append(f"{candidate.name}: missing manifest") diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index ebffdbc2699..de41845cf10 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -71,11 +71,10 @@ ) from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( BOOTSTRAP_DIRNAME, - ROLLOUT_SNAPSHOTS_DIRNAME, - bootstrap_fingerprint, - ensure_bootstrap_anchor, - reset_bootstrap_anchor, + BootstrapCompatibilityIdentity, + bootstrap_compatibility_identity, resolve_latest_snapshot, + validate_bootstrap_anchor, ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn @@ -164,7 +163,7 @@ class SingleControllerActorArgs: finalizer_actors: list[Any] # Defaulted fields must follow the required ones above, so these stay last. data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None - bootstrap_fingerprint: Optional[str] = None + bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = None # None when async_rl.generation_fleet_health is disabled; the SingleController # drives the probe loop when it is present. fleet_monitor: Optional[GenerationFleetHealth] = None @@ -954,7 +953,7 @@ def setup_single_controller( rollout_checkpoint_cfg = master_config.rollout_checkpointing if ( master_config.checkpointing.get("save_data_plane") - or rollout_checkpoint_cfg.interval_s is not None + or rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None ) and not data_plane_checkpointing_supported: raise NotImplementedError( "SingleController data-plane checkpointing is not supported for " @@ -1014,7 +1013,7 @@ def setup_single_controller( # ray_actor_environment_registry.py), so nothing here needs to change the # worker's environment. token_capture_cfg = master_config.token_capture - if rollout_checkpoint_cfg.interval_s is not None: + if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None: if not master_config.checkpointing["enabled"]: raise ValueError( "rollout checkpointing requires checkpointing.enabled=true" @@ -1115,10 +1114,15 @@ def setup_single_controller( bootstrap_anchor = checkpointer.checkpoint_dir / BOOTSTRAP_DIRNAME needs_bootstrap_identity = ( trainer_checkpoint_path is None - and rollout_checkpoint_cfg.interval_s is not None + and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None + ) + bootstrap_identity = ( + bootstrap_compatibility_identity(master_config) + if needs_bootstrap_identity + else None ) bootstrap_digest = ( - bootstrap_fingerprint(master_config) if needs_bootstrap_identity else None + bootstrap_identity.fingerprint() if bootstrap_identity is not None else None ) resolved_snapshot = None restored_trainer_version = ( @@ -1128,7 +1132,7 @@ def setup_single_controller( ) if ( trainer_checkpoint_path is not None - and rollout_checkpoint_cfg.interval_s is not None + and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None and restore_mode == "latest" ): resolved_snapshot = resolve_latest_snapshot( @@ -1138,30 +1142,24 @@ def setup_single_controller( expected_bootstrap_fingerprint=None, ) elif trainer_checkpoint_path is None: - if rollout_checkpoint_cfg.interval_s is not None: + if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None: assert bootstrap_digest is not None - if restore_mode == "latest": - bootstrap_anchor = ensure_bootstrap_anchor( - checkpointer.checkpoint_dir, - fingerprint=bootstrap_digest, - ) - else: - had_bootstrap_snapshots = ( - bootstrap_anchor / ROLLOUT_SNAPSHOTS_DIRNAME - ).is_dir() - bootstrap_anchor = reset_bootstrap_anchor( - checkpointer.checkpoint_dir, - fingerprint=bootstrap_digest, + assert bootstrap_identity is not None + if bootstrap_anchor.is_dir(): + validate_bootstrap_anchor( + bootstrap_anchor, + identity=bootstrap_identity, ) - if had_bootstrap_snapshots: - print( - "📦 Ignored existing bootstrap rollout snapshots and " - "started a new bootstrap lineage because " - f"rollout_checkpointing.restore_mode={restore_mode!r}.", - flush=True, + if restore_mode == "trainer_checkpoint": + raise ValueError( + "rollout_checkpointing.restore_mode='trainer_checkpoint' " + "cannot start a fresh bootstrap lineage because checkpoint " + f"state already exists at {bootstrap_anchor}. Use " + "restore_mode='latest' to recover it or choose a new " + "checkpoint_dir. Existing checkpoint state was not modified." ) if ( - rollout_checkpoint_cfg.interval_s is not None + rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None and restore_mode == "latest" and bootstrap_anchor.is_dir() ): @@ -1734,7 +1732,7 @@ def _build_generation_then_trainer( save_state=save_state, last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, - bootstrap_fingerprint=bootstrap_digest, + bootstrap_identity=bootstrap_identity, finalizer_actors=finalizer_actors, fleet_monitor=fleet_monitor, generation_router=generation_router, diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index 53ed80f1eac..69a07af0048 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -55,7 +55,7 @@ COMMON_OVERRIDES=( +checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling - ++rollout_checkpointing.interval_s="$SNAPSHOT_INTERVAL_S" + ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" ++rollout_checkpointing.keep_latest_k=8 ++rollout_checkpointing.restore_mode=latest async_rl.sampler.name=in_order @@ -102,7 +102,7 @@ deadline = time.monotonic() + float(sys.argv[6]) while time.monotonic() < deadline: for snapshot in sorted(root.glob("snapshot_*"), reverse=True): manifest_path = snapshot / "manifest.json" - if not (snapshot / "COMMITTED").is_file() or not manifest_path.is_file(): + if not manifest_path.is_file(): continue manifest = json.loads(manifest_path.read_text()) if ( diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 491be53a79e..f68b9db2e6c 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -82,10 +82,11 @@ from nemo_rl.algorithms.single_controller_utils.config import TokenCaptureConfig from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( BOOTSTRAP_DIRNAME, - ROLLOUT_SNAPSHOT_COMMITTED_FILENAME, ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + BootstrapCompatibilityIdentity, RolloutSnapshotManifest, + bootstrap_compatibility_identity, commit_snapshot, prepare_snapshot_paths, ) @@ -564,7 +565,7 @@ def _actor_master_config( max_num_epochs: int = 1, buffer_checkpoint: bool = False, data_plane_checkpoint: bool = True, - rollout_checkpoint_interval_s: Optional[float] = None, + rollout_checkpoint_attempt_interval_s: Optional[float] = None, token_capture_enabled: bool = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -626,7 +627,7 @@ def _actor_master_config( max_buffered_rollouts=4, ), rollout_checkpointing=RolloutCheckpointConfig( - interval_s=rollout_checkpoint_interval_s + snapshot_attempt_interval_s=rollout_checkpoint_attempt_interval_s ), token_capture=TokenCaptureConfig(enabled=token_capture_enabled), ) @@ -641,7 +642,7 @@ def _make_actor_args( dp_client: Optional[_FakeDPClient] = None, last_checkpoint_path: Optional[str] = None, data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None, - bootstrap_fingerprint: Optional[str] = None, + bootstrap_identity: Optional[BootstrapCompatibilityIdentity] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=_FakeGeneration(), @@ -663,7 +664,7 @@ def _make_actor_args( last_checkpoint_path=last_checkpoint_path, finalizer_actors=[], data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, - bootstrap_fingerprint=bootstrap_fingerprint, + bootstrap_identity=bootstrap_identity, ) @@ -1103,7 +1104,8 @@ def test_restore_mode_rejects_removed_none_value(self): @pytest.mark.parametrize( "config", [ - {"interval_s": 0}, + {"snapshot_attempt_interval_s": 0}, + {"interval_s": 1}, {"keep_latest_k": 0}, {"unknown_option": True}, ], @@ -1116,12 +1118,14 @@ def _actor(self, tmp_path: Path): config = _actor_master_config( tmp_path, buffer_checkpoint=True, - rollout_checkpoint_interval_s=120.0, + rollout_checkpoint_attempt_interval_s=120.0, token_capture_enabled=True, ) return _ACTOR_CLS( config, - _make_actor_args(bootstrap_fingerprint="bootstrap-digest"), + _make_actor_args( + bootstrap_identity=bootstrap_compatibility_identity(config) + ), SetupTimingMetrics(), ) @@ -1140,7 +1144,6 @@ def test_pre_step_snapshot_contains_only_rollout_state(self, tmp_path: Path): / "rollout_snapshots" / "snapshot_000001" ) - assert (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() assert (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).is_file() manifest = json.loads( (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).read_text() @@ -1213,7 +1216,7 @@ def test_periodic_pump_reports_each_consecutive_failure( capsys: pytest.CaptureFixture[str], ): actor = self._actor(tmp_path) - actor._master_config.rollout_checkpointing.interval_s = 0.001 + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 actor._train_steps = 1 async def _main() -> None: @@ -1250,7 +1253,7 @@ def test_periodic_pump_aborts_after_repeated_failures( capsys: pytest.CaptureFixture[str], ): actor = self._actor(tmp_path) - actor._master_config.rollout_checkpointing.interval_s = 0.001 + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 actor._train_steps = 1 async def _main() -> None: @@ -1276,6 +1279,32 @@ async def _failing_save(*, force: bool = False) -> bool: output = capsys.readouterr().out assert output.count("Periodic rollout checkpoint failed") == 3 + def test_periodic_pump_does_not_retry_invariant_failure(self, tmp_path: Path): + actor = self._actor(tmp_path) + actor._master_config.rollout_checkpointing.snapshot_attempt_interval_s = 0.001 + calls = 0 + + async def _main() -> None: + async def _failing_save(*, force: bool = False) -> bool: + nonlocal calls + del force + calls += 1 + raise RuntimeError("broken checkpoint invariant") + + actor._save_rollout_checkpoint = _failing_save + with pytest.raises(RuntimeError, match="broken checkpoint invariant"): + await asyncio.wait_for( + actor._rollout_checkpoint_pump(), + timeout=1.0, + ) + + try: + asyncio.run(_main()) + finally: + actor._checkpointer.shutdown() + + assert calls == 1 + class TestDataPlaneCheckpoint: def test_metadata_uses_explicit_snapshot_identity(self, tmp_path): @@ -1580,9 +1609,13 @@ async def _main() -> None: started = await asyncio.to_thread(dp_client.save_started.wait, 30.0) assert started - clear_task = asyncio.create_task( - actor._cleanup_consumed_metas([_consumed_meta("sample-0")]) - ) + async def _clear() -> None: + async with actor._data_plane_checkpoint_barrier.mutation() as cut: + await actor._cleanup_consumed_metas_unlocked( + cut, [_consumed_meta("sample-0")] + ) + + clear_task = asyncio.create_task(_clear()) await asyncio.sleep(0) assert dp_client.clear_calls == [] @@ -1603,7 +1636,10 @@ async def _main() -> int: mc, _make_actor_args(dp_client=dp_client), SetupTimingMetrics() ) event_loop_thread_id = threading.get_ident() - await actor._cleanup_consumed_metas([_consumed_meta("sample-0")]) + async with actor._data_plane_checkpoint_barrier.mutation() as cut: + await actor._cleanup_consumed_metas_unlocked( + cut, [_consumed_meta("sample-0")] + ) actor._checkpointer.shutdown() return event_loop_thread_id @@ -2068,7 +2104,9 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( final_snapshot = _write_periodic_snapshot(step_3) mc = _setup_master_config(str(ckpt_dir)) mc.checkpointing["save_period"] = 1 - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=120.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=120.0 + ) mc.token_capture = TokenCaptureConfig(enabled=True) mc.policy["generation"].update( { @@ -2079,7 +2117,7 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( "vllm_cfg": {"async_engine": True}, } ) - mc.logger = {"log_dir": str(tmp_path / "logs")} + mc.logger["log_dir"] = str(tmp_path / "logs") patched_factories["setup_response_data"].return_value = ( list(range(8)), None, @@ -2099,7 +2137,8 @@ def test_periodic_snapshot_restores_exact_dispatch_cursor( return_value=False, ), patch( - "nemo_rl.experience.finalizer_actor.create_finalizer_actors", + "nemo_rl.experience.rollout_reassembler_actor." + "create_rollout_reassembler_actors", return_value=[MagicMock(name="finalizer")], ), ): @@ -2124,7 +2163,7 @@ def test_disabled_periodic_checkpointing_uses_trainer_anchor( _write_periodic_snapshot(step_3) mc = _setup_master_config(str(ckpt_dir)) mc.rollout_checkpointing = RolloutCheckpointConfig( - interval_s=None, + snapshot_attempt_interval_s=None, restore_mode="latest", ) diff --git a/tests/unit/single_controller/test_finalizer_lifecycle.py b/tests/unit/single_controller/test_finalizer_lifecycle.py index 9cbfb2f661b..049e68855d4 100644 --- a/tests/unit/single_controller/test_finalizer_lifecycle.py +++ b/tests/unit/single_controller/test_finalizer_lifecycle.py @@ -260,7 +260,11 @@ def test_post_train_cleanup_clears_canonical_rows_and_route_plan_staging_keys() ) ctrl = _controller(SimpleNamespace()) - asyncio.run(ctrl._cleanup_consumed_metas([meta])) + async def _cleanup() -> None: + async with ctrl._data_plane_checkpoint_barrier.mutation() as cut: + await ctrl._cleanup_consumed_metas_unlocked(cut, [meta]) + + asyncio.run(_cleanup()) assert ctrl._dp_client.clear_calls == [ { diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 8d6cad80d7f..985bb020bdc 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -187,7 +187,8 @@ def _resolved(name: str) -> dict: ids=["grpo", "ppo"], ) def test_the_exemplars_still_validate(self, name): - MasterConfig(**self._resolved(name)) + config = MasterConfig(**self._resolved(name)) + assert config.checkpointing["save_period"] == 1 def test_rejects_a_config_with_no_algorithm_block(self): resolved = self._resolved("grpo_math_1B_megatron_single_controller.yaml") diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index a8dfd05c780..9c7cf81e31b 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -19,12 +19,19 @@ import pytest -from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint -from nemo_rl.algorithms.single_controller_utils.config import TokenCaptureConfig +from nemo_rl.algorithms.single_controller_utils.config import ( + AsyncRLConfig, + MasterConfig, + RolloutCheckpointConfig, + TokenCaptureConfig, +) from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + BOOTSTRAP_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + BootstrapCompatibilityIdentity, RolloutSnapshotManifest, bootstrap_compatibility_identity, bootstrap_fingerprint, @@ -32,12 +39,11 @@ ensure_bootstrap_anchor, prepare_snapshot_paths, prune_bootstrap_snapshots, - reset_bootstrap_anchor, resolve_latest_snapshot, + validate_bootstrap_anchor, ) from nemo_rl.data import DataConfig -from nemo_rl.models.generation.interfaces import GenerationConfig -from nemo_rl.models.generation.vllm.config import VllmSpecificArgs +from nemo_rl.models.generation.vllm.config import VllmConfig, VllmSpecificArgs from nemo_rl.models.policy import PolicyConfig @@ -50,179 +56,23 @@ def model_dump(self, *, mode: str) -> dict[str, Any]: return self._dumped -_BOOTSTRAP_PROJECTION_CASES = ( - ( - rollout_checkpoint._BOOTSTRAP_POLICY_FIELDS, - frozenset( - { - "disable_modelopt_layer_spec", - "draft", - "dtensor_cfg", - "dynamic_batching", - "generation", - "generation_batch_size", - "is_vlm", - "logprob_batch_size", - "logprob_chunk_size", - "make_sequence_length_divisible_by", - "max_grad_norm", - "megatron_cfg", - "optimizer", - "precision", - "quant_batch_size", - "quant_calib_data", - "quant_calib_size", - "quant_cfg", - "quant_sequence_length", - "refit_buffer_size_gb", - "reward_model_cfg", - "router_replay", - "scheduler", - "sequence_packing", - "train_global_batch_size", - "train_micro_batch_size", - } - ), - PolicyConfig, - ), - ( - rollout_checkpoint._BOOTSTRAP_GENERATION_FIELDS, - frozenset( - { - "_debug_payload_metrics", - "_mtp_weights_from_refit", - "_pad_token_id", - "bad_words", - "colocated", - "model_name", - "port_range_high", - "port_range_low", - "use_async_rollouts", - "val_temperature", - "val_top_k", - "val_top_p", - } - ), - GenerationConfig, - ), - ( - rollout_checkpoint._BOOTSTRAP_VLLM_FIELDS, - frozenset( - { - "async_engine", - "cap_max_tokens_to_context", - "enable_return_routed_experts", - "enforce_eager", - "env_vars", - "expert_parallel_size", - "expose_http_server", - "gpu_memory_utilization", - "http_refit_api_key_env_var", - "http_refit_server_port", - "is_mx", - "kv_cache_dtype", - "load_format", - "logprobs_mode", - "pipeline_parallel_size", - "precision", - "quantization_ignore_patterns", - "quantization_ignored_layer_kws", - "reset_encoder_cache_after_weight_update", - "skip_tokenizer_init", - "tensor_parallel_size", - "tool_parser_plugin", - "use_tqdm", - "video", - "zmq_refit_server_port", - } - ), - VllmSpecificArgs, - ), - ( - rollout_checkpoint._BOOTSTRAP_GRPO_FIELDS, - frozenset( - { - "adv_estimator", - "advantage_clip_high", - "advantage_clip_low", - "async_grpo", - "batch_multiplier", - "calculate_advantages_on_gpu", - "debug_payload_metrics", - "deduplicate_multimodal_data", - "dynamic_sampling_max_gen_batches", - "invalid_tool_call_advantage", - "malformed_thinking_advantage", - "max_num_epochs", - "max_num_steps", - "max_val_samples", - "normalize_rewards", - "overlong_filtering", - "reward_scaling", - "reward_shaping", - "seq_logprob_error_threshold", - "skip_reference_policy_logprobs_calculation", - "stop_at_validation_metric", - "stop_at_validation_threshold", - "use_dynamic_sampling", - "use_leave_one_out_baseline", - "val_at_end", - "val_at_start", - "val_batch_size", - "val_num_generations_per_prompt", - "val_period", - "val_start_at", - } - ), - GRPOConfig, - ), - ( - rollout_checkpoint._BOOTSTRAP_DATA_FIELDS, - frozenset( - { - "add_bos", - "add_eos", - "add_generation_prompt", - "add_system_prompt", - "custom_dataloader", - "num_prompts_per_dataloader", - "num_workers", - "use_multiple_dataloader", - "validation", - } - ), - DataConfig, - ), - ( - rollout_checkpoint._BOOTSTRAP_TOKEN_CAPTURE_FIELDS, - frozenset( - { - "capture_dir", - "control_auth_token", - "control_timeout_s", - "num_finalizer_workers", - } - ), - TokenCaptureConfig, - ), -) +def _test_bootstrap_identity(label: str = "v1") -> BootstrapCompatibilityIdentity: + return BootstrapCompatibilityIdentity( + schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + excluded_paths=(), + config={"test_identity": label}, + ) -@pytest.mark.parametrize( - ("declared", "ignored", "schema"), - [ - pytest.param(declared, ignored, schema, id=schema.__name__) - for declared, ignored, schema in _BOOTSTRAP_PROJECTION_CASES - ], -) -def test_bootstrap_projection_covers_config_schema(declared, ignored, schema): - fields = ( - set(schema.model_fields) - if hasattr(schema, "model_fields") - else set(schema.__annotations__) +def _test_bootstrap_fingerprint(label: str = "v1") -> str: + return _test_bootstrap_identity(label).fingerprint() + + +def _ensure_test_bootstrap_anchor(tmp_path: Path, label: str = "v1") -> Path: + return ensure_bootstrap_anchor( + tmp_path, + identity=_test_bootstrap_identity(label), ) - assert declared.isdisjoint(ignored) - assert declared | ignored == fields def _commit_snapshot( @@ -230,8 +80,10 @@ def _commit_snapshot( *, mutation_version: int, trainer_version: int = 0, - fingerprint: str | None = "fingerprint-v1", + fingerprint: str | None = None, ): + if fingerprint is None: + fingerprint = _test_bootstrap_fingerprint() tmp_path, final_path, _ = prepare_snapshot_paths(anchor) manifest = RolloutSnapshotManifest( schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, @@ -251,14 +103,80 @@ def _commit_snapshot( def test_bootstrap_anchor_rejects_different_initial_state(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") - assert ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") == anchor + anchor = _ensure_test_bootstrap_anchor(tmp_path) + assert _ensure_test_bootstrap_anchor(tmp_path) == anchor + manifest = json.loads((anchor / BOOTSTRAP_MANIFEST_FILENAME).read_text()) + assert manifest["bootstrap_identity"] == _test_bootstrap_identity().to_dict() + + with pytest.raises( + ValueError, + match="Existing checkpoint state was not modified", + ) as error: + _ensure_test_bootstrap_anchor(tmp_path, "v2") + + assert "test_identity: 'v1' -> 'v2'" in str(error.value) + assert "bootstrap_fingerprint" not in str(error.value) + + +def test_validate_bootstrap_anchor_is_read_only(tmp_path: Path) -> None: + identity = _test_bootstrap_identity() + anchor = ensure_bootstrap_anchor(tmp_path, identity=identity) + snapshot = anchor / "rollout_snapshots" / "snapshot_000001" + snapshot.mkdir(parents=True) + payload = snapshot / "payload" + payload.write_text("preserve me") + + before = { + path.relative_to(anchor): path.read_bytes() + for path in anchor.rglob("*") + if path.is_file() + } + validate_bootstrap_anchor(anchor, identity=identity) + after = { + path.relative_to(anchor): path.read_bytes() + for path in anchor.rglob("*") + if path.is_file() + } + + assert after == before + + +def test_bootstrap_anchor_mismatch_names_redacted_config_paths(tmp_path: Path) -> None: + original = _DumpedConfig( + { + "policy": {"generation": {"temperature": 1.0}}, + "env": {"service": {"api_key": "secret-one"}}, + "custom_algo": {"service_token": "custom-secret-one"}, + } + ) + changed = _DumpedConfig( + { + "policy": {"generation": {"temperature": 0.7}}, + "env": {"service": {"api_key": "secret-two"}}, + "custom_algo": {"service_token": "custom-secret-two"}, + } + ) + anchor = ensure_bootstrap_anchor( + tmp_path, + identity=bootstrap_compatibility_identity(cast(Any, original)), + ) + + with pytest.raises(ValueError) as error: + validate_bootstrap_anchor( + anchor, + identity=bootstrap_compatibility_identity(cast(Any, changed)), + ) - with pytest.raises(ValueError, match="does not match"): - ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v2") + message = str(error.value) + assert "policy.generation.temperature: 1.0 -> 0.7" in message + assert "secret-one" not in message + assert "secret-two" not in message + manifest = (anchor / BOOTSTRAP_MANIFEST_FILENAME).read_text() + assert "secret-one" not in manifest + assert "custom-secret-one" not in manifest -def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: +def test_bootstrap_fingerprint_ignores_default_operational_paths() -> None: base = { "policy": { "model_name": "model-a", @@ -267,41 +185,34 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: "backend": "vllm", "temperature": 1.0, "colocated": {"enabled": True, "resources": {"gpus": 8}}, - "vllm_cfg": { - "kv_cache_dtype": "auto", - "precision": "bfloat16", - "skip_tokenizer_init": False, - }, + "port_range_low": 3000, + "port_range_high": 4000, }, }, "data": { "train": [{"data_path": "/datasets/train.jsonl"}], "num_workers": 4, }, - "grpo": { - "num_generations_per_prompt": 4, - "max_num_steps": 10, - "batch_multiplier": 1, - "use_dynamic_sampling": False, - "reward_shaping": {"enabled": False}, - "reward_scaling": {"enabled": False}, - }, - "loss_fn": { - "reference_policy_kl_penalty": 0.01, - "use_kl_in_reward": False, - }, - "reward_penalties": {"penalize_unwanted_tokens": False}, + "grpo": {"num_generations_per_prompt": 4}, "token_capture": { + "capture_dir": "/run/one/capture", + "control_auth_token": "secret-one", "control_timeout_s": 60.0, "enabled": True, - "num_finalizer_workers": 2, + "num_reassembler_workers": 2, "on_capture_failure": "continue", "staging_partition": "rollout_staging", }, + "async_rl": { + "sampler": {"name": "windowed", "max_staleness_versions": 1}, + "stall_watchdog": {"interval_s": 30}, + }, + "checkpointing": {"checkpoint_dir": "/run/one/checkpoints"}, + "rollout_checkpointing": {"snapshot_attempt_interval_s": 120}, "cluster": {"num_nodes": 2}, "logger": {"log_dir": "/run/one"}, } - compatible_changed = { + operationally_changed = { **base, "policy": { **base["policy"], @@ -309,42 +220,150 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: "generation": { **base["policy"]["generation"], "colocated": {"enabled": False, "resources": {"gpus": 16}}, - "vllm_cfg": { - "kv_cache_dtype": "fp8", - "precision": "float16", - "skip_tokenizer_init": True, - }, + "port_range_low": 5000, + "port_range_high": 6000, }, }, "data": {**base["data"], "num_workers": 16}, - "grpo": { - **base["grpo"], - "max_num_steps": 100, - "batch_multiplier": 2, - "use_dynamic_sampling": True, - "reward_shaping": {"enabled": True}, - "reward_scaling": {"enabled": True}, - }, - "loss_fn": { - "reference_policy_kl_penalty": 0.1, - "use_kl_in_reward": True, - }, - "reward_penalties": {"penalize_unwanted_tokens": True}, "token_capture": { **base["token_capture"], + "capture_dir": "/run/two/capture", + "control_auth_token": "secret-two", "control_timeout_s": 15.0, - "num_finalizer_workers": 8, + "num_reassembler_workers": 8, }, + "async_rl": { + **base["async_rl"], + "stall_watchdog": {"interval_s": 5}, + }, + "checkpointing": {"checkpoint_dir": "/run/two/checkpoints"}, + "rollout_checkpointing": {"snapshot_attempt_interval_s": 300}, "cluster": {"num_nodes": 8}, "logger": {"log_dir": "/run/two"}, } fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) assert fingerprint == bootstrap_fingerprint( - cast(Any, _DumpedConfig(compatible_changed)) + cast(Any, _DumpedConfig(operationally_changed)) + ) + + +def test_bootstrap_fingerprint_includes_unknown_config_by_default() -> None: + base = {"custom_algo": {"semantic_setting": "one"}} + changed = {"custom_algo": {"semantic_setting": "two"}} + + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(changed))) + ) + + +def test_bootstrap_fingerprint_honors_extra_excluded_dotpaths() -> None: + base = { + "custom_algo": { + "semantic_setting": "same", + "runtime": {"endpoint": "host-one", "port": 1234}, + }, + "rollout_checkpointing": { + "extra_fingerprint_excluded_paths": ["custom_algo.runtime"] + }, + } + runtime_changed = { + **base, + "custom_algo": { + **base["custom_algo"], + "runtime": {"endpoint": "host-two", "port": 5678}, + }, + } + semantic_changed = { + **base, + "custom_algo": { + **base["custom_algo"], + "semantic_setting": "different", + }, + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(runtime_changed)) + ) + assert fingerprint != bootstrap_fingerprint( + cast(Any, _DumpedConfig(semantic_changed)) + ) + identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) + assert identity.config["custom_algo"] == {"semantic_setting": "same"} + assert "custom_algo.runtime" in identity.excluded_paths + + +def test_bootstrap_fingerprint_extra_excluded_dotpaths_support_lists() -> None: + base = { + "custom_algo": { + "workers": [ + {"name": "one", "log_dir": "/run/one"}, + {"name": "two", "log_dir": "/run/two"}, + ] + }, + "rollout_checkpointing": { + "extra_fingerprint_excluded_paths": ["custom_algo.workers.*.log_dir"] + }, + } + changed = { + **base, + "custom_algo": { + "workers": [ + {"name": "one", "log_dir": "/other/one"}, + {"name": "two", "log_dir": "/other/two"}, + ] + }, + } + + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) == ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(changed))) ) +@pytest.mark.parametrize( + "path", + ["", " custom.path", "custom.path ", ".x", "x.", "x..y", "*", "**"], +) +def test_rollout_checkpoint_config_rejects_invalid_extra_excluded_path( + path: str, +) -> None: + with pytest.raises(ValueError, match="extra_fingerprint_excluded_paths"): + RolloutCheckpointConfig(extra_fingerprint_excluded_paths=[path]) + + +def test_builtin_fingerprint_exclusions_reference_declared_config_fields() -> None: + """Keep typed portions of the built-in denylist from silently going stale.""" + + def _fields(schema: Any) -> set[str]: + model_fields = getattr(schema, "model_fields", None) + if model_fields is not None: + return set(model_fields) + return set(schema.__annotations__) + + schemas_by_prefix = { + (): _fields(MasterConfig), + ("async_rl",): _fields(AsyncRLConfig), + ("data",): _fields(DataConfig), + ("policy",): _fields(PolicyConfig), + ("policy", "generation"): _fields(VllmConfig), + ("policy", "generation", "vllm_cfg"): _fields(VllmSpecificArgs), + ("rollout_checkpointing",): _fields(RolloutCheckpointConfig), + ("token_capture",): _fields(TokenCaptureConfig), + } + for excluded_path in rollout_checkpoint._BOOTSTRAP_FINGERPRINT_EXCLUDED_PATHS: + segments = tuple(excluded_path.split(".")) + for prefix, fields in schemas_by_prefix.items(): + if segments[: len(prefix)] != prefix or len(segments) == len(prefix): + continue + next_segment = segments[len(prefix)] + if not any(character in next_segment for character in "*?["): + assert next_segment in fields, ( + f"bootstrap fingerprint exclusion {excluded_path!r} refers to " + f"unknown config field {'.'.join((*prefix, next_segment))!r}" + ) + + @pytest.mark.parametrize( ("section", "changed"), [ @@ -358,6 +377,7 @@ def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: "async_rl", {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, ), + ("rollout_recovery", {"default_granularity": "prompt_group"}), ], ) def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( @@ -377,6 +397,7 @@ def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( "mixed_weight_version_policy": "allow", }, "async_rl": {"sampler": {"name": "windowed", "max_staleness_versions": 1}}, + "rollout_recovery": {"default_granularity": "sibling"}, } modified = {**base, section: {**base[section], **changed}} @@ -410,6 +431,20 @@ def test_bootstrap_fingerprint_rejects_generation_semantic_changes() -> None: bootstrap_fingerprint(cast(Any, _DumpedConfig(sampling_changed))) ) + blocked_tokens_changed = { + **base, + "policy": { + **base["policy"], + "generation": { + **base["policy"]["generation"], + "bad_words": ["forbidden"], + }, + }, + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(blocked_tokens_changed))) + ) + context_changed = { **base, "policy": { @@ -540,13 +575,13 @@ def test_bootstrap_fingerprint_ignores_nemo_gym_service_routing() -> None: identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) assert ( "base_url" - not in identity.environment["nemo_gym"]["genrm_model"]["responses_api_models"][ - "genrm_model" - ] + not in identity.config["env"]["nemo_gym"]["genrm_model"][ + "responses_api_models" + ]["genrm_model"] ) assert ( "base_url" - not in identity.environment["nemo_gym"]["nl2bash_judge_model"][ + not in identity.config["env"]["nemo_gym"]["nl2bash_judge_model"][ "responses_api_models" ]["local_vllm_model"] ) @@ -554,9 +589,9 @@ def test_bootstrap_fingerprint_ignores_nemo_gym_service_routing() -> None: @pytest.mark.parametrize( "field", - sorted(rollout_checkpoint._BOOTSTRAP_ENV_RUNTIME_FIELDS), + ["concurrency", "nemo_gym_log_dir", "num_processes", "verbose"], ) -def test_bootstrap_fingerprint_ignores_declared_environment_runtime_fields( +def test_bootstrap_fingerprint_ignores_environment_runtime_fields( field: str, ) -> None: base = { @@ -614,7 +649,7 @@ def test_bootstrap_fingerprint_ignores_nested_environment_credentials( } identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) - assert field not in identity.environment["nemo_gym"]["service"] + assert field not in identity.config["env"]["nemo_gym"]["service"] assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) == ( bootstrap_fingerprint(cast(Any, _DumpedConfig(credential_changed))) ) @@ -644,7 +679,7 @@ def test_bootstrap_fingerprint_keeps_environment_token_semantics() -> None: } identity = bootstrap_compatibility_identity(cast(Any, _DumpedConfig(base))) - assert identity.environment["nemo_gym"]["service"] == { + assert identity.config["env"]["nemo_gym"]["service"] == { "max_tokens": 1024, "model": "model-a", "tokenizer": "tokenizer-a", @@ -675,22 +710,8 @@ def test_prune_bootstrap_snapshots_requires_durable_trainer_checkpoint(tmp_path) assert not snapshot_root.exists() -def test_reset_bootstrap_anchor_discards_skipped_snapshot_lineage( - tmp_path: Path, -) -> None: - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="old-fingerprint") - snapshot_root = anchor / "rollout_snapshots" - (snapshot_root / "snapshot_000001").mkdir(parents=True) - - reset = reset_bootstrap_anchor(tmp_path, fingerprint="new-fingerprint") - - assert reset == anchor - assert not snapshot_root.exists() - assert ensure_bootstrap_anchor(tmp_path, fingerprint="new-fingerprint") == anchor - - def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + anchor = _ensure_test_bootstrap_anchor(tmp_path) first = _commit_snapshot(anchor, mutation_version=1) second = _commit_snapshot(anchor, mutation_version=2) @@ -698,7 +719,7 @@ def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): anchor, expected_train_step=0, expected_trainer_version=0, - expected_bootstrap_fingerprint="fingerprint-v1", + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), ) assert resolved is not None @@ -708,7 +729,7 @@ def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): def test_resolver_falls_back_from_corrupt_newest_snapshot(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + anchor = _ensure_test_bootstrap_anchor(tmp_path) first = _commit_snapshot(anchor, mutation_version=1) second = _commit_snapshot(anchor, mutation_version=2) (second / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text("not-json") @@ -717,17 +738,17 @@ def test_resolver_falls_back_from_corrupt_newest_snapshot(tmp_path): anchor, expected_train_step=0, expected_trainer_version=0, - expected_bootstrap_fingerprint="fingerprint-v1", + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), ) assert resolved is not None assert resolved.path == first -def test_resolver_ignores_snapshot_without_commit_marker(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") - committed = _commit_snapshot(anchor, mutation_version=1) - incomplete = anchor / "rollout_snapshots" / "snapshot_000002" +def test_resolver_ignores_unpublished_temporary_snapshot(tmp_path): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + published = _commit_snapshot(anchor, mutation_version=1) + incomplete = anchor / "rollout_snapshots" / "tmp_snapshot_000002" incomplete.mkdir() (incomplete / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text("{}") @@ -735,23 +756,27 @@ def test_resolver_ignores_snapshot_without_commit_marker(tmp_path): anchor, expected_train_step=0, expected_trainer_version=0, - expected_bootstrap_fingerprint="fingerprint-v1", + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), ) assert resolved is not None - assert resolved.path == committed + assert resolved.path == published def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") - _commit_snapshot(anchor, mutation_version=1, fingerprint="different") + anchor = _ensure_test_bootstrap_anchor(tmp_path) + _commit_snapshot( + anchor, + mutation_version=1, + fingerprint=_test_bootstrap_fingerprint("different"), + ) with pytest.raises(ValueError, match="trainer-anchor mismatch"): resolve_latest_snapshot( anchor, expected_train_step=0, expected_trainer_version=0, - expected_bootstrap_fingerprint="fingerprint-v1", + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), ) @@ -761,27 +786,20 @@ def test_commit_snapshot_flushes_payload_before_publication(tmp_path, monkeypatc tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(anchor) (tmp_snapshot / "payload").write_text("payload") fsync_tree = Mock() - fsync_file = Mock() fsync_directory = Mock() monkeypatch.setattr(rollout_checkpoint, "_fsync_tree", fsync_tree) - monkeypatch.setattr(rollout_checkpoint, "_fsync_file", fsync_file) monkeypatch.setattr(rollout_checkpoint, "_fsync_directory", fsync_directory) commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=1) fsync_tree.assert_called_once_with(tmp_snapshot) - assert fsync_file.call_args_list == [ - call(tmp_snapshot / "COMMITTED"), - ] - assert fsync_directory.call_args_list[:2] == [ - call(tmp_snapshot), - call(anchor / "rollout_snapshots"), - ] - assert (final_snapshot / "COMMITTED").is_file() + assert fsync_directory.call_args_list == [call(anchor / "rollout_snapshots")] + assert final_snapshot.is_dir() + assert not tmp_snapshot.exists() def test_commit_snapshot_prunes_oldest_committed_snapshot(tmp_path): - anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + anchor = _ensure_test_bootstrap_anchor(tmp_path) first = _commit_snapshot(anchor, mutation_version=1) second = _commit_snapshot(anchor, mutation_version=2) third_tmp, third, _ = prepare_snapshot_paths(anchor) @@ -793,7 +811,7 @@ def test_commit_snapshot_prunes_oldest_committed_snapshot(tmp_path): sampler_dispatch_index=2, mutation_version=3, rolled_back_train_group_count=0, - bootstrap_fingerprint="fingerprint-v1", + bootstrap_fingerprint=_test_bootstrap_fingerprint(), ) (third_tmp / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( json.dumps(manifest.to_dict()) @@ -806,6 +824,72 @@ def test_commit_snapshot_prunes_oldest_committed_snapshot(tmp_path): assert third.is_dir() +def test_commit_snapshot_removes_stale_snapshot_from_live_namespace_before_delete( + tmp_path, + monkeypatch, +): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + first = _commit_snapshot(anchor, mutation_version=1) + _commit_snapshot(anchor, mutation_version=2) + third_tmp, third, _ = prepare_snapshot_paths(anchor) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=0, + trainer_version=0, + current_epoch=2, + sampler_dispatch_index=2, + mutation_version=3, + rolled_back_train_group_count=0, + bootstrap_fingerprint=_test_bootstrap_fingerprint(), + ) + (third_tmp / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + real_rmtree = rollout_checkpoint.shutil.rmtree + + def fail_trash_delete(path: Path) -> None: + path = Path(path) + if path.name.startswith("trash_snapshot_"): + raise OSError("simulated delete failure") + real_rmtree(path) + + monkeypatch.setattr(rollout_checkpoint.shutil, "rmtree", fail_trash_delete) + + with pytest.raises(OSError, match="simulated delete failure"): + commit_snapshot(third_tmp, third, keep_latest_k=2) + + assert not first.exists() + assert (first.parent / f"trash_{first.name}").is_dir() + assert third.is_dir() + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), + ) + assert resolved is not None + assert resolved.path == third + + +def test_prepare_snapshot_paths_sweeps_interrupted_snapshot_garbage(tmp_path): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + root = anchor / rollout_checkpoint.ROLLOUT_SNAPSHOTS_DIRNAME + stale_tmp = root / "tmp_snapshot_000001" + stale_trash = root / "trash_snapshot_000002" + stale_tmp.mkdir(parents=True) + stale_trash.mkdir() + (stale_tmp / "stale").write_text("stale") + (stale_trash / "stale").write_text("stale") + + tmp_path, final_path, sequence = prepare_snapshot_paths(anchor) + + assert not (root / "trash_snapshot_000002").exists() + assert tmp_path == root / "tmp_snapshot_000001" + assert not (tmp_path / "stale").exists() + assert final_path == root / "snapshot_000001" + assert sequence == 1 + + def test_manifest_rejects_bool_for_integer_field(): raw = { "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index f0f90001f76..b517fdd10b7 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -61,6 +61,10 @@ TokenCaptureConfig, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + bootstrap_compatibility_identity, + ensure_bootstrap_anchor, +) from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS @@ -699,7 +703,9 @@ def test_periodic_checkpointing_requires_trainer_checkpointing(self): mc = _make_master_config() mc.checkpointing["enabled"] = False mc.checkpointing["save_data_plane"] = True - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) with pytest.raises(ValueError, match="requires checkpointing.enabled=true"): setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -712,7 +718,9 @@ def test_periodic_checkpointing_requires_data_plane_save(self): ) mc.checkpointing["enabled"] = True mc.checkpointing["save_data_plane"] = False - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) with ( pytest.warns(UserWarning, match="cannot recover completed buffered"), @@ -726,7 +734,9 @@ def test_periodic_checkpointing_requires_token_capture(self): mc = _make_master_config() mc.checkpointing["enabled"] = True mc.checkpointing["save_data_plane"] = True - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) with pytest.raises(ValueError, match="requires token_capture.enabled=true"): setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -740,7 +750,9 @@ def test_periodic_checkpointing_requires_replay_capable_sampler(self): mc.checkpointing["enabled"] = True mc.checkpointing["save_data_plane"] = True mc.token_capture = TokenCaptureConfig(enabled=True) - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) with ( pytest.warns(UserWarning, match="cannot recover completed buffered"), @@ -757,7 +769,9 @@ def test_periodic_checkpointing_requires_claim_aware_custom_sampler(self): mc.checkpointing["enabled"] = True mc.checkpointing["save_data_plane"] = True mc.token_capture = TokenCaptureConfig(enabled=True) - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) with pytest.raises(ValueError, match="supports training-claim ownership"): setup_single_controller(mc, MagicMock(pad_token_id=0)) @@ -785,9 +799,11 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( "vllm_cfg": {"async_engine": True}, } ) - mc.logger = {"log_dir": str(tmp_path / "logs")} + mc.logger["log_dir"] = str(tmp_path / "logs") mc.token_capture.enabled = True - mc.rollout_checkpointing = RolloutCheckpointConfig(interval_s=1.0) + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0 + ) fake_finalizers = [MagicMock(name="finalizer")] patched_factories["setup_response_data"].return_value = ( list(range(8)), @@ -802,7 +818,8 @@ def test_periodic_checkpointing_warns_without_per_step_trainer_anchors( ), patch.object(sc_setup_mod, "router_replay_enabled", return_value=False), patch( - "nemo_rl.experience.finalizer_actor.create_finalizer_actors", + "nemo_rl.experience.rollout_reassembler_actor." + "create_rollout_reassembler_actors", return_value=fake_finalizers, ), ): @@ -822,7 +839,7 @@ def test_disabled_periodic_checkpointing_ignores_existing_snapshots( checkpoint_dir = tmp_path / "checkpoints" mc.checkpointing["checkpoint_dir"] = str(checkpoint_dir) mc.rollout_checkpointing = RolloutCheckpointConfig( - interval_s=None, + snapshot_attempt_interval_s=None, restore_mode="latest", ) (checkpoint_dir / "bootstrap" / "rollout_snapshots").mkdir(parents=True) @@ -836,6 +853,66 @@ def test_disabled_periodic_checkpointing_ignores_existing_snapshots( resolve.assert_not_called() assert actor_args.last_checkpoint_path is None + def test_trainer_checkpoint_restore_preserves_bootstrap_state( + self, + tmp_path: Path, + patched_factories, + ): + mc = _make_master_config(colocated=False, backend="vllm") + checkpoint_dir = tmp_path / "checkpoints" + mc.checkpointing.update( + { + "checkpoint_dir": str(checkpoint_dir), + "enabled": True, + "save_data_plane": True, + "save_period": 1, + } + ) + mc.policy["generation"].update( + { + "model_name": "test-model", + "stop_strings": None, + "stop_token_ids": None, + "top_k": None, + "vllm_cfg": {"async_engine": True}, + } + ) + mc.logger = {"log_dir": str(tmp_path / "logs")} + mc.token_capture.enabled = True + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_interval_s=1.0, + restore_mode="trainer_checkpoint", + ) + anchor = ensure_bootstrap_anchor( + checkpoint_dir, + identity=bootstrap_compatibility_identity(mc), + ) + payload = anchor / "rollout_snapshots" / "snapshot_000001" / "payload" + payload.parent.mkdir(parents=True) + payload.write_text("preserve me") + before = { + path.relative_to(anchor): path.read_bytes() + for path in anchor.rglob("*") + if path.is_file() + } + + with ( + patch.object(sc_setup_mod, "should_use_nemo_gym", return_value=True), + pytest.raises( + ValueError, + match="Existing checkpoint state was not modified", + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + after = { + path.relative_to(anchor): path.read_bytes() + for path in anchor.rglob("*") + if path.is_file() + } + assert after == before + patched_factories["setup_response_data"].assert_not_called() + def test_rejects_windowed_checkpointing_without_native_tq(self): mc = _make_master_config() mc.checkpointing["enabled"] = True diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 5d61313b21e..8423d0d033a 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -1207,6 +1207,9 @@ class _EmptyBuffer: def __len__(self) -> int: return 0 + def training_owned_group_ids(self) -> set[str]: + return set() + class _NoOpTrainer: def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: From 17918b512ce46de0f6348f6df65c4f90dbe879e6 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 6 Sep 2026 19:27:40 -0400 Subject: [PATCH 12/14] fix(sc): address periodic checkpoint review Signed-off-by: Anish Mahishi --- .../rollout_checkpoint.py | 16 ++++++++-- .../test_rollout_checkpoint.py | 29 ++++++++++++++++++- .../test_single_controller_actor.py | 6 ++++ 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py index c7dec0c191c..539aa7d86c4 100644 --- a/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -606,9 +606,19 @@ def resolve_latest_snapshot( if ( manifest.base_train_step != expected_train_step or manifest.trainer_version != expected_trainer_version - or manifest.bootstrap_fingerprint != expected_bootstrap_fingerprint ): - errors.append(f"{candidate.name}: trainer-anchor mismatch") + errors.append( + f"{candidate.name}: belongs to a different trainer lineage " + f"(step={manifest.base_train_step}, " + f"version={manifest.trainer_version}); expected " + f"step={expected_train_step}, version={expected_trainer_version}" + ) + continue + if manifest.bootstrap_fingerprint != expected_bootstrap_fingerprint: + errors.append( + f"{candidate.name}: bootstrap lineage fingerprint does not " + "match the selected trainer anchor" + ) continue return ResolvedRolloutCheckpoint(candidate, manifest) @@ -616,5 +626,7 @@ def resolve_latest_snapshot( raise ValueError( "no committed rollout snapshot matches the selected trainer anchor: " + "; ".join(errors) + + ". These snapshot directories contain stale or corrupted state; " + "inspect and remove them, or use a fresh checkpointing.checkpoint_dir." ) return None diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py index 9c7cf81e31b..a7974b62f15 100644 --- a/tests/unit/single_controller/test_rollout_checkpoint.py +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -771,7 +771,12 @@ def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): fingerprint=_test_bootstrap_fingerprint("different"), ) - with pytest.raises(ValueError, match="trainer-anchor mismatch"): + with pytest.raises( + ValueError, + match=( + "bootstrap lineage fingerprint does not match the selected trainer anchor" + ), + ): resolve_latest_snapshot( anchor, expected_train_step=0, @@ -780,6 +785,28 @@ def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): ) +def test_resolver_reports_trainer_lineage_mismatch(tmp_path): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + _commit_snapshot( + anchor, + mutation_version=1, + trainer_version=4, + ) + + with pytest.raises(ValueError) as error: + resolve_latest_snapshot( + anchor, + expected_train_step=1, + expected_trainer_version=2, + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), + ) + + message = str(error.value) + assert "step=4, version=4" in message + assert "expected step=1, version=2" in message + assert "fresh checkpointing.checkpoint_dir" in message + + def test_commit_snapshot_flushes_payload_before_publication(tmp_path, monkeypatch): anchor = tmp_path / "step_1" anchor.mkdir() diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 8423d0d033a..649dd0ecc5f 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -136,6 +136,7 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: last_checkpoint_path=None, finalizer_actors=[], data_plane_checkpoint_metadata=None, + bootstrap_identity=None, ) args.update(overrides) return SimpleNamespace(**args) @@ -708,6 +709,7 @@ def test_advantage_stage_writes_each_sample_filter_without_seq_threshold( ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -838,6 +840,7 @@ def test_advantage_stage_clips_training_values_and_metrics() -> None: ctrl._dp_client = data_plane ctrl._advantage_cfg = AdvantageConfig() ctrl._advantage_estimator = estimator + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._policy_logprobs_required = False ctrl._reference_logprobs_required = False ctrl._teacher_logprobs_required = False @@ -1210,6 +1213,9 @@ def __len__(self) -> int: def training_owned_group_ids(self) -> set[str]: return set() + def release_training_claims(self, group_ids: list[str]) -> None: + assert not group_ids + class _NoOpTrainer: def prepare_for_lp_inference(self, keep_train_buffers: bool = False) -> None: From 17b86fb10f5004d7a90e58401c75626ca46e41fa Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 8 Sep 2026 14:39:02 -0400 Subject: [PATCH 13/14] test(sc): run rollout recovery in fast CI Signed-off-by: Anish Mahishi --- tests/functional/L1_Functional_Tests_SingleController.sh | 4 ++-- .../grpo_async_gym_single_controller_sibling_recovery.sh | 2 +- .../grpo_async_gym_single_controller_streaming_recovery.sh | 2 +- tests/functional/ppo_async_single_controller.sh | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index b991087eaf5..04599bacac6 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -184,10 +184,10 @@ run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh ++token_capture.enabled=true # Two-process token-capture recovery: preserve one sealed sibling in TQ and # redispatch only its unfinished peer after restoring the step checkpoint. -run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh # Periodic native-TQ snapshot while a streamed step owns only part of its # rollout batch, followed by SIGKILL and rollback to the durable trainer anchor. -run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh index 1d2182f25c2..9d0c54d213f 100755 --- a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -23,7 +23,7 @@ COMMON_OVERRIDES=( checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.metric_name=null checkpointing.save_period=1 - +checkpointing.save_data_plane=true + checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling async_rl.sampler.name=in_order diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index 69a07af0048..f23327b58a8 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -52,7 +52,7 @@ COMMON_OVERRIDES=( checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.save_period=1 checkpointing.metric_name=null - +checkpointing.save_data_plane=true + checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" diff --git a/tests/functional/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh index 2861a038f60..b3df65ab72c 100755 --- a/tests/functional/ppo_async_single_controller.sh +++ b/tests/functional/ppo_async_single_controller.sh @@ -65,7 +65,7 @@ TRAIN_CMD=( checkpointing.checkpoint_dir="${CKPT_DIR}" checkpointing.metric_name=null checkpointing.save_period=1 - +checkpointing.save_data_plane=true + checkpointing.save_data_plane=true ) cd "${PROJECT_ROOT}" From 2835680edfc1035177197b13b3e54756a755868d Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 8 Sep 2026 16:23:17 -0400 Subject: [PATCH 14/14] test(sc): fix rollout recovery checkpoint overrides Signed-off-by: Anish Mahishi --- .../grpo_async_gym_single_controller_sibling_recovery.sh | 2 +- .../grpo_async_gym_single_controller_streaming_recovery.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh index 9d0c54d213f..1d2182f25c2 100755 --- a/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_sibling_recovery.sh @@ -23,7 +23,7 @@ COMMON_OVERRIDES=( checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.metric_name=null checkpointing.save_period=1 - checkpointing.save_data_plane=true + +checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling async_rl.sampler.name=in_order diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index f23327b58a8..69a07af0048 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -52,7 +52,7 @@ COMMON_OVERRIDES=( checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.save_period=1 checkpointing.metric_name=null - checkpointing.save_data_plane=true + +checkpointing.save_data_plane=true ++token_capture.enabled=true ++rollout_recovery.default_granularity=sibling ++rollout_checkpointing.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S"