diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index b3405c29cb7..c32dd476b28 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -71,6 +71,13 @@ checkpointing: checkpoint_dir: results/grpo-single-controller metric_name: null +# Frequent rollout-only snapshots are disabled unless interval_s is set. They +# reuse the latest durable trainer checkpoint and native TQ checkpoint support. +rollout_checkpointing: + interval_s: null + keep_latest_k: 2 + restore_mode: latest + policy: dtensor_cfg: enabled: false diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 5463e2fa179..ef028bee0f1 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -159,6 +159,12 @@ def __init__(self) -> None: self._checkpoint_active = False self._active_mutations = 0 self._mutation_depth_by_task: dict[asyncio.Task[Any], int] = {} + self._mutation_version = 0 + + @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[None]: @@ -187,6 +193,10 @@ async def mutation(self) -> AsyncIterator[None]: del self._mutation_depth_by_task[task] async with self._condition: self._active_mutations -= 1 + # Count completed mutation sections even when their body raised. + # A redundant periodic snapshot is safe; missing a mutation that + # partially changed TQ or controller metadata is not. + self._mutation_version += 1 if self._active_mutations == 0: self._condition.notify_all() @@ -1225,7 +1235,12 @@ async def _remove_unlocked( return len(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 @@ -1237,7 +1252,10 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: across the complete publish/index or clear/remove transition. This includes future finalizer paths; canonical writes are not required to originate specifically from :meth:`commit`. - In-flight reservations are intentionally omitted. + In-flight reservations are intentionally omitted. ``additional_groups`` + is used only by periodic snapshots to re-index canonical groups already + claimed by the current, uncheckpointed optimizer step. Their TQ rows are + still present, but normal sampler selection has removed their live slots. """ groups: list[TQReplayGroupMetadata] = [] for i, ready in enumerate(self.ready_list): @@ -1254,6 +1272,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(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, diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 31792919732..863d2c3deaf 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -136,6 +136,17 @@ def restore_dispatch_state( ) -> None: ... +@runtime_checkable +class TwoPhaseAdmissionSampler(Protocol): + """Sampler that separates a blocking gate from its admission commit.""" + + async def wait_until_admissible( + self, *, trainer_version_fn: Callable[[], int] + ) -> None: ... + + def commit_admission(self, *, trainer_version: int) -> Optional[int]: ... + + class BaseSampler(abc.ABC): """Shared machinery for the built-in policies. @@ -170,10 +181,19 @@ def set_dispatch_index(self, resume_from_trainer_version: int) -> None: self._dispatch_index = resume_from_trainer_version - 1 # ── rollout-pump side ──────────────────────────────────────────────── - @abc.abstractmethod - async def admit( + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + """Wait for and commit one prompt-batch admission.""" + await self.wait_until_admissible(trainer_version_fn=trainer_version_fn) + return self.commit_admission(trainer_version=trainer_version_fn()) + + async def wait_until_admissible( self, *, trainer_version_fn: Callable[[], int] - ) -> Optional[int]: ... + ) -> None: + """Wait until one batch may be admitted; ungated by default.""" + + def commit_admission(self, *, trainer_version: int) -> Optional[int]: + """Commit one already-admissible batch; unstamped by default.""" + return None # ── train-pump side ────────────────────────────────────────────────── @abc.abstractmethod @@ -377,9 +397,20 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: gate_window=self._gate_window, ) - async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + async def wait_until_admissible( + self, *, trainer_version_fn: Callable[[], int] + ) -> None: while self._dispatch_index >= trainer_version_fn() + self._gate_window: await asyncio.sleep(_GATE_POLL_SECONDS) + + def commit_admission(self, *, trainer_version: int) -> Optional[int]: + if self._dispatch_index >= trainer_version + self._gate_window: + raise RuntimeError( + "sampler admission was committed before its gate opened: " + f"dispatch_index={self._dispatch_index}, " + f"trainer_version={trainer_version}, " + f"gate_window={self._gate_window}" + ) self._dispatch_index += 1 return self._stamp() diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 35905a8d471..61d604c78ce 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -37,15 +37,20 @@ import asyncio import hashlib import io +import json import os +import shutil import statistics import time +import warnings +from dataclasses import dataclass from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Optional, Union, cast import ray import torch +import yaml from nemo_rl.algorithms.async_utils.replay_buffer import ( DATA_PLANE_CHECKPOINT_DIR, @@ -53,10 +58,12 @@ REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, DataPlaneCheckpointBarrier, + TQReplayGroupMetadata, TQReplayMetadataState, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( CheckpointDispatchSampler, + TwoPhaseAdmissionSampler, create_sampler, ) from nemo_rl.algorithms.grpo import GRPOSaveState, _write_latest_checkpoint_status @@ -64,9 +71,19 @@ from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, MasterConfig, + required_rollout_recovery_capacity, validate_sampler_buffer_capacity, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + RolloutSnapshotManifest, + commit_snapshot, + ensure_bootstrap_anchor, + prepare_snapshot_paths, + prune_bootstrap_snapshots, +) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.algorithms.single_controller_utils.utils import ( aggregate_step_metrics, @@ -103,6 +120,27 @@ Generation = Union[VllmGeneration, SGLangGeneration] +@dataclass(frozen=True) +class _RolloutWorkItem: + """One capacity-owned prompt group awaiting generation.""" + + prompt: DatumSpec + target_step: Optional[int] + group_id: Optional[str] + + +@dataclass(frozen=True) +class _RolloutCheckpointCut: + """Controller sidecars captured with one native TQ snapshot.""" + + dataloader_state: dict[str, Any] + replay_metadata: Optional[TQReplayMetadataState] + rollout_recovery_payload: Optional[bytes] + rollout_recovery_group_count: Optional[int] + rolled_back_train_group_count: int + mutation_version: int + + @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: """CPU-only Ray actor that orchestrates the RL training loop. @@ -247,6 +285,17 @@ def __init__( # A future staging/finalizer path must join the same barrier before # native restore can be authoritative. self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + # Full trainer checkpoints and lightweight rollout snapshots share a + # filesystem namespace and must never publish concurrently. + self._checkpoint_save_lock = asyncio.Lock() + self._last_rollout_snapshot_mutation_version: Optional[int] = None + self._last_missing_rollout_snapshot_anchor: Optional[tuple[int, int]] = None + self._bootstrap_fingerprint = actor_args.bootstrap_fingerprint + self._rollout_checkpoint_stop_requested = asyncio.Event() + # Narrow unsafe window after optimizer apply and before SC publishes + # the matching trainer counters. Streaming accumulation itself remains + # snapshot-safe and is represented by the ledger's open train step. + self._optimizer_commit_in_progress = False if self._buffer is not None: self._buffer.set_data_plane_checkpoint_barrier( self._data_plane_checkpoint_barrier @@ -331,9 +380,16 @@ async def run(self) -> dict[str, Any]: rollout_task = asyncio.create_task(self._rollout_pump()) train_task = asyncio.create_task(self._train_pump()) watchdog_task = asyncio.create_task(self._watchdog_pump()) + rollout_checkpoint_task = ( + asyncio.create_task(self._rollout_checkpoint_pump()) + if self._master_config.rollout_checkpointing.interval_s is not None + else None + ) tasks = {rollout_task, train_task, watchdog_task} if recovery_task is not None: tasks.add(recovery_task) + if rollout_checkpoint_task is not None: + tasks.add(rollout_checkpoint_task) try: pending = set(tasks) while pending: @@ -352,6 +408,15 @@ async def run(self) -> dict[str, Any]: rollout_task.cancel() await recovery_task break + if ( + rollout_checkpoint_task is not None + and rollout_checkpoint_task in done + ): + if self._rollout_checkpoint_stop_requested.is_set(): + break + raise RuntimeError( + "rollout checkpoint pump exited without requesting stop" + ) finally: for task in tasks: task.cancel() @@ -908,6 +973,22 @@ async def _prepare_rollout_recovery( f"unfinished={len(unfinished)}, " f"capacity={self._async_cfg.max_buffered_rollouts}" ) + restored_group_count = restored_replay_groups + len(unfinished) + min_streaming_groups = self._async_cfg.min_groups_for_streaming_train + free_capacity = self._async_cfg.max_buffered_rollouts - restored_group_count + if ( + 0 < restored_group_count < min_streaming_groups + and free_capacity < self._master_config.grpo.num_prompts_per_step + ): + required_capacity = required_rollout_recovery_capacity( + num_prompts_per_step=(self._master_config.grpo.num_prompts_per_step), + min_groups_for_streaming_train=min_streaming_groups, + ) + raise RuntimeError( + "restored rollout state cannot fill a streaming batch with the " + "configured capacity; increase async_rl.max_buffered_rollouts " + f"to at least {required_capacity}" + ) # Reclaim all saved permits before fresh dataloader work can run. The # capacity check above makes these acquisitions immediate and prevents @@ -948,6 +1029,11 @@ async def _recover_one_prepared_group( buffer_permit_owned = True rollout_slot_owned = False counted_inflight = False + ledger = self._rollout_recovery_ledger + assert ledger is not None + recovery_record = ledger.get_group(group.group_id) + reused_siblings = len(recovery_record.sealed_generation_indices) + redispatched_siblings = recovery_record.expected_generations - reused_siblings try: await self._rollout_slots.acquire() rollout_slot_owned = True @@ -970,6 +1056,12 @@ async def _recover_one_prepared_group( if request is None: return False await self._finalize_with_actor(request) + print( + "rollout recovery finalized group: " + f"group={group.group_id} reused={reused_siblings} " + f"redispatched={redispatched_siblings}", + flush=True, + ) buffer_permit_owned = False return True finally: @@ -1160,6 +1252,273 @@ async def _save_data_plane_checkpoint( flush=True, ) + async def _capture_rollout_checkpoint_cut( + self, + checkpoint_path: PathLike, + *, + periodic: bool, + ) -> _RolloutCheckpointCut: + """Save TQ and capture matching controller sidecars under the barrier. + + For a periodic cut, groups already claimed by an open optimizer step + are added back to the persisted replay index and rolled back to + ``FINALIZED`` in the persisted lineage. The live ledger and trainer are + not modified; only restart semantics discard those uncommitted gradients. + """ + dataloader_state = self._dataloader.state_dict() + replay_metadata: Optional[TQReplayMetadataState] = None + rollout_recovery_payload: Optional[bytes] = None + rollout_recovery_digest: Optional[str] = None + rollout_recovery_group_count: Optional[int] = None + rollout_recovery_state: Optional[dict[str, Any]] = None + ledger = self._rollout_recovery_ledger + + additional_groups: list[TQReplayGroupMetadata] = [] + if ledger is not None: + if periodic: + additional_groups = ledger.training_owned_replay_groups() + rollout_recovery_state = ledger.periodic_snapshot_state_dict() + else: + ledger.assert_full_step_checkpoint_safe() + rollout_recovery_state = ledger.state_dict() + payload_buffer = io.BytesIO() + torch.save(rollout_recovery_state, payload_buffer) + rollout_recovery_payload = payload_buffer.getvalue() + rollout_recovery_digest = hashlib.sha256( + rollout_recovery_payload + ).hexdigest() + rollout_recovery_group_count = len(rollout_recovery_state["groups"]) + + if self._sampler.supports_buffer_checkpoint: + replay_metadata = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts, + additional_groups=additional_groups, + ) + if replay_metadata is not None and rollout_recovery_state is not None: + replay_group_ids = { + group["group_id"] for group in replay_metadata["groups"] + } + finalized_group_ids = { + group["group_id"] + for group in rollout_recovery_state["groups"] + if group["status"] == PromptGroupStatus.FINALIZED.value + } + if replay_group_ids != finalized_group_ids: + raise RuntimeError( + "rollout snapshot replay metadata does not match finalized " + "lineage ownership: " + f"replay_only={sorted(replay_group_ids - finalized_group_ids)!r}, " + f"lineage_only={sorted(finalized_group_ids - replay_group_ids)!r}" + ) + if replay_metadata is not None: + await self._validate_replay_inventory(replay_metadata) + if rollout_recovery_payload is not None: + await self._validate_rollout_recovery_inventory(clear_unreferenced=False) + await self._save_data_plane_checkpoint( + checkpoint_path, + replay_metadata=replay_metadata, + rollout_recovery_payload_sha256=rollout_recovery_digest, + rollout_recovery_group_count=rollout_recovery_group_count, + ) + return _RolloutCheckpointCut( + dataloader_state=dataloader_state, + replay_metadata=replay_metadata, + rollout_recovery_payload=rollout_recovery_payload, + rollout_recovery_group_count=rollout_recovery_group_count, + rolled_back_train_group_count=len(additional_groups), + mutation_version=self._data_plane_checkpoint_barrier.mutation_version, + ) + + async def _write_rollout_checkpoint_sidecars( + self, + checkpoint_path: PathLike, + cut: _RolloutCheckpointCut, + *, + include_config: bool, + ) -> None: + """Write controller files that belong to a captured native TQ cut.""" + checkpoint_path = Path(checkpoint_path) + await asyncio.to_thread( + torch.save, + cut.dataloader_state, + checkpoint_path / "train_dataloader.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, + ) + if include_config: + dumped_config = self._master_config.model_dump(mode="json") + + def _write_config() -> None: + with (checkpoint_path / "config.yaml").open("w") as config_file: + yaml.safe_dump(dumped_config, config_file) + + await asyncio.to_thread(_write_config) + + async def _save_rollout_checkpoint(self, *, force: bool = False) -> bool: + """Publish one lightweight 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_fingerprint is None: + raise RuntimeError( + "rollout snapshotting requires a bootstrap fingerprint" + ) + anchor = await asyncio.to_thread( + ensure_bootstrap_anchor, + self._checkpointer.checkpoint_dir, + fingerprint=self._bootstrap_fingerprint, + ) + snapshot_fingerprint = self._bootstrap_fingerprint + else: + if self._trainer_version != self._train_steps: + raise RuntimeError( + "rollout snapshot trainer identity is ambiguous: " + f"step={self._train_steps}, " + f"trainer_version={self._trainer_version}" + ) + anchor = self._checkpointer.checkpoint_dir / f"step_{self._train_steps}" + if not anchor.is_dir(): + skip_key = (self._train_steps, self._trainer_version) + if self._last_missing_rollout_snapshot_anchor != skip_key: + print( + "rollout checkpoint skipped: matching trainer " + f"checkpoint is not durable yet: {anchor}", + flush=True, + ) + self._last_missing_rollout_snapshot_anchor = skip_key + return False + try: + await asyncio.to_thread( + prune_bootstrap_snapshots, + self._checkpointer.checkpoint_dir, + durable_trainer_checkpoint=anchor, + ) + except OSError as error: + warnings.warn( + "Failed to prune obsolete bootstrap rollout snapshots: " + f"{type(error).__name__}: {error}", + stacklevel=2, + ) + snapshot_fingerprint = None + + expected_train_step = self._train_steps + expected_trainer_version = self._trainer_version + tmp_path, final_path, _ = await asyncio.to_thread( + prepare_snapshot_paths, anchor + ) + try: + async with self._data_plane_checkpoint_barrier.checkpoint(): + if ( + self._optimizer_commit_in_progress + or self._train_steps != expected_train_step + or self._trainer_version != expected_trainer_version + ): + await asyncio.to_thread(shutil.rmtree, tmp_path) + return False + snapshot_epoch = self._current_epoch + cut = await self._capture_rollout_checkpoint_cut( + tmp_path, + periodic=True, + ) + + await self._write_rollout_checkpoint_sidecars( + tmp_path, + cut, + include_config=True, + ) + manifest = RolloutSnapshotManifest( + schema_version=ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + base_train_step=expected_train_step, + trainer_version=expected_trainer_version, + current_epoch=snapshot_epoch, + mutation_version=cut.mutation_version, + rolled_back_train_group_count=(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(shutil.rmtree, tmp_path) + raise + + self._last_rollout_snapshot_mutation_version = cut.mutation_version + self._last_missing_rollout_snapshot_anchor = None + print( + "rollout checkpoint save completed: " + f"{final_path} (step={expected_train_step}, " + f"trainer_version={expected_trainer_version}, " + f"ledger_groups={cut.rollout_recovery_group_count or 0})", + flush=True, + ) + return True + + async def _rollout_checkpoint_pump(self) -> None: + """Persist rollout state periodically, including during streamed train.""" + interval_s = self._master_config.rollout_checkpointing.interval_s + if interval_s is None: + raise RuntimeError("rollout checkpoint pump started while disabled") + while True: + await asyncio.sleep(interval_s) + deadline_due = self._train_steps == 0 and self._timeout.would_save() + try: + saved = await self._save_rollout_checkpoint(force=deadline_due) + except Exception as error: + if deadline_due: + raise RuntimeError( + "failed to save the required pre-step rollout checkpoint" + ) from error + warnings.warn( + "Periodic rollout checkpoint failed; retaining the previous " + f"committed snapshot: {type(error).__name__}: {error}", + stacklevel=2, + ) + continue + if deadline_due: + if not saved: + continue + if not self._timeout.check_save(): + continue + print( + "Checkpoint deadline reached before the first train step; " + "stopping after a durable rollout snapshot", + flush=True, + ) + self._rollout_checkpoint_stop_requested.set() + return + @staticmethod def _group_ids_from_meta(meta: KVBatchMeta) -> list[str]: """Return stable prompt-group IDs in their canonical sample order.""" @@ -1230,6 +1589,7 @@ async def _rollout_pump(self) -> None: async def _dispatch_one_prompt( prompt: DatumSpec, target_step: Optional[int], + recovery_group_id: Optional[str], task_started_event: asyncio.Event, dispatch_admitted_event: asyncio.Event, ) -> None: @@ -1247,10 +1607,15 @@ async def _dispatch_one_prompt( self._inflight_rollouts += 1 counted_inflight = True if self._finalizer_actors: + finalization_kwargs: dict[str, Any] = { + "target_step": target_step, + "inflight_registry": self._inflight_by_group_id, + } + if recovery_group_id is not None: + finalization_kwargs["recovery_group_id"] = recovery_group_id request = await self._rollout_manager.generate_for_finalization( prompt, - target_step=target_step, - inflight_registry=self._inflight_by_group_id, + **finalization_kwargs, ) if request is None: outcome = RolloutOutcome.SKIPPED @@ -1309,60 +1674,175 @@ def _release_buffer_if_task_not_started( self._buffer_capacity.release() dispatch_admitted_event.set() - max_epochs = self._master_config.grpo.max_num_epochs - async with asyncio.TaskGroup() as rollout_tasks: - while max_epochs is None or self._current_epoch < max_epochs: - for prompt_batch in self._dataloader: - target_step = await self._sampler.admit( - trainer_version_fn=lambda: self._trainer_version - ) + def _release_capacity(permits: int) -> None: + for _ in range(permits): + self._buffer_capacity.release() - num_prompts = prompt_batch.size - if target_step is not None: - buffered = self._buffer.count_for_target_step(target_step) - if buffered: - num_prompts = max(0, prompt_batch.size - buffered) - print( - f" target_step={target_step}: {buffered} group(s) " - f"already buffered; dispatching {num_prompts} of " - f"{prompt_batch.size} prompt(s), dropping the rest", - flush=True, - ) + async def _acquire_capacity(permits: int) -> None: + acquired = 0 + try: + for _ in range(permits): + await self._buffer_capacity.acquire() + acquired += 1 + except BaseException: + _release_capacity(acquired) + raise - for prompt_idx in range(num_prompts): - prompt: DatumSpec = { # type: ignore - k: v[prompt_idx] for k, v in prompt_batch.items() - } + def _num_prompts_to_dispatch( + batch_size: int, target_step: Optional[int] + ) -> int: + """Return only the target-step shortfall after replay restoration.""" + if target_step is None: + return batch_size + buffered = self._buffer.count_for_target_step(target_step) + return max(0, batch_size - buffered) - # check if buffer is full - await self._buffer_capacity.acquire() - - task_started_event = asyncio.Event() - dispatch_admitted_event = asyncio.Event() - # dispatch rollout - task = rollout_tasks.create_task( - _dispatch_one_prompt( - prompt, - target_step, - task_started_event, - dispatch_admitted_event, - ) + max_epochs = self._master_config.grpo.max_num_epochs + token_capture_config = getattr(self._master_config, "token_capture", None) + token_capture_enabled = bool( + token_capture_config is not None and token_capture_config.enabled + ) + if token_capture_enabled and not isinstance( + self._sampler, TwoPhaseAdmissionSampler + ): + raise RuntimeError( + f"token-capture sampler {type(self._sampler).__name__} must " + "support two-phase admission" + ) + + async with asyncio.TaskGroup() as rollout_tasks: + while max_epochs is None or self._current_epoch < max_epochs: + dataloader_iterator = iter(self._dataloader) + while True: + preacquired_capacity = 0 + if token_capture_enabled: + await self._sampler.wait_until_admissible( + trainer_version_fn=lambda: self._trainer_version ) - self._dispatched_rollouts.add(task) - task.add_done_callback(self._dispatched_rollouts.discard) - task.add_done_callback( - partial( - _release_buffer_if_task_not_started, - task_started_event=task_started_event, - dispatch_admitted_event=dispatch_admitted_event, - ) + configured_batch_size = ( + self._master_config.grpo.num_prompts_per_step + ) + await _acquire_capacity(configured_batch_size) + preacquired_capacity = configured_batch_size + end_of_epoch = False + try: + # The cursor and all durable prompt reservations move + # together. A snapshot therefore cannot skip a batch + # whose lineage was not yet made recoverable. + async with self._data_plane_checkpoint_barrier.mutation(): + try: + prompt_batch = next(dataloader_iterator) + except StopIteration: + end_of_epoch = True + if end_of_epoch: + # Keep the outer epoch counter in the same + # checkpoint cut as the exhausted loader + # cursor. Otherwise a snapshot can observe + # end-of-epoch data with the previous epoch + # number and replay the dataset on restart. + self._current_epoch += 1 + else: + if prompt_batch.size > preacquired_capacity: + raise RuntimeError( + "dataloader batch exceeds reserved " + "rollout capacity" + ) + target_step = self._sampler.commit_admission( + trainer_version=self._trainer_version + ) + num_prompts = _num_prompts_to_dispatch( + prompt_batch.size, target_step + ) + unused_capacity = preacquired_capacity - num_prompts + _release_capacity(unused_capacity) + preacquired_capacity = num_prompts + rollout_work_items = [] + for prompt_idx in range(num_prompts): + prompt: DatumSpec = { # type: ignore + key: value[prompt_idx] + for key, value in prompt_batch.items() + } + group_id = ( + self._rollout_manager.reserve_prompt_group( + prompt, + target_step=target_step, + ) + ) + rollout_work_items.append( + _RolloutWorkItem( + prompt=prompt, + target_step=target_step, + group_id=group_id, + ) + ) + except BaseException: + _release_capacity(preacquired_capacity) + raise + if end_of_epoch: + _release_capacity(preacquired_capacity) + break + else: + try: + prompt_batch = next(dataloader_iterator) + except StopIteration: + self._current_epoch += 1 + break + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version ) - # Keep dataloader production bounded by actual worker - # admission instead of filling the buffer with tasks - # queued behind recovery. - await dispatch_admitted_event.wait() + num_prompts = _num_prompts_to_dispatch( + prompt_batch.size, target_step + ) + rollout_work_items = [] + for prompt_idx in range(num_prompts): + prompt = { # type: ignore + key: value[prompt_idx] + for key, value in prompt_batch.items() + } + rollout_work_items.append( + _RolloutWorkItem(prompt, target_step, None) + ) - self._current_epoch += 1 + pending_preacquired_capacity = preacquired_capacity + try: + for work_item in rollout_work_items: + capacity_acquired_here = False + if not token_capture_enabled: + await self._buffer_capacity.acquire() + capacity_acquired_here = True + task_started_event = asyncio.Event() + dispatch_admitted_event = asyncio.Event() + try: + task = rollout_tasks.create_task( + _dispatch_one_prompt( + work_item.prompt, + work_item.target_step, + work_item.group_id, + task_started_event, + dispatch_admitted_event, + ) + ) + except BaseException: + if capacity_acquired_here: + self._buffer_capacity.release() + elif token_capture_enabled: + pending_preacquired_capacity -= 1 + self._buffer_capacity.release() + raise + if token_capture_enabled: + pending_preacquired_capacity -= 1 + self._dispatched_rollouts.add(task) + task.add_done_callback(self._dispatched_rollouts.discard) + task.add_done_callback( + partial( + _release_buffer_if_task_not_started, + task_started_event=task_started_event, + dispatch_admitted_event=dispatch_admitted_event, + ) + ) + await dispatch_admitted_event.wait() + finally: + _release_capacity(pending_preacquired_capacity) # Drain in-flight so return implies "all rollouts in TQ". inflight = list(self._dispatched_rollouts) @@ -1572,6 +2052,7 @@ async def _train_pump(self) -> None: with self._timer.time("policy_training"): result = await asyncio.to_thread(self._trainer.finish_train_step) + self._optimizer_commit_in_progress = True async with self._data_plane_checkpoint_barrier.mutation(): if self._rollout_recovery_ledger is not None: self._rollout_recovery_ledger.mark_train_step_applied( @@ -1600,6 +2081,7 @@ async def _train_pump(self) -> None: self._trainer_version += 1 self._train_steps += 1 + self._optimizer_commit_in_progress = False with self._timer.time("weight_sync"): calibration_data = ( BatchedDataDict.from_batches(calibration_batches) @@ -1838,6 +2320,11 @@ async def _abort_stale_inflight(self) -> int: return len(stale_tasks) async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: + """Serialize full and rollout-only checkpoint publication.""" + async with self._checkpoint_save_lock: + await self._save_checkpoint_locked(step_metrics) + + async def _save_checkpoint_locked(self, step_metrics: dict[str, Any]) -> None: """Write a full checkpoint for the just-finished train step. Everything except the (possibly async) policy weight write must be @@ -1851,9 +2338,7 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: save_state.consumed_samples = self._consumed_samples save_state.total_valid_tokens = self._total_valid_tokens save_state.sampler_name = self._async_cfg.sampler.name - # Snapshot before any await so it can't interleave with - # _rollout_pump iterating this same dataloader. - dataloader_state = self._dataloader.state_dict() + dataloader_state: Optional[dict[str, Any]] = None # SC has no validation loop yet; drop the default sentinel instead of # persisting a bogus val_reward. if hasattr(save_state, "val_reward"): @@ -1873,28 +2358,7 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: await asyncio.to_thread(self._checkpointer.finalize_pending) print(f"Saving checkpoint for step {self._train_steps}...") - checkpoint_path: PathLike = await asyncio.to_thread( # pyrefly: ignore[bad-assignment] the PathLike alias resolves inconsistently under pyrefly's import-cycle breaking - self._checkpointer.init_tmp_checkpoint, - self._train_steps, - vars(save_state), - self._master_config, - ) - # With async_save this returns after D2H staging; disk writes finish - # in the background. - await asyncio.to_thread( - self._trainer.save_checkpoint, - weights_path=os.path.join(checkpoint_path, "policy", "weights"), - optimizer_path=os.path.join(checkpoint_path, "policy", "optimizer") - if self._checkpointer.save_optimizer - else None, - tokenizer_path=os.path.join(checkpoint_path, "policy", "tokenizer"), - checkpointing_cfg=self._master_config.checkpointing, - ) - await asyncio.to_thread( - torch.save, - dataloader_state, - os.path.join(checkpoint_path, "train_dataloader.pt"), - ) + checkpoint_path: PathLike replay_metadata: Optional[TQReplayMetadataState] = None rollout_recovery_payload: Optional[bytes] = None rollout_recovery_digest: Optional[str] = None @@ -1910,6 +2374,17 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: # wait at commit, so TQ and the metadata sidecar describe exactly # the same set of training-ready groups. async with self._data_plane_checkpoint_barrier.checkpoint(): + # Dataloader advancement and prompt-lineage reservation use the + # mutation side of this barrier, so this cursor belongs to the + # same exact cut as replay metadata, lineage, and native TQ. + save_state.current_epoch = self._current_epoch + checkpoint_path = await asyncio.to_thread( + self._checkpointer.init_tmp_checkpoint, + self._train_steps, + vars(save_state), + self._master_config, + ) + dataloader_state = self._dataloader.state_dict() if self._sampler.supports_buffer_checkpoint: replay_metadata = self._buffer.metadata_state_dict( saved_capacity=self._async_cfg.max_buffered_rollouts @@ -1936,6 +2411,35 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: rollout_recovery_payload_sha256=rollout_recovery_digest, rollout_recovery_group_count=rollout_recovery_group_count, ) + self._last_rollout_snapshot_mutation_version = ( + self._data_plane_checkpoint_barrier.mutation_version + ) + else: + save_state.current_epoch = self._current_epoch + checkpoint_path = await asyncio.to_thread( + self._checkpointer.init_tmp_checkpoint, + self._train_steps, + vars(save_state), + self._master_config, + ) + dataloader_state = self._dataloader.state_dict() + # With async_save this returns after D2H staging; disk writes finish + # in the background. + await asyncio.to_thread( + self._trainer.save_checkpoint, + weights_path=os.path.join(checkpoint_path, "policy", "weights"), + optimizer_path=os.path.join(checkpoint_path, "policy", "optimizer") + if self._checkpointer.save_optimizer + else None, + tokenizer_path=os.path.join(checkpoint_path, "policy", "tokenizer"), + checkpointing_cfg=self._master_config.checkpointing, + ) + assert dataloader_state is not None + await asyncio.to_thread( + torch.save, + dataloader_state, + os.path.join(checkpoint_path, "train_dataloader.pt"), + ) if replay_metadata is not None: await asyncio.to_thread( torch.save, diff --git a/nemo_rl/algorithms/single_controller_utils/__init__.py b/nemo_rl/algorithms/single_controller_utils/__init__.py index 198bea987a1..b5111d487d3 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, ) @@ -30,6 +31,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 e05d56ec794..db9e777763d 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -299,6 +299,29 @@ class TokenCaptureConfig(BaseModel, extra="allow"): num_finalizer_workers: PositiveInt = 2 +class RolloutCheckpointConfig(BaseModel, extra="allow"): + """Frequent TQ/lineage snapshots anchored to durable trainer state. + + ``interval_s=None`` disables periodic snapshots. A snapshot taken before + the first trainer checkpoint is anchored to the initial model and a + rollout-semantic configuration fingerprint. Later snapshots are anchored + to the most recent durable trainer checkpoint. + + ``restore_mode="latest"`` selects the newest compatible periodic snapshot. + ``trainer_checkpoint`` ignores newer periodic snapshots, while ``none`` + resumes trainer state without restoring rollout, replay, lineage, or the + dataloader cursor. + + SingleController has no validation loop, so checkpoint selection must use + ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited + ``val:`` settings are rejected during setup. + """ + + interval_s: Optional[float] = Field(default=None, gt=0) + keep_latest_k: int = Field(default=2, ge=1) + restore_mode: Literal["latest", "trainer_checkpoint", "none"] = "latest" + + class MasterConfig(BaseModel, extra="allow"): policy: PolicyConfig loss_fn: ClippedPGLossConfig @@ -311,6 +334,9 @@ class MasterConfig(BaseModel, extra="allow"): data_plane: DataPlaneConfig async_rl: AsyncRLConfig token_capture: TokenCaptureConfig = Field(default_factory=TokenCaptureConfig) + rollout_checkpointing: RolloutCheckpointConfig = Field( + default_factory=RolloutCheckpointConfig + ) def validate_sampler_buffer_capacity( @@ -332,6 +358,13 @@ def validate_sampler_buffer_capacity( ) +def required_rollout_recovery_capacity( + *, num_prompts_per_step: int, min_groups_for_streaming_train: int +) -> int: + """Capacity needed to fill a partially restored streaming batch.""" + return num_prompts_per_step + min_groups_for_streaming_train - 1 + + def validate_single_controller_config(master_config: MasterConfig) -> None: """Validate cross-section SingleController constraints before setup.""" async_config = master_config.async_rl @@ -350,6 +383,38 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: f"({async_config.min_groups_for_streaming_train}); otherwise the rollout " "pump can fill every buffer slot while the trainer waits for more groups." ) + if ( + master_config.token_capture.enabled + and async_config.max_buffered_rollouts < num_prompts_per_step + ): + raise ValueError( + "token capture requires async_rl.max_buffered_rollouts " + f"({async_config.max_buffered_rollouts}) to be >= " + f"grpo.num_prompts_per_step ({num_prompts_per_step}) so one " + "dataloader batch can reserve capacity atomically" + ) + if ( + master_config.token_capture.enabled + and master_config.checkpointing["enabled"] + and master_config.rollout_checkpointing.restore_mode != "none" + ): + recovery_capacity = required_rollout_recovery_capacity( + num_prompts_per_step=num_prompts_per_step, + min_groups_for_streaming_train=( + async_config.min_groups_for_streaming_train + ), + ) + if async_config.max_buffered_rollouts < recovery_capacity: + raise ValueError( + "token-capture rollout recovery requires " + "async_rl.max_buffered_rollouts " + f"({async_config.max_buffered_rollouts}) to be >= " + "grpo.num_prompts_per_step + " + "async_rl.min_groups_for_streaming_train - 1 " + f"({recovery_capacity}); otherwise restored groups can leave " + "the trainer below its streaming threshold while fresh batch " + "reservation waits for capacity" + ) rl_step_samples = ( num_prompts_per_step * master_config.grpo.num_generations_per_prompt 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..40b8fd191bc --- /dev/null +++ b/nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py @@ -0,0 +1,515 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Filesystem contract for frequent Single Controller rollout snapshots.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Mapping, Optional + +from nemo_rl.algorithms.single_controller_utils.config import MasterConfig + +ROLLOUT_SNAPSHOT_SCHEMA_VERSION = 1 +BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION = 2 +BOOTSTRAP_DIRNAME = "bootstrap" +BOOTSTRAP_MANIFEST_FILENAME = "manifest.json" +ROLLOUT_SNAPSHOTS_DIRNAME = "rollout_snapshots" +ROLLOUT_SNAPSHOT_MANIFEST_FILENAME = "manifest.json" +ROLLOUT_SNAPSHOT_COMMITTED_FILENAME = "COMMITTED" +ROLLOUT_SNAPSHOT_LATEST_FILENAME = "LATEST" + +_SNAPSHOT_RE = re.compile(r"snapshot_(\d+)") + +# Keep this projection limited to values needed to interpret persisted rollout +# state or execute missing siblings. Trainer-only and post-rollout settings do +# not belong in bootstrap compatibility. +_BOOTSTRAP_POLICY_FIELDS = frozenset( + { + "model_name", + "pretrained_checkpoint", + "hf_config_overrides", + "max_total_sequence_length", + "tokenizer", + } +) +_BOOTSTRAP_GENERATION_FIELDS = frozenset( + { + "backend", + "max_new_tokens", + "stop_strings", + "stop_token_ids", + "temperature", + "top_k", + "top_p", + } +) +_BOOTSTRAP_VLLM_FIELDS = frozenset( + { + "http_server_serving_chat_kwargs", + "max_model_len", + "reasoning_parser_plugin", + } +) +_BOOTSTRAP_GRPO_FIELDS = frozenset( + { + "max_rollout_turns", + "num_generations_per_prompt", + "num_prompts_per_step", + "seed", + } +) +_BOOTSTRAP_DATA_FIELDS = frozenset( + { + "default", + "max_input_seq_length", + "shuffle", + "train", + } +) +_BOOTSTRAP_TOKEN_CAPTURE_FIELDS = frozenset( + { + "enabled", + "min_valid_fraction_per_group", + "mixed_weight_version_policy", + "staging_partition", + } +) +_BOOTSTRAP_ENV_RUNTIME_FIELDS = frozenset( + { + "apptainer_memory_limit_mb", + "concurrency", + "nemo_gym_log_dir", + "num_gpu_nodes", + "port_range_high", + "port_range_low", + "should_log_nemo_gym_responses", + "skip_venv_if_present", + "use_absolute_ip", + } +) + + +def _select_fields( + mapping: Mapping[str, Any] | None, + fields: frozenset[str], +) -> dict[str, Any]: + """Select explicitly rollout-semantic fields from one config section.""" + if mapping is None: + return {} + return {key: mapping[key] for key in sorted(fields) if key in mapping} + + +def _drop_runtime_fields(value: Any, runtime_fields: frozenset[str]) -> Any: + """Recursively strip known operational leaves from an environment.""" + if isinstance(value, Mapping): + return { + key: _drop_runtime_fields(child, runtime_fields) + for key, child in value.items() + if key not in runtime_fields + } + if isinstance(value, list): + return [_drop_runtime_fields(child, runtime_fields) for child in value] + return value + + +def _fsync_file(path: Path) -> None: + """Flush one completed regular file to its backing filesystem.""" + with path.open("rb") as file_obj: + os.fsync(file_obj.fileno()) + + +def _fsync_directory(path: Path) -> None: + """Flush directory-entry updates such as rename and replace.""" + flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0) + directory_fd = os.open(path, flags) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + + +def _fsync_tree(root: Path) -> None: + """Flush every snapshot payload before publishing its commit marker.""" + for directory, _, filenames in os.walk(root, topdown=False): + directory_path = Path(directory) + for filename in filenames: + file_path = directory_path / filename + if not file_path.is_symlink() and file_path.is_file(): + _fsync_file(file_path) + _fsync_directory(directory_path) + + +def _snapshot_sequence(path: Path) -> int: + match = _SNAPSHOT_RE.fullmatch(path.name) + if match is None: + raise ValueError(f"not a rollout snapshot directory: {path}") + return int(match.group(1)) + + +@dataclass(frozen=True) +class RolloutSnapshotManifest: + """Identity binding one rollout-state cut to reconstructable trainer state.""" + + schema_version: int + base_train_step: int + trainer_version: int + current_epoch: int + mutation_version: int + rolled_back_train_group_count: int + bootstrap_fingerprint: Optional[str] + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any]) -> RolloutSnapshotManifest: + """Parse and validate a committed snapshot manifest.""" + required_ints = ( + "schema_version", + "base_train_step", + "trainer_version", + "current_epoch", + "mutation_version", + "rolled_back_train_group_count", + ) + for key in required_ints: + value = raw.get(key) + if not isinstance(value, int) or isinstance(value, bool): + raise ValueError( + f"rollout snapshot manifest {key!r} must be an integer" + ) + fingerprint = raw.get("bootstrap_fingerprint") + if fingerprint is not None and not isinstance(fingerprint, str): + raise ValueError( + "rollout snapshot bootstrap_fingerprint must be a string or null" + ) + manifest = cls( + schema_version=raw["schema_version"], + base_train_step=raw["base_train_step"], + trainer_version=raw["trainer_version"], + current_epoch=raw["current_epoch"], + mutation_version=raw["mutation_version"], + rolled_back_train_group_count=raw["rolled_back_train_group_count"], + bootstrap_fingerprint=fingerprint, + ) + if manifest.schema_version != ROLLOUT_SNAPSHOT_SCHEMA_VERSION: + raise ValueError( + "unsupported rollout snapshot schema version: " + f"{manifest.schema_version}" + ) + if ( + min( + manifest.base_train_step, + manifest.trainer_version, + manifest.current_epoch, + manifest.mutation_version, + manifest.rolled_back_train_group_count, + ) + < 0 + ): + raise ValueError("rollout snapshot counters must be non-negative") + return manifest + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +@dataclass(frozen=True) +class ResolvedRolloutCheckpoint: + """A committed rollout snapshot selected for startup recovery.""" + + path: Path + manifest: RolloutSnapshotManifest + + +@dataclass(frozen=True) +class BootstrapCompatibilityIdentity: + """Rollout-semantic inputs that must match a trainer-version-zero cut.""" + + schema_version: int + model: Mapping[str, Any] + generation: Mapping[str, Any] + rollout: Mapping[str, Any] + dataset: Mapping[str, Any] + environment: Mapping[str, Any] + sampler: Mapping[str, Any] + token_capture: Mapping[str, Any] + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def bootstrap_compatibility_identity( + master_config: MasterConfig, +) -> BootstrapCompatibilityIdentity: + """Project a full run config onto inputs that affect recovered rollouts. + + Dataset identity is intentionally retained because a bootstrap snapshot + restores the dataloader cursor together with its unfinished prompt ledger. + Cluster shape, logging, checkpoint paths, worker counts, and other runtime + tuning are excluded so they may change across a restart. + """ + dumped = master_config.model_dump(mode="json") + policy = dumped.get("policy", {}) + generation = policy.get("generation", {}) + generation_identity = _select_fields( + generation, + _BOOTSTRAP_GENERATION_FIELDS, + ) + vllm_identity = _select_fields( + generation.get("vllm_cfg", {}), + _BOOTSTRAP_VLLM_FIELDS, + ) + if vllm_identity: + generation_identity["vllm_cfg"] = vllm_identity + + async_rl = dumped.get("async_rl", {}) + sampler = async_rl.get("sampler", {}) + if not isinstance(sampler, Mapping): + sampler = {} + + return BootstrapCompatibilityIdentity( + schema_version=BOOTSTRAP_COMPATIBILITY_SCHEMA_VERSION, + model=_select_fields(policy, _BOOTSTRAP_POLICY_FIELDS), + generation=generation_identity, + rollout=_select_fields(dumped.get("grpo", {}), _BOOTSTRAP_GRPO_FIELDS), + dataset=_select_fields(dumped.get("data", {}), _BOOTSTRAP_DATA_FIELDS), + environment=_drop_runtime_fields( + dumped.get("env", {}), + _BOOTSTRAP_ENV_RUNTIME_FIELDS, + ), + sampler=dict(sampler), + token_capture=_select_fields( + dumped.get("token_capture", {}), + _BOOTSTRAP_TOKEN_CAPTURE_FIELDS, + ), + ) + + +def bootstrap_fingerprint(master_config: MasterConfig) -> str: + """Hash rollout-semantic inputs needed to reuse a bootstrap snapshot. + + This is a compatibility guard, not a hash of the full training recipe. + Operational settings are deliberately excluded so a restart may use a + different cluster shape, checkpoint interval, or logging destination. + """ + payload = json.dumps( + bootstrap_compatibility_identity(master_config).to_dict(), + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def prune_bootstrap_snapshots( + checkpoint_dir: Path, + *, + durable_trainer_checkpoint: Path, +) -> bool: + """Remove trainer-version-zero snapshots once a trainer checkpoint exists.""" + if not durable_trainer_checkpoint.is_dir(): + raise FileNotFoundError( + "cannot prune bootstrap snapshots without a durable trainer " + f"checkpoint: {durable_trainer_checkpoint}" + ) + snapshot_root = checkpoint_dir / BOOTSTRAP_DIRNAME / ROLLOUT_SNAPSHOTS_DIRNAME + if not snapshot_root.is_dir(): + return False + shutil.rmtree(snapshot_root) + return True + + +def ensure_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: + """Create or validate the lightweight trainer-version-zero anchor.""" + anchor = checkpoint_dir / BOOTSTRAP_DIRNAME + anchor.mkdir(parents=True, exist_ok=True) + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + expected = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "bootstrap_fingerprint": fingerprint, + } + if manifest_path.is_file(): + raw = json.loads(manifest_path.read_text()) + if raw != expected: + raise ValueError( + "existing rollout bootstrap anchor does not match the current " + f"trainer configuration: checkpoint={raw!r}, expected={expected!r}" + ) + return anchor + + tmp_path = manifest_path.with_suffix(".json.tmp") + tmp_path.write_text(json.dumps(expected, sort_keys=True, indent=2) + "\n") + _fsync_file(tmp_path) + os.replace(tmp_path, manifest_path) + _fsync_directory(anchor) + _fsync_directory(anchor.parent) + return anchor + + +def reset_bootstrap_anchor(checkpoint_dir: Path, *, fingerprint: str) -> Path: + """Discard skipped pre-step snapshots and start a new bootstrap lineage. + + This is used only when restore mode deliberately skips partial-rollout + recovery and no trainer checkpoint exists. The user's restore choice makes + the previous state intentionally unreachable; removing it also prevents a + later periodic save from appending to an incompatible bootstrap anchor. + """ + anchor = checkpoint_dir / BOOTSTRAP_DIRNAME + snapshot_root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + if snapshot_root.exists(): + shutil.rmtree(snapshot_root) + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + if manifest_path.exists(): + manifest_path.unlink() + return ensure_bootstrap_anchor(checkpoint_dir, fingerprint=fingerprint) + + +def validate_bootstrap_anchor(anchor: Path, *, fingerprint: str) -> None: + """Fail loudly when bootstrap snapshots belong to different initial state.""" + manifest_path = anchor / BOOTSTRAP_MANIFEST_FILENAME + if not manifest_path.is_file(): + raise FileNotFoundError( + f"rollout bootstrap manifest is missing at {manifest_path}" + ) + raw = json.loads(manifest_path.read_text()) + expected = { + "schema_version": ROLLOUT_SNAPSHOT_SCHEMA_VERSION, + "base_train_step": 0, + "trainer_version": 0, + "bootstrap_fingerprint": fingerprint, + } + if raw != expected: + raise ValueError( + "rollout bootstrap anchor does not match the current trainer " + f"configuration: checkpoint={raw!r}, expected={expected!r}" + ) + + +def prepare_snapshot_paths(anchor: Path) -> tuple[Path, Path, int]: + """Allocate the next temporary/final snapshot directory pair.""" + root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + root.mkdir(parents=True, exist_ok=True) + sequences = [ + int(match.group(1)) + for child in root.iterdir() + if child.is_dir() and (match := _SNAPSHOT_RE.fullmatch(child.name)) + ] + sequence = max(sequences, default=0) + 1 + final_path = root / f"snapshot_{sequence:06d}" + tmp_path = root / f"tmp_snapshot_{sequence:06d}" + if tmp_path.exists(): + shutil.rmtree(tmp_path) + tmp_path.mkdir(parents=True) + return tmp_path, final_path, sequence + + +def commit_snapshot( + tmp_path: Path, + final_path: Path, + *, + keep_latest_k: int, +) -> None: + """Atomically publish one validated snapshot and retain recent fallbacks.""" + if keep_latest_k < 1: + raise ValueError("rollout snapshot retention must keep at least one snapshot") + _fsync_tree(tmp_path) + committed_path = tmp_path / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME + committed_path.write_text("committed\n") + _fsync_file(committed_path) + _fsync_directory(tmp_path) + os.rename(tmp_path, final_path) + + root = final_path.parent + _fsync_directory(root) + latest_path = root / ROLLOUT_SNAPSHOT_LATEST_FILENAME + latest_tmp = latest_path.with_suffix(".tmp") + latest_tmp.write_text(final_path.name + "\n") + _fsync_file(latest_tmp) + os.replace(latest_tmp, latest_path) + _fsync_directory(root) + + committed = sorted( + ( + child + for child in root.iterdir() + if child.is_dir() + and _SNAPSHOT_RE.fullmatch(child.name) + and (child / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() + ), + key=_snapshot_sequence, + reverse=True, + ) + stale_snapshots = committed[keep_latest_k:] + for stale in stale_snapshots: + shutil.rmtree(stale) + if stale_snapshots: + _fsync_directory(root) + + +def resolve_latest_snapshot( + anchor: Path, + *, + expected_train_step: int, + expected_trainer_version: int, + expected_bootstrap_fingerprint: Optional[str], +) -> Optional[ResolvedRolloutCheckpoint]: + """Select the newest complete snapshot compatible with its trainer anchor.""" + root = anchor / ROLLOUT_SNAPSHOTS_DIRNAME + if not root.is_dir(): + return None + + candidates = sorted( + ( + child + for child in root.iterdir() + if child.is_dir() and _SNAPSHOT_RE.fullmatch(child.name) + ), + key=_snapshot_sequence, + reverse=True, + ) + errors: list[str] = [] + for candidate in candidates: + if not (candidate / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file(): + continue + manifest_path = candidate / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME + if not manifest_path.is_file(): + errors.append(f"{candidate.name}: missing manifest") + continue + try: + raw = json.loads(manifest_path.read_text()) + manifest = RolloutSnapshotManifest.from_mapping(raw) + except (json.JSONDecodeError, OSError, ValueError) as error: + errors.append(f"{candidate.name}: {error}") + continue + if ( + manifest.base_train_step != expected_train_step + or manifest.trainer_version != expected_trainer_version + or manifest.bootstrap_fingerprint != expected_bootstrap_fingerprint + ): + errors.append(f"{candidate.name}: trainer-anchor mismatch") + continue + return ResolvedRolloutCheckpoint(candidate, manifest) + + if errors: + raise ValueError( + "no committed rollout snapshot matches the selected trainer anchor: " + + "; ".join(errors) + ) + return None diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index d795622760d..a3706e1a8c0 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -63,6 +63,15 @@ MasterConfig, validate_single_controller_config, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_DIRNAME, + ROLLOUT_SNAPSHOTS_DIRNAME, + bootstrap_fingerprint, + ensure_bootstrap_anchor, + reset_bootstrap_anchor, + resolve_latest_snapshot, + validate_bootstrap_anchor, +) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.utils import load_dataloader_state, setup_response_data @@ -127,6 +136,7 @@ class SingleControllerActorArgs: save_state: GRPOSaveState last_checkpoint_path: Optional[str] data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None + bootstrap_fingerprint: Optional[str] = None def _maybe_restore_native_data_plane_checkpoint( @@ -672,16 +682,20 @@ def setup_single_controller( if ( master_config.checkpointing["enabled"] and sampler_supports_buffer_checkpoint(master_config.async_rl.sampler) + and master_config.rollout_checkpointing.restore_mode != "none" and not dp_config.get("checkpointing_enabled") ): raise ValueError( "SingleController checkpointing with a replay-checkpoint-capable " "sampler requires data_plane.checkpointing_enabled=true so " - "completed, unconsumed rollouts are recoverable." + "completed, unconsumed rollouts are recoverable. Set " + "rollout_checkpointing.restore_mode='none' to explicitly resume " + "trainer state without rollout recovery." ) if ( master_config.checkpointing["enabled"] and master_config.token_capture.enabled + and master_config.rollout_checkpointing.restore_mode != "none" and not dp_config.get("checkpointing_enabled") ): raise ValueError( @@ -692,6 +706,7 @@ def setup_single_controller( if ( master_config.checkpointing["enabled"] and master_config.token_capture.enabled + and master_config.rollout_checkpointing.restore_mode != "none" and not sampler_supports_buffer_checkpoint(master_config.async_rl.sampler) ): raise ValueError( @@ -713,6 +728,25 @@ def setup_single_controller( # give capture-enabled vLLM workers a venv that carries nemo_gym (the # worker hosts Gym's capture core + adapter in-process). token_capture_cfg = master_config.token_capture + rollout_checkpoint_cfg = master_config.rollout_checkpointing + if rollout_checkpoint_cfg.interval_s is not None: + if not master_config.checkpointing["enabled"]: + raise ValueError( + "rollout checkpointing requires checkpointing.enabled=true" + ) + if not dp_config.get("checkpointing_enabled"): + raise ValueError( + "rollout checkpointing requires data_plane.checkpointing_enabled=true" + ) + if not token_capture_cfg.enabled: + raise ValueError( + "rollout checkpointing currently requires token_capture.enabled=true" + ) + if not sampler_supports_buffer_checkpoint(master_config.async_rl.sampler): + raise ValueError( + "rollout checkpointing requires a sampler that supports " + "replay-buffer recovery" + ) if token_capture_cfg.enabled: if not _should_use_nemo_gym(master_config): raise ValueError( @@ -766,12 +800,97 @@ 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 + ) + + restore_mode = rollout_checkpoint_cfg.restore_mode + restore_latest_rollout = restore_mode == "latest" + recovery_checkpoint_path = ( + trainer_checkpoint_path if restore_mode != "none" else None + ) + bootstrap_anchor = checkpointer.checkpoint_dir / BOOTSTRAP_DIRNAME + needs_bootstrap_identity = trainer_checkpoint_path is None and ( + rollout_checkpoint_cfg.interval_s is not None + or (restore_latest_rollout and bootstrap_anchor.is_dir()) + ) + bootstrap_digest = ( + bootstrap_fingerprint(master_config) if needs_bootstrap_identity else None + ) + resolved_snapshot = None + restored_trainer_version = ( + save_state.trainer_version + if save_state.trainer_version is not None + else save_state.current_step + ) + if trainer_checkpoint_path is not None and restore_latest_rollout: + resolved_snapshot = resolve_latest_snapshot( + Path(trainer_checkpoint_path), + expected_train_step=save_state.current_step, + expected_trainer_version=restored_trainer_version, + expected_bootstrap_fingerprint=None, + ) + elif trainer_checkpoint_path is None: + if rollout_checkpoint_cfg.interval_s is not None: + assert bootstrap_digest is not None + if restore_latest_rollout: + bootstrap_anchor = ensure_bootstrap_anchor( + checkpointer.checkpoint_dir, + fingerprint=bootstrap_digest, + ) + else: + had_bootstrap_snapshots = ( + bootstrap_anchor / ROLLOUT_SNAPSHOTS_DIRNAME + ).is_dir() + bootstrap_anchor = reset_bootstrap_anchor( + checkpointer.checkpoint_dir, + fingerprint=bootstrap_digest, + ) + if had_bootstrap_snapshots: + print( + "📦 Ignored existing bootstrap rollout snapshots and " + "started a new bootstrap lineage because " + f"rollout_checkpointing.restore_mode={restore_mode!r}.", + flush=True, + ) + elif restore_latest_rollout and bootstrap_anchor.is_dir(): + assert bootstrap_digest is not None + validate_bootstrap_anchor( + bootstrap_anchor, + fingerprint=bootstrap_digest, + ) + if restore_latest_rollout and bootstrap_anchor.is_dir(): + resolved_snapshot = resolve_latest_snapshot( + bootstrap_anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint=bootstrap_digest, + ) + if resolved_snapshot is not None: + recovery_checkpoint_path = str(resolved_snapshot.path) + save_state.current_epoch = resolved_snapshot.manifest.current_epoch + print( + f"📦 Selected rollout recovery snapshot: {recovery_checkpoint_path}", + flush=True, + ) + elif restore_mode == "trainer_checkpoint" and trainer_checkpoint_path: + print( + "📦 Restoring rollout state from the durable trainer checkpoint " + f"without considering newer periodic snapshots: {trainer_checkpoint_path}", + flush=True, + ) + elif restore_mode == "none" and trainer_checkpoint_path: + print( + "📦 Resuming trainer state without rollout, replay, lineage, or " + "dataloader recovery.", + flush=True, + ) # ========================== # Setup Dataset & Environments @@ -804,9 +923,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) @@ -930,13 +1051,13 @@ 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, ) recovery_ledger = _maybe_restore_rollout_recovery_ledger( - last_checkpoint_path=last_checkpoint_path, + last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, token_capture_enabled=token_capture_cfg.enabled, ) @@ -1118,7 +1239,8 @@ def _build_generation_then_trainer( partition_id=partition_id, finalizer_actors=finalizer_actors, save_state=save_state, - last_checkpoint_path=last_checkpoint_path, + last_checkpoint_path=recovery_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + bootstrap_fingerprint=bootstrap_digest, ) return actor_args, setup_timing_metrics diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 7810bc7fb1b..4456c733ac0 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -1315,6 +1315,29 @@ def set_weight_version(self, version: int) -> None: """ self._weight_version = int(version) + def reserve_prompt_group( + self, input_sample: DatumSpec, *, target_step: Optional[int] = None + ) -> str: + """Reserve durable lineage before the dataloader advances past a batch. + + The caller owns the SC data-plane mutation barrier, making dataloader + advancement and every reservation in the batch one checkpoint cut. + """ + if self._recovery_ledger is None: + raise RuntimeError("prompt-group reservation requires token capture") + group = self._recovery_ledger.reserve_group( + prompt_id=str(input_sample["idx"]), + prompt_ref=PromptRef( + sample_id=str(input_sample["idx"]), + task_name=input_sample.get("task_name"), + ), + prompt_payload=input_sample, + expected_generations=self._num_generations_per_prompt, + target_step=target_step, + start_weight_version=self._weight_version, + ) + return group.group_id + async def run_rollout( self, input_sample: DatumSpec, @@ -1498,6 +1521,7 @@ async def generate_for_finalization( input_sample: DatumSpec, *, target_step: Optional[int] = None, + recovery_group_id: Optional[str] = None, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, ) -> Optional["FinalizationRequest"]: """Run capture generation and return a metadata-only actor request. @@ -1519,19 +1543,25 @@ async def generate_for_finalization( raise RuntimeError( "generate_for_finalization requires a rollout recovery ledger" ) - async with self._recovery_mutation(): - recovery_group = self._recovery_ledger.reserve_group( - prompt_id=str(input_sample["idx"]), - prompt_ref=PromptRef( - sample_id=str(input_sample["idx"]), - task_name=input_sample.get("task_name"), - ), - prompt_payload=input_sample, - expected_generations=self._num_generations_per_prompt, - target_step=target_step, - start_weight_version=self._weight_version, - ) - recovery_group_id = recovery_group.group_id + if recovery_group_id is None: + async with self._recovery_mutation(): + recovery_group_id = self.reserve_prompt_group( + input_sample, target_step=target_step + ) + else: + recovery_group = self._recovery_ledger.get_group(recovery_group_id) + if recovery_group.prompt_id != str(input_sample["idx"]): + raise ValueError( + "reserved rollout group belongs to a different prompt: " + f"group={recovery_group.prompt_id!r}, " + f"input={str(input_sample['idx'])!r}" + ) + if recovery_group.target_step != target_step: + raise ValueError( + "reserved rollout group target step does not match dispatch: " + f"group={recovery_group.target_step!r}, " + f"dispatch={target_step!r}" + ) return await self._run_finalization_with_retries( input_sample, recovery_group_id=recovery_group_id, diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index d7001537cfe..e7215f6ccfd 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -32,6 +32,7 @@ from nemo_rl.experience.route_plan import decode_route_plan if TYPE_CHECKING: + from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayGroupMetadata from nemo_rl.data.interfaces import DatumSpec ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 @@ -678,6 +679,65 @@ def rollback_open_train_step(self, train_step: int) -> None: record.claimed_train_step = None self._open_train_step = None + def training_owned_replay_groups(self) -> list[TQReplayGroupMetadata]: + """Return replay metadata for rows claimed by the open train step. + + Sampler selection removes these groups from the live replay index, but + their canonical rows intentionally remain in TQ until optimizer-step + completion. A periodic snapshot includes this metadata so restart can + discard uncommitted gradients and train those groups again. + """ + open_step = self._open_train_step + if open_step is None: + return [] + if open_step.status != TrainStepStatus.OPEN: + raise RuntimeError( + "cannot snapshot an optimizer step after it was applied but " + "before its data-plane cleanup completed" + ) + groups: list[TQReplayGroupMetadata] = [] + for group_id in open_step.group_ids: + record = self._require_group(group_id) + if ( + record.status != PromptGroupStatus.CLAIMED_FOR_TRAINING + or record.canonical_meta is None + or record.group_min_weight_version is None + or record.group_max_weight_version is None + ): + raise RuntimeError( + f"open train step has incomplete replay metadata for {group_id!r}" + ) + groups.append( + { + "meta": copy.deepcopy(record.canonical_meta), + "start_weight": record.group_min_weight_version, + "end_weight": record.group_max_weight_version, + "target_step": record.target_step, + "group_id": group_id, + } + ) + return groups + + def periodic_snapshot_state_dict(self) -> dict[str, Any]: + """Return restart-ready lineage without mutating the live train step. + + The trainer state is not part of a lightweight rollout snapshot. On + restart it comes from the snapshot's durable trainer anchor, so every + group claimed by the current optimizer step must become ``FINALIZED`` + and replayable again in the persisted lineage. + """ + snapshot = type(self).from_state_dict(self.state_dict()) + open_step = snapshot.open_train_step + if open_step is not None: + if open_step.status != TrainStepStatus.OPEN: + raise RuntimeError( + "cannot create a periodic snapshot after optimizer apply " + "but before data-plane cleanup" + ) + snapshot.rollback_open_train_step(open_step.train_step) + snapshot.prepare_for_restart() + return snapshot.state_dict() + def discard_group(self, group_id: str) -> None: """Drop a group only after its external TQ/Gate ownership is cleaned.""" record = self._require_group(group_id) 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 8f03d7a500d..b5d32bc0006 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -62,6 +62,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 727738c2ba0..5ffdf42244a 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -62,6 +62,17 @@ run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_r # custodying token lineage and the finalizer publishing training rows. run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh ++token_capture.enabled=true +# Two-process token-capture recovery. The first run is killed after a durable +# cut with a partially completed sibling group; the restart must reuse sealed +# siblings and redispatch only the missing ones. +run_test uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_tq_recovery.sh + +# Periodic snapshots are allowed while one streamed optimizer step is open. +# Exercise both an early and a late cut: restart from the immutable step_1 +# trainer anchor, replay the rolled-back groups, and apply step 2 exactly once. +run_test env SC_TQ_STREAMING_RECOVERY_CLAIMED_GROUPS=2 uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +run_test env SC_TQ_STREAMING_RECOVERY_CLAIMED_GROUPS=6 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 coverage combine .coverage* diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh index 178eb202029..722b5ba3ec6 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -10,20 +10,24 @@ git config --global --add safe.directory $PROJECT_ROOT set -eou pipefail -EXP_NAME=$(basename $0 .sh) -EXP_DIR=$SCRIPT_DIR/$EXP_NAME -LOG_DIR=$EXP_DIR/logs -JSON_METRICS=$EXP_DIR/metrics.json -RUN_LOG=$EXP_DIR/run.log -CHECKPOINT_DIR=$EXP_DIR/checkpoints -DATA_DIR=$EXP_DIR/data +EXP_NAME=${SC_GYM_EXP_NAME:-$(basename "$0" .sh)} +EXP_DIR=${SC_GYM_EXP_DIR:-$SCRIPT_DIR/$EXP_NAME} +LOG_DIR=${SC_GYM_LOG_DIR:-$EXP_DIR/logs} +JSON_METRICS=${SC_GYM_JSON_METRICS:-$EXP_DIR/metrics.json} +RUN_LOG=${SC_GYM_RUN_LOG:-$EXP_DIR/run.log} +CHECKPOINT_DIR=${SC_GYM_CHECKPOINT_DIR:-$EXP_DIR/checkpoints} +DATA_DIR=${SC_GYM_DATA_DIR:-$EXP_DIR/data} export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} -rm -rf $EXP_DIR $LOG_DIR -mkdir -p $EXP_DIR $LOG_DIR $CHECKPOINT_DIR $DATA_DIR +if [[ "${SC_GYM_CLEAN_EXP_DIR:-1}" == "1" ]]; then + rm -rf "$EXP_DIR" +fi +mkdir -p "$EXP_DIR" "$LOG_DIR" "$CHECKPOINT_DIR" "$DATA_DIR" -# clean up checkpoint directory on exit -trap "rm -rf $CHECKPOINT_DIR" EXIT +# Preserve checkpoints only when a caller is coordinating a restart test. +if [[ "${SC_GYM_KEEP_CHECKPOINTS:-0}" != "1" ]]; then + trap 'rm -rf "$CHECKPOINT_DIR"' EXIT +fi cd $PROJECT_ROOT @@ -107,9 +111,11 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE $@ \ 2>&1 | tee $RUN_LOG -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS +if [[ "${SC_GYM_RUN_CONVERGENCE_CHECKS:-1}" == "1" ]]; then + uv run tests/json_dump_tb_logs.py "$LOG_DIR" --output_path "$JSON_METRICS" -# Observed to be between 0.8-1.3 -uv run tests/check_metrics.py $JSON_METRICS \ - 'median(data["train/gen_kl_error"]) < 1.3' \ - 'max(data["train/reward"]) > 0' + # Observed to be between 0.8-1.3 + uv run tests/check_metrics.py "$JSON_METRICS" \ + 'median(data["train/gen_kl_error"]) < 1.3' \ + 'max(data["train/reward"]) > 0' +fi 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..2ca919373b7 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -0,0 +1,261 @@ +#!/bin/bash +# Two-process NeMo-Gym test for a periodic snapshot taken during a streamed +# optimizer step. The snapshot intentionally omits uncommitted gradients and +# makes every already-claimed group replayable from its durable trainer anchor. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +BASE_TEST=$SCRIPT_DIR/grpo_async_gym_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_streaming_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +DATA_DIR=$TEST_DIR/data +PHASE1_DIR=$TEST_DIR/phase1 +PHASE2_DIR=$TEST_DIR/phase2 +PHASE1_LOG=$PHASE1_DIR/run.log +PHASE2_LOG=$PHASE2_DIR/run.log +SELECTION_FILE=$TEST_DIR/selected_snapshot +PHASE1_PID="" + +NUM_PROMPTS=${SC_TQ_STREAMING_RECOVERY_NUM_PROMPTS:-8} +NUM_GENERATIONS=${SC_TQ_STREAMING_RECOVERY_NUM_GENERATIONS:-2} +MIN_STREAMING_GROUPS=${SC_TQ_STREAMING_RECOVERY_MIN_GROUPS:-2} +CLAIMED_GROUPS=${SC_TQ_STREAMING_RECOVERY_CLAIMED_GROUPS:-2} +MAX_INFLIGHT_PROMPTS=${SC_TQ_STREAMING_RECOVERY_MAX_INFLIGHT_PROMPTS:-2} +MAX_BUFFERED_ROLLOUTS=${SC_TQ_STREAMING_RECOVERY_MAX_BUFFERED_ROLLOUTS:-16} +MAX_NUM_STEPS=${SC_TQ_STREAMING_RECOVERY_MAX_NUM_STEPS:-3} +SNAPSHOT_INTERVAL_S=${SC_TQ_STREAMING_RECOVERY_INTERVAL_S:-0.05} +SNAPSHOT_TIMEOUT_S=${SC_TQ_STREAMING_RECOVERY_SNAPSHOT_TIMEOUT_S:-2400} +TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) +REQUIRED_RECOVERY_CAPACITY=$((NUM_PROMPTS + MIN_STREAMING_GROUPS - 1)) + +if (( NUM_PROMPTS < 2 || NUM_GENERATIONS < 1 )); then + echo "NUM_PROMPTS must be at least 2 and NUM_GENERATIONS must be positive" + exit 2 +fi +if (( MIN_STREAMING_GROUPS < 1 || MIN_STREAMING_GROUPS >= NUM_PROMPTS )); then + echo "MIN_STREAMING_GROUPS must be between 1 and NUM_PROMPTS - 1" + exit 2 +fi +if (( CLAIMED_GROUPS < MIN_STREAMING_GROUPS || CLAIMED_GROUPS >= NUM_PROMPTS )); then + echo "CLAIMED_GROUPS must be between MIN_STREAMING_GROUPS and NUM_PROMPTS - 1" + exit 2 +fi +if (( CLAIMED_GROUPS % MIN_STREAMING_GROUPS != 0 )); then + echo "CLAIMED_GROUPS must be a multiple of MIN_STREAMING_GROUPS" + exit 2 +fi +if (( MAX_NUM_STEPS < 2 )); then + echo "MAX_NUM_STEPS must be at least 2 so step_1 can anchor the snapshot" + exit 2 +fi +if (( MAX_BUFFERED_ROLLOUTS < REQUIRED_RECOVERY_CAPACITY )); then + echo "MAX_BUFFERED_ROLLOUTS must be at least $REQUIRED_RECOVERY_CAPACITY" + exit 2 +fi + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +stop_phase1() { + if [[ -z "$PHASE1_PID" ]]; then + return + fi + # Use an abrupt whole-process-group failure. A graceful delay could let + # step 2 commit after the exact streamed-step cut selected above. + 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=( + ++token_capture.enabled=true + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + # The inherited Gym recipe tracks val:accuracy, but SC has no validation + # loop. This test exercises checkpoint mechanics, not top-k retention. + checkpointing.metric_name=null + ++data_plane.checkpointing_enabled=true + 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="$MAX_INFLIGHT_PROMPTS" + async_rl.max_buffered_rollouts="$MAX_BUFFERED_ROLLOUTS" + ++rollout_checkpointing.interval_s="$SNAPSHOT_INTERVAL_S" + ++rollout_checkpointing.keep_latest_k=256 + ++rollout_checkpointing.restore_mode=latest + grpo.num_prompts_per_step="$NUM_PROMPTS" + grpo.num_generations_per_prompt="$NUM_GENERATIONS" + policy.train_global_batch_size="$TRAIN_GLOBAL_BATCH_SIZE" + grpo.max_num_steps="$MAX_NUM_STEPS" +) + +echo "=== Phase 1: crash with $CLAIMED_GROUPS/$NUM_PROMPTS streamed groups claimed ===" +command -v setsid >/dev/null +setsid env \ + SC_GYM_EXP_DIR="$PHASE1_DIR" \ + SC_GYM_RUN_LOG="$PHASE1_LOG" \ + SC_GYM_CHECKPOINT_DIR="$CHECKPOINT_DIR" \ + SC_GYM_DATA_DIR="$DATA_DIR" \ + SC_GYM_KEEP_CHECKPOINTS=1 \ + SC_GYM_RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" & +PHASE1_PID=$! + +uv run --no-sync python - \ + "$CHECKPOINT_DIR/step_1/rollout_snapshots" \ + "$SELECTION_FILE" \ + "$PHASE1_PID" \ + "$PHASE1_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]) + + +def phase_log_tail() -> str: + if not phase_log.is_file(): + return f"phase-one log was not created at {phase_log}" + lines = phase_log.read_text(errors="replace").splitlines() + return "\n".join(lines[-40:]) + +while time.monotonic() < deadline: + for snapshot in sorted(root.glob("snapshot_*"), reverse=True): + manifest_path = snapshot / "manifest.json" + if not (snapshot / "COMMITTED").is_file() or not manifest_path.is_file(): + continue + manifest = json.loads(manifest_path.read_text()) + if ( + manifest["base_train_step"] == 1 + and manifest["trainer_version"] == 1 + and manifest["rolled_back_train_group_count"] == expected_claimed + ): + selection.write_text(snapshot.name + "\n") + print( + f"selected {snapshot.name}: " + f"rolled_back_train_group_count={expected_claimed}", + flush=True, + ) + raise SystemExit(0) + try: + os.kill(phase_pid, 0) + except ProcessLookupError as error: + raise RuntimeError( + "phase one exited before producing the requested streamed-step cut; " + f"last log lines from {phase_log}:\n{phase_log_tail()}" + ) from error + time.sleep(0.1) + +raise TimeoutError( + f"no committed snapshot captured {expected_claimed} claimed groups" +) +PY + +stop_phase1 +SNAPSHOT_NAME=$(tr -d '\n' < "$SELECTION_FILE") +SNAPSHOT_ROOT=$CHECKPOINT_DIR/step_1/rollout_snapshots +SNAPSHOT_DIR=$SNAPSHOT_ROOT/$SNAPSHOT_NAME + +# The resolver scans immutable directories, not only LATEST. Retain exactly +# the fault-injection cut so phase two cannot silently choose a newer snapshot. +for candidate in "$SNAPSHOT_ROOT"/snapshot_*; do + if [[ -d "$candidate" && "$(basename "$candidate")" != "$SNAPSHOT_NAME" ]]; then + rm -rf "$candidate" + fi +done +printf '%s\n' "$SNAPSHOT_NAME" > "$SNAPSHOT_ROOT/LATEST" + +test -f "$CHECKPOINT_DIR/step_1/training_info.json" +test -d "$CHECKPOINT_DIR/step_1/policy" +test -f "$SNAPSHOT_DIR/COMMITTED" +test -d "$SNAPSHOT_DIR/data_plane" +test -f "$SNAPSHOT_DIR/replay_buffer_metadata.pt" +test -f "$SNAPSHOT_DIR/rollout_recovery.pt" +test -f "$SNAPSHOT_DIR/train_dataloader.pt" +test ! -d "$SNAPSHOT_DIR/policy" + +uv run --no-sync python - \ + "$SNAPSHOT_DIR/manifest.json" \ + "$SNAPSHOT_DIR/rollout_recovery.pt" \ + "$CLAIMED_GROUPS" <<'PY' +import json +import sys + +import torch + +manifest = json.load(open(sys.argv[1])) +lineage = torch.load(sys.argv[2], weights_only=False) +expected_claimed = int(sys.argv[3]) + +assert manifest["base_train_step"] == 1, manifest +assert manifest["trainer_version"] == 1, manifest +assert manifest["rolled_back_train_group_count"] == expected_claimed, manifest +assert lineage["open_train_step"] is None, lineage["open_train_step"] +assert sum(group["status"] == "finalized" for group in lineage["groups"]) >= expected_claimed +print( + "validated restart-ready streamed-step snapshot: " + f"claimed_groups={expected_claimed}" +) +PY + +echo "=== Phase 2: replay the rolled-back groups and finish each step once ===" +SC_GYM_EXP_DIR="$PHASE2_DIR" \ +SC_GYM_RUN_LOG="$PHASE2_LOG" \ +SC_GYM_CHECKPOINT_DIR="$CHECKPOINT_DIR" \ +SC_GYM_DATA_DIR="$DATA_DIR" \ +SC_GYM_KEEP_CHECKPOINTS=1 \ +SC_GYM_RUN_CONVERGENCE_CHECKS=0 \ +bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" + +grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Native TQ replay inventory validated" "$PHASE2_LOG" +grep -Eq "Restored [1-9][0-9]* replay group" "$PHASE2_LOG" +grep -q "Restored sampler dispatch state" "$PHASE2_LOG" +grep -q "train step $MAX_NUM_STEPS/$MAX_NUM_STEPS" "$PHASE2_LOG" + +uv run --no-sync python - \ + "$PHASE2_LOG" \ + "$CHECKPOINT_DIR/step_$MAX_NUM_STEPS/training_info.json" \ + "$MAX_NUM_STEPS" <<'PY' +import json +import re +import sys +from pathlib import Path + +log = Path(sys.argv[1]).read_text() +training_info_path = Path(sys.argv[2]) +max_steps = int(sys.argv[3]) + +assert training_info_path.is_file(), training_info_path +training_info = json.loads(training_info_path.read_text()) +assert training_info["current_step"] == max_steps, training_info +assert training_info["trainer_version"] == max_steps, training_info + +# The durable anchor already contains step 1. Recovery must resume at step 2, +# replay its rolled-back groups, and never apply that optimizer step twice. +assert not re.search(r"train step 1/", log), "phase two repeated anchored 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, f"step {step} completed {len(matches)} times" +print(f"validated exactly-once resumed steps 2..{max_steps}") +PY + +echo "Streamed-step periodic recovery functional test passed." diff --git a/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh b/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh new file mode 100755 index 00000000000..54d3c41b1a4 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh @@ -0,0 +1,234 @@ +#!/bin/bash +# Two-process NeMo-Gym test for pre-step TQ + partial-sibling recovery. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +BASE_TEST=$SCRIPT_DIR/grpo_async_gym_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_async_gym_single_controller_tq_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +DATA_DIR=$TEST_DIR/data +PHASE1_DIR=$TEST_DIR/phase1 +PHASE2_DIR=$TEST_DIR/phase2 +PHASE1_LOG=$PHASE1_DIR/run.log +PHASE2_LOG=$PHASE2_DIR/run.log +SELECTION_FILE=$TEST_DIR/selected_snapshot +EXPECTED_STATE=$TEST_DIR/expected_rollout_recovery.pt +PHASE1_PID="" + +MODEL_NAME=${SC_TQ_RECOVERY_MODEL_NAME:-Qwen/Qwen3-0.6B} +NUM_PROMPTS=${SC_TQ_RECOVERY_NUM_PROMPTS:-4} +NUM_GENERATIONS=${SC_TQ_RECOVERY_NUM_GENERATIONS:-4} +MIN_SEALED=${SC_TQ_RECOVERY_MIN_SEALED:-1} +MAX_INFLIGHT_PROMPTS=${SC_TQ_RECOVERY_MAX_INFLIGHT_PROMPTS:-8} +MAX_BUFFERED_ROLLOUTS=${SC_TQ_RECOVERY_MAX_BUFFERED_ROLLOUTS:-8} +MAX_NUM_STEPS=${SC_TQ_RECOVERY_MAX_NUM_STEPS:-6} +SNAPSHOT_TIMEOUT_S=${SC_TQ_RECOVERY_SNAPSHOT_TIMEOUT_S:-1800} +TRAIN_GLOBAL_BATCH_SIZE=$((NUM_PROMPTS * NUM_GENERATIONS)) + +if (( NUM_PROMPTS < 1 || NUM_GENERATIONS < 2 )); then + echo "NUM_PROMPTS must be positive and NUM_GENERATIONS must be at least 2" + exit 2 +fi +if (( MIN_SEALED < 1 || MIN_SEALED >= NUM_GENERATIONS )); then + echo "MIN_SEALED must be between 1 and NUM_GENERATIONS - 1" + exit 2 +fi +if (( MAX_BUFFERED_ROLLOUTS < (2 * NUM_PROMPTS - 1) )); then + echo "MAX_BUFFERED_ROLLOUTS must be at least 2 * NUM_PROMPTS - 1" + exit 2 +fi + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +stop_phase1() { + if [[ -z "$PHASE1_PID" ]]; then + return + fi + # Use an abrupt whole-process-group failure. A graceful delay could let + # training publish a newer trainer checkpoint after the selected cut. + 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=( + ++token_capture.enabled=true + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + # The inherited Gym recipe tracks val:accuracy, but SC has no validation + # loop. This test exercises checkpoint mechanics, not top-k retention. + checkpointing.metric_name=null + ++data_plane.checkpointing_enabled=true + async_rl.sampler.name=windowed + '~async_rl.sampler.max_lookahead_versions' + '+async_rl.sampler.max_staleness_versions=1' + async_rl.min_groups_for_streaming_train="$NUM_PROMPTS" + async_rl.max_inflight_prompts="$MAX_INFLIGHT_PROMPTS" + async_rl.max_buffered_rollouts="$MAX_BUFFERED_ROLLOUTS" + ++rollout_checkpointing.interval_s=0.25 + ++rollout_checkpointing.keep_latest_k=128 + policy.model_name="$MODEL_NAME" + grpo.num_prompts_per_step="$NUM_PROMPTS" + grpo.num_generations_per_prompt="$NUM_GENERATIONS" + policy.train_global_batch_size="$TRAIN_GLOBAL_BATCH_SIZE" + grpo.max_num_steps="$MAX_NUM_STEPS" +) + +echo "=== Phase 1: crash after a committed partial-sibling snapshot ===" +command -v setsid >/dev/null +setsid env \ + SC_GYM_EXP_DIR="$PHASE1_DIR" \ + SC_GYM_RUN_LOG="$PHASE1_LOG" \ + SC_GYM_CHECKPOINT_DIR="$CHECKPOINT_DIR" \ + SC_GYM_DATA_DIR="$DATA_DIR" \ + SC_GYM_KEEP_CHECKPOINTS=1 \ + SC_GYM_RUN_CONVERGENCE_CHECKS=0 \ + bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" & +PHASE1_PID=$! + +uv run --no-sync python - \ + "$CHECKPOINT_DIR/bootstrap/rollout_snapshots" \ + "$SELECTION_FILE" \ + "$PHASE1_PID" \ + "$PHASE1_LOG" \ + "$NUM_GENERATIONS" \ + "$MIN_SEALED" \ + "$SNAPSHOT_TIMEOUT_S" <<'PY' +import os +import sys +import time +from pathlib import Path + +import torch + +root = Path(sys.argv[1]) +selection = Path(sys.argv[2]) +phase_pid = int(sys.argv[3]) +phase_log = Path(sys.argv[4]) +expected_generations = int(sys.argv[5]) +min_sealed = int(sys.argv[6]) +deadline = time.monotonic() + float(sys.argv[7]) + + +def phase_log_tail() -> str: + if not phase_log.is_file(): + return f"phase-one log was not created at {phase_log}" + lines = phase_log.read_text(errors="replace").splitlines() + return "\n".join(lines[-40:]) + +while time.monotonic() < deadline: + for snapshot in sorted(root.glob("snapshot_*"), reverse=True): + ledger_path = snapshot / "rollout_recovery.pt" + if not (snapshot / "COMMITTED").is_file() or not ledger_path.is_file(): + continue + state = torch.load(ledger_path, weights_only=False) + partial = [] + for group in state["groups"]: + statuses = [ + sibling["attempts"][-1]["status"] + for sibling in group["siblings"] + ] + if len(statuses) != expected_generations: + raise RuntimeError( + f"group {group['group_id']} has {len(statuses)} siblings; " + f"expected {expected_generations}" + ) + sealed = statuses.count("sealed") + if min_sealed <= sealed < expected_generations: + partial.append((group["group_id"], sealed, expected_generations)) + if partial: + selection.write_text(snapshot.name + "\n") + print(f"selected {snapshot.name}: partial_groups={partial}", flush=True) + raise SystemExit(0) + try: + os.kill(phase_pid, 0) + except ProcessLookupError as error: + raise RuntimeError( + "phase one exited before producing a partial-sibling snapshot; " + f"last log lines from {phase_log}:\n{phase_log_tail()}" + ) from error + time.sleep(0.25) + +raise TimeoutError("no committed partial-sibling snapshot was produced") +PY + +stop_phase1 +SNAPSHOT_NAME=$(tr -d '\n' < "$SELECTION_FILE") +SNAPSHOT_ROOT=$CHECKPOINT_DIR/bootstrap/rollout_snapshots +SNAPSHOT_DIR=$SNAPSHOT_ROOT/$SNAPSHOT_NAME + +# resolve_latest_snapshot scans immutable directories rather than trusting only +# LATEST. Retain exactly the cut selected by the fault-injection predicate. +for candidate in "$SNAPSHOT_ROOT"/snapshot_*; do + if [[ -d "$candidate" && "$(basename "$candidate")" != "$SNAPSHOT_NAME" ]]; then + rm -rf "$candidate" + fi +done +printf '%s\n' "$SNAPSHOT_NAME" > "$SNAPSHOT_ROOT/LATEST" + +test -f "$CHECKPOINT_DIR/bootstrap/manifest.json" +test -f "$SNAPSHOT_DIR/COMMITTED" +test -d "$SNAPSHOT_DIR/data_plane" +test -f "$SNAPSHOT_DIR/replay_buffer_metadata.pt" +test -f "$SNAPSHOT_DIR/rollout_recovery.pt" +test -f "$SNAPSHOT_DIR/train_dataloader.pt" +test ! -d "$SNAPSHOT_DIR/policy" +cp "$SNAPSHOT_DIR/rollout_recovery.pt" "$EXPECTED_STATE" + +echo "=== Phase 2: restore sealed siblings and redispatch only missing ones ===" +SC_GYM_EXP_DIR="$PHASE2_DIR" \ +SC_GYM_RUN_LOG="$PHASE2_LOG" \ +SC_GYM_CHECKPOINT_DIR="$CHECKPOINT_DIR" \ +SC_GYM_DATA_DIR="$DATA_DIR" \ +SC_GYM_KEEP_CHECKPOINTS=1 \ +SC_GYM_RUN_CONVERGENCE_CHECKS=1 \ +bash "$BASE_TEST" "${COMMON_OVERRIDES[@]}" + +grep -Fq "Selected rollout recovery snapshot: $SNAPSHOT_DIR" "$PHASE2_LOG" +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Rollout recovery completed" "$PHASE2_LOG" +grep -q "train step $MAX_NUM_STEPS/$MAX_NUM_STEPS" "$PHASE2_LOG" + +uv run --no-sync python - "$EXPECTED_STATE" "$PHASE2_LOG" <<'PY' +import re +import sys +from pathlib import Path + +import torch + +state = torch.load(sys.argv[1], weights_only=False) +log = Path(sys.argv[2]).read_text() +expected = {} +for group in state["groups"]: + if group["status"] not in {"generating", "ready_to_finalize"}: + continue + total = len(group["siblings"]) + sealed = sum( + sibling["attempts"][-1]["status"] == "sealed" + for sibling in group["siblings"] + ) + expected[group["group_id"]] = (sealed, total - sealed) + +pattern = re.compile( + r"rollout recovery finalized group: group=(\S+) " + r"reused=(\d+) redispatched=(\d+)" +) +observed = { + group_id: (int(reused), int(redispatched)) + for group_id, reused, redispatched in pattern.findall(log) +} +assert observed == expected, f"recovery mismatch: {observed=} {expected=}" +assert any(reused > 0 and redispatched > 0 for reused, redispatched in observed.values()) +print(f"validated partial-sibling reuse: {observed}") +PY + +echo "Partial-sibling TQ recovery functional test passed." diff --git a/tests/functional/grpo_dp_single_controller_tq_recovery.sh b/tests/functional/grpo_dp_single_controller_tq_recovery.sh index 671b67dd4ec..a660375f34c 100755 --- a/tests/functional/grpo_dp_single_controller_tq_recovery.sh +++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh @@ -16,6 +16,7 @@ COMMON_OVERRIDES=( checkpointing.enabled=true checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.save_period=1 + checkpointing.metric_name=null data_plane.checkpointing_enabled=true async_rl.sampler.name=windowed '~async_rl.sampler.max_lookahead_versions' diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index e8f28efcd3f..98a60596232 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -1118,6 +1118,26 @@ def test_mints_ids_and_returns_metadata_request(self): # actor-pool path; the manager leaves the reservation unready. assert buf.commit_calls == [] + def test_dispatch_uses_pre_reserved_dataloader_lineage(self): + buf = _FakeCaptureBuffer() + mgr = _make_capture_manager(buf) + sample = {"prompt": "p", "idx": 0} + group_id = mgr.reserve_prompt_group(sample, target_step=5) + + request = _run( + mgr.generate_for_finalization( + sample, + target_step=5, + recovery_group_id=group_id, + ) + ) + + assert request is not None + assert request.group_id == group_id + assert mgr.recovery_ledger is not None + assert [group.group_id for group in mgr.recovery_ledger.groups()] == [group_id] + assert buf._slots == [group_id] + def test_failed_dispatch_aborts_replay_reservation(self): buf = _FakeCaptureBuffer() diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 0e981d534fe..e18fca1cc92 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -202,6 +202,44 @@ def test_open_train_step_rolls_back_as_one_unit() -> None: ) +def test_periodic_snapshot_rolls_back_claimed_groups_without_mutating_live_ledger() -> ( + None +): + ledger = RolloutRecoveryLedger() + group = _reserve(ledger) + ledger.mark_group_dispatched(group.group_id) + _seal(ledger, group.group_id, 0) + _seal(ledger, group.group_id, 1) + meta = _finalize(ledger, group.group_id) + ledger.claim_groups_for_training( + [group.group_id], + train_step=4, + trainer_version=4, + expected_group_count=2, + ) + + replay_groups = ledger.training_owned_replay_groups() + snapshot_state = ledger.periodic_snapshot_state_dict() + restored = RolloutRecoveryLedger.from_state_dict(snapshot_state) + + assert replay_groups == [ + { + "meta": meta, + "start_weight": 7, + "end_weight": 8, + "target_step": 8, + "group_id": group.group_id, + } + ] + assert restored.open_train_step is None + assert restored.get_group(group.group_id).status == PromptGroupStatus.FINALIZED + assert ledger.open_train_step is not None + assert ( + ledger.get_group(group.group_id).status + == PromptGroupStatus.CLAIMED_FOR_TRAINING + ) + + def test_state_dict_round_trip_preserves_partial_and_claimed_lineage() -> None: ledger = RolloutRecoveryLedger() partial = _reserve(ledger, group_id="partial") diff --git a/tests/unit/single_controller/test_rollout_checkpoint.py b/tests/unit/single_controller/test_rollout_checkpoint.py new file mode 100644 index 00000000000..1226071e741 --- /dev/null +++ b/tests/unit/single_controller/test_rollout_checkpoint.py @@ -0,0 +1,395 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from pathlib import Path +from typing import Any, cast +from unittest.mock import Mock, call + +import pytest + +from nemo_rl.algorithms.single_controller_utils import rollout_checkpoint +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, + RolloutSnapshotManifest, + bootstrap_fingerprint, + commit_snapshot, + ensure_bootstrap_anchor, + prepare_snapshot_paths, + prune_bootstrap_snapshots, + reset_bootstrap_anchor, + resolve_latest_snapshot, + validate_bootstrap_anchor, +) + + +class _DumpedConfig: + def __init__(self, dumped: dict[str, Any]): + self._dumped = dumped + + def model_dump(self, *, mode: str) -> dict[str, Any]: + assert mode == "json" + return self._dumped + + +def _commit_snapshot( + anchor, + *, + mutation_version: int, + trainer_version: int = 0, + fingerprint: str | None = "fingerprint-v1", +): + tmp_path, final_path, _ = prepare_snapshot_paths(anchor) + manifest = RolloutSnapshotManifest( + schema_version=1, + base_train_step=trainer_version, + trainer_version=trainer_version, + current_epoch=2, + mutation_version=mutation_version, + rolled_back_train_group_count=0, + bootstrap_fingerprint=fingerprint, + ) + (tmp_path / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps(manifest.to_dict()) + ) + commit_snapshot(tmp_path, final_path, keep_latest_k=3) + return final_path + + +def test_bootstrap_anchor_rejects_different_initial_state(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v1") + + with pytest.raises(ValueError, match="does not match"): + validate_bootstrap_anchor(anchor, fingerprint="fingerprint-v2") + + +def test_bootstrap_fingerprint_ignores_non_recovery_configuration() -> None: + base = { + "policy": { + "model_name": "model-a", + "optimizer": {"lr": 1.0e-6}, + "generation": { + "backend": "vllm", + "temperature": 1.0, + "colocated": {"enabled": True, "resources": {"gpus": 8}}, + "vllm_cfg": { + "kv_cache_dtype": "auto", + "precision": "bfloat16", + "skip_tokenizer_init": False, + }, + }, + }, + "data": { + "train": [{"data_path": "/datasets/train.jsonl"}], + "num_workers": 4, + }, + "grpo": { + "num_generations_per_prompt": 4, + "max_num_steps": 10, + "batch_multiplier": 1, + "use_dynamic_sampling": False, + "reward_shaping": {"enabled": False}, + "reward_scaling": {"enabled": False}, + }, + "loss_fn": { + "reference_policy_kl_penalty": 0.01, + "use_kl_in_reward": False, + }, + "reward_penalties": {"penalize_unwanted_tokens": False}, + "token_capture": { + "enabled": True, + "staging_partition": "rollout_staging", + "on_capture_failure": "continue", + }, + "cluster": {"num_nodes": 2}, + "logger": {"log_dir": "/run/one"}, + } + compatible_changed = { + **base, + "policy": { + **base["policy"], + "optimizer": {"lr": 5.0e-7}, + "generation": { + **base["policy"]["generation"], + "colocated": {"enabled": False, "resources": {"gpus": 16}}, + "vllm_cfg": { + "kv_cache_dtype": "fp8", + "precision": "float16", + "skip_tokenizer_init": True, + }, + }, + }, + "data": {**base["data"], "num_workers": 16}, + "grpo": { + **base["grpo"], + "max_num_steps": 100, + "batch_multiplier": 2, + "use_dynamic_sampling": True, + "reward_shaping": {"enabled": True}, + "reward_scaling": {"enabled": True}, + }, + "loss_fn": { + "reference_policy_kl_penalty": 0.1, + "use_kl_in_reward": True, + }, + "reward_penalties": {"penalize_unwanted_tokens": True}, + "token_capture": { + **base["token_capture"], + "on_capture_failure": "abort", + }, + "cluster": {"num_nodes": 8}, + "logger": {"log_dir": "/run/two"}, + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(compatible_changed)) + ) + + +@pytest.mark.parametrize( + ("section", "changed"), + [ + ("policy", {"model_name": "model-b"}), + ("data", {"train": [{"data_path": "/datasets/other.jsonl"}]}), + ("grpo", {"num_generations_per_prompt": 8}), + ("token_capture", {"mixed_weight_version_policy": "reject"}), + ( + "async_rl", + {"sampler": {"name": "windowed", "max_staleness_versions": 2}}, + ), + ], +) +def test_bootstrap_fingerprint_rejects_rollout_semantic_changes( + section: str, + changed: dict[str, Any], +) -> None: + base = { + "policy": { + "model_name": "model-a", + "tokenizer": {"name": "tokenizer-a"}, + "generation": {"backend": "vllm", "temperature": 1.0}, + }, + "data": {"train": [{"data_path": "/datasets/train.jsonl"}]}, + "grpo": {"num_generations_per_prompt": 4}, + "token_capture": { + "enabled": True, + "mixed_weight_version_policy": "allow", + }, + "async_rl": {"sampler": {"name": "windowed", "max_staleness_versions": 1}}, + } + modified = {**base, section: {**base[section], **changed}} + + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(modified))) + ) + + +def test_bootstrap_fingerprint_rejects_generation_semantic_changes() -> None: + base = { + "policy": { + "model_name": "model-a", + "generation": { + "backend": "vllm", + "temperature": 1.0, + "vllm_cfg": {"max_model_len": 4096}, + }, + } + } + sampling_changed = { + **base, + "policy": { + **base["policy"], + "generation": { + **base["policy"]["generation"], + "temperature": 0.5, + }, + }, + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(sampling_changed))) + ) + + context_changed = { + **base, + "policy": { + **base["policy"], + "generation": { + **base["policy"]["generation"], + "vllm_cfg": {"max_model_len": 8192}, + }, + }, + } + assert bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) != ( + bootstrap_fingerprint(cast(Any, _DumpedConfig(context_changed))) + ) + + +def test_bootstrap_fingerprint_ignores_nested_gym_log_directory() -> None: + base = { + "policy": {"model": "model-a"}, + "env": { + "should_use_nemo_gym": True, + "nemo_gym": { + "nemo_gym_log_dir": "/run/one/nemo_gym", + "should_log_nemo_gym_responses": True, + "policy_model": {"temperature": 1.0}, + "agent": {"concurrency": 16, "max_turns": 20}, + }, + }, + } + runtime_changed = { + **base, + "env": { + **base["env"], + "nemo_gym": { + **base["env"]["nemo_gym"], + "nemo_gym_log_dir": "/run/two/nemo_gym", + "should_log_nemo_gym_responses": False, + "agent": {"concurrency": 64, "max_turns": 20}, + }, + }, + } + semantic_changed = { + **base, + "env": { + **base["env"], + "nemo_gym": { + **base["env"]["nemo_gym"], + "policy_model": {"temperature": 0.5}, + }, + }, + } + + fingerprint = bootstrap_fingerprint(cast(Any, _DumpedConfig(base))) + assert fingerprint == bootstrap_fingerprint( + cast(Any, _DumpedConfig(runtime_changed)) + ) + assert fingerprint != bootstrap_fingerprint( + cast(Any, _DumpedConfig(semantic_changed)) + ) + assert base["env"]["nemo_gym"]["nemo_gym_log_dir"] == "/run/one/nemo_gym" + + +def test_prune_bootstrap_snapshots_requires_durable_trainer_checkpoint(tmp_path): + snapshot_root = tmp_path / "bootstrap" / "rollout_snapshots" + snapshot_root.mkdir(parents=True) + (snapshot_root / "snapshot_000001").mkdir() + durable_anchor = tmp_path / "step_1" + + with pytest.raises(FileNotFoundError, match="durable trainer checkpoint"): + prune_bootstrap_snapshots( + tmp_path, + durable_trainer_checkpoint=durable_anchor, + ) + + assert snapshot_root.is_dir() + durable_anchor.mkdir() + assert prune_bootstrap_snapshots( + tmp_path, + durable_trainer_checkpoint=durable_anchor, + ) + assert not snapshot_root.exists() + + +def test_reset_bootstrap_anchor_discards_skipped_snapshot_lineage( + tmp_path: Path, +) -> None: + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="old-fingerprint") + snapshot_root = anchor / "rollout_snapshots" + (snapshot_root / "snapshot_000001").mkdir(parents=True) + + reset = reset_bootstrap_anchor(tmp_path, fingerprint="new-fingerprint") + + assert reset == anchor + assert not snapshot_root.exists() + validate_bootstrap_anchor(anchor, fingerprint="new-fingerprint") + + +def test_resolver_selects_latest_compatible_committed_snapshot(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + first = _commit_snapshot(anchor, mutation_version=1) + second = _commit_snapshot(anchor, mutation_version=2) + + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + assert resolved is not None + assert resolved.path == second + assert resolved.manifest.mutation_version == 2 + assert first.is_dir() + + +def test_resolver_falls_back_from_corrupt_newest_snapshot(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + first = _commit_snapshot(anchor, mutation_version=1) + second = _commit_snapshot(anchor, mutation_version=2) + (second / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text("not-json") + + resolved = resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + assert resolved is not None + assert resolved.path == first + + +def test_resolver_fails_when_no_committed_snapshot_matches_anchor(tmp_path): + anchor = ensure_bootstrap_anchor(tmp_path, fingerprint="fingerprint-v1") + _commit_snapshot(anchor, mutation_version=1, fingerprint="different") + + with pytest.raises(ValueError, match="trainer-anchor mismatch"): + resolve_latest_snapshot( + anchor, + expected_train_step=0, + expected_trainer_version=0, + expected_bootstrap_fingerprint="fingerprint-v1", + ) + + +def test_commit_snapshot_flushes_payload_before_publication(tmp_path, monkeypatch): + anchor = tmp_path / "step_1" + anchor.mkdir() + tmp_snapshot, final_snapshot, _ = prepare_snapshot_paths(anchor) + (tmp_snapshot / "payload").write_text("payload") + fsync_tree = Mock() + fsync_file = Mock() + fsync_directory = Mock() + monkeypatch.setattr(rollout_checkpoint, "_fsync_tree", fsync_tree) + monkeypatch.setattr(rollout_checkpoint, "_fsync_file", fsync_file) + monkeypatch.setattr(rollout_checkpoint, "_fsync_directory", fsync_directory) + + commit_snapshot(tmp_snapshot, final_snapshot, keep_latest_k=1) + + fsync_tree.assert_called_once_with(tmp_snapshot) + assert fsync_file.call_args_list == [ + call(tmp_snapshot / "COMMITTED"), + call(anchor / "rollout_snapshots" / "LATEST.tmp"), + ] + assert fsync_directory.call_args_list[:2] == [ + call(tmp_snapshot), + call(anchor / "rollout_snapshots"), + ] + assert (final_snapshot / "COMMITTED").is_file() + assert ( + anchor / "rollout_snapshots" / "LATEST" + ).read_text().strip() == final_snapshot.name diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index afdbb961e2e..919878efdd7 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -24,7 +24,10 @@ import ray import torch -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + TQReplayBuffer, +) from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSampler, WeightFifoSampler, @@ -94,6 +97,96 @@ async def generate_and_push( self._buffer.reserve(target_step=target_step) +def test_token_capture_reserves_batch_atomically_with_dataloader_advance() -> None: + events: list[str] = [] + barrier = DataPlaneCheckpointBarrier() + + class _RecordingSampler(InOrderSampler): + async def wait_until_admissible(self, *, trainer_version_fn) -> None: + assert barrier._active_mutations == 0 + events.append("wait") + await super().wait_until_admissible(trainer_version_fn=trainer_version_fn) + + def commit_admission(self, *, trainer_version: int) -> int | None: + assert barrier._active_mutations == 1 + events.append("commit") + return super().commit_admission(trainer_version=trainer_version) + + class _CaptureManager: + def __init__(self) -> None: + self.stats = RolloutStats() + + def reserve_prompt_group( + self, prompt: Any, *, target_step: int | None = None + ) -> str: + del target_step + events.append(f"reserve-{prompt['idx']}") + return f"group-{prompt['idx']}" + + async def generate_for_finalization( + self, + prompt: Any, + *, + target_step: int | None = None, + recovery_group_id: str, + inflight_registry: dict[str, Any] | None = None, + ) -> Any: + del target_step, inflight_registry + events.append(f"dispatch-{prompt['idx']}-{recovery_group_id}") + return SimpleNamespace(group_id=recovery_group_id) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._async_cfg = SimpleNamespace(max_inflight_prompts=2, diagnostics=False) + ctrl._master_config = SimpleNamespace( + grpo=GRPOConfig.model_construct( + max_num_epochs=1, + num_prompts_per_step=2, + ), + token_capture=SimpleNamespace(enabled=True), + ) + ctrl._rollout_manager = _CaptureManager() + ctrl._finalizer_actors = [object()] + ctrl._finalize_with_actor = lambda request: asyncio.sleep(0) + ctrl._buffer = _RecordingBuffer() + ctrl._sampler = _RecordingSampler(ctrl._buffer, max_lookahead_versions=1) + ctrl._dataloader = [ + BatchedDataDict( + { + "idx": [10, 11], + "message_log": [ + [{"role": "user", "content": "a"}], + [{"role": "user", "content": "b"}], + ], + } + ) + ] + ctrl._data_plane_checkpoint_barrier = barrier + ctrl._rollout_permitted = asyncio.Event() + ctrl._rollout_permitted.set() + ctrl._rollout_exhausted = asyncio.Event() + ctrl._buffer_capacity = asyncio.Semaphore(4) + ctrl._rollout_slots = asyncio.Semaphore(2) + ctrl._inflight_rollouts = 0 + ctrl._inflight_by_group_id = {} + ctrl._dispatched_rollouts = set() + ctrl._trainer_version = 0 + ctrl._current_epoch = 0 + + asyncio.run(ctrl._rollout_pump()) + + assert events == [ + "wait", + "commit", + "reserve-10", + "reserve-11", + "dispatch-10-group-10", + "dispatch-11-group-11", + "wait", + ] + assert ctrl._buffer_capacity._value == 2 + + @pytest.mark.parametrize( ("make_sampler", "expected_target_steps"), [ diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index be2a75dc83e..21981269454 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -41,7 +41,7 @@ import threading from pathlib import Path from typing import Any, Optional, Union -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest import torch @@ -54,6 +54,7 @@ REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, DataPlaneCheckpointBarrier, + replay_manifest_digest, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( WeightFifoSamplerConfig, @@ -71,8 +72,14 @@ from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, MasterConfig, + RolloutCheckpointConfig, setup_single_controller, ) +from nemo_rl.algorithms.single_controller_utils.rollout_checkpoint import ( + BOOTSTRAP_DIRNAME, + ROLLOUT_SNAPSHOT_COMMITTED_FILENAME, + ROLLOUT_SNAPSHOT_MANIFEST_FILENAME, +) from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.data.utils import load_dataloader_state from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta @@ -446,9 +453,18 @@ 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) + if additional_groups: + state["groups"] = [*state["groups"], *additional_groups] + state["manifest_digest"] = replay_manifest_digest(state["groups"]) + return state async def load_state_dict( self, @@ -515,6 +531,9 @@ def _actor_master_config( max_num_epochs: int = 1, buffer_checkpoint: bool = False, data_plane_checkpoint: bool = False, + token_capture: bool = False, + rollout_checkpoint_interval_s: Optional[float] = None, + rollout_checkpoint_keep_latest_k: int = 2, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -526,7 +545,7 @@ def _actor_master_config( if buffer_checkpoint else WeightFifoSamplerConfig(max_staleness_versions=1) ) - return MasterConfig.model_construct( + config = MasterConfig.model_construct( policy={ # One optimizer.step per RL step: prompts * generations == gbs. "train_global_batch_size": num_prompts_per_step * 2, @@ -573,7 +592,13 @@ def _actor_master_config( max_inflight_prompts=4, max_buffered_rollouts=4, ), + rollout_checkpointing=RolloutCheckpointConfig( + interval_s=rollout_checkpoint_interval_s, + keep_latest_k=rollout_checkpoint_keep_latest_k, + ), ) + config.token_capture.enabled = token_capture + return config def _make_actor_args( @@ -586,6 +611,7 @@ def _make_actor_args( last_checkpoint_path: Optional[str] = None, data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None, rollout_manager: Optional[_FakeRolloutManager] = None, + bootstrap_fingerprint: Optional[str] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=_FakeGeneration(), @@ -609,6 +635,7 @@ def _make_actor_args( ), last_checkpoint_path=last_checkpoint_path, data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + bootstrap_fingerprint=bootstrap_fingerprint, ) @@ -938,6 +965,269 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} +class TestPeriodicRolloutCheckpoint: + @staticmethod + def _make_actor( + tmp_path: Path, + *, + ledger: Optional[RolloutRecoveryLedger] = None, + dp_client: Optional[_FakeDPClient] = None, + tq_buffer: Optional[_FakeTQBuffer] = None, + ) -> Any: + mc = _actor_master_config( + tmp_path, + buffer_checkpoint=True, + data_plane_checkpoint=True, + token_capture=True, + rollout_checkpoint_interval_s=120.0, + ) + ledger = ledger or RolloutRecoveryLedger() + return _ACTOR_CLS( + mc, + _make_actor_args( + rollout_manager=_FakeRolloutManager(ledger), + dp_client=dp_client, + tq_buffer=tq_buffer, + bootstrap_fingerprint="bootstrap-v1", + ), + SetupTimingMetrics(), + ) + + def test_pre_step_snapshot_omits_trainer_payload(self, tmp_path: Path) -> None: + actor = self._make_actor(tmp_path) + + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + actor._checkpointer.shutdown() + + snapshot = ( + tmp_path + / "checkpoints" + / BOOTSTRAP_DIRNAME + / "rollout_snapshots" + / "snapshot_000001" + ) + assert (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).is_file() + assert (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).is_file() + assert (snapshot / "data_plane" / "metadata.json").is_file() + assert (snapshot / "train_dataloader.pt").is_file() + assert (snapshot / REPLAY_BUFFER_METADATA_FILENAME).is_file() + assert (snapshot / ROLLOUT_RECOVERY_STATE_FILENAME).is_file() + assert (snapshot / "config.yaml").is_file() + assert not (snapshot / "policy").exists() + + def test_unchanged_state_is_not_saved_twice(self, tmp_path: Path) -> None: + actor = self._make_actor(tmp_path) + + async def _save_twice() -> tuple[bool, bool]: + return ( + await actor._save_rollout_checkpoint(), + await actor._save_rollout_checkpoint(), + ) + + assert asyncio.run(_save_twice()) == (True, False) + actor._checkpointer.shutdown() + snapshot_root = ( + tmp_path / "checkpoints" / BOOTSTRAP_DIRNAME / "rollout_snapshots" + ) + assert sorted(path.name for path in snapshot_root.glob("snapshot_*")) == [ + "snapshot_000001" + ] + + def test_failed_save_retains_previous_committed_snapshot( + self, tmp_path: Path + ) -> None: + dp_client = _FakeDPClient() + actor = self._make_actor(tmp_path, dp_client=dp_client) + + async def _save_then_fail() -> None: + assert await actor._save_rollout_checkpoint(force=True) + dp_client.save_error = RuntimeError("injected periodic save failure") + await actor._save_rollout_checkpoint(force=True) + + with pytest.raises(RuntimeError, match="injected periodic save failure"): + asyncio.run(_save_then_fail()) + actor._checkpointer.shutdown() + + snapshot_root = ( + tmp_path / "checkpoints" / BOOTSTRAP_DIRNAME / "rollout_snapshots" + ) + assert (snapshot_root / "snapshot_000001" / "COMMITTED").is_file() + assert (snapshot_root / "LATEST").read_text().strip() == "snapshot_000001" + assert not (snapshot_root / "snapshot_000002").exists() + assert not (snapshot_root / "tmp_snapshot_000002").exists() + + def test_retention_keeps_recent_fallbacks(self, tmp_path: Path) -> None: + mc = _actor_master_config( + tmp_path, + buffer_checkpoint=True, + data_plane_checkpoint=True, + token_capture=True, + rollout_checkpoint_interval_s=120.0, + rollout_checkpoint_keep_latest_k=2, + ) + ledger = RolloutRecoveryLedger() + actor = _ACTOR_CLS( + mc, + _make_actor_args( + rollout_manager=_FakeRolloutManager(ledger), + bootstrap_fingerprint="bootstrap-v1", + ), + SetupTimingMetrics(), + ) + + async def _save_three() -> None: + for _ in range(3): + assert await actor._save_rollout_checkpoint(force=True) + + asyncio.run(_save_three()) + actor._checkpointer.shutdown() + snapshot_root = ( + tmp_path / "checkpoints" / BOOTSTRAP_DIRNAME / "rollout_snapshots" + ) + assert sorted(path.name for path in snapshot_root.glob("snapshot_*")) == [ + "snapshot_000002", + "snapshot_000003", + ] + assert (snapshot_root / "LATEST").read_text().strip() == "snapshot_000003" + + def test_post_step_snapshot_requires_matching_trainer_anchor( + self, tmp_path: Path, capsys: Any + ) -> None: + actor = self._make_actor(tmp_path) + actor._train_steps = 1 + actor._trainer_version = 1 + step_dir = tmp_path / "checkpoints" / "step_1" + + async def _save_without_then_with_anchor() -> tuple[bool, bool, bool]: + first = await actor._save_rollout_checkpoint(force=True) + second = await actor._save_rollout_checkpoint(force=True) + step_dir.mkdir(parents=True) + third = await actor._save_rollout_checkpoint(force=True) + return first, second, third + + assert asyncio.run(_save_without_then_with_anchor()) == (False, False, True) + actor._checkpointer.shutdown() + assert ( + capsys.readouterr().out.count( + "rollout checkpoint skipped: matching trainer checkpoint" + ) + == 1 + ) + manifest = json.loads( + ( + step_dir + / "rollout_snapshots" + / "snapshot_000001" + / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME + ).read_text() + ) + assert manifest["base_train_step"] == 1 + assert manifest["trainer_version"] == 1 + assert manifest["bootstrap_fingerprint"] is None + + def test_pre_step_deadline_retries_guarded_cut_before_stopping( + self, tmp_path: Path + ) -> None: + actor = self._make_actor(tmp_path) + actor._master_config.rollout_checkpointing.interval_s = 0.001 + actor._timeout.last_save_time = 0 + save = AsyncMock(side_effect=[False, True]) + + with patch.object(actor, "_save_rollout_checkpoint", new=save): + asyncio.run(asyncio.wait_for(actor._rollout_checkpoint_pump(), timeout=5)) + actor._checkpointer.shutdown() + + assert save.await_count == 2 + assert all(call.kwargs == {"force": True} for call in save.await_args_list) + assert actor._timeout.last_saved is True + assert actor._rollout_checkpoint_stop_requested.is_set() + + def test_active_streamed_step_is_persisted_as_replayable( + self, tmp_path: Path + ) -> None: + ledger = RolloutRecoveryLedger() + group = ledger.reserve_group( + group_id="claimed", + prompt_id="17", + prompt_ref=PromptRef(sample_id="17", task_name="nemo_gym"), + prompt_payload={"idx": 17, "task_name": "nemo_gym"}, + expected_generations=2, + target_step=0, + start_weight_version=0, + ) + ledger.mark_group_dispatched(group.group_id) + for generation_index, gate_id in enumerate(group.gate_rollout_ids): + ledger.mark_sibling_sealed( + group.group_id, + generation_index=generation_index, + gate_rollout_id=gate_id, + receipt={ + "rollout_id": gate_id, + "manifest": [{"staging_key": f"{gate_id}/call"}], + }, + reward=float(generation_index), + ) + canonical_meta = KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=group.logical_rollout_ids, + fields=["input_ids"], + sequence_lengths=[16, 16], + tags=[{"weight_version": 0}, {"weight_version": 0}], + ) + ledger.mark_finalization_started(group.group_id) + ledger.mark_group_finalized( + group.group_id, + meta=canonical_meta, + group_min_weight_version=0, + group_max_weight_version=0, + ) + ledger.claim_groups_for_training( + [group.group_id], + train_step=0, + trainer_version=0, + expected_group_count=2, + ) + dp_client = _FakeDPClient( + sample_ids=list(canonical_meta.sample_ids), + staging_sample_ids=sorted(ledger.expected_staging_keys()), + ) + actor = self._make_actor( + tmp_path, + ledger=ledger, + dp_client=dp_client, + tq_buffer=_FakeTQBuffer(), + ) + + assert asyncio.run(actor._save_rollout_checkpoint(force=True)) + actor._checkpointer.shutdown() + + snapshot = ( + tmp_path + / "checkpoints" + / BOOTSTRAP_DIRNAME + / "rollout_snapshots" + / "snapshot_000001" + ) + replay_state = torch.load( + snapshot / REPLAY_BUFFER_METADATA_FILENAME, + weights_only=False, + ) + lineage_state = torch.load( + snapshot / ROLLOUT_RECOVERY_STATE_FILENAME, + weights_only=False, + ) + restored = RolloutRecoveryLedger.from_state_dict(lineage_state) + assert [item["group_id"] for item in replay_state["groups"]] == ["claimed"] + assert restored.open_train_step is None + assert restored.get_group("claimed").status.value == "finalized" + assert ledger.open_train_step is not None + manifest = json.loads( + (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).read_text() + ) + assert manifest["rolled_back_train_group_count"] == 1 + + class TestDataPlaneCheckpoint: def test_saves_digest_bound_rollout_lineage_without_prompt_payload(self, tmp_path): mc = _actor_master_config( @@ -1485,6 +1775,68 @@ def test_setup_fresh_start_passes_none_paths( assert actor_args.save_state == _initial_grpo_save_state() assert actor_args.last_checkpoint_path is None + def test_latest_restore_mode_selects_newer_rollout_snapshot( + self, + patched_factories, # noqa: F811 + tmp_path: Path, + ) -> None: + ckpt_dir = tmp_path / "ckpts" + step_3 = _write_checkpoint( + ckpt_dir, + 3, + _STEP_3_SAVE_STATE, + dataloader_state={"fake_position": 3}, + ) + snapshot = step_3 / "rollout_snapshots" / "snapshot_000001" + snapshot.mkdir(parents=True) + (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).write_text("committed\n") + (snapshot / ROLLOUT_SNAPSHOT_MANIFEST_FILENAME).write_text( + json.dumps( + { + "schema_version": 1, + "base_train_step": 3, + "trainer_version": 3, + "current_epoch": 7, + "mutation_version": 1, + "rolled_back_train_group_count": 0, + "bootstrap_fingerprint": None, + } + ) + ) + torch.save({"fake_position": 7}, snapshot / "train_dataloader.pt") + mc = _setup_master_config(str(ckpt_dir)) + + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert actor_args.last_checkpoint_path == str(snapshot) + assert actor_args.save_state.current_step == 3 + assert actor_args.save_state.current_epoch == 7 + trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs + assert trainer_kwargs["weights_path"] == step_3 / "policy" / "weights" + + def test_trainer_restore_mode_ignores_periodic_snapshot( + self, + patched_factories, # noqa: F811 + tmp_path: Path, + ) -> None: + ckpt_dir = tmp_path / "ckpts" + step_3 = _write_checkpoint( + ckpt_dir, + 3, + _STEP_3_SAVE_STATE, + dataloader_state={"fake_position": 3}, + ) + snapshot = step_3 / "rollout_snapshots" / "snapshot_000001" + snapshot.mkdir(parents=True) + (snapshot / ROLLOUT_SNAPSHOT_COMMITTED_FILENAME).write_text("committed\n") + mc = _setup_master_config(str(ckpt_dir)) + mc.rollout_checkpointing.restore_mode = "trainer_checkpoint" + + actor_args, _ = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert actor_args.last_checkpoint_path == str(step_3) + assert actor_args.save_state.current_epoch == 1 + def test_setup_forwards_pretrained_checkpoint( self, patched_factories, # noqa: F811 diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index c5e1c23cf10..f1246cd5e73 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -358,6 +358,10 @@ def test_multiple_dataloader_not_supported(self): "deferred_routes_without_capture", "defer_routed_experts_to_policy requires", ), + ( + "recovery_capacity", + "num_prompts_per_step.*min_groups_for_streaming_train - 1", + ), ], ) def test_invalid_config_fails_before_setup_factories( @@ -381,6 +385,16 @@ def test_invalid_config_fails_before_setup_factories( mc.async_rl.max_buffered_rollouts = 3 elif invalid_case == "deferred_routes_without_capture": mc.token_capture.defer_routed_experts_to_policy = True + elif invalid_case == "recovery_capacity": + # Four restored groups may leave one group below the streaming + # threshold while the producer atomically reserves the next full + # four-prompt batch. Capacity 4 cannot make progress; 4 + 2 - 1 can. + mc.async_rl.sampler = WindowedSamplerConfig(max_staleness_versions=1) + mc.async_rl.min_groups_for_streaming_train = 2 + mc.async_rl.max_buffered_rollouts = 4 + mc.token_capture.enabled = True + mc.checkpointing["enabled"] = True + mc.data_plane["checkpointing_enabled"] = True else: # pragma: no cover raise AssertionError(f"unknown test case {invalid_case}") diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 61aa1adcc07..36f8eb5367d 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -208,6 +208,21 @@ def _add_group( class TestDataPlaneCheckpointBarrier: + def test_mutation_version_advances_once_per_outer_section(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + assert barrier.mutation_version == 0 + async with barrier.mutation(): + async with barrier.mutation(): + assert barrier.mutation_version == 0 + assert barrier.mutation_version == 1 + with pytest.raises(RuntimeError, match="injected"): + async with barrier.mutation(): + raise RuntimeError("injected") + assert barrier.mutation_version == 2 + + asyncio.run(exercise()) + def test_mutations_run_concurrently_without_checkpoint(self): async def exercise() -> None: barrier = DataPlaneCheckpointBarrier() @@ -1048,6 +1063,23 @@ def test_metadata_state_dict_omits_tensors_and_data_plane_reads(self): assert state["manifest_digest"] == replay_manifest_digest(state["groups"]) assert dp.get_calls == [] + def test_metadata_state_dict_adds_training_claimed_groups(self): + buf = _make_buffer(FakeDataPlaneClient()) + live_meta = _add_group(buf, weight=1) + claimed = _make_group_entry("claimed", weight=2) + + state = buf.metadata_state_dict( + saved_capacity=8, + additional_groups=[claimed], + ) + + assert [group["group_id"] for group in state["groups"]] == [ + _group_id_of(live_meta), + "claimed", + ] + assert state["groups"][1] == claimed + assert state["manifest_digest"] == replay_manifest_digest(state["groups"]) + def test_native_tq_round_trip_restores_index_without_reputting_rows(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)