diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 8c45f089b9..ee38971cfb 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -496,6 +496,10 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + # SingleController only: save native TQ state. Required when trainer + # checkpointing is enabled with a replay-checkpoint-capable sampler; + # supported samplers restore from metadata-only replay indexes. + checkpointing_enabled: false global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index fe4a6b3f31..ba841363c7 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -21,7 +21,6 @@ async_rl: # Enable per-rollout diagnostic prints (prompt content / completion previews). diagnostics: false -# SC does not support checkpointing yet. checkpointing: enabled: false checkpoint_dir: results/grpo-single-controller @@ -55,6 +54,8 @@ logger: data_plane: enabled: true + # Required shadow snapshot: a save failure aborts checkpoint finalization. + checkpointing_enabled: true cluster: gpus_per_node: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 1bf9c5ae93..61da9684d9 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -13,21 +13,130 @@ # limitations under the License. import asyncio +import hashlib +import json import statistics import threading as _threading import uuid from collections import Counter -from collections.abc import Mapping -from typing import Any, Iterable, Optional +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager +from typing import Any, Iterable, Literal, Optional, TypedDict import ray from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.async_utils import call_data_plane from nemo_rl.experience.interfaces import PromptGroupRecord from nemo_rl.experience.payload import pack_payload, record_to_train_batch +DATA_PLANE_CHECKPOINT_DIR = "data_plane" +DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 +REPLAY_BUFFER_METADATA_FILENAME = "replay_buffer_metadata.pt" +LEGACY_REPLAY_BUFFER_FILENAME = "replay_buffer.pt" +REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1 +REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" + + +class TQReplayGroupMetadata(TypedDict): + """Controller-local index for one training-ready group stored in TQ.""" + + meta: KVBatchMeta + start_weight: int + end_weight: int + target_step: Optional[int] + group_id: str + + +class TQReplayMetadataState(TypedDict): + """Versioned metadata-only replay sidecar paired with a TQ snapshot.""" + + schema_version: int + storage: Literal["tq_checkpoint"] + partition_id: str + saved_capacity: int + manifest_digest: str + groups: list[TQReplayGroupMetadata] + + +def replay_manifest_digest(groups: list[TQReplayGroupMetadata]) -> str: + """Return a stable digest binding replay metadata to a TQ checkpoint.""" + digest_input = [ + { + "group_id": group["group_id"], + "start_weight": group["start_weight"], + "end_weight": group["end_weight"], + "target_step": group["target_step"], + "meta": { + "partition_id": group["meta"].partition_id, + "task_name": group["meta"].task_name, + "sample_ids": list(group["meta"].sample_ids), + "fields": ( + list(group["meta"].fields) + if group["meta"].fields is not None + else None + ), + "sequence_lengths": ( + list(group["meta"].sequence_lengths) + if group["meta"].sequence_lengths is not None + else None + ), + "tags": group["meta"].tags, + "extra_info": group["meta"].extra_info, + }, + } + for group in groups + ] + encoded = json.dumps( + digest_input, sort_keys=True, separators=(",", ":") + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +class DataPlaneCheckpointBarrier: + """Allow concurrent mutations while giving checkpoints exclusive access.""" + + def __init__(self) -> None: + self._condition = asyncio.Condition() + self._checkpoint_active = False + self._active_mutations = 0 + + @asynccontextmanager + async def mutation(self) -> AsyncIterator[None]: + """Enter a commit/clear section, waiting only for an active checkpoint.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._checkpoint_active) + self._active_mutations += 1 + try: + yield + finally: + async with self._condition: + self._active_mutations -= 1 + if self._active_mutations == 0: + self._condition.notify_all() + + @asynccontextmanager + async def checkpoint(self) -> AsyncIterator[None]: + """Block new mutations and wait for active ones before snapshotting.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._checkpoint_active) + self._checkpoint_active = True + try: + await self._condition.wait_for(lambda: self._active_mutations == 0) + except BaseException: + self._checkpoint_active = False + self._condition.notify_all() + raise + try: + yield + finally: + async with self._condition: + self._checkpoint_active = False + self._condition.notify_all() + + # Classes with @ray.remote can't be inherited from, so we split the implementation out. class ReplayBufferImpl(ReplayBufferProtocol): """Replay buffer storing per-prompt groups. @@ -659,6 +768,22 @@ def __init__( self.target_step_list: list[Optional[int]] = [] self.ready_list: list[bool] = [] self._group_ids: list[str] = [] + self._data_plane_checkpoint_barrier: Optional[ + DataPlaneCheckpointBarrier + ] = None + + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + """Bind the controller's shared checkpoint/mutation barrier once. + + A private fallback barrier would not coordinate with controller-owned + saves and clears, so destructive operations fail loudly until the SC + actor supplies its barrier. + """ + if self._data_plane_checkpoint_barrier is not None: + raise RuntimeError("data-plane checkpoint barrier is already configured") + self._data_plane_checkpoint_barrier = barrier def reserve( self, @@ -716,52 +841,57 @@ async def commit( f"commit called with unknown group_id={group_id!r}; " f"reserve() must precede commit() (or the slot was already removed)" ) + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before committing samples" + ) train_batch = record_to_train_batch(record, pad_value_dict=self._pad_value_dict) sample_ids, fields, tags = pack_payload( train_batch, weight_version=start_weight_version, group_id=group_id ) - try: - await self._call_dp( - "put_samples", - sample_ids=sample_ids, - partition_id=self._partition_id, - fields=fields, - tags=tags, - ) - - # mirrors kv_first_write - lengths = train_batch["input_lengths"] - meta = KVBatchMeta( - partition_id=self._partition_id, - task_name="train", - sample_ids=list(sample_ids), - fields=list(fields.keys()), - sequence_lengths=[int(s) for s in lengths.tolist()], - tags=[dict(t) for t in tags], - ) - - idx = self._group_ids.index(group_id) - self.meta_list[idx] = meta - self.end_weight_list[idx] = end_weight_version - self.ready_list[idx] = True - return meta - except BaseException as commit_error: - # put_samples may have written rows before raising. Roll back by the - # deterministic IDs known here; the caller removes the reserved slot. + async with self._data_plane_checkpoint_barrier.mutation(): try: - await self._call_dp( - "clear_samples", - sample_ids=list(sample_ids), + await call_data_plane( + self._dp_client, + "put_samples", + sample_ids=sample_ids, partition_id=self._partition_id, + fields=fields, + tags=tags, ) - except BaseException as rollback_error: - if isinstance(commit_error, asyncio.CancelledError): - raise commit_error from rollback_error - raise BaseExceptionGroup( - f"commit and rollback both failed for group_id={group_id!r}", - [commit_error, rollback_error], + + # mirrors kv_first_write + lengths = train_batch["input_lengths"] + meta = KVBatchMeta( + partition_id=self._partition_id, + task_name="train", + sample_ids=list(sample_ids), + fields=list(fields.keys()), + sequence_lengths=[int(s) for s in lengths.tolist()], + tags=[dict(t) for t in tags], ) - raise + + idx = self._group_ids.index(group_id) + self.meta_list[idx] = meta + self.end_weight_list[idx] = end_weight_version + self.ready_list[idx] = True + return meta + except BaseException as commit_error: + # put_samples may have written rows before raising. Roll back by the + # deterministic IDs while retaining the barrier mutation slot. + try: + await self._clear_samples_unlocked( + sample_ids=list(sample_ids), + ) + except BaseException as rollback_error: + if isinstance(commit_error, asyncio.CancelledError): + raise commit_error from rollback_error + raise BaseExceptionGroup( + f"commit and rollback both failed for group_id={group_id!r}", + [commit_error, rollback_error], + ) + raise async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> int: """Remove the live slot identified by ``group_id``. @@ -776,11 +906,19 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in Raises: ValueError: ``group_id`` has no live slot. """ - try: - idx = self._group_ids.index(group_id) - except ValueError as error: - raise ValueError(f"unknown group_id={group_id!r}") from error - return await self.remove([idx], remove_in_dp=remove_in_dp) + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before removing a group" + ) + async with self._data_plane_checkpoint_barrier.mutation(): + try: + idx = self._group_ids.index(group_id) + except ValueError as error: + raise ValueError(f"unknown group_id={group_id!r}") from error + return await self._remove_unlocked( + [idx], clear_data_plane=remove_in_dp + ) async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: """Drop entries at the given indices and optionally clear them from DataPlane. @@ -794,14 +932,26 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: """ if len(idxs) == 0: return 0 - - drop_idxs = sorted(idxs, reverse=True) - if drop_idxs[0] >= len(self.meta_list): - raise IndexError( - f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; " - f"size={len(self.meta_list)}" + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before removing groups" + ) + async with self._data_plane_checkpoint_barrier.mutation(): + drop_idxs = sorted(idxs, reverse=True) + if drop_idxs[0] >= len(self.meta_list): + raise IndexError( + f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; " + f"size={len(self.meta_list)}" + ) + return await self._remove_unlocked( + drop_idxs, clear_data_plane=remove_in_dp ) + async def _remove_unlocked( + self, drop_idxs: list[int], *, clear_data_plane: bool + ) -> int: + """Remove validated indices while the caller owns any required lock.""" dropped_sample_ids: list[str] = [] for i in drop_idxs: meta = self.meta_list[i] @@ -814,15 +964,207 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: del self.ready_list[i] del self._group_ids[i] - if remove_in_dp: - await self._call_dp( - "clear_samples", + if clear_data_plane: + await self._clear_samples_unlocked( sample_ids=dropped_sample_ids, - partition_id=self._partition_id, ) return len(drop_idxs) + def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: + """Capture the controller index for ready groups without tensor payloads. + + The caller must hold the exclusive side of the shared data-plane + checkpoint barrier through this capture and the matching TQ save. + Commits and destructive clears use shared mutation slots, so the sidecar + and native snapshot describe one exact set of training-ready groups. + Every operation that mutates the canonical rollout partition or its + controller-local replay membership must participate in that barrier + 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. + """ + groups: list[TQReplayGroupMetadata] = [] + for i, ready in enumerate(self.ready_list): + if not ready: + continue + meta = self.meta_list[i] + assert meta is not None # commit sets meta before ready=True + groups.append( + { + "meta": meta, + "start_weight": self.start_weight_list[i], + "end_weight": self.end_weight_list[i], + "target_step": self.target_step_list[i], + "group_id": self._group_ids[i], + } + ) + return { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": self._partition_id, + "saved_capacity": saved_capacity, + "manifest_digest": replay_manifest_digest(groups), + "groups": groups, + } + + async def load_state_dict( + self, + state: dict[str, Any], + *, + max_groups: int, + expected_partition_id: str, + expected_group_size: int, + expected_manifest_digest: str, + ) -> int: + """Restore the local replay index for an already-restored TQ snapshot. + + The sidecar never contains tensor payloads and this method never writes + to the DataPlane. TQ must be restored first; the caller binds the two + artifacts by passing the manifest digest returned by TQ checkpoint + loading. + + Staleness is intentionally NOT handled here — load only loads. The + train pump's first ``sampler.evict`` drops any restored group that is + outside the staleness window and releases its capacity permit, keeping + eviction in one place. + + Args: + state: Envelope produced by ``metadata_state_dict``. + max_groups: Current max_buffered_rollouts; the restored count + never exceeds it. + expected_partition_id: Partition this buffer writes to; must + match the envelope. + expected_group_size: num_generations_per_prompt; every group must + hold exactly this many rows (a changed group size silently + breaks the group-relative baseline). + + Returns: + Number of groups restored into the buffer. + + Raises: + ValueError: If the envelope is malformed (missing keys, partition + mismatch, misaligned or wrongly sized groups, duplicate + sample_ids). + """ + if self.meta_list or self._group_ids: + raise RuntimeError( + "Replay-buffer checkpoint loading requires an empty local buffer" + ) + required_keys = { + "schema_version", + "storage", + "partition_id", + "saved_capacity", + "manifest_digest", + "groups", + } + missing_keys = required_keys - set(state) + if missing_keys: + raise ValueError( + f"Replay buffer checkpoint missing required keys: {missing_keys}" + ) + if state["schema_version"] != REPLAY_BUFFER_METADATA_SCHEMA_VERSION: + raise ValueError( + "Unsupported replay-buffer metadata schema version: " + f"{state['schema_version']!r}" + ) + if state["storage"] != REPLAY_BUFFER_METADATA_STORAGE: + raise ValueError( + "Replay-buffer metadata has unsupported storage: " + f"{state['storage']!r}" + ) + if state["partition_id"] != expected_partition_id: + raise ValueError( + "Replay buffer checkpoint partition_id mismatch: " + f"checkpoint={state['partition_id']!r}, " + f"expected={expected_partition_id!r}" + ) + + groups = list(state["groups"]) + group_keys = { + "meta", + "start_weight", + "end_weight", + "target_step", + "group_id", + } + seen_sample_ids: set[str] = set() + for group in groups: + if "fields_data" in group: + raise ValueError( + "Metadata-only replay checkpoint must not contain fields_data" + ) + missing_group_keys = group_keys - set(group) + if missing_group_keys: + raise ValueError( + f"Replay buffer checkpoint group missing keys: {missing_group_keys}" + ) + meta = group["meta"] + if meta.partition_id != expected_partition_id: + raise ValueError( + "Replay buffer checkpoint group partition_id mismatch: " + f"checkpoint={meta.partition_id!r}, " + f"expected={expected_partition_id!r}" + ) + num_tags = len(meta.tags) if meta.tags is not None else -1 + num_lengths = ( + len(meta.sequence_lengths) if meta.sequence_lengths is not None else -1 + ) + if not ( + len(meta.sample_ids) == num_tags == num_lengths == expected_group_size + ): + raise ValueError( + "Replay buffer checkpoint group misaligned: " + f"sample_ids={len(meta.sample_ids)}, tags={num_tags}, " + f"sequence_lengths={num_lengths}, " + f"expected_group_size={expected_group_size}" + ) + for sid in meta.sample_ids: + if sid in seen_sample_ids: + raise ValueError( + f"Replay buffer checkpoint has duplicate sample_id: {sid!r}" + ) + seen_sample_ids.add(sid) + + actual_digest = replay_manifest_digest(groups) + if state["manifest_digest"] != actual_digest: + raise ValueError( + "Replay-buffer metadata digest does not match its contents" + ) + if expected_manifest_digest != actual_digest: + raise ValueError( + "Replay-buffer metadata does not match the loaded TQ checkpoint" + ) + + if state["saved_capacity"] != max_groups: + print( + "TQReplayBuffer capacity changed: " + f"checkpoint={state['saved_capacity']}, current={max_groups}. " + "Using current config value." + ) + if len(groups) > max_groups: + raise ValueError( + "Native TQ checkpoint contains more replay groups than the current " + f"buffer capacity: checkpoint={len(groups)}, current={max_groups}" + ) + + for group in groups: + meta = group["meta"] + self.meta_list.append(meta) + self.start_weight_list.append(group["start_weight"]) + self.end_weight_list.append(group["end_weight"]) + self.target_step_list.append(group["target_step"]) + self.ready_list.append(True) + self._group_ids.append(group["group_id"]) + + print( + f"📦 Restored {len(groups)} replay group(s) from checkpoint", + flush=True, + ) + return len(groups) + def size(self) -> int: """Return the number of prompt-group entries currently held.""" return len(self.meta_list) @@ -830,13 +1172,22 @@ def size(self) -> int: def __len__(self) -> int: return len(self.meta_list) - async def _call_dp(self, method_name: str, **kwargs: Any) -> Any: - """Call a DataPlaneClient method, awaiting Ray remotes if needed.""" - method = getattr(self._dp_client, method_name) - remote = getattr(method, "remote", None) - if remote is not None: - return await remote(**kwargs) - result = method(**kwargs) - if asyncio.iscoroutine(result): - return await result - return result + async def _clear_samples(self, *, sample_ids: list[str]) -> None: + """Clear rows without overlapping a bound data-plane checkpoint.""" + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint barrier before clearing samples" + ) + async with self._data_plane_checkpoint_barrier.mutation(): + await self._clear_samples_unlocked(sample_ids=sample_ids) + + async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None: + """Clear rows while the caller holds a barrier mutation slot.""" + await call_data_plane( + self._dp_client, + "clear_samples", + offload_sync=True, + sample_ids=sample_ids, + partition_id=self._partition_id, + ) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 69959de6e8..2896192f22 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -44,6 +44,7 @@ from typing import ( Annotated, Callable, + ClassVar, Literal, Optional, Protocol, @@ -100,6 +101,16 @@ def is_on_policy(self) -> bool: """True when the policy admits zero staleness (sync mode).""" ... + @property + def supports_buffer_checkpoint(self) -> bool: + """True when buffered groups may be checkpointed and restored. + + Gated policies dispatch a fixed quota per trainer step, so groups + restored from a checkpoint can never complete an already-consumed + quota window — only ungated (over-sampled) policies can consume them. + """ + ... + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" ... @@ -113,11 +124,26 @@ class BaseSampler(abc.ABC): select-finalize / weight-window-evict helpers. """ - def __init__(self, buffer: TQReplayBuffer) -> None: + def __init__(self, buffer: TQReplayBuffer, *, resume_from_step: int = 0) -> None: + """Initialize shared sampler state. + + Args: + buffer: Shared TQReplayBuffer holding the candidate slots. + resume_from_step: Trainer step this run starts from — 0 for a + fresh run, the restored ``current_step`` when resuming from a + checkpoint. Seeds the dispatch cursor so gated ``admit`` and + ``InOrderSampler``'s target_step stamps line up with the + restored trainer version exactly as they would at step 0 of + a fresh run. + """ + if resume_from_step < 0: + raise ValueError( + f"resume_from_step must be non-negative, got {resume_from_step}" + ) self._buffer = buffer - # Pre-incremented before each admitted batch, so -1 lets the first - # batch through a zero-staleness gate. - self._dispatch_index: int = -1 + # Pre-incremented before each admitted batch, so the cursor trails + # the run's starting step by one. + self._dispatch_index: int = resume_from_step - 1 # ── rollout-pump side ──────────────────────────────────────────────── @abc.abstractmethod @@ -158,6 +184,10 @@ async def evict(self, *, current_train_weight: int) -> int: def is_on_policy(self) -> bool: return self._eviction_window() == 0 + @property + def supports_buffer_checkpoint(self) -> bool: + return False + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: return None @@ -215,8 +245,9 @@ def __init__( *, max_staleness_versions: int, sample_freshest_first: bool = False, + resume_from_step: int = 0, ) -> None: - super().__init__(buffer) + super().__init__(buffer, resume_from_step=resume_from_step) if max_staleness_versions < 0: raise ValueError( f"max_staleness_versions must be non-negative, got " @@ -228,6 +259,12 @@ def __init__( def _eviction_window(self) -> int: return self.max_staleness_versions + @property + def supports_buffer_checkpoint(self) -> bool: + # Ungated: restored groups are ordinary in-window candidates, so the + # buffer can round-trip through a checkpoint. + return True + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: # Over-sampled: dispatch is bounded by buffer capacity, not by version. return None @@ -275,8 +312,14 @@ class _GatedSampler(BaseSampler): (``gate_window`` versions of lookahead). """ - def __init__(self, buffer: TQReplayBuffer, *, gate_window: int) -> None: - super().__init__(buffer) + def __init__( + self, + buffer: TQReplayBuffer, + *, + gate_window: int, + resume_from_step: int = 0, + ) -> None: + super().__init__(buffer, resume_from_step=resume_from_step) if gate_window < 0: raise ValueError(f"gate_window must be non-negative, got {gate_window}") self._gate_window = gate_window @@ -308,8 +351,18 @@ class WeightFifoSampler(_GatedSampler): that weight's batch to fill. Evict uses the weight window (default). """ - def __init__(self, buffer: TQReplayBuffer, *, max_staleness_versions: int) -> None: - super().__init__(buffer, gate_window=max_staleness_versions) + def __init__( + self, + buffer: TQReplayBuffer, + *, + max_staleness_versions: int, + resume_from_step: int = 0, + ) -> None: + super().__init__( + buffer, + gate_window=max_staleness_versions, + resume_from_step=resume_from_step, + ) self.max_staleness_versions = max_staleness_versions async def select( @@ -349,8 +402,18 @@ class InOrderSampler(_GatedSampler): upcoming is never dropped early, and evict/select can't disagree. """ - def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> None: - super().__init__(buffer, gate_window=max_lookahead_versions) + def __init__( + self, + buffer: TQReplayBuffer, + *, + max_lookahead_versions: int, + resume_from_step: int = 0, + ) -> None: + super().__init__( + buffer, + gate_window=max_lookahead_versions, + resume_from_step=resume_from_step, + ) self.max_lookahead_versions = max_lookahead_versions def _stamp(self) -> Optional[int]: @@ -393,6 +456,7 @@ async def evict(self, *, current_train_weight: int) -> int: class WindowedSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = True name: Literal["windowed"] = "windowed" # Max weight-version gap a selected group may have from the trainer. max_staleness_versions: NonNegativeInt = 1 @@ -401,18 +465,22 @@ class WindowedSamplerConfig(BaseModel, extra="allow"): class WeightFifoSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["weight_fifo"] = "weight_fifo" # Lookahead + selectable weight window, in trainer versions. max_staleness_versions: NonNegativeInt = 1 class InOrderSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["in_order"] = "in_order" # How far generation may run ahead of the trainer, in dispatch batches. max_lookahead_versions: NonNegativeInt = 1 class CustomSamplerConfig(BaseModel, extra="allow"): + # A custom implementation's capability is known only after construction. + supports_buffer_checkpoint: ClassVar[Optional[bool]] = None name: Literal["custom"] = "custom" # "module:ClassName" of a PromptGroupSampler defined outside this repo. # Extra keys are forwarded to the constructor (after ``buffer``). @@ -454,32 +522,69 @@ def required_buffer_capacity_for_config( def create_sampler( buffer: TQReplayBuffer, cfg: SamplerConfig, + *, + resume_from_step: int = 0, ) -> PromptGroupSampler: - """Build a sampler from its config (or import one by FQN).""" + """Build a sampler from its config (or import one by FQN). + + Args: + buffer: Shared TQReplayBuffer holding the candidate slots. + cfg: Discriminated sampler config selecting the policy. + resume_from_step: Trainer step this run starts from; 0 for a fresh + run, the restored ``current_step`` on checkpoint resume (see + ``BaseSampler.__init__``). Custom samplers receive it only on + resume, so fresh starts don't constrain their constructor + signature — a custom class that doesn't accept the kwarg fails + loudly the first time a run actually resumes with it. + """ + sampler: PromptGroupSampler if isinstance(cfg, WindowedSamplerConfig): - return WindowedSampler( + sampler = WindowedSampler( buffer, max_staleness_versions=cfg.max_staleness_versions, sample_freshest_first=cfg.sample_freshest_first, + resume_from_step=resume_from_step, ) - if isinstance(cfg, WeightFifoSamplerConfig): - return WeightFifoSampler( - buffer, max_staleness_versions=cfg.max_staleness_versions + elif isinstance(cfg, WeightFifoSamplerConfig): + sampler = WeightFifoSampler( + buffer, + max_staleness_versions=cfg.max_staleness_versions, + resume_from_step=resume_from_step, ) - if isinstance(cfg, InOrderSamplerConfig): - return InOrderSampler(buffer, max_lookahead_versions=cfg.max_lookahead_versions) - if isinstance(cfg, CustomSamplerConfig): + elif isinstance(cfg, InOrderSamplerConfig): + sampler = InOrderSampler( + buffer, + max_lookahead_versions=cfg.max_lookahead_versions, + resume_from_step=resume_from_step, + ) + elif isinstance(cfg, CustomSamplerConfig): module_name, sep, class_name = cfg.target.partition(":") if not sep: raise ValueError( f"custom sampler target must be 'module:ClassName', got {cfg.target!r}" ) sampler_cls = getattr(importlib.import_module(module_name), class_name) - sampler = sampler_cls(buffer, **(cfg.model_extra or {})) + custom_kwargs = dict(cfg.model_extra or {}) + if resume_from_step != 0: + custom_kwargs["resume_from_step"] = resume_from_step + sampler = sampler_cls(buffer, **custom_kwargs) if not isinstance(sampler, PromptGroupSampler): raise TypeError( f"{cfg.target} does not implement the PromptGroupSampler " - f"interface (needs admit/select/evict)" + f"interface (needs admit/select/evict, is_on_policy, " + f"supports_buffer_checkpoint, required_buffer_capacity)" ) - return sampler - raise ValueError(f"unknown sampler config {type(cfg).__name__}") + else: + raise ValueError(f"unknown sampler config {type(cfg).__name__}") + + expected_capability = cfg.supports_buffer_checkpoint + if ( + expected_capability is not None + and sampler.supports_buffer_checkpoint != expected_capability + ): + raise RuntimeError( + f"{type(cfg).__name__}.supports_buffer_checkpoint=" + f"{expected_capability} disagrees with {type(sampler).__name__}." + f"supports_buffer_checkpoint={sampler.supports_buffer_checkpoint}" + ) + return sampler diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 85bac9836c..0b3228f61c 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -295,6 +295,7 @@ class GRPOSaveState(TypedDict): current_epoch: int total_steps: int total_valid_tokens: int # Track total number of non-padding tokens during training + trainer_version: NotRequired[int] val_reward: NotRequired[ float ] # Optional field - may not be present during training diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 9effa5168c..6350739296 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -35,14 +35,27 @@ from __future__ import annotations import asyncio +import os +import sys import time +import warnings from functools import partial from typing import Any, Optional, Union import ray import torch +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + TQReplayMetadataState, +) from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler +from nemo_rl.algorithms.grpo import GRPOSaveState, _write_latest_checkpoint_status from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, MasterConfig, @@ -59,17 +72,18 @@ ) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.async_utils import call_data_plane from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmGeneration from nemo_rl.models.policy.tq_policy import TQPolicy +from nemo_rl.utils.checkpoint import CheckpointManager from nemo_rl.utils.logger import Logger -from nemo_rl.utils.timer import Timer +from nemo_rl.utils.timer import TimeoutChecker, Timer Generation = Union[VllmGeneration, SGLangGeneration] - @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: """CPU-only Ray actor that orchestrates the RL training loop. @@ -126,6 +140,29 @@ def __init__( self._logger = Logger(master_config.logger) # type: ignore self._timer = Timer() + # Also built here, not on the driver: TimeoutChecker must capture + # wall-clock start times inside the actor, not at driver setup time. + # actor_args only carries the driver-side restore products + # (save_state, last_checkpoint_path). + self._checkpointer = CheckpointManager(master_config.checkpointing) + self._timeout = TimeoutChecker( + timeout=master_config.checkpointing["checkpoint_must_save_by"], + fit_last_save_time=True, + ) + self._timeout.start_iterations() + + # Loaded (or default) GRPOSaveState; keys SC does not own + # (val_reward, ...) pass through to saved checkpoints untouched. + self._save_state: GRPOSaveState = actor_args.save_state + self._last_checkpoint_path: Optional[str] = actor_args.last_checkpoint_path + self._data_plane_checkpoint_metadata: Optional[dict[str, Any]] = ( + actor_args.data_plane_checkpoint_metadata + ) + self._consumed_samples: int = actor_args.save_state["consumed_samples"] + self._total_valid_tokens: int = actor_args.save_state.get( + "total_valid_tokens", 0 + ) # Default to 0 for backward compatibility with older checkpoints + # Pin clusters so RayVirtualCluster.__del__ doesn't remove the PGs. self._train_cluster = actor_args.train_cluster self._inference_cluster = actor_args.inference_cluster @@ -134,7 +171,20 @@ def __init__( self._sampler = create_sampler( self._buffer, self._async_cfg.sampler, + resume_from_step=actor_args.save_state.get( + "trainer_version", actor_args.save_state["current_step"] + ), ) + if ( + self._master_config.checkpointing["enabled"] + and self._sampler.supports_buffer_checkpoint + and not self._master_config.data_plane.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." + ) required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step) validate_sampler_buffer_capacity( self._async_cfg, @@ -143,6 +193,19 @@ def __init__( ) # ── asyncio state ────────────────────────────────────────────────── + # Commits and destructive clears use this lock with TQ snapshots. This + # makes the native snapshot match the controller's metadata-only replay + # index exactly. Generation may continue, but completed rollouts wait at + # commit; _buffer_capacity bounds reservations and eventually stalls + # dispatch instead of allowing unbounded TQ growth. + # A future staging/finalizer path must join the same barrier before + # native restore can be authoritative. + self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + if self._buffer is not None: + self._buffer.set_data_plane_checkpoint_barrier( + self._data_plane_checkpoint_barrier + ) + # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() self._rollout_permitted.set() @@ -165,9 +228,11 @@ def __init__( self._async_cfg.max_buffered_rollouts ) - self._trainer_version: int = 0 - self._train_steps: int = 0 - self._current_epoch: int = 0 + self._trainer_version: int = actor_args.save_state.get( + "trainer_version", actor_args.save_state["current_step"] + ) + self._train_steps: int = actor_args.save_state["current_step"] + self._current_epoch: int = actor_args.save_state["current_epoch"] self._step_log_dict: dict[str, list] = { "rewards": [], "masked_advantages": [], @@ -190,6 +255,8 @@ async def run(self) -> dict[str, Any]: # Synchronize weights before starting the pumps await self._sync_weights() + await self._maybe_restore_replay_buffer() + # Start the rollout and train pumps rollout_task = asyncio.create_task(self._rollout_pump()) train_task = asyncio.create_task(self._train_pump()) @@ -206,7 +273,21 @@ async def run(self) -> dict[str, Any]: rollout_task.cancel() train_task.cancel() await asyncio.gather(rollout_task, train_task, return_exceptions=True) - self._logger.finish() + # Flush the last checkpoint's background finalization; a failure + # raises on a clean exit but never masks a propagating exception. + propagating = sys.exc_info()[0] is not None + try: + await asyncio.to_thread(self._checkpointer.shutdown) + except Exception: + if not propagating: + raise + warnings.warn( + "Checkpoint finalization failed while handling an " + "exception; the original exception will be re-raised.", + stacklevel=2, + ) + finally: + self._logger.finish() return { "train_steps": self._train_steps, @@ -226,20 +307,195 @@ async def ping(self) -> dict[str, Any]: # ── internal helpers ─────────────────────────────────────────────────── - async def _ray_get(self, obj_ref: Any) -> Any: - """Await a Ray ObjectRef without blocking the asyncio event loop.""" - return await obj_ref - - async def _call_dp(self, method_name: str, **kwargs) -> Any: - """Call a DataPlaneClient method or a Ray actor exposing that method.""" - method = getattr(self._dp_client, method_name) - remote = getattr(method, "remote", None) - if remote is not None: - return await self._ray_get(remote(**kwargs)) - result = method(**kwargs) - if asyncio.iscoroutine(result): - return await result - return result + async def _maybe_restore_replay_buffer(self) -> None: + """Restore replay-buffer groups from the previous run's checkpoint. + + No-op unless the sampler supports_buffer_checkpoint (ungated only). + """ + if self._last_checkpoint_path is None: + return + metadata_path = os.path.join( + self._last_checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME + ) + if ( + os.path.exists(metadata_path) + and not self._sampler.supports_buffer_checkpoint + ): + raise RuntimeError( + "The checkpoint contains native replay state, but the configured " + f"sampler {self._async_cfg.sampler.name!r} does not support " + "replay-buffer recovery" + ) + if not self._sampler.supports_buffer_checkpoint: + return + if not os.path.exists(metadata_path): + legacy_path = os.path.join( + self._last_checkpoint_path, LEGACY_REPLAY_BUFFER_FILENAME + ) + if os.path.exists(legacy_path): + raise RuntimeError( + "Checkpoint contains legacy replay_buffer.pt state, which " + "predates authoritative native TQ replay recovery. Resume it " + "with the older implementation or explicitly start without " + "restoring buffered rollouts." + ) + print( + f"⚠️ No native replay metadata found at {metadata_path}. " + "Starting with an empty replay buffer.", + flush=True, + ) + return + print(f"📦 Restoring replay buffer metadata: {metadata_path}") + # weights_only=False: the metadata sidecar contains pickled KVBatchMeta + # objects but no rollout tensor payloads. It is a trusted same-job artifact. + buffer_state = await asyncio.to_thread( + torch.load, metadata_path, weights_only=False + ) + if self._data_plane_checkpoint_metadata is None: + raise RuntimeError( + "Found metadata-only replay checkpoint, but the matching " + "native TQ checkpoint was not restored during setup" + ) + expected_manifest_digest_value = self._data_plane_checkpoint_metadata.get( + "replay_manifest_digest" + ) + if not isinstance(expected_manifest_digest_value, str): + raise ValueError( + "Restored TQ checkpoint metadata is missing a replay manifest digest" + ) + expected_group_count = self._data_plane_checkpoint_metadata.get( + "replay_group_count" + ) + groups = buffer_state.get("groups") + if ( + not isinstance(expected_group_count, int) + or not isinstance(groups, list) + or len(groups) != expected_group_count + ): + raise ValueError( + "Replay-buffer metadata group count does not match the " + "loaded TQ checkpoint metadata" + ) + restored = await self._buffer.load_state_dict( + buffer_state, + max_groups=self._async_cfg.max_buffered_rollouts, + expected_partition_id=self._partition_id, + expected_group_size=self._master_config.grpo["num_generations_per_prompt"], + expected_manifest_digest=expected_manifest_digest_value, + ) + await self._validate_replay_inventory(buffer_state) + + # Each buffered group holds one _buffer_capacity permit. Restore fails + # above if the saved group count exceeds current capacity. + assert restored <= self._async_cfg.max_buffered_rollouts + for _ in range(restored): + await self._buffer_capacity.acquire() + + async def _validate_replay_inventory( + self, replay_metadata: TQReplayMetadataState + ) -> None: + """Require the canonical TQ keys to match the SC replay index exactly.""" + expected_sample_ids = { + sample_id + for group in replay_metadata["groups"] + for sample_id in group["meta"].sample_ids + } + actual_sample_ids = set( + await call_data_plane( + self._dp_client, + "list_sample_ids", + offload_sync=True, + partition_id=self._partition_id, + ) + ) + missing_sample_ids = sorted(expected_sample_ids - actual_sample_ids) + unexpected_sample_ids = sorted(actual_sample_ids - expected_sample_ids) + if missing_sample_ids or unexpected_sample_ids: + raise RuntimeError( + "Native TQ checkpoint inventory does not match the replay " + "metadata sidecar: " + f"missing={missing_sample_ids[:10]!r} " + f"(total={len(missing_sample_ids)}), " + f"unexpected={unexpected_sample_ids[:10]!r} " + f"(total={len(unexpected_sample_ids)})" + ) + print( + "📦 Native TQ replay inventory validated: " + f"samples={len(actual_sample_ids)}", + flush=True, + ) + + async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: + """Clear consumed rows without overlapping a data-plane checkpoint.""" + async with self._data_plane_checkpoint_barrier.mutation(): + await call_data_plane( + self._dp_client, + "clear_samples", + offload_sync=True, + sample_ids=sample_ids, + partition_id=self._partition_id, + ) + + async def _save_data_plane_checkpoint( + self, + checkpoint_path: str, + replay_metadata: Optional[TQReplayMetadataState] = None, + ) -> None: + """Save a required TQ snapshot inside an SC checkpoint bundle. + + A sampler with replay-buffer recovery writes an authoritative native + TQ snapshot bound to its metadata-only sidecar by a digest. Other + samplers retain shadow-mode snapshots until their recovery contract is + defined. Failures propagate so a finalized bundle never silently omits + the advertised data-plane component. + """ + checkpoint_dir = os.path.join( + checkpoint_path, + DATA_PLANE_CHECKPOINT_DIR, + ) + metadata = { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": self._train_steps, + "single_controller_trainer_version": self._trainer_version, + "single_controller_epoch": self._current_epoch, + "partition_id": self._partition_id, + "sampler_name": self._async_cfg.sampler.name, + "mode": "authoritative" if replay_metadata is not None else "shadow", + } + if replay_metadata is not None: + metadata.update( + { + "replay_metadata_schema_version": ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION + ), + "replay_manifest_digest": replay_metadata["manifest_digest"], + "replay_group_count": len(replay_metadata["groups"]), + } + ) + started = time.monotonic() + print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) + try: + await call_data_plane( + self._dp_client, + "save_checkpoint", + offload_sync=True, + checkpoint_dir=checkpoint_dir, + metadata=metadata, + ) + except Exception as error: + print( + "data-plane checkpoint save failed: " + f"{checkpoint_dir} ({type(error).__name__}: {error})", + flush=True, + ) + raise + print( + "data-plane checkpoint save completed: " + f"{checkpoint_dir} ({time.monotonic() - started:.2f}s)", + flush=True, + ) # ── the three pumps + the inline advantage stage ─────────────────────── @@ -497,11 +753,7 @@ async def _train_pump(self) -> None: min_sample_version = curr_min_sample_version # Remove consumed sample_ids from the buffer - await self._call_dp( - "clear_samples", - sample_ids=list(train_meta.sample_ids), - partition_id=self._partition_id, - ) + await self._clear_data_plane_samples(list(train_meta.sample_ids)) groups_dispatched += num_groups @@ -524,6 +776,28 @@ async def _train_pump(self) -> None: ) await self._sync_weights(calibration_data=calibration_data) + # Checkpointing (mirrors async_grpo_train's save block). + self._consumed_samples += grpo_cfg["num_prompts_per_step"] + self._total_valid_tokens += step_metrics.get("global_valid_toks", 0) + self._timeout.mark_iteration() + + is_last_step = self._train_steps >= grpo_cfg["max_num_steps"] + # _train_steps was already incremented above, so it equals + # the legacy loop's 1-indexed `step + 1`. + should_save_by_step = ( + is_last_step + or self._train_steps + % self._master_config.checkpointing["save_period"] + == 0 + ) + should_save_by_timeout = self._timeout.check_save() + + if self._master_config.checkpointing["enabled"] and ( + should_save_by_step or should_save_by_timeout + ): + with self._timer.time("checkpointing"): + await self._save_checkpoint(step_metrics) + timing_metrics: dict[str, float] = self._timer.get_timing_metrics( reduction_op="sum" ) # type: ignore @@ -549,8 +823,6 @@ async def _train_pump(self) -> None: percent = (v / total_time * 100) if total_time > 0 else 0.0 print(f" • {k}: {v:.2f}s ({percent:.1f}%)") - # TODO: checkpointing (save_period/top-k metric_name, - # policy.save_checkpoint, dataloader state, TQReplayBuffer state). # TODO: per-step train_data jsonl dump, vllm metrics logger, # histogram log, rollout_metrics, seq_logprob_error_metrics, # pretty-print "Training Results" block, print_performance_metrics. @@ -573,6 +845,120 @@ async def _train_pump(self) -> None: flush=True, ) + if should_save_by_timeout: + print("Timeout has been reached, stopping training early", flush=True) + break + + async def _save_checkpoint( + self, + step_metrics: dict[str, Any], + val_metrics: Optional[dict[str, Any]] = None, + ) -> None: + """Write a full checkpoint for the just-finished train step. + + Everything except the (possibly async) policy weight write must be + on disk before begin_finalization; rollouts keep running throughout. + val_metrics stays None until SC grows a validation loop. + """ + save_state = self._save_state + save_state["current_step"] = self._train_steps + save_state["total_steps"] = self._train_steps + save_state["trainer_version"] = self._trainer_version + save_state["current_epoch"] = self._current_epoch + save_state["consumed_samples"] = self._consumed_samples + save_state["total_valid_tokens"] = self._total_valid_tokens + # Snapshot before any await so it can't interleave with + # _rollout_pump iterating this same dataloader. + dataloader_state = self._dataloader.state_dict() + if val_metrics is not None: + save_state["val_reward"] = val_metrics["accuracy"] + elif "val_reward" in save_state: + del save_state["val_reward"] + + full_metric_name = self._master_config.checkpointing["metric_name"] + if full_metric_name is not None: + assert full_metric_name.startswith("train:") or full_metric_name.startswith( + "val:" + ), ( + f"metric_name={full_metric_name} must start with 'val:' or 'train:',\n" + f'followed by the corresponding name in the "val" or "train" metrics dictionary.' + f" If you are using an old config, please updated checkpointing.metric_name to the new format, " + f" e.g. 'val_reward --> 'val:accuracy'" + ) + prefix, metric_name = full_metric_name.split(":", 1) + metrics_source = step_metrics if prefix == "train" else val_metrics + if not metrics_source: + warnings.warn( + f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. " + "This checkpoint will not be saved as top-k.", + stacklevel=2, + ) + if full_metric_name in save_state: + del save_state[full_metric_name] + elif metric_name not in metrics_source: + raise ValueError(f"Metric {metric_name} not found in {prefix} metrics") + else: + save_state[full_metric_name] = metrics_source[metric_name] + + # Flush the previous checkpoint's background finalization first; + # re-raises a failure from it. + await asyncio.to_thread(self._checkpointer.finalize_pending) + + print(f"Saving checkpoint for step {self._train_steps}...") + checkpoint_path = await asyncio.to_thread( + self._checkpointer.init_tmp_checkpoint, + self._train_steps, + 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"), + ) + replay_metadata: Optional[TQReplayMetadataState] = None + if self._master_config.data_plane.get("checkpointing_enabled"): + # Commits and destructive clears take the same barrier. Generation + # may continue while a snapshot is written, but completed groups + # 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(): + if self._sampler.supports_buffer_checkpoint: + replay_metadata = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + await self._save_data_plane_checkpoint( + checkpoint_path, replay_metadata=replay_metadata + ) + if replay_metadata is not None: + await self._validate_replay_inventory(replay_metadata) + if replay_metadata is not None: + await asyncio.to_thread( + torch.save, + replay_metadata, + os.path.join(checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME), + ) + # Rename happens in the background once the async weight writes + # finish; flushed at the next save or on exit. + self._checkpointer.begin_finalization( + checkpoint_path, + wait_fn=self._trainer.finalize_async_save, + ) + _write_latest_checkpoint_status( + self._checkpointer, last_checkpoint_step=self._train_steps + ) + async def _sync_weights( self, *, @@ -633,7 +1019,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: return meta adv_cfg = self._advantage_cfg - data = await self._call_dp( + data = await call_data_plane( + self._dp_client, "get_samples", sample_ids=meta.sample_ids, partition_id=meta.partition_id, @@ -683,7 +1070,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: response_advantages.detach().cpu() ) - await self._call_dp( + await call_data_plane( + self._dp_client, "put_samples", sample_ids=meta.sample_ids, partition_id=meta.partition_id, diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 6ad8d107a3..e9de1ba7d4 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -21,18 +21,29 @@ from __future__ import annotations +import os from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass +from pathlib import Path from typing import Any, Optional, cast from torchdata.stateful_dataloader import StatefulDataLoader from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + TQReplayBuffer, +) from nemo_rl.algorithms.grpo import MasterConfig as GrpoMasterConfig from nemo_rl.algorithms.grpo import ( + GRPOSaveState, _create_advantage_estimator, + _default_grpo_save_state, _should_use_nemo_gym, ) from nemo_rl.algorithms.loss import ClippedPGLossFn @@ -43,7 +54,7 @@ ) from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn -from nemo_rl.data.utils import setup_response_data +from nemo_rl.data.utils import load_dataloader_state, setup_response_data from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.environments.interfaces import EnvironmentInterface @@ -61,6 +72,7 @@ router_replay_enabled, ) from nemo_rl.models.policy.tq_policy import TQPolicy +from nemo_rl.utils.checkpoint import CheckpointManager from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer @@ -85,6 +97,97 @@ class SingleControllerActorArgs: rollout_manager: RolloutManager tq_buffer: TQReplayBuffer partition_id: str + save_state: GRPOSaveState + last_checkpoint_path: Optional[str] + data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None + + +def _maybe_restore_native_data_plane_checkpoint( + policy: TQPolicy, + *, + last_checkpoint_path: Optional[str], + save_state: GRPOSaveState, + partition_id: str, + sampler_name: str, +) -> Optional[dict[str, Any]]: + """Load and validate an authoritative native TQ checkpoint when present. + + The metadata-only replay sidecar is the format marker. Checkpoints without + any replay artifact resume trainer state with an empty replay buffer; + legacy tensor-bearing replay files are rejected rather than silently + ignored. Rollout tensors are never serialized into a controller-side + replay checkpoint. + """ + if last_checkpoint_path is None: + return None + checkpoint_path = Path(last_checkpoint_path) + replay_metadata_path = checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME + if not replay_metadata_path.is_file(): + legacy_replay_path = checkpoint_path / LEGACY_REPLAY_BUFFER_FILENAME + if legacy_replay_path.is_file(): + raise RuntimeError( + "Checkpoint contains legacy replay_buffer.pt state, which " + "predates authoritative native TQ replay recovery. Resume it " + "with the older implementation or explicitly start without " + "restoring buffered rollouts." + ) + return None + + data_plane_path = checkpoint_path / DATA_PLANE_CHECKPOINT_DIR + if not data_plane_path.is_dir(): + raise FileNotFoundError( + "Metadata-only replay checkpoint requires a matching native TQ " + f"checkpoint at {data_plane_path}" + ) + + print(f"📦 Restoring native TQ checkpoint: {data_plane_path}", flush=True) + metadata = policy.load_data_plane_checkpoint(data_plane_path) + if not isinstance(metadata, dict): + raise TypeError( + "Native TQ checkpoint load must return a metadata dictionary, " + f"got {type(metadata).__name__}" + ) + expected_values: dict[str, Any] = { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": save_state["current_step"], + "single_controller_trainer_version": save_state.get( + "trainer_version", save_state["current_step"] + ), + "single_controller_epoch": save_state["current_epoch"], + "partition_id": partition_id, + "sampler_name": sampler_name, + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + } + mismatches = { + key: {"checkpoint": metadata.get(key), "expected": expected} + for key, expected in expected_values.items() + if metadata.get(key) != expected + } + if mismatches: + raise ValueError( + "Native TQ checkpoint metadata does not match the trainer " + f"checkpoint: {mismatches}" + ) + manifest_digest = metadata.get("replay_manifest_digest") + if not isinstance(manifest_digest, str) or not manifest_digest: + raise ValueError( + "Native TQ checkpoint metadata is missing replay_manifest_digest" + ) + group_count = metadata.get("replay_group_count") + if not isinstance(group_count, int) or group_count < 0: + raise ValueError( + "Native TQ checkpoint metadata has invalid replay_group_count: " + f"{group_count!r}" + ) + print( + "📦 Native TQ checkpoint restored and validated: " + f"groups={group_count}", + flush=True, + ) + return metadata def _build_clusters( @@ -197,6 +300,9 @@ def _build_trainer( master_config: MasterConfig, tokenizer, processor, + *, + weights_path: Optional[Path], + optimizer_path: Optional[Path], ): """Build the TQ-mediated trainer (driver-side TQPolicy). @@ -212,8 +318,8 @@ def _build_trainer( config=master_config.policy, tokenizer=tokenizer, processor=processor, - weights_path=None, - optimizer_path=None, + weights_path=weights_path, + optimizer_path=optimizer_path, init_optimizer=True, init_reference_model=init_reference_model, dp_cfg=master_config.data_plane, @@ -296,18 +402,28 @@ def setup_single_controller( "SingleController doesn't support validation now, will support " "later. Set grpo.val_period=0, val_at_start=false, val_at_end=false." ) - if master_config.checkpointing["enabled"]: - raise NotImplementedError( - "SingleController doesn't support checkpointing now, will support " - "later. Set checkpointing.enabled=false." - ) - if dp_config is None or not dp_config.get("enabled", False): raise ValueError( "single_controller_utils.setup requires " "master_config.data_plane.enabled=True. The async-RL " "SingleController path is built on the TransferQueue data plane." ) + if dp_config.get("checkpointing_enabled") and dp_config["backend"] != "simple": + raise NotImplementedError( + "SingleController data-plane checkpointing currently requires " + "data_plane.backend='simple'; Mooncake storage cannot be restored " + "by TQ v0.1.9." + ) + if ( + master_config.checkpointing["enabled"] + and master_config.async_rl.sampler.supports_buffer_checkpoint is True + 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." + ) assert generation_config is not None, ( "single_controller_utils.setup requires policy.generation in master_config" @@ -319,8 +435,24 @@ def setup_single_controller( "data.use_multiple_dataloader=True yet." ) + checkpointing_pretrained = master_config.checkpointing.get("pretrained_checkpoint") + if checkpointing_pretrained is not None: + policy_config["pretrained_checkpoint"] = checkpointing_pretrained + set_seed(grpo_config["seed"]) + # ========================== + # Checkpointing + # ========================== + checkpointer = CheckpointManager(master_config.checkpointing) + last_checkpoint_path = checkpointer.get_latest_checkpoint_path() + save_state = cast( + Optional[GRPOSaveState], checkpointer.load_training_info(last_checkpoint_path) + ) + if save_state is None: + save_state = _default_grpo_save_state() + weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path) + # ========================== # Setup Dataset & Environments # ========================== @@ -352,6 +484,20 @@ def setup_single_controller( drop_last=True, num_workers=data_config["num_workers"], ) + if last_checkpoint_path is not None: + dataloader_state_path = os.path.join( + last_checkpoint_path, "train_dataloader.pt" + ) + if os.path.exists(dataloader_state_path): + print( + f"📦 Restoring dataloader state from checkpoint: {dataloader_state_path}" + ) + load_dataloader_state(dataloader, last_checkpoint_path, data_config) + else: + print( + f"⚠️ No dataloader state found at {dataloader_state_path}. " + "Starting with a fresh dataloader position." + ) _clamp_max_num_steps(master_config, dataloader) _maybe_inject_megatron_train_iters(master_config) @@ -365,7 +511,14 @@ def setup_single_controller( # Colocated: vLLM prefers a clean GPU at load time, so generation # comes up before the policy. generation = _build_generation(inference_cluster, master_config) - policy = _build_trainer(train_cluster, master_config, tokenizer, processor) + policy = _build_trainer( + train_cluster, + master_config, + tokenizer, + processor, + weights_path=weights_path, + optimizer_path=optimizer_path, + ) else: # Non-colocated: generation + policy run on disjoint GPUs, so # bring them up in parallel. @@ -374,11 +527,27 @@ def setup_single_controller( _build_generation, inference_cluster, master_config ) policy_future = executor.submit( - _build_trainer, train_cluster, master_config, tokenizer, processor + _build_trainer, + train_cluster, + master_config, + tokenizer, + processor, + weights_path=weights_path, + optimizer_path=optimizer_path, ) generation = gen_future.result() policy = policy_future.result() + # Native TQ restore must run through the bootstrap client before the SC + # client is created and before any rollout/train data-plane operation. + data_plane_checkpoint_metadata = _maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=last_checkpoint_path, + save_state=save_state, + partition_id=partition_id, + sampler_name=master_config.async_rl.sampler.name, + ) + # ========================== # NeMo-Gym actor (after generation is up so OpenAI URLs are available) # ========================== @@ -458,4 +627,7 @@ def setup_single_controller( rollout_manager=rollout_manager, tq_buffer=tq_buffer, partition_id=partition_id, + save_state=save_state, + last_checkpoint_path=last_checkpoint_path, + data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, ) diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py index 1c5b00a5e4..34d2f76c7d 100644 --- a/nemo_rl/data_plane/adapters/noop.py +++ b/nemo_rl/data_plane/adapters/noop.py @@ -25,7 +25,10 @@ from __future__ import annotations +import pickle +import shutil from dataclasses import dataclass, field +from pathlib import Path from typing import Any import torch @@ -220,6 +223,11 @@ def get_samples( stacked = {f: _stack_or_nest(out[f]) for f in select_fields} return TensorDict(stacked, batch_size=(len(sample_ids),)) + def list_sample_ids(self, partition_id: str) -> list[str]: + """List stored sample IDs without reading their tensor payloads.""" + rec = self._partitions.get(partition_id) + return sorted(rec.rows) if rec is not None else [] + def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: rec = self._partitions.get(partition_id) if rec is None: @@ -237,6 +245,59 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None for s in rec.consumed.values(): s.discard(sid) + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """Persist the trusted in-memory fixture for adapter contract tests. + + This test-only adapter uses pickle; callers must not load checkpoints + from untrusted paths. + """ + checkpoint_dir = Path(checkpoint_dir) + tmp_dir = checkpoint_dir.with_name(f"{checkpoint_dir.name}.tmp") + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + tmp_dir.mkdir(parents=True) + try: + with (tmp_dir / "noop_state.pkl").open("wb") as checkpoint_file: + pickle.dump( + { + "partitions": self._partitions, + "metadata": metadata or {}, + }, + checkpoint_file, + protocol=pickle.HIGHEST_PROTOCOL, + ) + if checkpoint_dir.exists(): + shutil.rmtree(checkpoint_dir) + tmp_dir.rename(checkpoint_dir) + except Exception: + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + raise + + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + """Restore the in-memory fixture into a clean client.""" + if self._partitions: + raise RuntimeError( + "load_checkpoint requires a clean data-plane client with no " + "registered partitions" + ) + checkpoint_file = Path(checkpoint_dir) / "noop_state.pkl" + if not checkpoint_file.is_file(): + raise FileNotFoundError(f"NoOp checkpoint not found: {checkpoint_file}") + with checkpoint_file.open("rb") as state_file: + state = pickle.load(state_file) + metadata = state.get("metadata", {}) + if not isinstance(metadata, dict): + raise ValueError("NoOp checkpoint metadata must be a dictionary") + self._partitions = state["partitions"] + self._closed = False + return dict(metadata) + def close(self) -> None: if self._closed: return diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index ce293b4e26..d2d8038919 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -23,12 +23,15 @@ from __future__ import annotations import ipaddress +import json import os import socket import subprocess import time +import warnings from importlib import resources -from typing import Any +from pathlib import Path +from typing import Any, cast import torch import transfer_queue as tq @@ -171,27 +174,28 @@ def patched(*args, **kwargs): cls.options = patched # type: ignore[method-assign] return True - patched_any = False + unpatched_classes: list[str] = [] try: - from transfer_queue.storage.simple_backend import SimpleStorageUnit + from transfer_queue.storage.simple_storage import SimpleStorageUnit - patched_any |= _install(SimpleStorageUnit) + if not _install(SimpleStorageUnit): + unpatched_classes.append("SimpleStorageUnit") except ImportError: - pass + unpatched_classes.append("SimpleStorageUnit") try: from transfer_queue.controller import TransferQueueController - patched_any |= _install(TransferQueueController) + if not _install(TransferQueueController): + unpatched_classes.append("TransferQueueController") except ImportError: - pass + unpatched_classes.append("TransferQueueController") - if not patched_any: + if unpatched_classes: # Soft-fail: TQ may have moved its actor classes. The driver will # still work; multi-node TQ may need the per-node `uv sync` workaround. - import warnings - warnings.warn( - "Could not patch TQ actor classes for runtime_env injection. " + "Could not patch every TQ actor class for runtime_env injection: " + f"unpatched={unpatched_classes}. " "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' " "on worker nodes. Workaround: run `uv sync` inside each node's " "container before the driver runs.", @@ -306,6 +310,77 @@ def _init_tq(cfg: DataPlaneConfig) -> None: # ────────────────────────────────────────────────────────────────────────── +_PROMOTED_1D_TAG_PREFIX = "__nemo_rl_promoted_1d__:" +_WIRE_SHAPE_VERSION_TAG = "__nemo_rl_wire_shape_version__" +_WIRE_SHAPE_VERSION = 1 + + +def _add_wire_shape_tags( + tags: list[dict[str, Any]], td: TensorDict +) -> list[dict[str, Any]]: + """Record which tensor fields the Mooncake workaround will promote. + + TQ metadata is durable and visible to readers in other processes, unlike + client-local state. The version marker distinguishes new unpromoted fields + from legacy rows; promoted fields get an additional boolean marker. + """ + field_markers: dict[str, int | bool] = { + _WIRE_SHAPE_VERSION_TAG: _WIRE_SHAPE_VERSION + } + for k in td.keys(): + value = td.get(k) + if isinstance(value, torch.Tensor) and not value.is_nested and value.dim() == 1: + field_markers[f"{_PROMOTED_1D_TAG_PREFIX}{k}"] = True + return [{**tag, **field_markers} for tag in tags] + + +def _promoted_1d_fields_from_tags( + tags: list[dict[str, Any]], field_names: list[str] +) -> set[str] | None: + """Recover promoted fields, or ``None`` for legacy unmarked rows. + + Every selected tensor field must have the same marker on every row. Mixed + metadata means rows were written by incompatible codecs; fail instead of + silently returning tensors with the wrong rank. + """ + versions = [tag.get(_WIRE_SHAPE_VERSION_TAG) for tag in tags] + if not any(version is not None for version in versions): + return None + if not versions or any(version != _WIRE_SHAPE_VERSION for version in versions): + raise RuntimeError( + "Inconsistent Mooncake shape metadata: expected wire-shape " + f"version {_WIRE_SHAPE_VERSION} on every sample." + ) + + promoted: set[str] = set() + for field in field_names: + marker = f"{_PROMOTED_1D_TAG_PREFIX}{field}" + values = [tag.get(marker) for tag in tags] + if any(value is not None for value in values) and any( + value is not True for value in values + ): + raise RuntimeError( + "Inconsistent Mooncake shape metadata for " + f"field {field!r}: expected the promotion marker on every sample." + ) + if values and values[0] is True: + promoted.add(field) + return promoted + + +def _strip_wire_shape_tags(tags: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Remove adapter-private shape metadata from user-visible tags.""" + return [ + { + key: value + for key, value in tag.items() + if key != _WIRE_SHAPE_VERSION_TAG + and not key.startswith(_PROMOTED_1D_TAG_PREFIX) + } + for tag in tags + ] + + def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None: """Guard against silent leaf drops through TensorDict constructor rebuild. @@ -355,19 +430,36 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict: return new_td -def _from_wire(td: TensorDict) -> TensorDict: - """Inverse of `_promote_1d_leaves`: squeeze trailing 1 back to (N,).""" +def _from_wire( + td: TensorDict, promoted_1d_fields: set[str] | None = None +) -> TensorDict: + """Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed. + + Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a + nested tensor, including fields whose rows all have the same shape. + Densify those uniform nested tensors first so regular batched inputs retain + their dense representation. Truly ragged fields remain nested. Finally, + squeeze only singleton dimensions known to have been introduced by + :func:`_promote_1d_leaves`. ``None`` retains the old shape-only behavior + for checkpoints written before shape tags existed. + """ # Same top-level iteration as `_promote_1d_leaves`: NonTensorData / # NonTensorStack leaves are only visible via td.keys(), not leaves_only. new_dict: dict[str, Any] = {} changed = False for k in td.keys(): v = td.get(k) + if isinstance(v, torch.Tensor) and v.is_nested: + rows = list(v.unbind()) + if rows and all(row.shape == rows[0].shape for row in rows[1:]): + v = torch.stack(rows) + changed = True if ( isinstance(v, torch.Tensor) and not v.is_nested and v.dim() >= 2 and v.shape[-1] == 1 + and (promoted_1d_fields is None or str(k) in promoted_1d_fields) ): new_dict[str(k)] = v.squeeze(-1).contiguous() changed = True @@ -426,6 +518,7 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: # is unaffected). Writer unsqueezes 1D → (N, 1) on put; reader # squeezes the trailing 1 back on get. Drop when upstream TQ # unifies the schema/data shapes for 1D fields. + self._backend = cfg["backend"] self._promote_1d = cfg["backend"] == "mooncake_cpu" if bootstrap: @@ -434,6 +527,22 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: _connect_existing() self._poll_interval_s = cfg["claim_meta_poll_interval_s"] self._closed = False + # TQ restore is non-transactional and requires a globally clean system. + # This process-local guard catches incorrect ordering through this + # adapter; setup must still ensure no other client has touched TQ. + self._data_operations_started = False + + def _mark_data_operation_started(self) -> None: + """Make a later checkpoint load fail instead of mixing TQ states.""" + self._data_operations_started = True + + def _require_clean_for_load(self) -> None: + """Reject restore after this client has performed a data operation.""" + if self._data_operations_started: + raise RuntimeError( + "load_checkpoint requires a clean TQ client before any " + "register, claim, get, list, put, clear, or consumption operation" + ) # ── (A) task-mediated ─────────────────────────────────────────────── @@ -466,6 +575,7 @@ def register_partition( # stale metadata from a previous registration. if not fields: return + self._mark_data_operation_started() schema_key = ( f"__schema__:{partition_id}:{os.getpid()}:{id(self)}:{time.time_ns()}" ) @@ -491,6 +601,7 @@ def claim_meta( blocking: bool = True, timeout_s: float = 60.0, ) -> KVBatchMeta: + self._mark_data_operation_started() client = tq.get_client() deadline = time.time() + max(0.0, timeout_s) sampling_config: dict[str, Any] = {} @@ -532,7 +643,10 @@ def claim_meta( # because shard_meta_for_dp reads it directly), but the rest # of the tag dict travels through unchanged so consumers can # filter on it without fetching data. - tags = list(tq_meta.custom_meta) if tq_meta.custom_meta else [{} for _ in keys] + wire_tags = ( + list(tq_meta.custom_meta) if tq_meta.custom_meta else [{} for _ in keys] + ) + tags = _strip_wire_shape_tags(wire_tags) seqlens: list[int] | None = None if tags and any("input_lengths" in t for t in tags): seqlens = [int(t.get("input_lengths", 0)) for t in tags] @@ -562,6 +676,7 @@ def get_data( def check_consumption_status( self, partition_id: str, task_names: list[str] ) -> bool: + self._mark_data_operation_started() client = tq.get_client() for t in task_names: if not client.check_consumption_status( @@ -583,8 +698,10 @@ def put_samples( return KVBatchMeta( partition_id=partition_id, task_name=None, sample_ids=[], fields=None ) - if tags is None: - tags = [{} for _ in sample_ids] + user_tags = ( + [{} for _ in sample_ids] if tags is None else [dict(tag) for tag in tags] + ) + wire_tags = user_tags wire_fields: TensorDict | None = None field_names: list[str] | None = None @@ -594,17 +711,23 @@ def put_samples( # TDs. TQ's encoder forces ``.contiguous()`` per tensor leaf # itself, so the call here was redundant for tensors and # destructive for non-tensors. - wire_fields = fields.detach() # type: ignore[bad-assignment,missing-argument] + detached_fields = cast( + TensorDict, + fields.detach(), # type: ignore[missing-argument] + ) if self._promote_1d: - wire_fields = _promote_1d_leaves(wire_fields) # type: ignore[bad-argument-type] - field_names = list(wire_fields.keys()) + wire_tags = _add_wire_shape_tags(wire_tags, detached_fields) + detached_fields = _promote_1d_leaves(detached_fields) + wire_fields = detached_fields + field_names = [str(key) for key in detached_fields.keys()] + self._mark_data_operation_started() # TQ's wire vocabulary is `keys=` — translation point. tq.kv_batch_put( keys=list(sample_ids), partition_id=partition_id, fields=wire_fields, - tags=tags, + tags=wire_tags, ) return KVBatchMeta( @@ -612,7 +735,7 @@ def put_samples( task_name=None, sample_ids=list(sample_ids), fields=field_names, - tags=[dict(t) for t in tags] if tags else None, + tags=user_tags if user_tags else None, ) def get_samples( @@ -623,19 +746,44 @@ def get_samples( ) -> TensorDict: if not sample_ids: return TensorDict({}, batch_size=(0,)) - # TQ's wire vocabulary is `keys=` — translation point. - td = tq.kv_batch_get( - keys=list(sample_ids), - partition_id=partition_id, - select_fields=select_fields, + self._mark_data_operation_started() + if not self._promote_1d: + td = tq.kv_batch_get( + keys=list(sample_ids), + partition_id=partition_id, + select_fields=select_fields, + ) + return _from_wire(td, promoted_1d_fields=set()) + + # Inline TQ's public kv_batch_get flow so the adapter can also inspect + # the durable per-row shape markers carried by BatchMeta.custom_meta. + client = tq.get_client() + tq_meta = client.kv_retrieve_meta( + keys=list(sample_ids), partition_id=partition_id, create=False ) - if self._promote_1d: - td = _from_wire(td) - return td + if tq_meta.size == 0: + raise ValueError("keys or partition were not found!") + tq_meta = tq_meta.select_fields(list(select_fields)) + if not tq_meta.is_ready: + raise ValueError("Some fields are not ready in all the requested keys!") + td = client.get_data(tq_meta) + tensor_fields = [ + str(k) for k in td.keys() if isinstance(td.get(k), torch.Tensor) + ] + wire_tags = list(tq_meta.custom_meta) if tq_meta.custom_meta else [] + promoted_1d_fields = _promoted_1d_fields_from_tags(wire_tags, tensor_fields) + return _from_wire(td, promoted_1d_fields) + + def list_sample_ids(self, partition_id: str) -> list[str]: + """List TQ keys in ``partition_id`` without fetching tensor payloads.""" + self._mark_data_operation_started() + listing = tq.kv_list(partition_id=partition_id) + return sorted(listing.get(partition_id, {}).keys()) def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: cleared_via_none = sample_ids is None if sample_ids is None: + self._mark_data_operation_started() # No local state — ask TQ's controller for the current key # set in this partition. ``kv_list`` errors propagate; we # don't want a network blip to silently turn into "cleared @@ -657,11 +805,56 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None stacklevel=2, ) return + self._mark_data_operation_started() # TQ's wire vocabulary is `keys=` — translation point. tq.kv_clear(keys=list(sample_ids), partition_id=partition_id) # ── (C) lifecycle ────────────────────────────────────────────────── + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """Save TQ controller metadata and storage data.""" + if self._backend == "mooncake_cpu": + raise NotImplementedError( + "TQ checkpointing is not supported for the mooncake_cpu " + "backend: MooncakeStorageManager cannot persist its in-memory " + "rows, so TQ would silently create a metadata-only checkpoint." + ) + _connect_existing() + tq.save_checkpoint(checkpoint_dir, metadata=metadata) + + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + """Restore TQ state after initialization and before data operations. + + The local lifecycle guard cannot observe operations issued by another + TQ client, so the recovery coordinator must also guarantee globally + clean setup ordering. + """ + if self._backend == "mooncake_cpu": + raise NotImplementedError( + "TQ checkpoint restore is not supported for the mooncake_cpu " + "backend because its in-memory rows cannot be restored." + ) + self._require_clean_for_load() + # Validate the adapter-owned metadata before starting TQ's + # non-transactional storage/controller restore. + metadata_path = Path(checkpoint_dir) / "metadata.json" + with metadata_path.open() as metadata_file: + checkpoint_metadata = json.load(metadata_file) + user_metadata = checkpoint_metadata.get("user_metadata", {}) + if not isinstance(user_metadata, dict): + raise ValueError("TQ checkpoint user_metadata must be a dictionary") + _connect_existing() + # A failed TQ load may have partially modified distributed storage, so + # this client is no longer safe for a retry even when an error escapes. + self._mark_data_operation_started() + tq.load_checkpoint(checkpoint_dir) + return dict(user_metadata) + def close(self) -> None: if self._closed: return diff --git a/nemo_rl/data_plane/async_utils.py b/nemo_rl/data_plane/async_utils.py new file mode 100644 index 0000000000..aee3ef9b53 --- /dev/null +++ b/nemo_rl/data_plane/async_utils.py @@ -0,0 +1,56 @@ +# 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. + +"""Async dispatch helpers for local and Ray data-plane clients.""" + +from __future__ import annotations + +import asyncio +from typing import Any + + +async def call_data_plane( + client: Any, + method_name: str, + *, + offload_sync: bool = False, + **kwargs: Any, +) -> Any: + """Call a local data-plane client or a Ray actor exposing its methods. + + Synchronous offloading is opt-in because it allows the actor event loop to + issue other calls while this one is running. Callers should enable it only + when that concurrency is supported or externally serialized. + + Args: + client: Local ``DataPlaneClient`` or Ray actor handle. + method_name: Data-plane method to invoke. + offload_sync: Run a synchronous local implementation in a worker + thread. Ray methods are already asynchronous and ignore this flag. + **kwargs: Keyword arguments forwarded to the data-plane method. + + Returns: + The method result after awaiting Ray or coroutine results. + """ + method = getattr(client, method_name) + remote = getattr(method, "remote", None) + if remote is not None: + return await remote(**kwargs) + if offload_sync: + result = await asyncio.to_thread(method, **kwargs) + else: + result = method(**kwargs) + if asyncio.iscoroutine(result): + return await result + return result diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 41a98f0c0e..465f42b706 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -37,6 +37,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict from tensordict import TensorDict @@ -58,6 +59,13 @@ class DataPlaneConfig(TypedDict): ``backend == "mooncake_cpu"``; the simple backend ignores them. They are required (not NotRequired) so the YAML carries the full schema and there are no hidden Python defaults. + + ``checkpointing_enabled`` opts SingleController into saving required TQ + state inside its checkpoint bundle. Samplers that support replay-buffer + recovery pair the native snapshot with a metadata-only local index; other + samplers save it in shadow mode. Other algorithm entrypoints do not consume + this field. It is optional because existing configs predate data-plane + checkpointing; exemplar configs carry the recommended default explicitly. """ enabled: bool @@ -68,6 +76,7 @@ class DataPlaneConfig(TypedDict): claim_meta_poll_interval_s: float global_segment_size: int local_buffer_size: int + checkpointing_enabled: NotRequired[bool] controller_address: NotRequired[str] ack_timeout_ms: NotRequired[int] observability: NotRequired["ObservabilityConfig"] @@ -259,7 +268,8 @@ class DataPlaneClient(ABC): B. *Direct-by-key* — used by stages that already know the exact uids (e.g. driver-side fan-out to DP ranks): :meth:`put_samples`, :meth:`get_samples`, :meth:`clear_samples`. - C. *Lifecycle* — :meth:`close`. + C. *Lifecycle* — :meth:`save_checkpoint`, :meth:`load_checkpoint`, and + :meth:`close`. Stage-completion signal: there is intentionally no ``mark_consumed``. The authoritative signal in TransferQueue is *field production* — @@ -411,6 +421,22 @@ def get_samples( ``TensorDict`` keyed by field name, batched along ``sample_ids``. """ + @abstractmethod + def list_sample_ids(self, partition_id: str) -> list[str]: + """List the sample IDs currently stored in a partition. + + This metadata-only operation is intended for recovery validation and + reconciliation. It must not fetch tensor payloads or advance consumer + cursors. + + Args: + partition_id: Partition whose stored keys should be listed. + + Returns: + Stable, sorted sample IDs. An unknown or empty partition returns + an empty list. + """ + @abstractmethod def clear_samples( self, @@ -442,6 +468,44 @@ def clear_samples( # ── (C) lifecycle ────────────────────────────────────────────────── + @abstractmethod + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """Persist the complete data-plane state to ``checkpoint_dir``. + + The checkpoint must include both data and the implementation's + scheduling/consumption metadata. Callers must serialize checkpoint + saves and prevent destructive operations such as clears until this + method returns. + + Args: + checkpoint_dir: New durable directory for this checkpoint. + metadata: Optional JSON-compatible recovery metadata. + """ + + @abstractmethod + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + """Restore a complete data-plane checkpoint. + + The data-plane implementation must already be initialized, but no data + operations may have run before restore. Implementations must reject a + load after operations through the same client; callers must also ensure + that no other client has modified shared data-plane state. + + Args: + checkpoint_dir: Directory previously written by + :meth:`save_checkpoint`. + + Returns: + User metadata supplied to :meth:`save_checkpoint`. The caller may + validate this metadata, but restoring data-plane state does not + restore the surrounding controller or trainer state. + """ + @abstractmethod def close(self) -> None: """Release controller / storage handles. Idempotent.""" diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 63e551dc20..0e66d4c1d9 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -29,6 +29,7 @@ import logging from dataclasses import asdict, dataclass +from pathlib import Path from time import monotonic from typing import Any, Callable, Literal, TypedDict @@ -322,6 +323,13 @@ def get_samples(self, sample_ids, partition_id, select_fields): n_keys=len(sample_ids), ) + def list_sample_ids(self, partition_id): + return self._run( + "list_sample_ids", + partition_id, + lambda: self._inner.list_sample_ids(partition_id), + ) + def clear_samples(self, sample_ids, partition_id): sample_ids_list = ( sample_ids @@ -337,6 +345,28 @@ def clear_samples(self, sample_ids, partition_id): ) self._record_clear(partition_id, sample_ids_list) + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + self._run( + "save_checkpoint", + "", + lambda: self._inner.save_checkpoint( + checkpoint_dir, + metadata=metadata, + ), + ) + + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + return self._run( + "load_checkpoint", + "", + lambda: self._inner.load_checkpoint(checkpoint_dir), + ) + def close(self) -> None: self._run( "close", diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 61da43511c..43ca603cdb 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -33,6 +33,7 @@ from collections import defaultdict from contextlib import nullcontext from dataclasses import replace +from pathlib import Path from typing import Any, Optional import ray @@ -134,6 +135,12 @@ def __init__( # ── lifecycle ────────────────────────────────────────────────────── + def load_data_plane_checkpoint( + self, checkpoint_dir: str | Path + ) -> dict[str, Any]: + """Restore TQ through the clean bootstrap client during SC setup.""" + return self.dp_client.load_checkpoint(checkpoint_dir) + def shutdown(self) -> bool: # type: ignore[override] """Close the TQ client before shutting down the worker group.""" try: diff --git a/pyproject.toml b/pyproject.toml index 5fd74de631..3b2da988c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,11 +84,10 @@ dependencies = [ # automatically include them. Removes the need for a `[data-plane]` # extra and the corresponding plumbing in the per-worker venv builder. "tensordict", - # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore - # refactor: `clear` switched from unanchored `remove_by_regex` to - # exact-key `batch_remove`, which fixes a collateral-key-deletion bug - # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released. - "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39", + # TransferQueue v0.1.9 adds full-system checkpoint save/load APIs and + # retains the exact-key MooncakeStore clear behavior required by DAPO. + # Pin the immutable release commit rather than the mutable version tag. + "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@c51614308b68c8d7a87c9b3ef62d59e14c69bde2", # Backs data_plane.backend="mooncake_cpu". Default backend is "simple" # (in-process), but the mooncake_cpu path needs the `mooncake_master` # binary that ships in this wheel at /mooncake/. Bundled @@ -379,10 +378,9 @@ override-dependencies = [ "pytest>=9.0.3", "langchain>=0.3.28", # Address CVE-2025-65106 "langchain-core>=0.3.80", # Address CVE-2025-65106 - # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs - # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane - # extra composes with mcore/automodel without version-mirroring TQ's - # requirements.txt. Forward-compatible across TQ minor bumps. + # Forces numpy past tensorrt-llm's `numpy>=2.0.0,<2.4` cap (resolves to + # 2.5.x). The original driver — TransferQueue pinning `numpy<2.0.0` — is + # gone as of TQ v0.1.9; drop this override once the trtllm cap lifts. "numpy>=2.1.0", # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame). # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it diff --git a/pyrefly.toml b/pyrefly.toml index 65929bb208..1b7fc28002 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -114,6 +114,7 @@ project-includes = [ "nemo_rl/data_plane/adapters/__init__.py", "nemo_rl/data_plane/adapters/noop.py", "nemo_rl/data_plane/adapters/transfer_queue.py", + "nemo_rl/data_plane/async_utils.py", "nemo_rl/data_plane/codec.py", "nemo_rl/data_plane/column_io.py", "nemo_rl/data_plane/factory.py", @@ -244,6 +245,7 @@ project-includes = [ "tools/model_diagnostics/5.prefix_caching_nan.py", "tools/model_diagnostics/6.vllm_routed_experts_completeness.py", "tools/refit_bandwidth_calculator.py", + "tools/verify_tq_data_plane_checkpoint.py", "tools/x_token/__init__.py", "tools/x_token/reapply_exact_map.py", "tools/x_token/sort_and_cut_projection_matrix.py", diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 2a6fad63ee..ea53d7fe6b 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -36,6 +36,7 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh +run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_dp_single_controller.sh b/tests/functional/grpo_dp_single_controller.sh index 855760445d..f62ff6aac0 100755 --- a/tests/functional/grpo_dp_single_controller.sh +++ b/tests/functional/grpo_dp_single_controller.sh @@ -22,7 +22,7 @@ rm -rf $EXP_DIR $LOG_DIR mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ +uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo_single_controller.py \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ @@ -47,11 +47,13 @@ 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 [[ "${RUN_CONVERGENCE_CHECKS:-1}" == "1" ]]; then + uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/gen_kl_error"]) < 0.002' \ - 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ - 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_max"]) < 1.21' + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/gen_kl_error"]) < 0.002' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.21' +fi diff --git a/tests/functional/grpo_dp_single_controller_tq_recovery.sh b/tests/functional/grpo_dp_single_controller_tq_recovery.sh new file mode 100755 index 0000000000..671b67dd4e --- /dev/null +++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Two-process functional test for native TQ + metadata-only replay recovery. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +BASE_TEST=$SCRIPT_DIR/grpo_dp_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_dp_single_controller_tq_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +BASE_RUN_LOG=$SCRIPT_DIR/grpo_dp_single_controller/run.log + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +COMMON_OVERRIDES=( + 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.max_inflight_prompts=8 + async_rl.max_buffered_rollouts=8 +) + +echo "=== Phase 1: save an authoritative native TQ checkpoint ===" +# Keep the two-step training horizon identical across both processes so the +# Megatron optimizer scheduler can be restored. The timeout makes phase 1 save +# after its first completed step and exit early, simulating an interrupted job. +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" \ + grpo.max_num_steps=2 \ + checkpointing.checkpoint_must_save_by=0:0:0:1 + +test -d "$CHECKPOINT_DIR/step_1/data_plane" +test -f "$CHECKPOINT_DIR/step_1/replay_buffer_metadata.pt" +test ! -f "$CHECKPOINT_DIR/step_1/replay_buffer.pt" +uv run --no-sync python -c \ + 'import json, sys; metadata = json.load(open(sys.argv[1]))["user_metadata"]; assert metadata["mode"] == "authoritative"; assert metadata["replay_group_count"] > 0, metadata' \ + "$CHECKPOINT_DIR/step_1/data_plane/metadata.json" + +echo "=== Phase 2: start a fresh process, restore TQ, and train one more step ===" +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" grpo.max_num_steps=2 + +grep -q "Native TQ checkpoint restored and validated" "$BASE_RUN_LOG" +grep -q "Native TQ replay inventory validated" "$BASE_RUN_LOG" +grep -Eq "Restored [1-9][0-9]* replay group" "$BASE_RUN_LOG" +test -d "$CHECKPOINT_DIR/step_2/data_plane" +test -f "$CHECKPOINT_DIR/step_2/replay_buffer_metadata.pt" + +echo "Native TQ recovery functional test passed." diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py index a105892bc0..a5167f053f 100644 --- a/tests/unit/data_plane/test_architecture_invariants.py +++ b/tests/unit/data_plane/test_architecture_invariants.py @@ -78,8 +78,11 @@ def test_sync_trainer_rejects_message_level_advantage_penalties(): "get_data", "put_samples", "get_samples", + "list_sample_ids", "clear_samples", "check_consumption_status", + "save_checkpoint", + "load_checkpoint", "close", ], ) diff --git a/tests/unit/data_plane/test_async_utils.py b/tests/unit/data_plane/test_async_utils.py new file mode 100644 index 0000000000..25d29b5d71 --- /dev/null +++ b/tests/unit/data_plane/test_async_utils.py @@ -0,0 +1,77 @@ +# 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. + +"""Tests for local and Ray-style async data-plane dispatch.""" + +import asyncio +import threading + +from nemo_rl.data_plane.async_utils import call_data_plane + + +class _LocalClient: + def thread_id(self) -> int: + return threading.get_ident() + + async def async_value(self, *, value: int) -> int: + return value + + +class _RemoteMethod: + def __init__(self) -> None: + self.calls: list[int] = [] + + async def remote(self, *, value: int) -> int: + self.calls.append(value) + return value + + +class _RemoteClient: + def __init__(self) -> None: + self.value = _RemoteMethod() + + +def test_sync_call_stays_inline_by_default() -> None: + caller_thread_id = threading.get_ident() + + result = asyncio.run(call_data_plane(_LocalClient(), "thread_id")) + + assert result == caller_thread_id + + +def test_sync_call_can_be_offloaded() -> None: + caller_thread_id = threading.get_ident() + + result = asyncio.run( + call_data_plane(_LocalClient(), "thread_id", offload_sync=True) + ) + + assert result != caller_thread_id + + +def test_local_coroutine_result_is_awaited() -> None: + result = asyncio.run( + call_data_plane(_LocalClient(), "async_value", value=7) + ) + + assert result == 7 + + +def test_ray_style_remote_result_is_awaited() -> None: + client = _RemoteClient() + + result = asyncio.run(call_data_plane(client, "value", value=11)) + + assert result == 11 + assert client.value.calls == [11] diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py index c9b71820d7..19907c79a5 100644 --- a/tests/unit/data_plane/test_codec_mooncake.py +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -22,6 +22,7 @@ from __future__ import annotations +import pytest import torch from nemo_rl.data_plane.codec import pack_per_token_field, to_nested_by_length @@ -73,6 +74,242 @@ def test_promote_1d_roundtrip_via_from_wire() -> None: assert torch.equal(back["reward"], original) +def test_from_wire_densifies_uniform_nested_rows() -> None: + """TQ v0.1.9's uniform nested reads are restored to dense tensors.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + rows = [torch.tensor([i], dtype=torch.float32) for i in range(4)] + wire = TensorDict( + {"reward": torch.nested.as_nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + back = _from_wire(wire) + + assert not back["reward"].is_nested + assert back["reward"].shape == (len(rows),) + assert torch.equal(back["reward"], torch.arange(len(rows), dtype=torch.float32)) + + +def test_from_wire_preserves_genuine_length_one_token_column() -> None: + """Only fields promoted from ``(N,)`` are squeezed after a TQ read.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + n = 4 + wire = TensorDict( + { + "reward": torch.nested.as_nested_tensor( + [torch.tensor([float(i)]) for i in range(n)], layout=torch.jagged + ), + "input_ids": torch.nested.as_nested_tensor( + [torch.tensor([i]) for i in range(n)], layout=torch.jagged + ), + }, + batch_size=[n], + ) + + back = _from_wire(wire, promoted_1d_fields={"reward"}) + + assert back["reward"].shape == (n,) + assert back["input_ids"].shape == (n, 1) + assert torch.equal(back["input_ids"], torch.arange(n).unsqueeze(-1)) + + +def test_wire_shape_tags_roundtrip_promoted_field_names() -> None: + """Durable tags distinguish scalar and length-one tensor columns.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import ( + _add_wire_shape_tags, + _promoted_1d_fields_from_tags, + _strip_wire_shape_tags, + ) + + n = 3 + fields = TensorDict( + { + "reward": torch.arange(n, dtype=torch.float32), + "input_ids": torch.arange(n).unsqueeze(-1), + }, + batch_size=[n], + ) + user_tags = [{"weight_version": 7} for _ in range(n)] + + wire_tags = _add_wire_shape_tags(user_tags, fields) + + assert _promoted_1d_fields_from_tags(wire_tags, ["reward", "input_ids"]) == { + "reward" + } + assert _strip_wire_shape_tags(wire_tags) == user_tags + assert all( + not any(key.startswith("__nemo_rl_promoted_1d__:") for key in tag) + for tag in user_tags + ) + + +def test_wire_shape_tags_legacy_rows_have_no_promotion_provenance() -> None: + """Rows written before wire-shape tags use the legacy decode path.""" + from nemo_rl.data_plane.adapters.transfer_queue import ( + _promoted_1d_fields_from_tags, + ) + + assert _promoted_1d_fields_from_tags([], ["reward"]) is None + + +def test_wire_shape_tags_reject_mixed_codec_versions() -> None: + """A batch cannot safely mix tagged and legacy rows.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import ( + _WIRE_SHAPE_VERSION_TAG, + _add_wire_shape_tags, + _promoted_1d_fields_from_tags, + ) + + fields = TensorDict({"input_ids": torch.arange(4).reshape(2, 2)}, batch_size=[2]) + wire_tags = _add_wire_shape_tags([{}, {}], fields) + wire_tags[1].pop(_WIRE_SHAPE_VERSION_TAG) + + with pytest.raises(RuntimeError, match="wire-shape version"): + _promoted_1d_fields_from_tags(wire_tags, ["input_ids"]) + + +def test_wire_shape_tags_reject_partial_promotion_marker() -> None: + """All rows must agree that a field was promoted from one dimension.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import ( + _PROMOTED_1D_TAG_PREFIX, + _add_wire_shape_tags, + _promoted_1d_fields_from_tags, + ) + + fields = TensorDict({"reward": torch.arange(2)}, batch_size=[2]) + wire_tags = _add_wire_shape_tags([{}, {}], fields) + wire_tags[1].pop(f"{_PROMOTED_1D_TAG_PREFIX}reward") + + with pytest.raises(RuntimeError, match="promotion marker"): + _promoted_1d_fields_from_tags(wire_tags, ["reward"]) + + +def test_get_samples_uses_persisted_shape_tags(monkeypatch) -> None: + """The Mooncake adapter restores ranks using TQ's durable row metadata.""" + from tensordict import TensorDict + + import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter + + n = 3 + original = TensorDict( + { + "reward": torch.arange(n, dtype=torch.float32), + "input_ids": torch.arange(n).unsqueeze(-1), + }, + batch_size=[n], + ) + wire_tags = tq_adapter._add_wire_shape_tags([{} for _ in range(n)], original) + wire_data = TensorDict( + { + "reward": torch.nested.as_nested_tensor( + [row for row in original["reward"].unsqueeze(-1)], + layout=torch.jagged, + ), + "input_ids": torch.nested.as_nested_tensor( + [row for row in original["input_ids"]], layout=torch.jagged + ), + }, + batch_size=[n], + ) + + class FakeMeta: + size = n + is_ready = True + custom_meta = wire_tags + + def select_fields(self, fields): + assert fields == ["reward", "input_ids"] + return self + + class FakeClient: + def kv_retrieve_meta(self, *, keys, partition_id, create): + assert keys == ["a", "b", "c"] + assert partition_id == "train" + assert create is False + return FakeMeta() + + def get_data(self, meta): + assert isinstance(meta, FakeMeta) + return wire_data + + monkeypatch.setattr( + tq_adapter.tq, "get_client", lambda: FakeClient(), raising=False + ) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._promote_1d = True + client._data_operations_started = False + + restored = client.get_samples(["a", "b", "c"], "train", ["reward", "input_ids"]) + + assert restored["reward"].shape == (n,) + assert restored["input_ids"].shape == (n, 1) + assert torch.equal(restored["reward"], original["reward"]) + assert torch.equal(restored["input_ids"], original["input_ids"]) + + +def test_get_samples_densifies_uniform_rows_without_1d_promotion(monkeypatch) -> None: + """The simple backend normalizes uniform nested rows without squeezing.""" + from tensordict import TensorDict + + import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter + + rows = [torch.tensor([1, 2]), torch.tensor([3, 4])] + wire_data = TensorDict( + {"input_ids": torch.nested.as_nested_tensor(rows, layout=torch.jagged)}, + batch_size=[len(rows)], + ) + + def fake_kv_batch_get( + *, keys: list[str], partition_id: str, select_fields: list[str] + ) -> TensorDict: + assert keys == ["a", "b"] + assert partition_id == "train" + assert select_fields == ["input_ids"] + return wire_data + + monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get, raising=False) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._promote_1d = False + client._data_operations_started = False + + restored = client.get_samples(["a", "b"], "train", ["input_ids"]) + + assert not restored["input_ids"].is_nested + assert restored["input_ids"].shape == (2, 2) + assert torch.equal(restored["input_ids"], torch.stack(rows)) + + +def test_from_wire_preserves_ragged_nested_rows() -> None: + """Variable-length rollout fields must remain nested.""" + from tensordict import TensorDict + + from nemo_rl.data_plane.adapters.transfer_queue import _from_wire + + rows = [torch.arange(i + 1) for i in range(3)] + nested = torch.nested.as_nested_tensor(rows, layout=torch.jagged) + wire = TensorDict({"token_ids": nested}, batch_size=[len(rows)]) + + back = _from_wire(wire) + + assert back["token_ids"].is_nested + assert all( + torch.equal(actual, expected) + for actual, expected in zip(back["token_ids"].unbind(), rows, strict=True) + ) + + # ── P2: pack_per_token_field — tolerates SP padding ────────────────────────── diff --git a/tests/unit/data_plane/test_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py index 3426c3b506..7c9c900259 100644 --- a/tests/unit/data_plane/test_interface_contract.py +++ b/tests/unit/data_plane/test_interface_contract.py @@ -65,8 +65,10 @@ def test_register_put_get_clear(client: DataPlaneClient): out = client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) assert torch.equal(out["x"], torch.arange(4)) + assert client.list_sample_ids("p") == keys client.clear_samples(sample_ids=None, partition_id="p") + assert client.list_sample_ids("p") == [] with pytest.raises(KeyError): client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) @@ -124,3 +126,65 @@ def test_kv_batch_put_rejects_non_tensor_leaves(client: DataPlaneClient): def test_close_is_idempotent(client: DataPlaneClient): client.close() client.close() + + +def test_checkpoint_round_trip_restores_data_and_consumption(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + source = NoOpDataPlaneClient() + source.register_partition( + partition_id="p", + fields=["x"], + num_samples=3, + consumer_tasks=["train"], + ) + source.put_samples( + sample_ids=["a", "b", "c"], + partition_id="p", + fields=TensorDict({"x": torch.tensor([10, 20, 30])}, batch_size=[3]), + ) + consumed = source.claim_meta( + partition_id="p", + task_name="train", + required_fields=["x"], + batch_size=1, + ) + source.save_checkpoint(checkpoint_dir, metadata={"step": 7}) + + restored = NoOpDataPlaneClient() + metadata = restored.load_checkpoint(checkpoint_dir) + assert metadata == {"step": 7} + data = restored.get_samples( + sample_ids=["a", "b", "c"], + partition_id="p", + select_fields=["x"], + ) + assert torch.equal(data["x"], torch.tensor([10, 20, 30])) + + remaining = restored.claim_meta( + partition_id="p", + task_name="train", + required_fields=["x"], + batch_size=3, + ) + assert consumed.sample_ids[0] not in remaining.sample_ids + assert set(consumed.sample_ids + remaining.sample_ids) == {"a", "b", "c"} + assert restored.check_consumption_status("p", ["train"]) + + source.close() + restored.close() + + +def test_checkpoint_load_requires_clean_client(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + source = NoOpDataPlaneClient() + source.save_checkpoint(checkpoint_dir) + + source.register_partition( + partition_id="already-used", + fields=["x"], + num_samples=1, + consumer_tasks=["train"], + ) + with pytest.raises(RuntimeError, match="clean data-plane client"): + source.load_checkpoint(checkpoint_dir) + source.close() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0d471bc266..13a5f59a38 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -89,6 +89,22 @@ def test_register_and_clear_recorded(wrapped_client): assert ops.count("clear") == 1 +def test_list_sample_ids_is_forwarded_and_recorded(wrapped_client): + client, events = wrapped_client + client.register_partition( + partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["r"] + ) + client.put_samples( + sample_ids=["b", "a"], + partition_id="p", + fields=TensorDict({"x": torch.ones(2)}, batch_size=[2]), + ) + + assert client.list_sample_ids("p") == ["a", "b"] + assert events[-1]["op"] == "list_sample_ids" + assert events[-1]["status"] == "ok" + + def test_error_status_recorded_and_reraised(wrapped_client): """Decorator does NOT swallow errors — re-raise after recording.""" client, events = wrapped_client @@ -131,6 +147,42 @@ def test_close_propagates(wrapped_client): client.close() +def test_checkpoint_lifecycle_is_forwarded_and_recorded(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + source_events: list[dict] = [] + source = MetricsDataPlaneClient( + NoOpDataPlaneClient(), + on_event=source_events.append, + ) + source.register_partition( + partition_id="p", + fields=["x"], + num_samples=1, + consumer_tasks=["train"], + ) + source.put_samples( + sample_ids=["a"], + partition_id="p", + fields=TensorDict({"x": torch.tensor([1])}, batch_size=[1]), + ) + source.save_checkpoint(checkpoint_dir, metadata={"step": 3}) + + restore_events: list[dict] = [] + restored = MetricsDataPlaneClient( + NoOpDataPlaneClient(), + on_event=restore_events.append, + ) + metadata = restored.load_checkpoint(checkpoint_dir) + + assert metadata == {"step": 3} + assert [event["op"] for event in source_events][-1] == "save_checkpoint" + assert source_events[-1]["status"] == "ok" + assert [event["op"] for event in restore_events] == ["load_checkpoint"] + assert restore_events[-1]["status"] == "ok" + source.close() + restored.close() + + def test_factory_wraps_when_observability_enabled(): """Programmatic wrap path; factory.py uses the same MetricsDataPlaneClient.""" inner = NoOpDataPlaneClient() diff --git a/tests/unit/data_plane/test_smoke.py b/tests/unit/data_plane/test_smoke.py index 25505d1053..e373adf076 100644 --- a/tests/unit/data_plane/test_smoke.py +++ b/tests/unit/data_plane/test_smoke.py @@ -89,6 +89,7 @@ def test_dataplane_client_abc_surface() -> None: # direct-by-key "put_samples", "get_samples", + "list_sample_ids", "clear_samples", # lifecycle "close", diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index a349c27c2e..bc0fa2dd22 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -22,6 +22,9 @@ from __future__ import annotations +import json +from unittest.mock import MagicMock + import numpy as np import pytest import torch @@ -51,6 +54,7 @@ def fake_clear(**kwargs): monkeypatch.setattr(tq_adapter.tq, "kv_clear", fake_clear) client = object.__new__(tq_adapter.TQDataPlaneClient) + client._data_operations_started = False client.register_partition( partition_id="obj-backend", fields=["msg_log"], @@ -82,6 +86,154 @@ def fake_clear(**kwargs): ] +def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect_calls = [] + save_calls = [] + load_calls = [] + monkeypatch.setattr( + tq_adapter, + "_connect_existing", + lambda: connect_calls.append(None), + ) + monkeypatch.setattr( + tq_adapter.tq, + "save_checkpoint", + lambda checkpoint_dir, *, metadata=None: save_calls.append( + (checkpoint_dir, metadata) + ), + ) + monkeypatch.setattr( + tq_adapter.tq, + "load_checkpoint", + lambda checkpoint_dir: load_calls.append(checkpoint_dir), + ) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + client._data_operations_started = False + checkpoint_dir = tmp_path / "step-7" + checkpoint_dir.mkdir() + (checkpoint_dir / "metadata.json").write_text( + json.dumps({"storage_saved": True, "user_metadata": {"step": 7}}) + ) + client.save_checkpoint(checkpoint_dir, metadata={"step": 7}) + metadata = client.load_checkpoint(checkpoint_dir) + + assert connect_calls == [None, None] + assert save_calls == [(checkpoint_dir, {"step": 7})] + assert load_calls == [checkpoint_dir] + assert metadata == {"step": 7} + assert client._data_operations_started + + +def test_list_sample_ids_uses_tq_partition_listing(monkeypatch) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + list_call = MagicMock( + return_value={"rollout_data": {"sample-b": {}, "sample-a": {}}} + ) + monkeypatch.setattr(tq_adapter.tq, "kv_list", list_call) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._data_operations_started = False + + sample_ids = client.list_sample_ids("rollout_data") + + assert sample_ids == ["sample-a", "sample-b"] + assert client._data_operations_started + list_call.assert_called_once_with(partition_id="rollout_data") + + +def test_checkpoint_load_rejects_client_after_data_operation( + monkeypatch, + tmp_path, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect = MagicMock() + load = MagicMock() + monkeypatch.setattr(tq_adapter, "_connect_existing", connect) + monkeypatch.setattr(tq_adapter.tq, "load_checkpoint", load) + monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", MagicMock()) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + client._promote_1d = False + client._data_operations_started = False + client.put_samples( + sample_ids=["sample-0"], + partition_id="rollout_data", + fields=TensorDict({"x": torch.tensor([1])}, batch_size=[1]), + ) + + with pytest.raises(RuntimeError, match="requires a clean TQ client"): + client.load_checkpoint(tmp_path / "data-plane") + + connect.assert_not_called() + load.assert_not_called() + + +def test_failed_checkpoint_load_leaves_client_in_dirty_state( + monkeypatch, + tmp_path, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect = MagicMock() + load = MagicMock(side_effect=RuntimeError("injected partial restore")) + monkeypatch.setattr(tq_adapter, "_connect_existing", connect) + monkeypatch.setattr(tq_adapter.tq, "load_checkpoint", load) + + checkpoint_dir = tmp_path / "data-plane" + checkpoint_dir.mkdir() + (checkpoint_dir / "metadata.json").write_text( + json.dumps({"storage_saved": True, "user_metadata": {"step": 7}}) + ) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + client._data_operations_started = False + + with pytest.raises(RuntimeError, match="injected partial restore"): + client.load_checkpoint(checkpoint_dir) + with pytest.raises(RuntimeError, match="requires a clean TQ client"): + client.load_checkpoint(checkpoint_dir) + + connect.assert_called_once_with() + load.assert_called_once_with(checkpoint_dir) + + +@pytest.mark.parametrize("operation", ["save", "load"]) +def test_mooncake_checkpoint_lifecycle_fails_loudly( + monkeypatch, + tmp_path, + operation: str, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect = MagicMock() + save = MagicMock() + load = MagicMock() + monkeypatch.setattr(tq_adapter, "_connect_existing", connect) + monkeypatch.setattr(tq_adapter.tq, "save_checkpoint", save) + monkeypatch.setattr(tq_adapter.tq, "load_checkpoint", load) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "mooncake_cpu" + client._data_operations_started = False + checkpoint_dir = tmp_path / "step-7" + + with pytest.raises(NotImplementedError, match="mooncake_cpu"): + if operation == "save": + client.save_checkpoint(checkpoint_dir) + else: + client.load_checkpoint(checkpoint_dir) + + connect.assert_not_called() + save.assert_not_called() + load.assert_not_called() + + # ``tq_client`` (simple) and ``tq_client_backends`` (parametrized over # simple + mooncake_cpu) are session-scoped fixtures provided by # ``tests/unit/data_plane/conftest.py``. See that file for the rationale. @@ -134,10 +286,11 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None: consumer_tasks=["read"], ) keys = ["a", "b", "c", "d"] + values = torch.arange(12).reshape(4, 3) client.put_samples( sample_ids=keys, partition_id="smoke-backend", - fields=TensorDict({"x": torch.arange(4)}, batch_size=[4]), + fields=TensorDict({"x": values}, batch_size=[4]), ) meta = client.claim_meta( @@ -150,7 +303,9 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None: assert meta.size == 4 data = client.get_data(meta) - expected = torch.tensor([keys.index(k) for k in meta.sample_ids]) + expected = torch.stack([values[keys.index(k)] for k in meta.sample_ids]) + assert not data["x"].is_nested + assert data["x"].shape == expected.shape assert torch.equal(data["x"], expected) client.clear_samples(sample_ids=None, partition_id="smoke-backend") diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 2124588dac..05ee6fee7d 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -485,6 +485,9 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + # SingleController only: save native TQ state. Supported samplers restore + # from metadata-only replay indexes. + checkpointing_enabled: false global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py index 1c5d9a8205..3d63f5915e 100644 --- a/tests/unit/single_controller/_dp_fakes.py +++ b/tests/unit/single_controller/_dp_fakes.py @@ -147,6 +147,9 @@ def clear_samples(self, sample_ids: list[str], partition_id: str) -> Any: ) ) + def list_sample_ids(self, partition_id: str) -> list[str]: + return ray.get(self._handle.list_sample_ids.remote(partition_id)) + @staticmethod def _padded(td: TensorDict) -> TensorDict: out: dict[str, torch.Tensor] = {} diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index cfeda7b1f5..ad337c75fb 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -31,6 +31,7 @@ WindowedSampler, WindowedSamplerConfig, ) +from nemo_rl.algorithms.grpo import _default_grpo_save_state from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AsyncRLConfig, @@ -309,6 +310,18 @@ def test_rollout_pump_writes_expected_tq_data( "mlflow_enabled": False, "monitor_gpus": False, }, + # Actor __init__ builds a CheckpointManager + TimeoutChecker from + # this block; enabled=False keeps the run write-free. + checkpointing={ + "enabled": False, + "checkpoint_dir": str(tmp_path / "checkpoints"), + "metric_name": None, + "higher_is_better": False, + "keep_top_k": None, + "save_period": 10_000, + "save_optimizer": False, + "checkpoint_must_save_by": None, + }, ) # Wrap each value in a single-element list so size==1 and v[0] returns the original field. batched_sample = BatchedDataDict({k: [v] for k, v in input_sample.items()}) @@ -343,6 +356,8 @@ def test_rollout_pump_writes_expected_tq_data( rollout_manager=rollout_manager, tq_buffer=tq_buffer, partition_id=_PARTITION_ID, + save_state=_default_grpo_save_state(), + last_checkpoint_path=None, ) ctrl = SingleControllerActor.remote( master_config=master_config, actor_args=actor_args diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 8e0ef3e20c..35f1ca0c10 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -27,10 +27,12 @@ import pytest from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, InOrderSampler, InOrderSamplerConfig, PromptGroupSampler, WeightFifoSampler, + WeightFifoSamplerConfig, WindowedSampler, WindowedSamplerConfig, create_sampler, @@ -151,6 +153,24 @@ def test_unready_slot_is_never_evicted(self): class TestFactory: + @pytest.mark.parametrize( + ("config", "expected"), + [ + (WindowedSamplerConfig(), True), + (WeightFifoSamplerConfig(), False), + (InOrderSamplerConfig(), False), + ( + CustomSamplerConfig(target=f"{__name__}:EchoSampler"), + None, + ), + ], + ) + def test_config_declares_checkpoint_capability_without_serializing_it( + self, config, expected + ): + assert config.supports_buffer_checkpoint is expected + assert "supports_buffer_checkpoint" not in config.model_dump() + def test_windowed_config_builds_windowed(self): s = create_sampler( FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3) @@ -164,24 +184,27 @@ def test_in_order_config_builds_in_order(self): assert s.max_lookahead_versions == 2 def test_weight_fifo_config_builds_weight_fifo(self): - from nemo_rl.algorithms.async_utils.staleness_sampler import ( - WeightFifoSamplerConfig, - ) - s = create_sampler( FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4) ) assert isinstance(s, WeightFifoSampler) assert s.max_staleness_versions == 4 + def test_factory_rejects_config_implementation_capability_drift( + self, monkeypatch + ): + monkeypatch.setattr( + WindowedSampler, + "supports_buffer_checkpoint", + property(lambda _self: False), + ) + with pytest.raises(RuntimeError, match="disagrees with"): + create_sampler(FakeBuffer(), WindowedSamplerConfig()) + class TestCustomFqnSampler: def test_custom_target_loads_out_of_repo_sampler(self): # A user sampler defined anywhere importable; here, this test module. - from nemo_rl.algorithms.async_utils.staleness_sampler import ( - CustomSamplerConfig, - ) - s = create_sampler( FakeBuffer(), CustomSamplerConfig( @@ -304,5 +327,93 @@ def test_windowed_evict_skips_unready_stale(self): assert _run(s.evict(current_train_weight=5)) == 0 +class TestSupportsBufferCheckpoint: + """Only ungated (over-sampled) policies may checkpoint buffered groups.""" + + def test_windowed_supports_buffer_checkpoint(self): + assert WindowedSampler( + FakeBuffer(), max_staleness_versions=1 + ).supports_buffer_checkpoint + + def test_gated_samplers_do_not(self): + assert not WeightFifoSampler( + FakeBuffer(), max_staleness_versions=1 + ).supports_buffer_checkpoint + assert not InOrderSampler( + FakeBuffer(), max_lookahead_versions=1 + ).supports_buffer_checkpoint + + +class TestDispatchCursorRestore: + """Checkpoint resume passes resume_from_step=current_step, restoring the + fresh-start invariant _dispatch_index == trainer_version - 1. Without it, + a restored InOrderSampler would stamp target_steps starting at 0 and every + dispatched batch would be instantly evicted (target < trainer_version).""" + + def test_resumed_in_order_stamps_from_trainer_version(self): + s = InOrderSampler(FakeBuffer(), max_lookahead_versions=1, resume_from_step=7) + assert _run(s.admit(trainer_version_fn=lambda: 7)) == 7 + assert _run(s.admit(trainer_version_fn=lambda: 8)) == 8 + + def test_resumed_gate_admits_window_then_blocks(self): + s = WeightFifoSampler( + FakeBuffer(), max_staleness_versions=0, resume_from_step=7 + ) + # Resumed at step 7, window 0: one batch admitted, then the gate + # closes exactly as it would on a fresh run at step 0. + assert _run(s.admit(trainer_version_fn=lambda: 7)) is None + with pytest.raises(asyncio.TimeoutError): + _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 7), timeout=0.05)) + + def test_negative_resume_step_rejected(self): + with pytest.raises(ValueError, match="resume_from_step"): + WindowedSampler(FakeBuffer(), max_staleness_versions=1, resume_from_step=-1) + + +class TestFactorySeeding: + def test_builtin_receives_resume_step(self): + s = create_sampler( + FakeBuffer(), + InOrderSamplerConfig(max_lookahead_versions=1), + resume_from_step=7, + ) + assert _run(s.admit(trainer_version_fn=lambda: 7)) == 7 + + def test_custom_receives_resume_step_on_restore(self): + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, + ) + + s = create_sampler( + FakeBuffer(), + CustomSamplerConfig( + target=f"{__name__}:EchoSampler", max_lookahead_versions=1 + ), + resume_from_step=6, + ) + assert _run(s.admit(trainer_version_fn=lambda: 6)) == 6 + + def test_custom_without_restore_support_fails_only_on_restore(self): + from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, + ) + + cfg = CustomSamplerConfig( + target=f"{__name__}:NoRestoreSampler", max_lookahead_versions=1 + ) + # Fresh start: the kwarg is withheld, existing custom samplers keep working. + assert isinstance(create_sampler(FakeBuffer(), cfg), NoRestoreSampler) + # Resume: loud TypeError instead of a silently unseeded cursor. + with pytest.raises(TypeError): + create_sampler(FakeBuffer(), cfg, resume_from_step=5) + + class EchoSampler(InOrderSampler): """Stand-in for a user-defined sampler loaded by FQN.""" + + +class NoRestoreSampler(InOrderSampler): + """User sampler whose constructor predates resume_from_step.""" + + def __init__(self, buffer, *, max_lookahead_versions: int) -> None: + super().__init__(buffer, max_lookahead_versions=max_lookahead_versions) diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py new file mode 100644 index 0000000000..6073f3a132 --- /dev/null +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -0,0 +1,1512 @@ +# 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. + +"""Unit tests for SC checkpointing. + +Covers: + - counter restore from save_state (train_steps / trainer_version / sampler + dispatch cursor, current_epoch); + - save trigger + write path through _train_pump with fakes (period + boundary, last step, timeout, disabled); + - async-save finalization (rename deferred until finalize_async_save, + background failure re-raised at the next save); + - metric_name behavior (val:* warn-and-save, train:* value recorded); + - dataloader state: train_dataloader.pt written at save, position + round-trip through a real StatefulDataLoader, dataset-swap guard, + setup restore wiring + missing-file fresh-position fallback; + - native replay persistence requires both sampler support and TQ checkpointing; + - setup_single_controller resume-path wiring (get_resume_paths forwarded + to the trainer factory, save_state loaded from training_info.json). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import threading +from pathlib import Path +from typing import Any, Optional +from unittest.mock import MagicMock, patch + +import pytest +import torch +import yaml +from torchdata.stateful_dataloader import StatefulDataLoader + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + WindowedSamplerConfig, +) +from nemo_rl.algorithms.grpo import _default_grpo_save_state +from nemo_rl.algorithms.loss import ClippedPGLossConfig +from nemo_rl.algorithms.single_controller import SingleControllerActor +from nemo_rl.algorithms.single_controller_utils import ( + AsyncRLConfig, + MasterConfig, + setup_single_controller, +) +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 KVBatchMeta +from nemo_rl.utils.checkpoint import CheckpointManager + +# Reuse the factory patches from the setup tests (same cross-module fixture +# import pattern as test_rollout_pump.py). +from tests.unit.single_controller.test_single_controller_setup import ( + patched_factories, # noqa: F401 +) + +# Instantiate the underlying class in-process (same pattern as +# tests/unit/algorithms/test_async_utils.py for AsyncTrajectoryCollector). +_ACTOR_CLS = SingleControllerActor.__ray_metadata__.modified_class + +_PARTITION_ID = "rollout_data" + + +# ── fakes ──────────────────────────────────────────────────────────────────── + + +class _FakeTrainer: + """TQPolicy stand-in: train methods are no-ops, save_checkpoint records calls.""" + + def __init__(self, step_metrics: Optional[dict[str, Any]] = None) -> None: + self._step_metrics = dict(step_metrics or {}) + self.save_calls: list[dict[str, Any]] = [] + self.finalize_calls: int = 0 + + def prepare_for_lp_inference(self) -> None: + pass + + def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None: + pass + + def get_reference_policy_logprobs_from_meta(self, meta: KVBatchMeta) -> None: + pass + + def prepare_for_training(self) -> None: + pass + + def begin_train_step(self, loss_fn: Any) -> None: + pass + + def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None: + pass + + def finish_train_step(self) -> dict[str, Any]: + return dict(self._step_metrics) + + def save_checkpoint( + self, + *, + weights_path: str, + optimizer_path: Optional[str], + tokenizer_path: str, + checkpointing_cfg: dict[str, Any], + ) -> None: + self.save_calls.append( + { + "weights_path": weights_path, + "optimizer_path": optimizer_path, + "tokenizer_path": tokenizer_path, + "checkpointing_cfg": checkpointing_cfg, + } + ) + # Mimic the real Policy: materialize the checkpoint subdirs. + os.makedirs(weights_path, exist_ok=True) + if optimizer_path is not None: + os.makedirs(optimizer_path, exist_ok=True) + os.makedirs(tokenizer_path, exist_ok=True) + + def finalize_async_save(self) -> None: + self.finalize_calls += 1 + + +class _GatedFinalizeTrainer(_FakeTrainer): + """Async-save stand-in: finalize_async_save blocks until released.""" + + def __init__(self) -> None: + super().__init__() + self.release = threading.Event() + + def finalize_async_save(self) -> None: + assert self.release.wait(timeout=30.0), "test never released the writer" + super().finalize_async_save() + + +class _FailingFinalizeTrainer(_FakeTrainer): + def finalize_async_save(self) -> None: + raise RuntimeError("injected async-writer failure") + + +class _FakeSampler: + """PromptGroupSampler stand-in: always returns a full, fresh batch.""" + + def __init__(self, supports_buffer_checkpoint: bool = True) -> None: + self._supports_buffer_checkpoint = supports_buffer_checkpoint + self._step = 0 + + async def admit(self, *, trainer_version_fn) -> Optional[int]: + return None + + async def evict(self, *, current_train_weight: int) -> int: + return 0 + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[KVBatchMeta, int]: + n = max_prompt_groups + sample_ids = [f"s{self._step}-{i}" for i in range(n)] + self._step += 1 + meta = KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=sample_ids, + sequence_lengths=[16] * n, + tags=[{"weight_version": current_train_weight}] * n, + ) + return meta, n + + @property + def is_on_policy(self) -> bool: + return False + + @property + def supports_buffer_checkpoint(self) -> bool: + return self._supports_buffer_checkpoint + + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: + return None + + +class _FakeDPClient: + def __init__( + self, + *, + save_error: Optional[Exception] = None, + sample_ids: Optional[list[str]] = None, + ) -> None: + self.clear_calls: list[tuple[list[str], str]] = [] + self.clear_thread_ids: list[int] = [] + self.save_calls: list[dict[str, Any]] = [] + self.save_error = save_error + self.sample_ids = list(sample_ids or []) + + def list_sample_ids(self, partition_id: str) -> list[str]: + assert partition_id == _PARTITION_ID + return sorted(self.sample_ids) + + def clear_samples(self, sample_ids: list[str], partition_id: str) -> None: + self.clear_thread_ids.append(threading.get_ident()) + self.clear_calls.append((list(sample_ids), partition_id)) + + def save_checkpoint( + self, + checkpoint_dir: str, + *, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + self.save_calls.append( + { + "checkpoint_dir": checkpoint_dir, + "metadata": dict(metadata or {}), + } + ) + if self.save_error is not None: + raise self.save_error + os.makedirs(checkpoint_dir, exist_ok=True) + with open(os.path.join(checkpoint_dir, "metadata.json"), "w") as f: + json.dump({"user_metadata": metadata or {}}, f) + + +class _BlockingDPClient(_FakeDPClient): + def __init__(self) -> None: + super().__init__() + self.save_started = threading.Event() + self.release_save = threading.Event() + + def save_checkpoint( + self, + checkpoint_dir: str, + *, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + self.save_started.set() + assert self.release_save.wait(timeout=30.0), "test never released TQ save" + super().save_checkpoint(checkpoint_dir, metadata=metadata) + + +class _FakeWeightSynchronizer: + def __init__(self) -> None: + self.sync_count = 0 + + def sync_weights(self, *, kv_scales: Any = None) -> None: + self.sync_count += 1 + + +class _FakeRolloutManager: + def __init__(self) -> None: + self.weight_versions: list[int] = [] + self._tq_buffer = None + + def set_weight_version(self, version: int) -> None: + self.weight_versions.append(version) + + +class _FakeTQBuffer: + """TQReplayBuffer stand-in for the SC save/restore integration tests.""" + + def __init__( + self, + metadata_state: Optional[dict[str, Any]] = None, + load_return: int = 0, + ) -> None: + self._metadata_state = metadata_state or { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": _PARTITION_ID, + "saved_capacity": 4, + "manifest_digest": "fake-manifest-digest", + "groups": [], + } + self.load_return = load_return + self.metadata_state_dict_calls: list[int] = [] + self.load_calls: list[dict[str, Any]] = [] + self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None + + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + self.checkpoint_barrier = barrier + + def metadata_state_dict(self, *, saved_capacity: int) -> dict[str, Any]: + self.metadata_state_dict_calls.append(saved_capacity) + return dict(self._metadata_state) + + async def load_state_dict( + self, + state: dict[str, Any], + *, + max_groups: int, + expected_partition_id: str, + expected_group_size: int, + expected_manifest_digest: str, + ) -> int: + self.load_calls.append( + { + "state": state, + "max_groups": max_groups, + "expected_partition_id": expected_partition_id, + "expected_group_size": expected_group_size, + "expected_manifest_digest": expected_manifest_digest, + } + ) + return self.load_return + + +# Default position sentinel the fake dataloader reports via state_dict(). +_SENTINEL_DL_STATE = {"fake_position": 42} + + +class _FakeDataloader(list): + """List-backed dataloader with the StatefulDataLoader state_dict surface. + + The save block snapshots ``self._dataloader.state_dict()``; return a + sentinel dict so tests can assert the exact object written to + train_dataloader.pt. + """ + + def __init__(self, batches: Any = (), state: Optional[dict[str, Any]] = None): + super().__init__(batches) + self._state = dict(state) if state is not None else dict(_SENTINEL_DL_STATE) + + def state_dict(self) -> dict[str, Any]: + return dict(self._state) + + +# ── builders ───────────────────────────────────────────────────────────────── + + +def _actor_master_config( + tmp_path: Path, + *, + max_num_steps: int = 4, + save_period: int = 2, + enabled: bool = True, + metric_name: Optional[str] = None, + save_optimizer: bool = True, + checkpoint_must_save_by: Optional[str] = None, + num_prompts_per_step: int = 2, + max_num_epochs: int = 1, + buffer_checkpoint: bool = False, + data_plane_checkpoint: bool = False, +) -> MasterConfig: + """MasterConfig for in-process SingleControllerActor tests. + + All fields are populated (init_tmp_checkpoint dumps the whole config to + config.yaml); values satisfy validate_single_controller_config. + buffer_checkpoint selects the sampler: windowed supports replay-buffer + checkpointing, the gated in_order sampler does not. + """ + sampler_cfg = ( + WindowedSamplerConfig(max_staleness_versions=1) + if buffer_checkpoint + else InOrderSamplerConfig(max_lookahead_versions=1) + ) + return MasterConfig.model_construct( + policy={ + # One optimizer.step per RL step: prompts * generations == gbs. + "train_global_batch_size": num_prompts_per_step * 2, + }, + loss_fn=ClippedPGLossConfig(), + env={}, + data={"shuffle": False, "num_workers": 0}, + grpo={ + "max_num_steps": max_num_steps, + "max_num_epochs": max_num_epochs, + "num_prompts_per_step": num_prompts_per_step, + "num_generations_per_prompt": 2, + "seed": 42, + }, + logger={ + "log_dir": str(tmp_path / "logs"), + "wandb_enabled": False, + "swanlab_enabled": False, + "tensorboard_enabled": False, + "mlflow_enabled": False, + "monitor_gpus": False, + }, + cluster={"num_nodes": 1, "gpus_per_node": 1}, + checkpointing={ + "enabled": enabled, + "checkpoint_dir": str(tmp_path / "checkpoints"), + "metric_name": metric_name, + "higher_is_better": True, + "keep_top_k": None, + "save_period": save_period, + "save_optimizer": save_optimizer, + "checkpoint_must_save_by": checkpoint_must_save_by, + }, + data_plane={ + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "checkpointing_enabled": data_plane_checkpoint, + }, + async_rl=AsyncRLConfig( + sampler=sampler_cfg, + min_groups_for_streaming_train=1, + max_inflight_prompts=4, + max_buffered_rollouts=4, + ), + ) + + +def _make_actor_args( + *, + trainer: Optional[_FakeTrainer] = None, + save_state: Optional[dict[str, Any]] = None, + dataloader: Optional[_FakeDataloader] = None, + tq_buffer: Optional[_FakeTQBuffer] = None, + dp_client: Optional[_FakeDPClient] = None, + last_checkpoint_path: Optional[str] = None, + data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None, +) -> SingleControllerActorArgs: + return SingleControllerActorArgs( + gen_handle=object(), + trainer_handle=trainer if trainer is not None else _FakeTrainer(), + env_handles={}, + train_cluster=None, # type: ignore[arg-type] + inference_cluster=None, # type: ignore[arg-type] + dp_client=dp_client if dp_client is not None else _FakeDPClient(), + dataloader=dataloader if dataloader is not None else _FakeDataloader(), + weight_synchronizer=_FakeWeightSynchronizer(), # type: ignore[arg-type] + advantage_estimator=None, + loss_fn=object(), # type: ignore[arg-type] + rollout_manager=_FakeRolloutManager(), # type: ignore[arg-type] + tq_buffer=tq_buffer if tq_buffer is not None else _FakeTQBuffer(), # type: ignore[arg-type] + partition_id=_PARTITION_ID, + save_state=( + save_state if save_state is not None else _default_grpo_save_state() + ), + last_checkpoint_path=last_checkpoint_path, + data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, + ) + + +def _run_train_pump( + mc: MasterConfig, + actor_args: SingleControllerActorArgs, + *, + flush: bool = True, +): + """Construct the actor in-process and drive _train_pump to completion. + + flush=True joins the (possibly async) checkpoint finalization afterwards, + like run()'s exit path does, so step_N dirs are visible to assertions. + """ + + async def _main(): + actor = _ACTOR_CLS(mc, actor_args) + actor._sampler = _FakeSampler( + supports_buffer_checkpoint=(mc.async_rl.sampler.name == "windowed") + ) + # In-process runs have no Ray runtime; the pump only reads the GPU + # count for a throughput metric. + with patch("ray.cluster_resources", return_value={"GPU": 0}): + await actor._train_pump() + if flush: + actor._checkpointer.shutdown() + return actor + + return asyncio.run(_main()) + + +def _run_actor_run(mc: MasterConfig, actor_args: SingleControllerActorArgs): + """Construct the actor in-process and drive run() to completion. + + max_num_steps=0 makes _train_pump exit immediately, so run() executes + only the restore block + pump startup/teardown. The wait_for bounds the + would-be-deadlock cases (an over-capacity permit acquisition would hang + run() forever). + """ + + async def _main(): + actor = _ACTOR_CLS(mc, actor_args) + result = await asyncio.wait_for(actor.run(), timeout=60.0) + return actor, result + + return asyncio.run(_main()) + + +def _step_dir_names(ckpt_dir: Path) -> set[str]: + if not ckpt_dir.exists(): + return set() + return { + p.name for p in ckpt_dir.iterdir() if p.name != "latest_checkpoint_status.json" + } + + +def _training_info(ckpt_dir: Path, step: int) -> dict[str, Any]: + with open(ckpt_dir / f"step_{step}" / "training_info.json") as f: + return json.load(f) + + +# ── counter restore ────────────────────────────────────────────────────────── + + +class TestCounterRestore: + def test_restore_from_step_n(self, tmp_path): + save_state = _default_grpo_save_state() + save_state["current_step"] = 7 + save_state["current_epoch"] = 2 + save_state["consumed_samples"] = 42 + save_state["total_valid_tokens"] = 1234 + + actor = _ACTOR_CLS( + _actor_master_config(tmp_path), _make_actor_args(save_state=save_state) + ) + + assert actor._train_steps == 7 + assert actor._trainer_version == 7 + # The sampler dispatch cursor is seeded to preserve the fresh-start + # invariant _dispatch_index == trainer_version - 1. + assert actor._sampler._dispatch_index == 6 + assert actor._consumed_samples == 42 + assert actor._current_epoch == 2 + assert actor._total_valid_tokens == 1234 + + def test_restores_trainer_version_independently_from_train_step(self, tmp_path): + save_state = _default_grpo_save_state() + save_state["current_step"] = 7 + save_state["trainer_version"] = 11 + + actor = _ACTOR_CLS( + _actor_master_config(tmp_path), _make_actor_args(save_state=save_state) + ) + + assert actor._train_steps == 7 + assert actor._trainer_version == 11 + assert actor._sampler._dispatch_index == 10 + + def test_fresh_start_defaults(self, tmp_path): + actor = _ACTOR_CLS(_actor_master_config(tmp_path), _make_actor_args()) + + assert actor._train_steps == 0 + assert actor._trainer_version == 0 + assert actor._sampler._dispatch_index == -1 + assert actor._consumed_samples == 0 + assert actor._current_epoch == 0 + assert actor._total_valid_tokens == 0 + + def test_old_checkpoint_without_total_valid_tokens(self, tmp_path): + # Older checkpoints may predate the total_valid_tokens key. + save_state = { + "consumed_samples": 10, + "current_step": 5, + "current_epoch": 0, + "total_steps": 5, + } + + actor = _ACTOR_CLS( + _actor_master_config(tmp_path), _make_actor_args(save_state=save_state) + ) + + assert actor._train_steps == 5 + assert actor._sampler._dispatch_index == 4 + assert actor._total_valid_tokens == 0 + + +# ── save trigger + write path ──────────────────────────────────────────────── + + +class TestSaveTrigger: + def test_saves_on_period_boundary_and_last_step(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=4, save_period=2) + trainer = _FakeTrainer() + + actor = _run_train_pump(mc, _make_actor_args(trainer=trainer)) + + assert actor._train_steps == 4 + ckpt_dir = tmp_path / "checkpoints" + # Finalized exactly at steps 2 and 4; no tmp_step_* leftovers. + assert _step_dir_names(ckpt_dir) == {"step_2", "step_4"} + + info_2 = _training_info(ckpt_dir, 2) + assert info_2["current_step"] == 2 + assert info_2["trainer_version"] == 2 + assert info_2["total_steps"] == 2 + assert info_2["consumed_samples"] == 4 # 2 prompts/step * 2 steps + # No validation ran, so the default val_reward is dropped. + assert "val_reward" not in info_2 + + info_4 = _training_info(ckpt_dir, 4) + assert info_4["current_step"] == 4 + assert info_4["consumed_samples"] == 8 + + # config.yaml is dumped next to training_info.json. + assert (ckpt_dir / "step_2" / "config.yaml").exists() + + # save_checkpoint was called into the tmp dir with all three paths. + assert len(trainer.save_calls) == 2 + first = trainer.save_calls[0] + assert first["weights_path"] == str( + ckpt_dir / "tmp_step_2" / "policy" / "weights" + ) + assert first["optimizer_path"] == str( + ckpt_dir / "tmp_step_2" / "policy" / "optimizer" + ) + assert first["tokenizer_path"] == str( + ckpt_dir / "tmp_step_2" / "policy" / "tokenizer" + ) + assert first["checkpointing_cfg"] is mc.checkpointing + assert trainer.save_calls[1]["weights_path"] == str( + ckpt_dir / "tmp_step_4" / "policy" / "weights" + ) + + # The async writers were waited on before each rename. + assert trainer.finalize_calls == 2 + + # The tmp dirs were finalized: policy/* survive under step_*. + assert (ckpt_dir / "step_2" / "policy" / "weights").is_dir() + assert (ckpt_dir / "step_2" / "policy" / "optimizer").is_dir() + assert (ckpt_dir / "step_4" / "policy" / "tokenizer").is_dir() + + def test_last_step_saves_off_period_boundary(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=3, save_period=2) + + _run_train_pump(mc, _make_actor_args()) + + # step 2 (boundary) + step 3 (last step), no step_1. + assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} + + def test_save_optimizer_false_gates_optimizer_path(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, save_optimizer=False + ) + trainer = _FakeTrainer() + + _run_train_pump(mc, _make_actor_args(trainer=trainer)) + + assert len(trainer.save_calls) == 1 + assert trainer.save_calls[0]["optimizer_path"] is None + ckpt_dir = tmp_path / "checkpoints" + assert (ckpt_dir / "step_2" / "policy" / "weights").is_dir() + assert not (ckpt_dir / "step_2" / "policy" / "optimizer").exists() + + def test_no_save_when_disabled(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=1, enabled=False + ) + trainer = _FakeTrainer() + + actor = _run_train_pump(mc, _make_actor_args(trainer=trainer)) + + assert actor._train_steps == 2 + assert trainer.save_calls == [] + assert _step_dir_names(tmp_path / "checkpoints") == set() + + def test_timeout_saves_and_stops_training_early(self, tmp_path): + # 0-second budget: the first check_save() fires; the pump must save + # at step 1 (off the period boundary) and break out of the loop. + mc = _actor_master_config( + tmp_path, + max_num_steps=4, + save_period=100, + checkpoint_must_save_by="00:00:00:00", + ) + trainer = _FakeTrainer() + + actor = _run_train_pump(mc, _make_actor_args(trainer=trainer)) + + assert actor._train_steps == 1 + assert len(trainer.save_calls) == 1 + assert _step_dir_names(tmp_path / "checkpoints") == {"step_1"} + + +class TestDataPlaneCheckpoint: + def test_saves_authoritative_tq_state_and_metadata_only_replay_index( + self, tmp_path + ): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + sample_ids = ["g0-0", "g0-1"] + dp_client = _FakeDPClient(sample_ids=sample_ids) + replay_metadata = { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": _PARTITION_ID, + "saved_capacity": 4, + "manifest_digest": "digest-1", + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=sample_ids, + fields=["input_ids"], + sequence_lengths=[16, 16], + tags=[ + {"weight_version": 0}, + {"weight_version": 0}, + ], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": None, + "group_id": "g0", + } + ], + } + buffer = _FakeTQBuffer(metadata_state=replay_metadata) + + _run_train_pump( + mc, + _make_actor_args(dp_client=dp_client, tq_buffer=buffer), + ) + + assert len(dp_client.save_calls) == 1 + save_call = dp_client.save_calls[0] + assert save_call["checkpoint_dir"] == str( + tmp_path / "checkpoints" / "tmp_step_1" / "data_plane" + ) + assert save_call["metadata"] == { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": 1, + "single_controller_trainer_version": 1, + "single_controller_epoch": 0, + "partition_id": _PARTITION_ID, + "sampler_name": "windowed", + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "replay_manifest_digest": "digest-1", + "replay_group_count": 1, + } + step_dir = tmp_path / "checkpoints" / "step_1" + assert (step_dir / "data_plane" / "metadata.json").is_file() + assert torch.load( + step_dir / REPLAY_BUFFER_METADATA_FILENAME, weights_only=False + ) == replay_metadata + assert not (step_dir / "replay_buffer.pt").exists() + assert buffer.metadata_state_dict_calls == [4] + + @pytest.mark.parametrize( + ("actual_sample_ids", "error_fragment"), + [ + (["g0-0"], r"missing=\['g0-1'\]"), + ( + ["g0-0", "g0-1", "orphan-0"], + r"unexpected=\['orphan-0'\]", + ), + ], + ) + def test_tq_save_rejects_inventory_mismatch( + self, tmp_path, actual_sample_ids, error_fragment + ): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + sample_ids = ["g0-0", "g0-1"] + replay_metadata = { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": _PARTITION_ID, + "saved_capacity": 4, + "manifest_digest": "digest-1", + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=sample_ids, + fields=["input_ids"], + sequence_lengths=[16, 16], + tags=[ + {"weight_version": 0}, + {"weight_version": 0}, + ], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": None, + "group_id": "g0", + } + ], + } + + with pytest.raises(RuntimeError, match=error_fragment): + _run_train_pump( + mc, + _make_actor_args( + dp_client=_FakeDPClient(sample_ids=actual_sample_ids), + tq_buffer=_FakeTQBuffer(metadata_state=replay_metadata), + ), + ) + + assert not (tmp_path / "checkpoints" / "step_1").exists() + + def test_gated_sampler_keeps_tq_checkpoint_in_shadow_mode(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + buffer_checkpoint=False, + data_plane_checkpoint=True, + ) + dp_client = _FakeDPClient() + buffer = _FakeTQBuffer() + + _run_train_pump( + mc, + _make_actor_args(dp_client=dp_client, tq_buffer=buffer), + ) + + assert dp_client.save_calls[0]["metadata"]["mode"] == "shadow" + step_dir = tmp_path / "checkpoints" / "step_1" + assert not (step_dir / REPLAY_BUFFER_METADATA_FILENAME).exists() + assert not (step_dir / "replay_buffer.pt").exists() + assert buffer.metadata_state_dict_calls == [] + + def test_tq_save_failure_aborts_checkpoint(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + data_plane_checkpoint=True, + ) + dp_client = _FakeDPClient(save_error=RuntimeError("injected TQ failure")) + + with pytest.raises(RuntimeError, match="injected TQ failure"): + _run_train_pump(mc, _make_actor_args(dp_client=dp_client)) + + assert not (tmp_path / "checkpoints" / "step_1").exists() + + def test_consumed_clear_waits_for_tq_save(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + data_plane_checkpoint=True, + ) + dp_client = _BlockingDPClient() + + async def _main() -> None: + actor = _ACTOR_CLS(mc, _make_actor_args(dp_client=dp_client)) + actor._train_steps = 1 + actor._trainer_version = 1 + save_task = asyncio.create_task(actor._save_checkpoint({"loss": 1.0})) + started = await asyncio.to_thread(dp_client.save_started.wait, 30.0) + assert started + + clear_task = asyncio.create_task( + actor._clear_data_plane_samples(["sample-0"]) + ) + await asyncio.sleep(0) + assert dp_client.clear_calls == [] + + dp_client.release_save.set() + await save_task + await clear_task + actor._checkpointer.shutdown() + + asyncio.run(_main()) + assert dp_client.clear_calls == [(["sample-0"], _PARTITION_ID)] + + def test_consumed_clear_does_not_block_actor_event_loop(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=1, save_period=1) + dp_client = _FakeDPClient() + + async def _main() -> int: + actor = _ACTOR_CLS(mc, _make_actor_args(dp_client=dp_client)) + event_loop_thread_id = threading.get_ident() + await actor._clear_data_plane_samples(["sample-0"]) + actor._checkpointer.shutdown() + return event_loop_thread_id + + event_loop_thread_id = asyncio.run(_main()) + assert dp_client.clear_thread_ids + assert dp_client.clear_thread_ids[0] != event_loop_thread_id + + +# ── async-save finalization ────────────────────────────────────────────────── + + +class TestAsyncSaveFinalization: + def test_rename_deferred_until_async_writes_finish(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2) + trainer = _GatedFinalizeTrainer() + + actor = _run_train_pump(mc, _make_actor_args(trainer=trainer), flush=False) + + # The async writer hasn't finished: the checkpoint must still be a + # tmp dir, invisible to resume lookups. + ckpt_dir = tmp_path / "checkpoints" + assert (ckpt_dir / "tmp_step_2").is_dir() + assert not (ckpt_dir / "step_2").exists() + + trainer.release.set() + actor._checkpointer.shutdown() + + assert (ckpt_dir / "step_2").is_dir() + assert not (ckpt_dir / "tmp_step_2").exists() + assert trainer.finalize_calls == 1 + + def test_failed_background_finalization_raises_at_next_save(self, tmp_path): + # The step-2 checkpoint's background finalization fails; the failure + # must surface at the step-4 save's finalize_pending, not vanish. + mc = _actor_master_config(tmp_path, max_num_steps=4, save_period=2) + trainer = _FailingFinalizeTrainer() + + with pytest.raises(RuntimeError, match="finalization failed"): + _run_train_pump(mc, _make_actor_args(trainer=trainer), flush=False) + + def test_save_checkpoint_records_val_metrics(self, tmp_path): + # Direct _save_checkpoint call: the val_metrics parameter feeds both + # val_reward and a val:* metric_name (wired by the future validation + # loop; the pump passes None today). + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, metric_name="val:accuracy" + ) + + async def _main(): + actor = _ACTOR_CLS(mc, _make_actor_args()) + actor._train_steps = 2 + await actor._save_checkpoint({"loss": 1.0}, val_metrics={"accuracy": 0.75}) + actor._checkpointer.shutdown() + return actor + + actor = asyncio.run(_main()) + + info = _training_info(tmp_path / "checkpoints", 2) + assert info["val_reward"] == 0.75 + assert info["val:accuracy"] == 0.75 + assert actor._save_state["val_reward"] == 0.75 + + +# ── metric_name behavior ───────────────────────────────────────────────────── + + +class TestMetricName: + def test_val_metric_without_validation_warns_and_still_saves(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, metric_name="val:accuracy" + ) + + with pytest.warns(UserWarning, match="no val metrics were collected"): + _run_train_pump(mc, _make_actor_args()) + + ckpt_dir = tmp_path / "checkpoints" + assert _step_dir_names(ckpt_dir) == {"step_2"} + assert "val:accuracy" not in _training_info(ckpt_dir, 2) + + def test_train_metric_lands_in_training_info(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, metric_name="train:loss" + ) + trainer = _FakeTrainer(step_metrics={"loss": 0.5}) + + _run_train_pump(mc, _make_actor_args(trainer=trainer)) + + info = _training_info(tmp_path / "checkpoints", 2) + assert info["train:loss"] == 0.5 + + def test_train_metric_missing_key_raises(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, metric_name="train:not_a_metric" + ) + + with pytest.raises(ValueError, match="not found in train metrics"): + _run_train_pump(mc, _make_actor_args()) + + def test_metric_name_requires_train_or_val_prefix(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, metric_name="reward" + ) + + with pytest.raises(AssertionError, match="must start with"): + _run_train_pump(mc, _make_actor_args()) + + +# ── setup resume-path wiring ───────────────────────────────────────────────── + + +_STEP_3_SAVE_STATE = { + "consumed_samples": 24, + "current_step": 3, + "current_epoch": 1, + "total_steps": 3, + "total_valid_tokens": 999, +} + + +def _write_checkpoint( + ckpt_dir: Path, + step: int, + save_state: dict[str, Any], + *, + with_optimizer: bool = True, + dataloader_state: Optional[dict[str, Any]] = None, + config: Optional[dict[str, Any]] = None, +) -> Path: + step_dir = ckpt_dir / f"step_{step}" + (step_dir / "policy" / "weights").mkdir(parents=True) + if with_optimizer: + (step_dir / "policy" / "optimizer").mkdir(parents=True) + with open(step_dir / "training_info.json", "w") as f: + json.dump(save_state, f) + if dataloader_state is not None: + torch.save(dataloader_state, step_dir / "train_dataloader.pt") + if config is not None: + with open(step_dir / "config.yaml", "w") as f: + yaml.safe_dump(config, f) + return step_dir + + +def _setup_master_config(checkpoint_dir: str) -> MasterConfig: + """Partially-populated MasterConfig for setup_single_controller tests. + + Same shape as test_single_controller_setup._make_master_config, plus the + checkpointing block setup now reads. + """ + return MasterConfig.model_construct( + data_plane={"enabled": True, "impl": "transfer_queue"}, + data={ + "use_multiple_dataloader": False, + "shuffle": False, + "num_workers": 0, + "train": [{"env_name": "math"}], + }, + grpo={ + "max_num_steps": 100, + "max_num_epochs": 1, + "num_prompts_per_step": 4, + "num_generations_per_prompt": 2, + "max_rollout_turns": 1, + "seed": 42, + "val_period": 0, + "val_at_start": False, + "val_at_end": False, + }, + policy={ + "train_global_batch_size": 8, + "max_total_sequence_length": 32, + "tokenizer": {"use_fastokens": False}, + "megatron_cfg": {"enabled": False}, + "generation": { + "backend": "vllm", + "colocated": {"enabled": True, "resources": {}}, + }, + }, + loss_fn=ClippedPGLossConfig(), + env={}, + async_rl=AsyncRLConfig( + min_groups_for_streaming_train=4, + max_buffered_rollouts=8, + ), + checkpointing={ + "enabled": True, + "checkpoint_dir": checkpoint_dir, + "metric_name": None, + "higher_is_better": True, + "keep_top_k": None, + "save_period": 2, + "save_optimizer": True, + "checkpoint_must_save_by": None, + }, + ) + + +class TestGetResumePaths: + def test_resume_paths_from_fixture_layout(self, tmp_path): + step_dir = _write_checkpoint(tmp_path, 3, _STEP_3_SAVE_STATE) + + weights_path, optimizer_path = CheckpointManager.get_resume_paths(str(step_dir)) + + assert weights_path == step_dir / "policy" / "weights" + assert optimizer_path == step_dir / "policy" / "optimizer" + + def test_resume_paths_without_optimizer_state(self, tmp_path): + step_dir = _write_checkpoint( + tmp_path, 3, _STEP_3_SAVE_STATE, with_optimizer=False + ) + + with pytest.warns(UserWarning, match="Optimizer state not found"): + weights_path, optimizer_path = CheckpointManager.get_resume_paths( + str(step_dir) + ) + + assert weights_path == step_dir / "policy" / "weights" + assert optimizer_path is None + + def test_no_checkpoint_gives_none(self): + assert CheckpointManager.get_resume_paths(None) == (None, None) + + +class TestSetupResumeWiring: + def test_setup_forwards_latest_resume_paths( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + ckpt_dir = tmp_path / "ckpts" + _write_checkpoint(ckpt_dir, 1, {**_STEP_3_SAVE_STATE, "current_step": 1}) + step_3 = _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE) + mc = _setup_master_config(str(ckpt_dir)) + + actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + # Latest checkpoint (step_3) wins; its paths reach the trainer factory. + trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs + assert trainer_kwargs["weights_path"] == step_3 / "policy" / "weights" + assert trainer_kwargs["optimizer_path"] == step_3 / "policy" / "optimizer" + # training_info.json is loaded into the actor args for the actor. + assert actor_args.save_state == _STEP_3_SAVE_STATE + assert actor_args.last_checkpoint_path == str(step_3) + + def test_setup_fresh_start_passes_none_paths( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + ckpt_dir = tmp_path / "empty_ckpts" + ckpt_dir.mkdir() + mc = _setup_master_config(str(ckpt_dir)) + + actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs + assert trainer_kwargs["weights_path"] is None + assert trainer_kwargs["optimizer_path"] is None + assert actor_args.save_state == _default_grpo_save_state() + assert actor_args.last_checkpoint_path is None + + def test_setup_forwards_pretrained_checkpoint( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + mc = _setup_master_config(str(tmp_path / "ckpts")) + pretrained = {"path": "/some/ckpt", "format": "megatron_bridge"} + mc.checkpointing["pretrained_checkpoint"] = pretrained + + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert mc.policy["pretrained_checkpoint"] == pretrained + + +def _make_int_dataloader() -> StatefulDataLoader: + """8 ints, batch_size=2 → batches [0,1], [2,3], [4,5], [6,7].""" + return StatefulDataLoader( + list(range(8)), + batch_size=2, + shuffle=False, + drop_last=True, + num_workers=0, + ) + + +class TestDataloaderState: + def test_save_writes_dataloader_state(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2) + save_state = _default_grpo_save_state() + save_state["current_epoch"] = 3 + dataloader = _FakeDataloader(state={"fake_position": 7}) + + _run_train_pump( + mc, _make_actor_args(save_state=save_state, dataloader=dataloader) + ) + + ckpt_dir = tmp_path / "checkpoints" + dl_state_path = ckpt_dir / "step_2" / "train_dataloader.pt" + assert dl_state_path.exists() + # The snapshot taken at save time round-trips through torch.save. + assert torch.load(dl_state_path) == {"fake_position": 7} + # current_epoch flows save_state → actor → training_info.json. + assert _training_info(ckpt_dir, 2)["current_epoch"] == 3 + + def test_stateful_dataloader_position_roundtrip(self, tmp_path): + data_config = {"train": [{"dataset_name": "math_train"}]} + dataloader = _make_int_dataloader() + it = iter(dataloader) + assert [next(it).tolist() for _ in range(2)] == [[0, 1], [2, 3]] + + step_dir = _write_checkpoint( + tmp_path, + 5, + _default_grpo_save_state(), + dataloader_state=dataloader.state_dict(), + config={"data": {"train": [{"dataset_name": "math_train"}]}}, + ) + + restored = _make_int_dataloader() + load_dataloader_state(restored, str(step_dir), data_config) + + # Resumes at batch k+1, not from the top. + assert next(iter(restored)).tolist() == [4, 5] + + def test_dataset_swap_skips_restore(self, tmp_path): + dataloader = _make_int_dataloader() + it = iter(dataloader) + assert [next(it).tolist() for _ in range(2)] == [[0, 1], [2, 3]] + + step_dir = _write_checkpoint( + tmp_path, + 5, + _default_grpo_save_state(), + dataloader_state=dataloader.state_dict(), + config={"data": {"train": [{"dataset_name": "old_dataset"}]}}, + ) + + restored = _make_int_dataloader() + load_dataloader_state( + restored, str(step_dir), {"train": [{"dataset_name": "new_dataset"}]} + ) + + # Restore skipped on dataset swap: the new dataset starts from index 0. + assert next(iter(restored)).tolist() == [0, 1] + + def test_setup_restores_dataloader_state( + self, + patched_factories, # noqa: F811 + tmp_path, + ): + ckpt_dir = tmp_path / "ckpts" + sentinel = {"fake_position": 123} + _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE, dataloader_state=sentinel) + mc = _setup_master_config(str(ckpt_dir)) + + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + fake_dataloader = patched_factories["dataloader"] + fake_dataloader.load_state_dict.assert_called_once() + assert fake_dataloader.load_state_dict.call_args.args[0] == sentinel + + def test_setup_missing_dataloader_state_starts_fresh( + self, + patched_factories, # noqa: F811 + tmp_path, + monkeypatch, + ): + # Checkpoint with training_info.json + policy/ but no + # train_dataloader.pt. Setup must warn and keep the fresh position + # while still forwarding the weight resume paths. + ckpt_dir = tmp_path / "ckpts" + step_3 = _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE) + mc = _setup_master_config(str(ckpt_dir)) + printed: list[str] = [] + # Repo addopts run pytest with -s, so capsys sees nothing; record + # print calls instead. + monkeypatch.setattr( + "builtins.print", + lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + ) + + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + patched_factories["dataloader"].load_state_dict.assert_not_called() + assert any("No dataloader state found" in line for line in printed) + trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs + assert trainer_kwargs["weights_path"] == step_3 / "policy" / "weights" + + +# ── replay buffer persistence ──────────────────────────────────────────────── + + +class TestReplayBufferPersistence: + def test_checkpoint_capable_sampler_without_native_tq_is_rejected(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=2, + save_period=2, + buffer_checkpoint=True, + ) + + with pytest.raises( + ValueError, + match="replay-checkpoint-capable sampler requires", + ): + _ACTOR_CLS(mc, _make_actor_args()) + + def test_no_replay_buffer_with_gated_sampler(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, buffer_checkpoint=False + ) + buffer = _FakeTQBuffer() + + _run_train_pump(mc, _make_actor_args(tq_buffer=buffer)) + + ckpt_dir = tmp_path / "checkpoints" + assert (ckpt_dir / "step_2" / "training_info.json").exists() + assert not (ckpt_dir / "step_2" / "replay_buffer.pt").exists() + assert not ( + ckpt_dir / "step_2" / REPLAY_BUFFER_METADATA_FILENAME + ).exists() + + def test_run_rejects_legacy_replay_file(self, tmp_path): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + torch.save( + {"groups": ["legacy"]}, ckpt_dir / LEGACY_REPLAY_BUFFER_FILENAME + ) + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + buffer = _FakeTQBuffer() + + with pytest.raises(RuntimeError, match="legacy replay_buffer.pt"): + _run_actor_run( + mc, + _make_actor_args( + tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir) + ), + ) + + assert buffer.load_calls == [] + + def test_run_restores_native_tq_replay_metadata_without_payload_reput( + self, tmp_path + ): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + sample_ids = ["g0-0", "g0-1", "g1-0", "g1-1"] + groups = [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=[f"g{i}-0", f"g{i}-1"], + sequence_lengths=[16, 16], + tags=[{"weight_version": 0}, {"weight_version": 0}], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": i, + "group_id": f"g{i}", + } + for i in range(2) + ] + envelope = {"groups": groups} + torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + tq_metadata = { + "replay_manifest_digest": "digest-1", + "replay_group_count": 2, + } + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + buffer = _FakeTQBuffer(load_return=2) + + actor, result = _run_actor_run( + mc, + _make_actor_args( + tq_buffer=buffer, + dp_client=_FakeDPClient(sample_ids=sample_ids), + last_checkpoint_path=str(ckpt_dir), + data_plane_checkpoint_metadata=tq_metadata, + ), + ) + + assert buffer.load_calls == [ + { + "state": envelope, + "max_groups": 4, + "expected_partition_id": _PARTITION_ID, + "expected_group_size": 2, + "expected_manifest_digest": "digest-1", + } + ] + assert actor._buffer_capacity._value == 2 + assert result["train_steps"] == 0 + + @pytest.mark.parametrize( + ("actual_sample_ids", "error_fragment"), + [ + (["g0-0"], r"missing=\['g0-1'\]"), + ( + ["g0-0", "g0-1", "orphan-0"], + r"unexpected=\['orphan-0'\]", + ), + ], + ) + def test_native_restore_rejects_tq_inventory_mismatch( + self, tmp_path, actual_sample_ids, error_fragment + ): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + envelope = { + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=["g0-0", "g0-1"], + sequence_lengths=[16, 16], + tags=[{"weight_version": 0}, {"weight_version": 0}], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": 0, + "group_id": "g0", + } + ] + } + torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + tq_metadata = { + "replay_manifest_digest": "digest-1", + "replay_group_count": 1, + } + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + + with pytest.raises(RuntimeError, match=error_fragment): + _run_actor_run( + mc, + _make_actor_args( + tq_buffer=_FakeTQBuffer(load_return=1), + dp_client=_FakeDPClient(sample_ids=actual_sample_ids), + last_checkpoint_path=str(ckpt_dir), + data_plane_checkpoint_metadata=tq_metadata, + ), + ) + + def test_native_replay_metadata_requires_setup_side_tq_restore(self, tmp_path): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + torch.save({"groups": []}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + + with pytest.raises(RuntimeError, match="native TQ checkpoint was not restored"): + _run_actor_run( + mc, + _make_actor_args(last_checkpoint_path=str(ckpt_dir)), + ) + + def test_run_missing_native_replay_metadata_starts_empty( + self, tmp_path, monkeypatch + ): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + buffer = _FakeTQBuffer() + printed: list[str] = [] + monkeypatch.setattr( + "builtins.print", + lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + ) + + actor, _ = _run_actor_run( + mc, + _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)), + ) + + assert buffer.load_calls == [] + assert actor._buffer_capacity._value == 4 # zero permits consumed + assert any("No native replay metadata found" in line for line in printed) + + def test_run_rejects_native_replay_state_with_gated_sampler(self, tmp_path): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + torch.save({"groups": []}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + mc = _actor_master_config( + tmp_path, max_num_steps=0, buffer_checkpoint=False + ) + + with pytest.raises(RuntimeError, match="does not support replay-buffer"): + _run_actor_run( + mc, + _make_actor_args(last_checkpoint_path=str(ckpt_dir)), + ) diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 3a71a23b6a..e36f724913 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -16,11 +16,21 @@ from __future__ import annotations +from typing import Any, Optional from unittest.mock import MagicMock, patch import pytest import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.grpo import _default_grpo_save_state from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, @@ -77,7 +87,18 @@ def _make_master_config( "colocated": {"enabled": colocated, "resources": {}}, }, }, - checkpointing={"enabled": False}, + # Full block: setup builds a CheckpointManager unconditionally (resume + # lookup), which indexes these keys directly. Nothing is written while + # enabled=False and the dir doesn't exist. + checkpointing={ + "enabled": False, + "checkpoint_dir": "results/_sc_setup_test_ckpt", + "metric_name": None, + "higher_is_better": False, + "keep_top_k": None, + "save_period": 10, + "save_optimizer": False, + }, loss_fn=ClippedPGLossConfig(), env=env if env is not None else {}, async_rl=AsyncRLConfig( @@ -87,6 +108,27 @@ def _make_master_config( ) +def _native_tq_metadata( + *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1 +) -> dict[str, Any]: + return { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": step, + "single_controller_trainer_version": ( + step if trainer_version is None else trainer_version + ), + "single_controller_epoch": epoch, + "partition_id": "rollout_data", + "sampler_name": "in_order", + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "replay_manifest_digest": "digest-1", + "replay_group_count": 2, + } + + @pytest.fixture def patched_factories(): """Patch every external factory setup calls. @@ -207,6 +249,37 @@ def test_raises_when_data_plane_disabled(self): with pytest.raises(ValueError, match="data_plane.enabled=True"): setup_single_controller(mc, MagicMock()) + def test_rejects_mooncake_data_plane_checkpointing(self): + mc = _make_master_config() + mc.data_plane.update( + { + "backend": "mooncake_cpu", + "checkpointing_enabled": True, + } + ) + with pytest.raises(NotImplementedError, match="backend='simple'"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_rejects_windowed_checkpointing_without_native_tq(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.async_rl.sampler = WindowedSamplerConfig(max_staleness_versions=1) + mc.data_plane.update( + { + "backend": "simple", + "checkpointing_enabled": False, + } + ) + + with pytest.raises( + ValueError, + match=( + "replay-checkpoint-capable sampler requires " + "data_plane.checkpointing_enabled=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_multiple_dataloader_not_supported(self): mc = _make_master_config(use_multiple_dataloader=True) with pytest.raises(NotImplementedError, match="use_multiple_dataloader"): @@ -443,3 +516,149 @@ def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend): ): setup_single_controller(mc, MagicMock(pad_token_id=0)) mock_spinup.assert_not_called() + + +class TestNativeTQRecoverySetup: + def test_setup_loads_tq_before_creating_single_controller_client( + self, tmp_path, patched_factories + ): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + save_state = _default_grpo_save_state() + save_state["current_step"] = 3 + save_state["current_epoch"] = 1 + policy = patched_factories["_build_trainer"].return_value + events: list[str] = [] + policy.load_data_plane_checkpoint.side_effect = ( + lambda checkpoint_dir: events.append("load") or _native_tq_metadata() + ) + patched_factories["build_data_plane_client"].side_effect = ( + lambda *args, **kwargs: events.append("build") + or MagicMock(name="dp_client") + ) + checkpointer = MagicMock() + checkpointer.get_latest_checkpoint_path.return_value = str(checkpoint_path) + checkpointer.load_training_info.return_value = save_state + checkpointer.get_resume_paths.return_value = (None, None) + mc = _make_master_config() + + with patch.object( + sc_setup_mod, "CheckpointManager", return_value=checkpointer + ): + actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0)) + + assert events == ["load", "build"] + assert actor_args.data_plane_checkpoint_metadata == _native_tq_metadata() + + def test_loads_authoritative_tq_checkpoint_when_metadata_sidecar_exists( + self, tmp_path + ): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + policy = MagicMock() + metadata = _native_tq_metadata() + policy.load_data_plane_checkpoint.return_value = metadata + save_state = { + "current_step": 3, + "current_epoch": 1, + } + + restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=str(checkpoint_path), + save_state=save_state, + partition_id="rollout_data", + sampler_name="in_order", + ) + + assert restored == metadata + policy.load_data_plane_checkpoint.assert_called_once_with( + checkpoint_path / DATA_PLANE_CHECKPOINT_DIR + ) + + def test_validates_trainer_version_independently_from_train_step(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + policy = MagicMock() + metadata = _native_tq_metadata(step=3, trainer_version=7) + policy.load_data_plane_checkpoint.return_value = metadata + + restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=str(checkpoint_path), + save_state={ + "current_step": 3, + "trainer_version": 7, + "current_epoch": 1, + }, + partition_id="rollout_data", + sampler_name="in_order", + ) + + assert restored == metadata + + def test_legacy_replay_checkpoint_is_rejected(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + checkpoint_path.mkdir() + (checkpoint_path / LEGACY_REPLAY_BUFFER_FILENAME).touch() + policy = MagicMock() + + with pytest.raises(RuntimeError, match="legacy replay_buffer.pt"): + sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=str(checkpoint_path), + save_state={"current_step": 3, "current_epoch": 1}, + partition_id="rollout_data", + sampler_name="in_order", + ) + + policy.load_data_plane_checkpoint.assert_not_called() + + def test_checkpoint_without_replay_artifacts_does_not_load_tq(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + checkpoint_path.mkdir() + policy = MagicMock() + + restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=str(checkpoint_path), + save_state={"current_step": 3, "current_epoch": 1}, + partition_id="rollout_data", + sampler_name="in_order", + ) + + assert restored is None + policy.load_data_plane_checkpoint.assert_not_called() + + def test_metadata_sidecar_requires_matching_tq_directory(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + checkpoint_path.mkdir() + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + + with pytest.raises(FileNotFoundError, match="matching native TQ checkpoint"): + sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + MagicMock(), + last_checkpoint_path=str(checkpoint_path), + save_state={"current_step": 3, "current_epoch": 1}, + partition_id="rollout_data", + sampler_name="in_order", + ) + + def test_rejects_tq_checkpoint_from_different_training_step(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + policy = MagicMock() + policy.load_data_plane_checkpoint.return_value = _native_tq_metadata(step=2) + + with pytest.raises(ValueError, match="does not match the trainer checkpoint"): + sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + policy, + last_checkpoint_path=str(checkpoint_path), + save_state={"current_step": 3, "current_epoch": 1}, + partition_id="rollout_data", + sampler_name="in_order", + ) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 76cb919849..1e1336b10f 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -17,13 +17,20 @@ from __future__ import annotations import asyncio +import threading from typing import Any import pytest import torch import nemo_rl.algorithms.async_utils.replay_buffer as _replay_buffer_module -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, + DataPlaneCheckpointBarrier, + TQReplayBuffer, + replay_manifest_digest, +) from nemo_rl.data_plane import KVBatchMeta from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord @@ -63,6 +70,8 @@ def __init__(self, partition_id: str = "rollout_data") -> None: self._rows: dict[str, dict[str, Any]] = {} self.put_calls: list[dict[str, Any]] = [] self.clear_calls: list[list[str]] = [] + self.clear_thread_ids: list[int] = [] + self.get_calls: list[dict[str, Any]] = [] def put_samples( self, @@ -93,11 +102,34 @@ def put_samples( def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: assert partition_id == self._partition_id + self.clear_thread_ids.append(threading.get_ident()) ids = list(sample_ids) if sample_ids is not None else list(self._rows) self.clear_calls.append(list(ids)) for sid in ids: self._rows.pop(sid, None) + def list_sample_ids(self, partition_id: str) -> list[str]: + assert partition_id == self._partition_id + return sorted(self._rows) + + def get_samples( + self, + sample_ids: list[str], + partition_id: str, + select_fields: list[str] | None = None, + ) -> dict[str, Any]: + assert partition_id == self._partition_id + self.get_calls.append( + { + "sample_ids": list(sample_ids), + "select_fields": ( + list(select_fields) if select_fields is not None else None + ), + } + ) + # Opaque payload used by tests that inspect direct DataPlane reads. + return {"payload_for": list(sample_ids)} + def depth(self) -> int: return len(self._rows) @@ -132,10 +164,18 @@ def _make_record() -> PromptGroupRecord: ) -def _make_buffer(dp: FakeDataPlaneClient) -> TQReplayBuffer: - return TQReplayBuffer( +def _make_buffer( + dp: FakeDataPlaneClient, + *, + checkpoint_barrier: DataPlaneCheckpointBarrier | None = None, +) -> TQReplayBuffer: + buffer = TQReplayBuffer( dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0} ) + buffer.set_data_plane_checkpoint_barrier( + checkpoint_barrier or DataPlaneCheckpointBarrier() + ) + return buffer def _add_group( @@ -154,7 +194,87 @@ def _add_group( ) +class TestDataPlaneCheckpointBarrier: + def test_mutations_run_concurrently_without_checkpoint(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + both_entered = asyncio.Event() + release = asyncio.Event() + active = 0 + + async def mutate() -> None: + nonlocal active + async with barrier.mutation(): + active += 1 + if active == 2: + both_entered.set() + await release.wait() + active -= 1 + + tasks = [asyncio.create_task(mutate()) for _ in range(2)] + await asyncio.wait_for(both_entered.wait(), timeout=5.0) + assert active == 2 + release.set() + await asyncio.gather(*tasks) + + asyncio.run(exercise()) + + def test_checkpoint_waits_for_active_mutation(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + mutation_entered = asyncio.Event() + release_mutation = asyncio.Event() + checkpoint_entered = asyncio.Event() + + async def mutate() -> None: + async with barrier.mutation(): + mutation_entered.set() + await release_mutation.wait() + + async def checkpoint() -> None: + async with barrier.checkpoint(): + checkpoint_entered.set() + + mutation_task = asyncio.create_task(mutate()) + await mutation_entered.wait() + checkpoint_task = asyncio.create_task(checkpoint()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + + release_mutation.set() + await asyncio.gather(mutation_task, checkpoint_task) + assert checkpoint_entered.is_set() + + asyncio.run(exercise()) + + class TestTQReplayBufferReserveCommit: + def test_commit_waits_for_active_checkpoint(self): + async def exercise() -> None: + dp = FakeDataPlaneClient() + checkpoint_barrier = DataPlaneCheckpointBarrier() + buf = _make_buffer(dp, checkpoint_barrier=checkpoint_barrier) + group_id = buf.reserve(weight_version=3) + + async with checkpoint_barrier.checkpoint(): + commit_task = asyncio.create_task( + buf.commit( + group_id, + _make_record(), + start_weight_version=3, + end_weight_version=3, + ) + ) + await asyncio.sleep(0) + assert dp.put_calls == [] + assert buf.ready_list == [False] + + await commit_task + assert len(dp.put_calls) == 1 + assert buf.ready_list == [True] + + asyncio.run(exercise()) + def test_commit_clears_rows_when_put_raises_after_writing(self): dp = FailAfterPutDataPlaneClient() buf = _make_buffer(dp) @@ -285,6 +405,56 @@ def test_commit_appends_multiple_records_in_order(self): class TestTQReplayBufferRemove: + def test_dp_clear_fails_without_bound_checkpoint_barrier(self): + dp = FakeDataPlaneClient() + buf = TQReplayBuffer( + dp, + partition_id="rollout_data", + pad_value_dict={"token_ids": 0}, + ) + + with pytest.raises(RuntimeError, match="must be bound"): + _run(buf._clear_samples(sample_ids=["sample-1"])) + + assert dp.clear_calls == [] + + def test_dp_clear_waits_for_active_checkpoint(self): + async def exercise() -> None: + dp = FakeDataPlaneClient() + checkpoint_barrier = DataPlaneCheckpointBarrier() + buf = _make_buffer(dp, checkpoint_barrier=checkpoint_barrier) + group_id = buf.reserve(weight_version=0) + await buf.commit( + group_id, + _make_record(), + start_weight_version=0, + end_weight_version=0, + ) + + async with checkpoint_barrier.checkpoint(): + remove_task = asyncio.create_task( + buf.remove([0], remove_in_dp=True) + ) + await asyncio.sleep(0) + assert dp.clear_calls == [] + + await remove_task + assert dp.clear_calls == [dp.put_calls[0]["sample_ids"]] + + asyncio.run(exercise()) + + def test_dp_clear_does_not_block_actor_event_loop(self): + async def exercise() -> tuple[FakeDataPlaneClient, int]: + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + event_loop_thread_id = threading.get_ident() + await buf._clear_samples(sample_ids=["sample-1"]) + return dp, event_loop_thread_id + + dp, event_loop_thread_id = asyncio.run(exercise()) + assert dp.clear_thread_ids + assert dp.clear_thread_ids[0] != event_loop_thread_id + def test_remove_drops_indices_and_clears_dp_when_requested(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) @@ -363,3 +533,195 @@ def test_size_and_len(self): _run(buf.remove([0], remove_in_dp=True)) assert buf.size() == 1 assert len(buf) == 1 + + +# ── state_dict / load_state_dict (checkpointing) ───────────────────────────── + + +def _group_id_of(meta: KVBatchMeta) -> str: + head, _, _ = meta.sample_ids[0].rpartition("_g") + return head + + +def _make_group_entry( + group_id: str, + weight: int, + *, + n: int = _N_GENS, + target_step: int | None = None, + sample_ids: list[str] | None = None, + sequence_lengths: list[int] | None = None, + partition_id: str = "rollout_data", +) -> dict[str, Any]: + """Hand-built envelope group (bypasses commit) for preflight tests.""" + sids = ( + list(sample_ids) + if sample_ids is not None + else [f"{group_id}_g{i}" for i in range(n)] + ) + meta = KVBatchMeta( + partition_id=partition_id, + task_name="train", + sample_ids=sids, + fields=["input_ids", "input_lengths", "total_reward"], + sequence_lengths=( + sequence_lengths if sequence_lengths is not None else [3] * len(sids) + ), + tags=[{"weight_version": weight}] * len(sids), + ) + return { + "meta": meta, + "start_weight": weight, + "end_weight": weight, + "target_step": target_step, + "group_id": group_id, + } + + +def _make_metadata_envelope( + groups: list[dict[str, Any]], + *, + partition_id: str = "rollout_data", + saved_capacity: int = 8, +) -> dict[str, Any]: + metadata_groups = [dict(group) for group in groups] + return { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": partition_id, + "saved_capacity": saved_capacity, + "manifest_digest": replay_manifest_digest(metadata_groups), + "groups": metadata_groups, + } + + +def _load( + buf: TQReplayBuffer, + state: dict[str, Any], + *, + max_groups: int = 8, + expected_partition_id: str = "rollout_data", + expected_group_size: int = _N_GENS, + expected_manifest_digest: str | None = None, +) -> int: + if expected_manifest_digest is None: + expected_manifest_digest = str(state.get("manifest_digest", "")) + return _run( + buf.load_state_dict( + state, + max_groups=max_groups, + expected_partition_id=expected_partition_id, + expected_group_size=expected_group_size, + expected_manifest_digest=expected_manifest_digest, + ) + ) + + +class TestTQReplayBufferStateDict: + def test_metadata_state_dict_omits_tensors_and_data_plane_reads(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + metas = [_add_group(buf, weight=w) for w in (1, 2)] + buf.reserve(weight_version=3) + + state = buf.metadata_state_dict(saved_capacity=8) + + assert state["schema_version"] == REPLAY_BUFFER_METADATA_SCHEMA_VERSION + assert state["storage"] == REPLAY_BUFFER_METADATA_STORAGE + assert len(state["groups"]) == 2 + assert all("fields_data" not in group for group in state["groups"]) + assert [group["meta"].sample_ids for group in state["groups"]] == [ + list(meta.sample_ids) for meta in metas + ] + assert state["manifest_digest"] == replay_manifest_digest(state["groups"]) + assert dp.get_calls == [] + + def test_native_tq_round_trip_restores_index_without_reputting_rows(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + metas = [_add_group(buf, weight=w) for w in (1, 2)] + state = buf.metadata_state_dict(saved_capacity=8) + + restored_dp = FakeDataPlaneClient() + restored_buf = _make_buffer(restored_dp) + restored = _load( + restored_buf, + state, + expected_manifest_digest=state["manifest_digest"], + ) + + assert restored == 2 + assert restored_buf.start_weight_list == [1, 2] + assert restored_buf.ready_list == [True, True] + assert [meta.sample_ids for meta in restored_buf.meta_list] == [ + list(meta.sample_ids) for meta in metas + ] + assert restored_dp.put_calls == [] + + +class TestTQReplayBufferLoadPreflight: + """Malformed envelopes raise ValueError before any DataPlane write.""" + + def _assert_rejected(self, state: dict[str, Any], match: str, **load_kwargs): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + with pytest.raises(ValueError, match=match): + _load(buf, state, **load_kwargs) + assert dp.put_calls == [] + assert buf.size() == 0 + + def test_missing_envelope_keys(self): + self._assert_rejected({"groups": []}, match="missing required keys") + + def test_partition_id_mismatch(self): + state = _make_metadata_envelope([], partition_id="other_partition") + self._assert_rejected(state, match="partition_id mismatch") + + def test_group_missing_keys(self): + state = _make_metadata_envelope([_make_group_entry("g0", weight=1)]) + del state["groups"][0]["group_id"] + self._assert_rejected(state, match="group missing keys") + + def test_group_with_tensor_payload_is_rejected(self): + group = _make_group_entry("g0", weight=1) + group["fields_data"] = {"input_ids": torch.ones(2, 3)} + state = _make_metadata_envelope([group]) + self._assert_rejected(state, match="must not contain fields_data") + + def test_group_misaligned_sequence_lengths(self): + group = _make_group_entry("g0", weight=1, sequence_lengths=[3]) + self._assert_rejected( + _make_metadata_envelope([group]), match="misaligned" + ) + + def test_group_size_mismatch(self): + state = _make_metadata_envelope([_make_group_entry("g0", weight=1, n=2)]) + self._assert_rejected(state, match="misaligned", expected_group_size=3) + + def test_duplicate_sample_ids_across_groups(self): + g0 = _make_group_entry("g0", weight=1) + g1 = _make_group_entry( + "g1", weight=2, sample_ids=["g0_g0", "g1_g1"] + ) # g0_g0 collides + self._assert_rejected( + _make_metadata_envelope([g0, g1]), match="duplicate sample_id" + ) + + def test_metadata_only_restore_rejects_tq_digest_mismatch(self): + state = _make_metadata_envelope([_make_group_entry("g0", weight=1)]) + self._assert_rejected( + state, + match="does not match the loaded TQ checkpoint", + expected_manifest_digest="wrong-digest", + ) + + def test_metadata_only_restore_rejects_capacity_truncation(self): + state = _make_metadata_envelope( + [_make_group_entry(f"g{w}", weight=w) for w in (1, 2, 3)] + ) + self._assert_rejected( + state, + match="more replay groups than the current buffer capacity", + max_groups=2, + expected_manifest_digest=state["manifest_digest"], + ) diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index fe2b59ebbd..2d0998fd85 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -28,6 +28,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.grpo import _default_grpo_save_state from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.algorithms.single_controller_utils.config import ( AsyncRLConfig, @@ -320,6 +321,18 @@ def test_train_pump_drives_mcore_training_step( "mlflow_enabled": False, "monitor_gpus": False, }, + # Actor __init__ builds a CheckpointManager + TimeoutChecker from + # this block; enabled=False keeps the run write-free. + checkpointing={ + "enabled": False, + "checkpoint_dir": str(tmp_path / "checkpoints"), + "metric_name": None, + "higher_is_better": False, + "keep_top_k": None, + "save_period": 10_000, + "save_optimizer": False, + "checkpoint_must_save_by": None, + }, ) actor_args = SingleControllerActorArgs( @@ -336,6 +349,8 @@ def test_train_pump_drives_mcore_training_step( rollout_manager=rollout_manager, tq_buffer=tq_buffer, partition_id=_PARTITION_ID, + save_state=_default_grpo_save_state(), + last_checkpoint_path=None, ) ctrl = _RecordingSingleControllerActor.remote( metric_log_handle=log, diff --git a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py new file mode 100644 index 0000000000..6359fc672f --- /dev/null +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -0,0 +1,70 @@ +# 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. + +from unittest.mock import MagicMock + +import pytest + +from tools import verify_tq_data_plane_checkpoint as verifier + + +def test_save_finalizes_by_renaming_parent_bundle(monkeypatch, tmp_path) -> None: + final_bundle = tmp_path / "step_7" + expected_staging_bundle = tmp_path / "tmp_step_7" + save_calls = [] + + def fake_save(checkpoint_dir, num_storage_units) -> None: + save_calls.append((checkpoint_dir, num_storage_units)) + assert checkpoint_dir.parent.is_dir() + checkpoint_dir.mkdir() + (checkpoint_dir / "marker").write_text("saved") + + monkeypatch.setattr(verifier, "_save", fake_save) + + verifier._save_and_finalize_bundle(final_bundle, num_storage_units=3) + + assert save_calls == [(expected_staging_bundle / "data_plane", 3)] + assert not expected_staging_bundle.exists() + assert (final_bundle / "data_plane" / "marker").read_text() == "saved" + + +def test_save_refuses_to_replace_final_bundle(monkeypatch, tmp_path) -> None: + final_bundle = tmp_path / "step_7" + final_bundle.mkdir() + save = MagicMock() + monkeypatch.setattr(verifier, "_save", save) + + with pytest.raises(FileExistsError, match=str(final_bundle)): + verifier._save_and_finalize_bundle(final_bundle, num_storage_units=1) + + save.assert_not_called() + + +def test_save_failure_removes_created_staging_bundle(monkeypatch, tmp_path) -> None: + final_bundle = tmp_path / "step_7" + staging_bundle = tmp_path / "tmp_step_7" + + def failing_save(checkpoint_dir, num_storage_units) -> None: + del num_storage_units + assert checkpoint_dir.parent == staging_bundle + assert staging_bundle.is_dir() + raise RuntimeError("injected TQ save failure") + + monkeypatch.setattr(verifier, "_save", failing_save) + + with pytest.raises(RuntimeError, match="injected TQ save failure"): + verifier._save_and_finalize_bundle(final_bundle, num_storage_units=1) + + assert not staging_bundle.exists() + assert not final_bundle.exists() diff --git a/tools/verify_tq_data_plane_checkpoint.py b/tools/verify_tq_data_plane_checkpoint.py new file mode 100644 index 0000000000..33c9f50ada --- /dev/null +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -0,0 +1,246 @@ +# 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. + +"""Verify TQ data-plane save/load across a fresh process and parent rename. + +The save phase writes sample tensors, tags, and partial consumer progress to +TQ under ``tmp_/data_plane``, then renames the parent bundle to its +final path just like ``CheckpointManager``. The load phase starts a fresh TQ +instance, restores from ``/data_plane`` before any partition operations, +and verifies both the tensors and the consumer cursor. + +Example: + uv run --no-sync python tools/verify_tq_data_plane_checkpoint.py \ + --checkpoint-dir /lustre/.../tq-data-plane-checkpoint-smoke +""" + +from __future__ import annotations + +import argparse +import shutil +import subprocess +import sys +from pathlib import Path +from typing import cast + +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane import DataPlaneConfig, build_data_plane_client + +PARTITION_ID = "tq_checkpoint_smoke" +TASK_NAME = "train" +SAMPLE_IDS = [f"prompt-0:generation-{index}" for index in range(4)] +SEQ_LEN = 16 +FIELDS = ["token_ids", "token_mask", "generation_logprobs"] +DATA_PLANE_DIR = "data_plane" + + +def _data_plane_config(num_storage_units: int) -> DataPlaneConfig: + return cast( + DataPlaneConfig, + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "checkpointing_enabled": True, + "storage_capacity": 1024, + "num_storage_units": num_storage_units, + "claim_meta_poll_interval_s": 0.05, + "global_segment_size": 8 * 1024**3, + "local_buffer_size": 1024**3, + }, + ) + + +def _expected_fields() -> TensorDict: + token_ids = torch.arange(len(SAMPLE_IDS) * SEQ_LEN, dtype=torch.int64).reshape( + len(SAMPLE_IDS), + SEQ_LEN, + ) + return TensorDict( + { + "token_ids": token_ids, + "token_mask": torch.ones_like(token_ids), + "generation_logprobs": -token_ids.to(torch.float32) / 100.0, + }, + batch_size=[len(SAMPLE_IDS)], + ) + + +def _save(checkpoint_dir: Path, num_storage_units: int) -> None: + dp_client = build_data_plane_client( + _data_plane_config(num_storage_units), + bootstrap=True, + ) + try: + dp_client.register_partition( + partition_id=PARTITION_ID, + fields=FIELDS, + num_samples=len(SAMPLE_IDS), + consumer_tasks=[TASK_NAME], + ) + dp_client.put_samples( + sample_ids=SAMPLE_IDS, + partition_id=PARTITION_ID, + fields=_expected_fields(), + tags=[{"policy_version": 3, "prompt_id": "prompt-0"} for _ in SAMPLE_IDS], + ) + + consumed = dp_client.claim_meta( + partition_id=PARTITION_ID, + task_name=TASK_NAME, + required_fields=FIELDS, + batch_size=1, + timeout_s=30.0, + ) + if consumed.size != 1: + raise AssertionError(f"Expected one consumed row, got {consumed.size}") + + dp_client.save_checkpoint( + checkpoint_dir, + metadata={ + "data_plane_checkpoint_schema_version": 1, + "expected_consumed_ids": consumed.sample_ids, + }, + ) + finally: + dp_client.close() + + +def _load(checkpoint_dir: Path, num_storage_units: int) -> None: + dp_client = build_data_plane_client( + _data_plane_config(num_storage_units), + bootstrap=True, + ) + try: + metadata = dp_client.load_checkpoint(checkpoint_dir) + + restored = dp_client.get_samples( + sample_ids=SAMPLE_IDS, + partition_id=PARTITION_ID, + select_fields=FIELDS, + ) + expected = _expected_fields() + for field in FIELDS: + if not torch.equal(restored[field], expected[field]): + raise AssertionError(f"Restored field differs: {field}") + + if metadata["data_plane_checkpoint_schema_version"] != 1: + raise AssertionError("Unexpected data-plane checkpoint schema") + consumed_ids = set(metadata["expected_consumed_ids"]) + expected_remaining_ids = set(SAMPLE_IDS) - consumed_ids + + if dp_client.check_consumption_status(PARTITION_ID, [TASK_NAME]): + raise AssertionError( + "Restored consumer cursor marked every row consumed before " + "the expected remaining rows were claimed" + ) + remaining = dp_client.claim_meta( + partition_id=PARTITION_ID, + task_name=TASK_NAME, + required_fields=FIELDS, + batch_size=len(expected_remaining_ids), + timeout_s=30.0, + ) + if consumed_ids.intersection(remaining.sample_ids): + raise AssertionError("A previously consumed row was claimed after restore") + if set(remaining.sample_ids) != expected_remaining_ids: + raise AssertionError("Restored consumption state lost or added rows") + if not dp_client.check_consumption_status(PARTITION_ID, [TASK_NAME]): + raise AssertionError("Restored consumer cursor did not reach completion") + finally: + dp_client.close() + + +def _save_and_finalize_bundle( + bundle_dir: Path, + num_storage_units: int, +) -> None: + """Save below a temporary parent, then rename it to ``bundle_dir``.""" + staging_dir = bundle_dir.with_name(f"tmp_{bundle_dir.name}") + if bundle_dir.exists(): + raise FileExistsError(f"Final checkpoint bundle already exists: {bundle_dir}") + if staging_dir.exists(): + raise FileExistsError( + f"Staging checkpoint bundle already exists: {staging_dir}" + ) + + # CheckpointManager creates tmp_step_N before component writers run. + # Mirror that precondition instead of relying on TQ to create the parent. + staging_dir.mkdir(parents=True) + try: + _save(staging_dir / DATA_PLANE_DIR, num_storage_units) + staging_dir.rename(bundle_dir) + except Exception: + if staging_dir.exists(): + shutil.rmtree(staging_dir) + raise + + +def _run_child( + phase: str, + checkpoint_dir: Path, + num_storage_units: int, +) -> None: + subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--phase", + phase, + "--checkpoint-dir", + str(checkpoint_dir), + "--num-storage-units", + str(num_storage_units), + ], + check=True, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--phase", + choices=("round-trip", "save", "load"), + default="round-trip", + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--checkpoint-dir", + type=Path, + required=True, + help="Final SC-like checkpoint bundle directory.", + ) + parser.add_argument("--num-storage-units", type=int, default=4) + args = parser.parse_args() + + checkpoint_dir = args.checkpoint_dir.expanduser().resolve() + if args.phase == "save": + _save_and_finalize_bundle(checkpoint_dir, args.num_storage_units) + return + if args.phase == "load": + _load(checkpoint_dir / DATA_PLANE_DIR, args.num_storage_units) + return + + _run_child("save", checkpoint_dir, args.num_storage_units) + _run_child("load", checkpoint_dir, args.num_storage_units) + print( + "PASS: TQ checkpoint survived a parent rename and fresh process", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index 9654360842..146f7104e6 100644 --- a/uv.lock +++ b/uv.lock @@ -4421,7 +4421,7 @@ requires-dist = [ { name = "torchdata" }, { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.26.0", index = "https://pypi.org/simple" }, - { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39" }, + { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2" }, { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformers", specifier = ">=5.5.0,<5.9.0" }, { name = "transformers", marker = "extra == 'automodel'", specifier = ">=5.5.0,<5.6.0" }, @@ -7874,14 +7874,15 @@ wheels = [ [[package]] name = "transferqueue" -version = "0.1.7.dev0" -source = { git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39#b266d39a15aae114730de36cf8317b6285436f7f" } +version = "0.1.9" +source = { git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2#c51614308b68c8d7a87c9b3ef62d59e14c69bde2" } dependencies = [ { name = "hydra-core", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "hydra-core", version = "1.3.4", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" }, { name = "msgspec" }, { name = "numpy" }, { name = "omegaconf" }, + { name = "prometheus-client" }, { name = "psutil" }, { name = "pyzmq" }, { name = "ray", extra = ["default"] },