From 1229d6813ef771db3e9079e6e51a35837b418aa9 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 17:03:36 -0400 Subject: [PATCH 1/4] feat(rollout): checkpoint active streamed steps Add periodic TQ, replay, lineage, and dataloader snapshots anchored to durable trainer state. Roll back groups claimed by an unfinished optimizer step into replayable state during persistence, restore them on restart, and cover partial-sibling plus 2/8 and 6/8 streamed recovery paths. (cherry picked from commit 98c196f4b0793541c4bde4e1a01ba0b5c7d21cb3) --- .../algorithms/async_utils/replay_buffer.py | 43 +- .../async_utils/staleness_sampler.py | 39 +- nemo_rl/algorithms/single_controller.py | 648 ++++++++++++++++-- .../single_controller_utils/__init__.py | 2 + .../single_controller_utils/config.py | 61 ++ .../rollout_checkpoint.py | 515 ++++++++++++++ .../single_controller_utils/setup.py | 142 +++- nemo_rl/experience/rollout_manager.py | 56 +- nemo_rl/experience/rollout_recovery.py | 60 ++ nemo_rl/utils/timer.py | 21 +- pyrefly.toml | 1 + .../L1_Functional_Tests_SingleController.sh | 11 + .../grpo_async_gym_single_controller.sh | 38 +- ...ym_single_controller_streaming_recovery.sh | 248 +++++++ ...async_gym_single_controller_tq_recovery.sh | 221 ++++++ tests/unit/experience/test_rollout_manager.py | 20 + .../unit/experience/test_rollout_recovery.py | 38 + .../test_rollout_checkpoint.py | 395 +++++++++++ .../single_controller/test_rollout_pump.py | 95 ++- .../test_sc_checkpointing.py | 360 +++++++++- .../test_single_controller_setup.py | 14 + .../test_tq_replay_buffer.py | 32 + tests/unit/utils/test_timer.py | 8 + 23 files changed, 2935 insertions(+), 133 deletions(-) create mode 100644 nemo_rl/algorithms/single_controller_utils/rollout_checkpoint.py create mode 100755 tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh create mode 100755 tests/functional/grpo_async_gym_single_controller_tq_recovery.sh create mode 100644 tests/unit/single_controller/test_rollout_checkpoint.py 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..fc87683dc17 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,169 @@ 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() - } + 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" + ) - # 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, - ) + 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 = 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 + ) + 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() + rollout_work_items = [] + for prompt_idx in range(prompt_batch.size): + 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 +2046,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 +2075,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 +2314,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 +2332,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 +2352,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 +2368,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 +2405,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..fb75cc8648e 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -299,6 +299,25 @@ 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. + """ + + 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 +330,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 +354,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 +379,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..28ea279721d --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -0,0 +1,248 @@ +#!/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 + ++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" \ + "$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]) +expected_claimed = int(sys.argv[4]) +deadline = time.monotonic() + float(sys.argv[5]) + +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" + ) 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..81eba536694 --- /dev/null +++ b/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh @@ -0,0 +1,221 @@ +#!/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 + ++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" \ + "$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]) +expected_generations = int(sys.argv[4]) +min_sealed = int(sys.argv[5]) +deadline = time.monotonic() + float(sys.argv[6]) + +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" + ) 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/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) From 16a9480a06f7d42d2b5fc8702157e09193efdb84 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 17:19:04 -0400 Subject: [PATCH 2/4] fix(test): neutralize SC checkpoint metric Set checkpointing.metric_name=null in every SC recovery harness that enables checkpointing, since SC has no validation loop. Include phase-one log tails when crash-driven recovery tests exit before producing their target snapshot. (cherry picked from commit a5e9a16a42f4286611e16a2723ae5b032c0d4687) --- ...ym_single_controller_streaming_recovery.sh | 19 ++++++++++++++--- ...async_gym_single_controller_tq_recovery.sh | 21 +++++++++++++++---- .../grpo_dp_single_controller_tq_recovery.sh | 1 + 3 files changed, 34 insertions(+), 7 deletions(-) diff --git a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh index 28ea279721d..2ca919373b7 100755 --- a/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_streaming_recovery.sh @@ -79,6 +79,9 @@ COMMON_OVERRIDES=( 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 @@ -110,6 +113,7 @@ 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 @@ -121,8 +125,16 @@ from pathlib import Path root = Path(sys.argv[1]) selection = Path(sys.argv[2]) phase_pid = int(sys.argv[3]) -expected_claimed = int(sys.argv[4]) -deadline = time.monotonic() + float(sys.argv[5]) +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): @@ -146,7 +158,8 @@ while time.monotonic() < deadline: os.kill(phase_pid, 0) except ProcessLookupError as error: raise RuntimeError( - "phase one exited before producing the requested streamed-step cut" + "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) diff --git a/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh b/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh index 81eba536694..54d3c41b1a4 100755 --- a/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh +++ b/tests/functional/grpo_async_gym_single_controller_tq_recovery.sh @@ -64,6 +64,9 @@ COMMON_OVERRIDES=( 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' @@ -96,6 +99,7 @@ uv run --no-sync python - \ "$CHECKPOINT_DIR/bootstrap/rollout_snapshots" \ "$SELECTION_FILE" \ "$PHASE1_PID" \ + "$PHASE1_LOG" \ "$NUM_GENERATIONS" \ "$MIN_SEALED" \ "$SNAPSHOT_TIMEOUT_S" <<'PY' @@ -109,9 +113,17 @@ import torch root = Path(sys.argv[1]) selection = Path(sys.argv[2]) phase_pid = int(sys.argv[3]) -expected_generations = int(sys.argv[4]) -min_sealed = int(sys.argv[5]) -deadline = time.monotonic() + float(sys.argv[6]) +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): @@ -141,7 +153,8 @@ while time.monotonic() < deadline: os.kill(phase_pid, 0) except ProcessLookupError as error: raise RuntimeError( - "phase one exited before producing a partial-sibling snapshot" + "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) 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' From 1926de57bae3f114c461d0a918d44c642426c836 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 20:52:27 -0400 Subject: [PATCH 3/4] docs(rollout): surface periodic checkpoint defaults (cherry picked from commit 4cd773d826a87c6ec8b444b9eec5e50064c06658) --- .../configs/grpo_math_1B_megatron_single_controller.yaml | 7 +++++++ nemo_rl/algorithms/single_controller_utils/config.py | 4 ++++ 2 files changed, 11 insertions(+) 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/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index fb75cc8648e..db9e777763d 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -311,6 +311,10 @@ class RolloutCheckpointConfig(BaseModel, extra="allow"): ``trainer_checkpoint`` ignores newer periodic snapshots, while ``none`` resumes trainer state without restoring rollout, replay, lineage, or the dataloader cursor. + + SingleController has no validation loop, so checkpoint selection must use + ``checkpointing.metric_name=None`` or a ``train:`` metric. Inherited + ``val:`` settings are rejected during setup. """ interval_s: Optional[float] = Field(default=None, gt=0) From 9355cc54785813bbdea09a7f6117dd5a50d28889 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 19 Aug 2026 20:54:07 -0400 Subject: [PATCH 4/4] fix(rollout): preserve restored target-step top-up (cherry picked from commit f1d8ad044804ac84f15f0cffa0bba4e6f37b8da1) --- nemo_rl/algorithms/single_controller.py | 26 +++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index fc87683dc17..61d604c78ce 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -1688,6 +1688,15 @@ async def _acquire_capacity(permits: int) -> None: _release_capacity(acquired) raise + 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) + max_epochs = self._master_config.grpo.max_num_epochs token_capture_config = getattr(self._master_config, "token_capture", None) token_capture_enabled = bool( @@ -1741,15 +1750,9 @@ async def _acquire_capacity(permits: int) -> None: target_step = self._sampler.commit_admission( trainer_version=self._trainer_version ) - 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 - ) + 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 @@ -1787,8 +1790,11 @@ async def _acquire_capacity(permits: int) -> None: target_step = await self._sampler.admit( trainer_version_fn=lambda: self._trainer_version ) + num_prompts = _num_prompts_to_dispatch( + prompt_batch.size, target_step + ) rollout_work_items = [] - for prompt_idx in range(prompt_batch.size): + for prompt_idx in range(num_prompts): prompt = { # type: ignore key: value[prompt_idx] for key, value in prompt_batch.items()