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 96f2d1ea69d..38cf8d4ece6 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -141,10 +141,28 @@ 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 +# 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: + 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: enabled: false diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index 0c414a80216..3a13053c7f7 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -146,6 +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 +# 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: + 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/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 60f94efd3eb..608da389533 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,75 @@ 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 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)): + 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 +1559,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 +1593,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 +1618,11 @@ 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. + 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. """ groups: list[TQReplayGroupMetadata] = [] for i, ready in enumerate(self.ready_list): @@ -1533,6 +1639,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 +1711,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..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]: @@ -93,7 +98,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: @@ -117,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.""" ... @@ -159,6 +172,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 @@ -287,7 +301,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) @@ -655,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 @@ -747,6 +762,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, @@ -790,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 af55ab9b493..cf1e039b31c 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,16 @@ 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, + BootstrapCompatibilityIdentity, + 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, @@ -149,6 +162,22 @@ # Logger this module also uses as `self._logger`. log = logging.getLogger(__name__) +_MAX_CONSECUTIVE_ROLLOUT_CHECKPOINT_FAILURES = 3 + + +@dataclass(frozen=True) +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] + rollout_recovery_group_count: Optional[int] + rolled_back_train_group_count: int + mutation_version: int + def _pooled_opd_metrics( stat_sum: float, stat_sumsq: float, count: int @@ -405,6 +434,20 @@ 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_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 + # 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() @@ -508,6 +551,14 @@ 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_task = ( + asyncio.create_task(self._rollout_checkpoint_pump()) + if self._master_config.rollout_checkpointing.snapshot_attempt_interval_s + 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 +573,39 @@ async def run(self) -> dict[str, Any]: done, _ = await asyncio.wait( set(tasks), return_when=asyncio.FIRST_COMPLETED ) - if probe_task is not None and probe_task in done: + 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 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 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() @@ -1135,6 +1206,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, @@ -1151,20 +1226,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", @@ -1458,10 +1526,21 @@ 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.""" + 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 ─────────────────────── @@ -2205,6 +2284,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_training_claim_ids: list[str] = [] + consumed_group_count = 0 step_finalizer_metrics: dict[str, list[float]] = {} with self._timer.time("total_step_time"): @@ -2248,11 +2329,48 @@ async def _train_pump(self) -> None: self._async_cfg.min_groups_for_streaming_train, max_prompt_groups, ) + selected_group_ids: list[str] = [] + selected_training_claim_ids: list[str] = [] + 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: @@ -2280,16 +2398,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_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( group_id, {} @@ -2378,6 +2488,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, @@ -2505,10 +2619,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 = {} @@ -2516,6 +2627,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_training_claim_ids) + for _ in range(consumed_group_count): + self._buffer_capacity.release() step_metrics.update( { name: statistics.fmean(values) @@ -2558,6 +2674,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 ) @@ -3275,11 +3392,287 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: ) return len(stale_tasks) + async def _capture_rollout_checkpoint_cut( + self, + cut: DataPlaneMutationCut, + checkpoint_path: PathLike, + ) -> _RolloutCheckpointCut: + """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. + """ + cut.require_live() + 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, + additional_groups=training_owned_groups, + ) + 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( + 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"]), + ) + 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, + 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, + ) + + 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.""" + async with self._checkpoint_save_lock: + if self._optimizer_commit_in_progress: + return False + 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_identity is None: + raise RuntimeError( + "rollout snapshotting requires a bootstrap identity" + ) + anchor = await asyncio.to_thread( + ensure_bootstrap_anchor, + self._checkpointer.checkpoint_dir, + identity=self._bootstrap_identity, + ) + snapshot_fingerprint = self._bootstrap_identity.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() as cut: + if ( + self._optimizer_commit_in_progress + or self._train_steps != expected_train_step + or self._trainer_version != expected_trainer_version + ): + await asyncio.to_thread(partial(shutil.rmtree, tmp_path)) + return False + snapshot_epoch = self._current_epoch + snapshot_cut = await self._capture_rollout_checkpoint_cut( + cut, tmp_path + ) + + 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=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( + (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(partial(shutil.rmtree, tmp_path)) + raise + + 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={snapshot_cut.rollout_recovery_group_count or 0})", + flush=True, + ) + return True + + async def _rollout_checkpoint_pump(self) -> None: + """Persist rollout state periodically, including during streamed train.""" + 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(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 (OSError, TimeoutError) as error: + if deadline_due: + raise RuntimeError( + "failed to save the required pre-step rollout checkpoint" + ) from error + consecutive_failures += 1 + print( + "Periodic rollout checkpoint failed; retaining the previous " + "committed snapshot: " + f"consecutive_failures={consecutive_failures}, " + 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(): + 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.""" + async with self._checkpoint_save_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. @@ -3350,6 +3743,12 @@ async def _save_checkpoint( ) 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 @@ -3392,6 +3791,9 @@ async def _save_checkpoint( 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=( @@ -3400,6 +3802,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: @@ -3921,16 +4326,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/__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..bc381273b09 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -714,6 +714,59 @@ 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. + + ``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; 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. 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 + ``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. + """ + + 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"): # algo configs grpo: Optional[GRPOConfig] = None @@ -734,6 +787,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..539aa7d86c4 --- /dev/null +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -0,0 +1,632 @@ +# 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 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 = 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" + +_SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") +_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( + { + "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 _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_excluded_paths( + value: Any, + *, + excluded_paths: tuple[tuple[str, ...], ...], + path: tuple[str, ...] = (), +) -> Any: + """Recursively remove denylisted mapping paths from a JSON config dump.""" + if isinstance(value, Mapping): + 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, + excluded_paths=excluded_paths, + path=child_path, + ) + return projected + if isinstance(value, list): + 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, + ) + ) + return projected_list + 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 directory.""" + 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 + sampler_dispatch_index: 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", + "sampler_dispatch_index", + "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"], + 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, + ) + 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") + 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]: + 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 + excluded_paths: tuple[str, ...] + config: Mapping[str, Any] + + def to_dict(self) -> dict[str, Any]: + 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: + """Remove explicitly operational paths from the fail-closed run identity.""" + dumped = master_config.model_dump(mode="json") + 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", [] + ) + 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 + ) + return BootstrapCompatibilityIdentity( + schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + excluded_paths=excluded_paths, + config=_drop_excluded_paths( + dumped, + excluded_paths=parsed_excluded_paths, + ), + ) + + +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. + """ + 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( + 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, + *, + 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 = _bootstrap_anchor_manifest(identity) + if manifest_path.is_file(): + validate_bootstrap_anchor(anchor, identity=identity) + 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 validate_bootstrap_anchor( + anchor: Path, + *, + identity: BootstrapCompatibilityIdentity, +) -> None: + """Validate a bootstrap anchor without modifying checkpoint 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()) + 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() + 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) + os.rename(tmp_path, final_path) + + root = final_path.parent + _fsync_directory(root) + + committed = sorted( + ( + child + for child in root.iterdir() + 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: + 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) + + +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: + 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 + ): + 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) + + if errors: + 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/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 41f53c89d3d..de41845cf10 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, @@ -68,6 +69,13 @@ is_ppo_run, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_DIRNAME, + 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 from nemo_rl.data.multimodal_utils import WIRE_MULTIMODAL_FIELDS @@ -155,8 +163,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_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 # None unless async_rl.generation_router is enabled. generation_router: Optional[ray.actor.ActorHandle[GenerationRouterImpl]] = None @@ -941,10 +950,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.snapshot_attempt_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}." @@ -1003,6 +1013,39 @@ 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.snapshot_attempt_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 not sampler_supports_training_claims(master_config.async_rl.sampler): + raise ValueError( + "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( @@ -1045,24 +1088,104 @@ 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 + bootstrap_anchor = checkpointer.checkpoint_dir / BOOTSTRAP_DIRNAME + needs_bootstrap_identity = ( + trainer_checkpoint_path is None + and rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None + ) + bootstrap_identity = ( + bootstrap_compatibility_identity(master_config) + if needs_bootstrap_identity + else None + ) + bootstrap_digest = ( + bootstrap_identity.fingerprint() if bootstrap_identity is not None 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 rollout_checkpoint_cfg.snapshot_attempt_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, + expected_trainer_version=restored_trainer_version, + expected_bootstrap_fingerprint=None, + ) + elif trainer_checkpoint_path is None: + if rollout_checkpoint_cfg.snapshot_attempt_interval_s is not None: + assert bootstrap_digest is not None + assert bootstrap_identity is not None + if bootstrap_anchor.is_dir(): + validate_bootstrap_anchor( + bootstrap_anchor, + identity=bootstrap_identity, + ) + 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.snapshot_attempt_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, + 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 + save_state.sampler_dispatch_index = ( + resolved_snapshot.manifest.sampler_dispatch_index + ) + 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, + ) + # ========================== # Setup Dataset & Environments # ========================== @@ -1103,9 +1226,11 @@ 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( + f"📦 Restoring dataloader state from checkpoint: {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 +1493,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 +1730,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_identity=bootstrap_identity, + 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/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index ded20da8876..04599bacac6 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -184,7 +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 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_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh new file mode 100755 index 00000000000..69a07af0048 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -0,0 +1,191 @@ +#!/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.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 + 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.snapshot_attempt_interval_s="$SNAPSHOT_INTERVAL_S" + ++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" + 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 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 at least {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 +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 "open_train_step" not in lineage, lineage +PY + +echo "=== Phase 2: restore claimed rows and finish without duplicate steps ===" +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" +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/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}" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index a813610185f..f68b9db2e6c 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 ( @@ -75,8 +76,20 @@ from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, MasterConfig, + 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_MANIFEST_FILENAME, + ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + BootstrapCompatibilityIdentity, + RolloutSnapshotManifest, + bootstrap_compatibility_identity, + commit_snapshot, + prepare_snapshot_paths, +) 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 @@ -104,6 +117,7 @@ _ACTOR_CLS = SingleControllerActor.__ray_metadata__.modified_class _PARTITION_ID = "rollout_data" +_STAGING_PARTITION_ID = "rollout_staging" def _consumed_meta(*sample_ids: str) -> KVBatchMeta: @@ -327,8 +341,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()) @@ -443,6 +460,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, ...]: @@ -453,9 +471,31 @@ 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 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} + 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.""" @@ -525,6 +565,8 @@ def _actor_master_config( max_num_epochs: int = 1, buffer_checkpoint: bool = False, data_plane_checkpoint: bool = True, + rollout_checkpoint_attempt_interval_s: Optional[float] = None, + token_capture_enabled: bool = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -584,6 +626,10 @@ def _actor_master_config( max_inflight_prompts=4, max_buffered_rollouts=4, ), + rollout_checkpointing=RolloutCheckpointConfig( + snapshot_attempt_interval_s=rollout_checkpoint_attempt_interval_s + ), + token_capture=TokenCaptureConfig(enabled=token_capture_enabled), ) @@ -596,6 +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_identity: Optional[BootstrapCompatibilityIdentity] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=_FakeGeneration(), @@ -617,6 +664,7 @@ def _make_actor_args( last_checkpoint_path=last_checkpoint_path, finalizer_actors=[], data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + bootstrap_identity=bootstrap_identity, ) @@ -1048,8 +1096,218 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} +class TestPeriodicRolloutCheckpoint: + 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", + [ + {"snapshot_attempt_interval_s": 0}, + {"interval_s": 1}, + {"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, + buffer_checkpoint=True, + rollout_checkpoint_attempt_interval_s=120.0, + token_capture_enabled=True, + ) + return _ACTOR_CLS( + config, + _make_actor_args( + bootstrap_identity=bootstrap_compatibility_identity(config) + ), + SetupTimingMetrics(), + ) + + 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() + + snapshot = ( + tmp_path + / "checkpoints" + / BOOTSTRAP_DIRNAME + / "rollout_snapshots" + / "snapshot_000001" + ) + 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() + assert (snapshot / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() + assert not (snapshot / "policy").exists() + + 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) + + try: + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + finally: + actor._checkpointer.shutdown() + + 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: + assert not asyncio.run(actor._save_rollout_checkpoint(force=True)) + assert actor._dp_client.save_calls == [] + 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.snapshot_attempt_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 + + 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.snapshot_attempt_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 + + 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_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, @@ -1067,12 +1325,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()) @@ -1347,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 == [] @@ -1370,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 @@ -1497,6 +1766,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 @@ -1680,6 +1950,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. @@ -1795,6 +2089,90 @@ 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}, + ) + final_snapshot = _write_periodic_snapshot(step_3) + mc = _setup_master_config(str(ckpt_dir)) + mc.checkpointing["save_period"] = 1 + mc.rollout_checkpointing = RolloutCheckpointConfig( + snapshot_attempt_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, + ) + + 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.rollout_reassembler_actor." + "create_rollout_reassembler_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( + snapshot_attempt_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_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 new file mode 100644 index 00000000000..a7974b62f15 --- /dev/null +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -0,0 +1,950 @@ +# 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.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, + commit_snapshot, + ensure_bootstrap_anchor, + prepare_snapshot_paths, + prune_bootstrap_snapshots, + resolve_latest_snapshot, + validate_bootstrap_anchor, +) +from nemo_rl.data import DataConfig +from nemo_rl.models.generation.vllm.config import VllmConfig, VllmSpecificArgs +from nemo_rl.models.policy import PolicyConfig + + +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 _test_bootstrap_identity(label: str = "v1") -> BootstrapCompatibilityIdentity: + return BootstrapCompatibilityIdentity( + schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + excluded_paths=(), + config={"test_identity": label}, + ) + + +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), + ) + + +def _commit_snapshot( + anchor, + *, + mutation_version: int, + trainer_version: int = 0, + 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, + 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, + ) + (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_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)), + ) + + 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_default_operational_paths() -> None: + base = { + "policy": { + "model_name": "model-a", + "optimizer": {"lr": 1.0e-6}, + "generation": { + "backend": "vllm", + "temperature": 1.0, + "colocated": {"enabled": True, "resources": {"gpus": 8}}, + "port_range_low": 3000, + "port_range_high": 4000, + }, + }, + "data": { + "train": [{"data_path": "/datasets/train.jsonl"}], + "num_workers": 4, + }, + "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_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"}, + } + operationally_changed = { + **base, + "policy": { + **base["policy"], + "optimizer": {"lr": 5.0e-7}, + "generation": { + **base["policy"]["generation"], + "colocated": {"enabled": False, "resources": {"gpus": 16}}, + "port_range_low": 5000, + "port_range_high": 6000, + }, + }, + "data": {**base["data"], "num_workers": 16}, + "token_capture": { + **base["token_capture"], + "capture_dir": "/run/two/capture", + "control_auth_token": "secret-two", + "control_timeout_s": 15.0, + "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(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"), + [ + ("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"}), + ("token_capture", {"defer_routed_experts_to_policy": True}), + ("token_capture", {"on_capture_failure": "abort"}), + ( + "async_rl", + {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, + ), + ("rollout_recovery", {"default_granularity": "prompt_group"}), + ], +) +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}}, + "rollout_recovery": {"default_granularity": "sibling"}, + } + 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))) + ) + + 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": { + **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_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.config["env"]["nemo_gym"]["genrm_model"][ + "responses_api_models" + ]["genrm_model"] + ) + assert ( + "base_url" + not in identity.config["env"]["nemo_gym"]["nl2bash_judge_model"][ + "responses_api_models" + ]["local_vllm_model"] + ) + + +@pytest.mark.parametrize( + "field", + ["concurrency", "nemo_gym_log_dir", "num_processes", "verbose"], +) +def test_bootstrap_fingerprint_ignores_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))) + ) + + +@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.config["env"]["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.config["env"]["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) + (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_resolver_selects_latest_compatible_committed_snapshot(tmp_path): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + 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=_test_bootstrap_fingerprint(), + ) + + 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_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") + + 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 == first + + +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("{}") + + 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 == published + + +def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): + anchor = _ensure_test_bootstrap_anchor(tmp_path) + _commit_snapshot( + anchor, + mutation_version=1, + fingerprint=_test_bootstrap_fingerprint("different"), + ) + + with pytest.raises( + ValueError, + match=( + "bootstrap lineage fingerprint does not match the selected trainer anchor" + ), + ): + resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint=_test_bootstrap_fingerprint(), + ) + + +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() + tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(anchor) + (tmp_snapshot / "payload").write_text("payload") + fsync_tree = Mock() + fsync_directory = Mock() + monkeypatch.setattr(rollout_checkpoint, "_fsync_tree", fsync_tree) + 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_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_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) + 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()) + ) + + commit_snapshot(third_tmp, third, keep_latest_k=2) + + assert not first.exists() + assert second.is_dir() + 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, + "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 28640eeae30..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 @@ -93,6 +94,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) @@ -272,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(), @@ -693,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..b517fdd10b7 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -57,8 +57,14 @@ setup_single_controller, ) from nemo_rl.algorithms.single_controller_utils.config import ( + RolloutCheckpointConfig, + 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 @@ -79,6 +85,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 +100,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 +699,220 @@ 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( + snapshot_attempt_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( + snapshot_attempt_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( + 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)) + + 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( + snapshot_attempt_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( + snapshot_attempt_interval_s=1.0 + ) + + 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( + snapshot_attempt_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.rollout_reassembler_actor." + "create_rollout_reassembler_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( + snapshot_attempt_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_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 9416e9989c1..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) @@ -608,6 +609,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 +670,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( @@ -706,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 @@ -770,6 +774,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 @@ -835,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 @@ -902,6 +908,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 +970,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 +1051,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 @@ -1201,6 +1210,12 @@ class _EmptyBuffer: def __len__(self) -> int: return 0 + 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: @@ -2264,6 +2279,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 1d2f8fa6ee1..46504442116 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -270,6 +270,20 @@ 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(): + 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 +1135,134 @@ 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_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) + + 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) 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)