diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index d84c2850840..edce5196e93 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -71,7 +71,45 @@ uv run examples/run_grpo_single_controller.py --config use_importance_sampling_correction: true ``` -5. **(PPO) Set `ppo:` instead of `grpo:`** — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs `value:`, `value_loss_fn:` and `ppo.adv_estimator.name: gae` (same schemas as legacy PPO), a Megatron critic, and `policy.offload_optimizer_for_logprob: true`, which is what keeps the policy optimizer off the GPU while the critic runs. `ppo.policy_training_start_step: N` gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step. `ppo.warm_start_value_checkpoint` seeds that critic from another run's checkpoint instead, so a fresh run can skip the online warmup entirely — see [Warm-Starting the Critic](./ppo.md#warm-starting-the-critic). +5. **Save the data plane for replay recovery.** When Single-Controller checkpointing is enabled, all built-in samplers require `checkpointing.save_data_plane: true` so completed, unconsumed rollout groups survive a restart. Native TQ checkpointing currently supports only the `simple` storage backend. For multi-node runs, `checkpoint_dir` must be on a durable filesystem visible at the same path from every node. + + ```yaml + checkpointing: + enabled: true + checkpoint_dir: /shared/checkpoints/my-run + save_data_plane: true + + data_plane: + enabled: true + backend: "simple" + ``` + +6. **(PPO) Set `ppo:` instead of `grpo:`** — the two algorithm blocks are mutually exclusive, and SC reads every step setting from whichever one is present. A PPO run also needs `value:`, `value_loss_fn:` and `ppo.adv_estimator.name: gae` (same schemas as legacy PPO), a Megatron critic, and `policy.offload_optimizer_for_logprob: true`, which is what keeps the policy optimizer off the GPU while the critic runs. `ppo.policy_training_start_step: N` gives the usual critic warmup: for the first N steps the policy is neither trained nor refit, while the critic trains every step. `ppo.warm_start_value_checkpoint` seeds that critic from another run's checkpoint instead, so a fresh run can skip the online warmup entirely — see [Warm-Starting the Critic](./ppo.md#warm-starting-the-critic). + +## Checkpointing and Replay Recovery + +With `checkpointing.save_data_plane: true`, each Single-Controller checkpoint contains: + +- The normal model, dataloader, and controller state, plus optimizer state when configured. +- A native TQ snapshot containing rollout tensor payloads and TQ state. +- A metadata-only replay index describing the completed rollout groups stored in TQ. +- A `rollout_recovery.pt` ownership ledger describing unfinished prompt groups that must be redispatched after a restart. +- A `replacement_reserve.pt` sidecar containing prompts held for dropped-rollout replacement, when applicable. +- The sampler dispatch position needed to continue scheduling from the correct point. + +The TQ snapshot and replay index are captured under the same checkpoint barrier. Generation may continue while the snapshot is written, but completed-group commits and destructive TQ clears wait at the barrier. This ensures that the TQ snapshot and replay index describe the same set of groups. + +On resume, Single-Controller validates the TQ snapshot against the trainer checkpoint, restores the replay index, and makes completed, committed, unconsumed groups available to the sampler before training resumes. + +Replay recovery is supported by all built-in samplers: `in_order`, `weight_fifo`, `ready_first`, and `windowed`. Custom samplers must explicitly declare `supports_buffer_checkpoint = True`. Otherwise, setup emits a warning and completed buffered groups are not restored. + +:::{note} +Completed groups are restored directly from the TQ snapshot. Prompt groups whose generations were still in flight at the checkpoint boundary are recovered by ownership: `rollout_recovery.pt` records them, and on resume they are redispatched and regenerated from the same dataset rows. Only rows already committed to TQ preserve their exact generated tokens; redispatched groups produce new samples from the same prompts. +::: + +When a sampler does not support replay recovery, a requested data-plane checkpoint is written in `shadow` mode. The TQ snapshot is retained, but no authoritative replay index is written and its rows are not restored into the training replay buffer. + +Native TQ save/load currently requires `data_plane.backend: "simple"`. Mooncake-backed storage is not recoverable through this mechanism. A failure while saving or validating the TQ snapshot prevents the incomplete checkpoint bundle from becoming the latest resumable checkpoint. ## Async-RL Knobs and Sampler Modes @@ -79,7 +117,7 @@ All SC async-RL runtime knobs live under `async_rl:` in the master config. The m ### Sampler modes -Pick one of four modes with `sampler.name`. Each mode takes its own knobs, listed below — a knob from one mode has no effect under another: +Pick one of five modes with `sampler.name`. Each mode takes its own knobs, listed below — a knob from one mode has no effect under another: ![Sampler modes: same buffer, four different training batches](../assets/sc-sampler-modes.png) @@ -90,13 +128,14 @@ Pick one of four modes with `sampler.name`. Each mode takes its own knobs, liste | -------------- | ----------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------ | | `in_order` | Dispatch may lead the trainer by up to `max_lookahead_versions` batches. Each dispatch is stamped with a `target_step`. | Consume the group whose `target_step == current_train_weight`. | Sync mode (`max_lookahead_versions=0`) and legacy-async exact-batch semantics (`max_lookahead_versions>=1`). The only mode supported on a PPO run. | | `weight_fifo` | Same gate as `in_order` (`max_staleness_versions` of lookahead). | Drain the oldest in-window `start_weight` first, waiting for that weight's batch to fill. | Strict weight-version FIFO under a bounded lookahead. | +| `ready_first` | Same gate as `weight_fifo` (`max_staleness_versions` of lookahead). | Take any ready group generated by a policy version no newer than the trainer, including late stragglers. | Completion-order streaming without stale-group eviction. | | `windowed` | Ungated — rollout keeps producing until the buffer fills. | Take any ready group with `start_weight` in `[train - max_staleness_versions, train]`, optionally freshest-first. | Over-sampled streaming; aged groups outside the window are evicted (wasted compute). | | `custom` | Determined by the imported class. | Determined by the imported class. | `target: "module:ClassName"` — bring your own `PromptGroupSampler`. | ### Config → behavior map -The shipped exemplars cover three of the four modes: +The shipped exemplars cover three of the five modes: | Mode | `sampler.name` | Sampler knob | `min_groups_for_streaming_train` | `max_buffered_rollouts` | Exemplar | @@ -104,6 +143,7 @@ The shipped exemplars cover three of the four modes: | Sync / on-policy | `in_order` | `max_lookahead_versions: 0` | `${grpo.num_prompts_per_step}` | `num_prompts_per_step × 1` | [`grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml`](../../examples/configs/recipes/llm/grpo-qwen2.5-math-1.5b-instruct-1n8g-megatron-single-controller-sync.yaml) | | Async, exact batch→step matching | `in_order` | `max_lookahead_versions: >= 1` | `x <= num_prompts_per_step` | `num_prompts_per_step × (max_lookahead_versions + 1)` | [`grpo_math_1B_megatron_single_controller.yaml`](../../examples/configs/grpo_math_1B_megatron_single_controller.yaml) | | Streaming, gated dispatch | `weight_fifo` | `max_staleness_versions: >= 1` | `x <= num_prompts_per_step` | `num_prompts_per_step × (max_staleness_versions + 1)` | — (none shipped) | +| Streaming, ready-first | `ready_first` | `max_staleness_versions: >= 1` | `x <= num_prompts_per_step` | `num_prompts_per_step × (max_staleness_versions + 1)` | — (none shipped) | | Streaming, over-sampled | `windowed` | `max_staleness_versions: >= 1` | `x <= num_prompts_per_step` | Larger than the gated capacity (dispatch is ungated) | [`grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml`](../../examples/configs/recipes/llm/grpo-llama3.1-8b-instruct-2n8g-async-1off-single-controller-streaming2.yaml) | @@ -144,7 +184,7 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th #### 4. Samplers (`nemo_rl/algorithms/async_utils/staleness_sampler.py`) - Filter-only prompt-group selector over `TQReplayBuffer`. The base `PromptGroupSampler` protocol defines `admit`, `select`, and `evict`. -- `WindowedSampler`, `WeightFifoSampler`, `InOrderSampler` are the built-in policies (one per row in the [Sampler modes](#sampler-modes) table). The `custom` mode (`CustomSamplerConfig.target`) makes `create_sampler` import a user-supplied class by FQN and type-check it against `PromptGroupSampler`. +- `WindowedSampler`, `ReadyFirstSampler`, `WeightFifoSampler`, and `InOrderSampler` are the built-in policies (one per row in the [Sampler modes](#sampler-modes) table). The `custom` mode (`CustomSamplerConfig.target`) makes `create_sampler` import a user-supplied class by FQN and type-check it against `PromptGroupSampler`. #### 5. `_rollout_pump` and `_train_pump` @@ -169,7 +209,7 @@ The [legacy async GRPO](./async-grpo.md) (`grpo.async_grpo.enabled: true` under | Entrypoint | `run_grpo.py` | `run_grpo_single_controller.py` | | Data-plane | Direct actor RPC | TransferQueue (`data_plane.enabled: true` required) | | Rollout batching | Full-batch `AsyncTrajectoryCollector` | Per-prompt `RolloutManager.generate_and_push` into a group-granular `TQReplayBuffer` | -| Staleness policy | Single knob (`max_trajectory_age_steps`) | Pluggable `StalenessSampler` (`in_order` / `weight_fifo` / `windowed` / `custom`) | +| Staleness policy | Single knob (`max_trajectory_age_steps`) | Pluggable `StalenessSampler` (`in_order` / `weight_fifo` / `ready_first` / `windowed` / `custom`) | | Batch boundary | Sampled by target weight | Sampler-defined; can decouple rollout dispatch from train batch (streaming) | diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d7f17aeac11..4e7d4914c74 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -125,6 +125,10 @@ checkpointing: model_save_format: "safetensors" save_consolidated: false save_optimizer: true + # SingleController only: include native TQ state and the metadata-only replay + # index. Required for recovery-capable samplers to preserve completed, + # unconsumed rollouts. + save_data_plane: false policy: model_name: "Qwen/Qwen2.5-1.5B" diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index e0ba69a6367..2b50750038e 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -128,6 +128,9 @@ checkpointing: enabled: false checkpoint_dir: results/grpo-single-controller metric_name: null + # Include native TQ state and the metadata-only replay index. A save failure + # aborts checkpoint finalization. + save_data_plane: true policy: dtensor_cfg: diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index f5061aa75ef..2cf2d09df08 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -14,18 +14,33 @@ import asyncio import gc +import hashlib +import json +import math import statistics import threading as _threading import uuid from collections import Counter -from collections.abc import Mapping -from typing import Any, Awaitable, Callable, Iterable, Optional +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager +from numbers import Integral, Real +from typing import ( + Any, + Awaitable, + Callable, + Iterable, + Literal, + NotRequired, + Optional, + TypedDict, +) import ray import torch 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.data_plane.schema import ROUTED_EXPERTS_FIELD from nemo_rl.experience.interfaces import ( NEMO_GYM_TASK_INDEX_KEY, @@ -36,6 +51,207 @@ from nemo_rl.experience.payload import pack_payload, record_to_train_batch from nemo_rl.utils.r3_trace import trace_rollout_payload +DATA_PLANE_CHECKPOINT_DIR = "data_plane" +REPLAY_BUFFER_METADATA_FILENAME = "replay_buffer_metadata.pt" +LEGACY_REPLAY_BUFFER_FILENAME = "replay_buffer.pt" +REPLACEMENT_RESERVE_FILENAME = "replacement_reserve.pt" +REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1 +REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" + +# These TypedDicts describe the versioned, plain-mapping checkpoint wire +# format. They are intentionally not dataclass instances: persisting a +# dataclass would couple recovery to its Python import path and class layout. +# Runtime objects such as KVBatchMeta remain explicitly represented as fields +# inside this schema. + + +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 index paired with a TQ snapshot.""" + + schema_version: int + storage: Literal["tq_checkpoint"] + partition_id: str + saved_capacity: int + manifest_digest: str + groups: list[TQReplayGroupMetadata] + + +class DataPlaneCheckpointMetadata(TypedDict): + """SC metadata envelope stored with a native data-plane checkpoint. + + The replay fields are present together in ``authoritative`` mode and + absent in ``shadow`` mode. + """ + + data_plane_checkpoint_schema_version: int + single_controller_train_steps: int + single_controller_trainer_version: int + single_controller_epoch: int + partition_id: str + sampler_name: str + mode: Literal["authoritative", "shadow"] + replay_metadata_schema_version: NotRequired[int] + replay_manifest_digest: NotRequired[str] + replay_group_count: NotRequired[int] + rollout_recovery_schema_version: NotRequired[int] + rollout_recovery_payload_sha256: NotRequired[str] + rollout_recovery_group_count: NotRequired[int] + + +def _canonical_manifest_value(value: Any, *, path: str) -> Any: + """Return a deterministic JSON value or reject unsupported metadata.""" + if value is None or isinstance(value, (bool, str)): + return value + if isinstance(value, Integral): + return int(value) + if isinstance(value, Real): + float_value = float(value) + if not math.isfinite(float_value): + raise TypeError(f"Replay metadata at {path} must be finite") + return float_value + if isinstance(value, Mapping): + canonical: dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise TypeError(f"Replay metadata at {path} has non-string key {key!r}") + canonical[key] = _canonical_manifest_value( + item, + path=f"{path}.{key}", + ) + return canonical + if isinstance(value, (list, tuple)): + return [ + _canonical_manifest_value(item, path=f"{path}[{index}]") + for index, item in enumerate(value) + ] + raise TypeError( + f"Replay metadata at {path} has unsupported type " + f"{type(value).__name__}; expected JSON-compatible primitive values" + ) + + +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": _canonical_manifest_value( + group["meta"].tags, + path=f"groups[{group_index}].meta.tags", + ), + "extra_info": _canonical_manifest_value( + group["meta"].extra_info, + path=f"groups[{group_index}].meta.extra_info", + ), + }, + } + for group_index, group in enumerate(groups) + ] + encoded = json.dumps( + digest_input, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +class DataPlaneMutationCut: + """Live capability proving code runs inside a data-plane barrier cut.""" + + __slots__ = ("_barrier", "_live") + + def __init__(self, barrier: "DataPlaneCheckpointBarrier") -> None: + self._barrier = barrier + self._live = True + + def require_live(self) -> None: + """Fail when a mutation tries to reuse an absent or expired cut.""" + if not self._live: + raise RuntimeError("data-plane mutation cut is no longer active") + + def _invalidate(self) -> None: + self._live = False + + +class DataPlaneCheckpointBarrier: + """Allow concurrent mutations while giving live checkpoints exclusivity. + + At most one checkpoint holder is active. New mutations queue behind it, + and a checkpoint waits for all active mutations before yielding. Every + live canonical TQ commit/clear and native save must use this barrier so the + snapshot and controller replay index describe the same rows. + """ + + def __init__(self) -> None: + self._condition = asyncio.Condition() + self._checkpoint_active = False + self._active_mutations = 0 + + @asynccontextmanager + async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: + """Yield a live mutation capability after any active checkpoint exits.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._checkpoint_active) + self._active_mutations += 1 + cut = DataPlaneMutationCut(self) + try: + yield cut + finally: + cut._invalidate() + async with self._condition: + self._active_mutations -= 1 + if self._active_mutations == 0: + self._condition.notify_all() + + @asynccontextmanager + async def checkpoint(self) -> AsyncIterator[DataPlaneMutationCut]: + """Yield a live capability after blocking and draining all mutations.""" + 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 + cut = DataPlaneMutationCut(self) + try: + yield cut + finally: + cut._invalidate() + async with self._condition: + self._checkpoint_active = False + self._condition.notify_all() + class PostWriteEnrichmentError(RuntimeError): """A rollout reached TQ but failed in required post-write processing.""" @@ -778,10 +994,31 @@ 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 self._post_write_enricher: Optional[ Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]] ] = 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 + + @property + def data_plane_checkpoint_barrier(self) -> DataPlaneCheckpointBarrier: + """Return the shared barrier used by controller and post-commit ownership.""" + if self._data_plane_checkpoint_barrier is None: + raise RuntimeError("data-plane checkpoint barrier is not configured") + return self._data_plane_checkpoint_barrier + def set_post_write_enricher( self, enricher: Callable[[KVBatchMeta, PromptGroupRecord], Awaitable[KVBatchMeta]], @@ -801,7 +1038,9 @@ def reserve( Args: weight_version: Weight version stamped on the slot. target_step: Training step this slot targets; only consulted by StalenessSampler.force_in_order. - group_id: Per-group sample_id prefix; defaults to a fresh uuid4. + group_id: Pre-minted logical group ID and sample-ID prefix. The + checkpoint-enabled lineage path always supplies this. ``None`` + creates a fresh UUID only for untracked callers. Returns: group_id used by the matching commit. @@ -846,9 +1085,17 @@ 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 + train_batch, + weight_version=start_weight_version, + group_id=group_id, + prompt_idx=record.prompt_idx, ) if self._require_routed_experts and ROUTED_EXPERTS_FIELD not in fields: raise RuntimeError( @@ -858,56 +1105,56 @@ async def commit( "the async message-log flattening path." ) trace_rollout_payload(keys=sample_ids, data=train_batch) - 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], - ) - - if self._post_write_enricher is not None: - try: - meta = await self._post_write_enricher(meta, record) - except Exception as error: - raise PostWriteEnrichmentError( - f"post-write enrichment failed for group_id={group_id!r}" - ) from error - - 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 + + if self._post_write_enricher is not None: + try: + meta = await self._post_write_enricher(meta, record) + except Exception as error: + raise PostWriteEnrichmentError( + f"post-write enrichment failed for group_id={group_id!r}" + ) from error + + 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``. @@ -922,11 +1169,17 @@ 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. @@ -940,14 +1193,24 @@ 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] @@ -960,72 +1223,52 @@ 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) - async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]: - """Serialize ready groups (meta + DataPlane payloads) for checkpointing. - - Snapshots the ready slots synchronously on the event loop first, then - fetches each group's rows from the DataPlane. Unready reservations are - in-flight rollouts and are dropped, matching legacy semantics. The - snapshot stays consistent during the async fetch: concurrent commits - only append/flip *other* slots, and the train pump — the only - remover — is the caller itself; groups committed mid-save land in the - next checkpoint. - - Args: - saved_capacity: max_buffered_rollouts at save time, recorded so - load_state_dict can report capacity changes across restarts. - - Returns: - Envelope: ``{"partition_id": ..., "saved_capacity": ..., - "groups": [{"meta", "start_weight", "end_weight", "target_step", - "group_id", "fields_data"}, ...]}``. + 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 replay + index 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 either participate in that + barrier across the complete publish/index or clear/remove transition, + or run in the same asyncio task as the checkpoint save. The advantage + stage relies on the latter: it and ``_save_checkpoint`` both live in + ``_train_pump``, so they cannot interleave. Any new writer outside + ``_train_pump`` -- including future finalizer paths -- must take a + mutation slot; canonical writes are not required to originate + specifically from :meth:`commit`. + In-flight reservations are intentionally omitted. """ - snapshot: list[tuple[KVBatchMeta, int, int, Optional[int], str]] = [] + 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 - snapshot.append( - ( - meta, - self.start_weight_list[i], - self.end_weight_list[i], - self.target_step_list[i], - self._group_ids[i], - ) - ) - - groups: list[dict[str, Any]] = [] - for meta, start_weight, end_weight, target_step, group_id in snapshot: - fields_data = await self._call_dp( - "get_samples", - sample_ids=meta.sample_ids, - partition_id=self._partition_id, - select_fields=meta.fields, - ) groups.append( { "meta": meta, - "start_weight": start_weight, - "end_weight": end_weight, - "target_step": target_step, - "group_id": group_id, - "fields_data": fields_data, + "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, } @@ -1036,16 +1279,14 @@ async def load_state_dict( max_groups: int, expected_partition_id: str, expected_group_size: int, + expected_manifest_digest: str, ) -> int: - """Validate and re-put checkpointed groups into the buffer. + """Restore the local replay index for an already-restored TQ snapshot. - The preflight runs entirely before any DataPlane write (legacy - precedent: validate, then truncate): - 1. Validate the envelope and raise ValueError on malformed state. - 2. Truncate to ``max_groups``, keeping the freshest groups, so the - restored count can never exceed the buffer's capacity. Groups - carrying a ``target_step`` are never truncated — an over-capacity - in-order checkpoint raises instead (see Raises). + The replay index 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 @@ -1053,7 +1294,7 @@ async def load_state_dict( eviction in one place. Args: - state: Envelope produced by ``state_dict``. + 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 @@ -1061,6 +1302,8 @@ async def load_state_dict( 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). + expected_manifest_digest: Digest returned by the matching native + TQ checkpoint load. It must match the replay metadata file. Returns: Number of groups restored into the buffer. @@ -1068,14 +1311,35 @@ async def load_state_dict( Raises: ValueError: If the envelope is malformed (missing keys, partition mismatch, misaligned or wrongly sized groups, duplicate - sample_ids), or if target-stamped groups exceed ``max_groups``. + sample_ids), disagrees with the native TQ snapshot, or exceeds + ``max_groups``. """ - required_keys = {"partition_id", "saved_capacity", "groups"} + 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( + f"Replay-buffer metadata has unsupported storage: {state['storage']!r}" + ) if state["partition_id"] != expected_partition_id: raise ValueError( "Replay buffer checkpoint partition_id mismatch: " @@ -1090,16 +1354,25 @@ async def load_state_dict( "end_weight", "target_step", "group_id", - "fields_data", } 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 @@ -1120,43 +1393,36 @@ async def load_state_dict( ) 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." ) - num_truncated = 0 if len(groups) > max_groups: - if any(group["target_step"] is not None for group in groups): - raise ValueError( - f"Replay buffer checkpoint holds {len(groups)} group(s) " - f"but async_rl.max_buffered_rollouts is {max_groups}. " - "These groups carry target_step stamps (in-order " - "sampling) and are selected as whole per-step batches, so " - "dropping any of them would deadlock the resumed run. " - "Resume with async_rl.max_buffered_rollouts >= " - f"{len(groups)}, or delete replay_buffer.pt from the " - "checkpoint to resume with an empty buffer." - ) - num_truncated = len(groups) - max_groups - # Keep the freshest max_groups groups, preserving original order. - prioritized = sorted( - range(len(groups)), - key=lambda i: (groups[i]["start_weight"], i), + raise ValueError( + "Native TQ checkpoint contains more replay groups than the current " + f"buffer capacity: checkpoint={len(groups)}, current={max_groups}. " + f"Resume with async_rl.max_buffered_rollouts >= {len(groups)} to " + f"keep them. Deleting {REPLAY_BUFFER_METADATA_FILENAME} from the " + "checkpoint directory also allows startup, but skips loading the " + "matching TQ checkpoint and discards these groups and the prompts " + "that produced them because the dataloader has already moved past " + "them." ) - indices_to_keep = sorted(prioritized[num_truncated:]) - groups = [groups[i] for i in indices_to_keep] for group in groups: meta = group["meta"] - await self._call_dp( - "put_samples", - sample_ids=list(meta.sample_ids), - partition_id=self._partition_id, - fields=group["fields_data"], - tags=[dict(t) for t in meta.tags], - ) self.meta_list.append(meta) self.start_weight_list.append(group["start_weight"]) self.end_weight_list.append(group["end_weight"]) @@ -1164,10 +1430,10 @@ async def load_state_dict( self.ready_list.append(True) self._group_ids.append(group["group_id"]) - summary = f"📦 Restored {len(groups)} replay group(s) from checkpoint" - if num_truncated: - summary += f"; truncated {num_truncated} group(s) over capacity" - print(summary, flush=True) + print( + f"📦 Restored {len(groups)} replay group(s) from checkpoint", + flush=True, + ) return len(groups) def count_for_target_step(self, target_step: int) -> int: @@ -1226,13 +1492,12 @@ 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_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 58fe5d43ebe..ff927ff3493 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, @@ -53,7 +54,10 @@ from pydantic import BaseModel, Field, NonNegativeInt, model_validator -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneMutationCut, + TQReplayBuffer, +) from nemo_rl.data_plane import KVBatchMeta # Poll interval for the rollout-pump admission gate. @@ -109,12 +113,39 @@ def is_on_policy(self) -> bool: """True when the policy admits zero staleness (sync mode).""" ... + supports_buffer_checkpoint: ClassVar[bool] + """Whether completed buffered groups can be restored safely.""" + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" ... - def set_dispatch_index(self, resume_from_step: int) -> None: - """Seed the dispatch cursor when resuming from a checkpoint.""" + @property + def dispatch_index(self) -> int: + """Last admitted dispatch batch index.""" + ... + + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: + """Seed the cursor for checkpoints that predate exact sampler state.""" + ... + + def restore_dispatch_index(self, dispatch_index: int) -> None: + """Restore the exact dispatch cursor from controller state.""" + ... + + +@runtime_checkable +class TransactionalAdmissionSampler(Protocol): + """Sampler whose blocking wait is separate from its cursor mutation.""" + + async def wait_until_admissible( + self, *, trainer_version_fn: Callable[[], int] + ) -> None: + """Wait until one admission can commit without mutating sampler state.""" + ... + + def commit_admission(self, cut: DataPlaneMutationCut) -> Optional[int]: + """Advance the cursor under a live data-plane mutation cut.""" ... @@ -126,28 +157,45 @@ class BaseSampler(abc.ABC): select-finalize / weight-window-evict helpers. """ + supports_buffer_checkpoint: ClassVar[bool] = False + def __init__(self, buffer: TQReplayBuffer) -> None: 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 - def set_dispatch_index(self, resume_from_step: int) -> None: - """Seed the dispatch cursor when resuming from a checkpoint. + @property + def dispatch_index(self) -> int: + """Return the last admitted dispatch batch index.""" + return self._dispatch_index + + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: + """Seed the cursor for checkpoints that predate exact sampler state. Args: - resume_from_step: Trainer step this run starts from — 0 for a - fresh run, the restored ``current_step`` when resuming. Sets - the cursor to ``resume_from_step - 1`` so gated ``admit`` and - ``InOrderSampler``'s target_step stamps line up with the - restored trainer version exactly as at step 0 of a fresh run. - Call before the first ``admit``. + resume_from_trainer_version: Trainer version from which the run + resumes. The next admitted batch receives that version. """ - if resume_from_step < 0: + if resume_from_trainer_version < 0: raise ValueError( - f"resume_from_step must be non-negative, got {resume_from_step}" + "resume_from_trainer_version must be non-negative, got " + f"{resume_from_trainer_version}" ) - self._dispatch_index = resume_from_step - 1 + self._dispatch_index = resume_from_trainer_version - 1 + + def restore_dispatch_index(self, dispatch_index: int) -> None: + """Restore the exact dispatch cursor. + + Args: + dispatch_index: Last admitted batch index, or ``-1`` when no batch + has been admitted. Call before the first ``admit``. + """ + if dispatch_index < -1: + raise ValueError( + f"dispatch_index must be at least -1, got {dispatch_index}" + ) + self._dispatch_index = dispatch_index # ── rollout-pump side ──────────────────────────────────────────────── @abc.abstractmethod @@ -247,6 +295,9 @@ class WindowedSampler(BaseSampler): freshest-first. """ + # Ungated restored groups are ordinary in-window candidates. + supports_buffer_checkpoint: ClassVar[bool] = True + def __init__( self, buffer: TQReplayBuffer, @@ -278,8 +329,19 @@ def should_abort_inflight( ) return start_weight_version < min_valid_version + async def wait_until_admissible( + self, *, trainer_version_fn: Callable[[], int] + ) -> None: + """Return immediately because buffer capacity is this policy's gate.""" + del trainer_version_fn + + def commit_admission(self, cut: DataPlaneMutationCut) -> Optional[int]: + """Return the unstamped admission result without changing a cursor.""" + cut.require_live() + return None + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: - # Over-sampled: dispatch is bounded by buffer capacity, not by version. + await self.wait_until_admissible(trainer_version_fn=trainer_version_fn) return None async def select( @@ -345,12 +407,27 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: gate_window=self._gate_window, ) - async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + async def wait_until_admissible( + self, *, trainer_version_fn: Callable[[], int] + ) -> None: + """Wait for the gate without advancing the durable dispatch cursor.""" while self._dispatch_index >= trainer_version_fn() + self._gate_window: await asyncio.sleep(_GATE_POLL_SECONDS) + + def commit_admission(self, cut: DataPlaneMutationCut) -> Optional[int]: + """Advance the cursor after the controller enters its mutation cut.""" + cut.require_live() + return self._commit_admission() + + def _commit_admission(self) -> Optional[int]: + """Advance admission for the legacy monolithic API.""" self._dispatch_index += 1 return self._stamp() + async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: + await self.wait_until_admissible(trainer_version_fn=trainer_version_fn) + return self._commit_admission() + def _stamp(self) -> Optional[int]: return None @@ -366,6 +443,9 @@ class ReadyFirstSampler(_GatedSampler): rollout is ever discarded. """ + # Committed groups retain start_weight, which is sufficient for selection. + supports_buffer_checkpoint: ClassVar[bool] = True + def __init__( self, buffer: TQReplayBuffer, @@ -404,6 +484,9 @@ class WeightFifoSampler(_GatedSampler): that weight's batch to fill. Evict uses the weight window (default). """ + # Committed groups retain start_weight, which is sufficient for selection. + supports_buffer_checkpoint: ClassVar[bool] = True + def __init__(self, buffer: TQReplayBuffer, *, max_staleness_versions: int) -> None: super().__init__(buffer, gate_window=max_staleness_versions) self.max_staleness_versions = max_staleness_versions @@ -451,6 +534,10 @@ class InOrderSampler(_GatedSampler): capacity is sized for the peak of the two. """ + # Committed groups retain target_step, which is sufficient for selection. + # The controller checkpoints the exact dispatch cursor separately. + supports_buffer_checkpoint: ClassVar[bool] = True + def __init__( self, buffer: TQReplayBuffer, @@ -560,7 +647,10 @@ def peak_lookahead_versions(self) -> int: class CustomSamplerConfig(BaseModel, extra="allow"): name: Literal["custom"] = "custom" # "module:ClassName" of a PromptGroupSampler defined outside this repo. - # Extra keys are forwarded to the constructor (after ``buffer``). + # Extra keys are forwarded to the constructor (after ``buffer``). The + # target class must declare a boolean ``supports_buffer_checkpoint`` class + # attribute so setup can validate recovery requirements before allocating + # cluster resources. target: str @@ -602,45 +692,101 @@ def required_buffer_capacity_for_config( return None +def _custom_sampler_class(cfg: CustomSamplerConfig) -> type: + """Import and return a custom sampler class without constructing it.""" + 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) + if not isinstance(sampler_cls, type): + raise TypeError(f"custom sampler target is not a class: {cfg.target!r}") + return sampler_cls + + +def _sampler_class_for_config(cfg: SamplerConfig) -> type: + """Return the sampler class selected by a built-in or custom config.""" + if isinstance(cfg, CustomSamplerConfig): + return _custom_sampler_class(cfg) + try: + return { + WindowedSamplerConfig: WindowedSampler, + ReadyFirstSamplerConfig: ReadyFirstSampler, + WeightFifoSamplerConfig: WeightFifoSampler, + InOrderSamplerConfig: InOrderSampler, + }[type(cfg)] + except KeyError: + raise ValueError(f"unknown sampler config {type(cfg).__name__}") from None + + +def sampler_supports_buffer_checkpoint(cfg: SamplerConfig) -> bool: + """Return a sampler class's static replay-checkpoint capability. + + Custom classes are imported but not instantiated, allowing setup to fail + before allocating cluster resources or triggering constructor side effects. + """ + sampler_cls = _sampler_class_for_config(cfg) + + if isinstance(cfg, CustomSamplerConfig): + # A custom subclass must opt in explicitly instead of inheriting a + # built-in sampler's capability declaration accidentally. + capability = sampler_cls.__dict__.get("supports_buffer_checkpoint", False) + else: + capability = getattr(sampler_cls, "supports_buffer_checkpoint", None) + if not isinstance(capability, bool): + raise TypeError( + f"{sampler_cls.__name__}.supports_buffer_checkpoint must be a " + f"boolean class attribute, got {capability!r}" + ) + return capability + + def create_sampler( buffer: TQReplayBuffer, cfg: SamplerConfig, ) -> 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. + """ + sampler_cls = _sampler_class_for_config(cfg) + sampler: PromptGroupSampler if isinstance(cfg, WindowedSamplerConfig): - return WindowedSampler( + sampler = sampler_cls( buffer, max_staleness_versions=cfg.max_staleness_versions, sample_freshest_first=cfg.sample_freshest_first, ) - if isinstance(cfg, ReadyFirstSamplerConfig): - return ReadyFirstSampler( + elif isinstance(cfg, ReadyFirstSamplerConfig): + sampler = sampler_cls( buffer, max_staleness_versions=cfg.max_staleness_versions, ) - if isinstance(cfg, WeightFifoSamplerConfig): - return WeightFifoSampler( - buffer, max_staleness_versions=cfg.max_staleness_versions + elif isinstance(cfg, WeightFifoSamplerConfig): + sampler = sampler_cls( + buffer, + max_staleness_versions=cfg.max_staleness_versions, ) - if isinstance(cfg, InOrderSamplerConfig): - return InOrderSampler( + elif isinstance(cfg, InOrderSamplerConfig): + sampler = sampler_cls( buffer, max_lookahead_versions=cfg.max_lookahead_versions, warmup_lookahead_versions=cfg.warmup_lookahead_versions, ) - if 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) + elif isinstance(cfg, CustomSamplerConfig): + sampler_supports_buffer_checkpoint(cfg) sampler = sampler_cls(buffer, **(cfg.model_extra or {})) if not isinstance(sampler, PromptGroupSampler): raise TypeError( f"{cfg.target} does not implement the PromptGroupSampler " f"interface (needs admit/select/evict/should_abort_inflight, " - f"set_dispatch_index, is_on_policy, required_buffer_capacity)" + f"dispatch_index, set_dispatch_index, restore_dispatch_index, " + f"is_on_policy, supports_buffer_checkpoint, " + f"required_buffer_capacity)" ) - return sampler - raise ValueError(f"unknown sampler config {type(cfg).__name__}") + else: + raise ValueError(f"unknown sampler config {type(cfg).__name__}") + return sampler diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index c84ed3a86b1..df2253070f1 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -386,10 +386,16 @@ class GRPOSaveState: total_steps: int total_valid_tokens: int # Track total number of non-padding tokens during training val_reward: float # May be removed when no validation metrics are available + # SC may advance the policy version independently from the optimizer-step + # counter. None preserves compatibility with checkpoints predating it. + trainer_version: Optional[int] = None # SingleController only: name of the sampler that wrote the replay buffer, # used to gate the SC buffer restore. None on checkpoints from the other # algorithms and from SC runs that predate this field. sampler_name: Optional[str] = None + # SingleController only: exact last admitted dispatch batch. None preserves + # compatibility with checkpoints that only recorded the trainer version. + sampler_dispatch_index: Optional[int] = None def _initial_grpo_save_state() -> GRPOSaveState: @@ -400,7 +406,9 @@ def _initial_grpo_save_state() -> GRPOSaveState: total_steps=0, total_valid_tokens=0, val_reward=-99999999.0, + trainer_version=None, sampler_name=None, + sampler_dispatch_index=None, ) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 4b345049707..405c9ec436a 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -40,23 +40,41 @@ import asyncio import contextlib +import hashlib +import io import logging import math import os import threading import time +import uuid import warnings from collections import deque from collections.abc import Iterator from functools import partial -from typing import Any, Awaitable, Callable, Optional, Union +from pathlib import Path +from typing import Any, Awaitable, Callable, Optional, Union, cast import ray import torch from ray.exceptions import RayActorError from nemo_rl.algorithms import opd as opd_module -from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLACEMENT_RESERVE_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, + DataPlaneCheckpointMetadata, + DataPlaneMutationCut, + TQReplayMetadataState, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + TransactionalAdmissionSampler, + create_sampler, +) from nemo_rl.algorithms.grpo import ( GRPOSaveState, _write_latest_checkpoint_status, @@ -81,7 +99,8 @@ tensor_field, ) from nemo_rl.data.interfaces import DatumSpec -from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, 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.distributed.refit_watchdog import RefitAborted, is_refit_context_lost @@ -89,6 +108,14 @@ from nemo_rl.experience.failures import RolloutStall from nemo_rl.experience.payload import VIOLATION_TAG_KEYS from nemo_rl.experience.rollout_manager import RolloutOutcome +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_SCHEMA_VERSION, + ROLLOUT_RECOVERY_STATE_FILENAME, + PromptGroupPhase, + RolloutRecoveryState, + build_rollout_recovery_state, + parse_rollout_recovery_state, +) from nemo_rl.models.generation.fleet_health import ShardState from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration @@ -251,6 +278,9 @@ def __init__( # already defaulted any fields missing from older checkpoints. 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[DataPlaneCheckpointMetadata] = ( + 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.total_valid_tokens @@ -258,9 +288,45 @@ def __init__( self._train_cluster = actor_args.train_cluster self._inference_cluster = actor_args.inference_cluster + restored_trainer_version = ( + actor_args.save_state.trainer_version + if actor_args.save_state.trainer_version is not None + else actor_args.save_state.current_step + ) num_prompts_per_step = self._algo_cfg.num_prompts_per_step self._sampler = create_sampler(self._buffer, self._async_cfg.sampler) - self._sampler.set_dispatch_index(actor_args.save_state.current_step) + restored_dispatch_index = actor_args.save_state.sampler_dispatch_index + if restored_dispatch_index is None: + # Checkpoints predating exact sampler state reconstruct the original + # fresh-step invariant from the restored trainer version. + self._sampler.set_dispatch_index(restored_trainer_version) + else: + self._sampler.restore_dispatch_index(restored_dispatch_index) + if ( + self._master_config.checkpointing["enabled"] + and self._sampler.supports_buffer_checkpoint + and not self._master_config.checkpointing.get("save_data_plane") + ): + raise ValueError( + "SingleController checkpointing with a replay-checkpoint-capable " + "sampler requires checkpointing.save_data_plane=true so " + "completed, unconsumed rollouts are recoverable." + ) + restoring_rollout_recovery = bool( + self._data_plane_checkpoint_metadata is not None + and self._data_plane_checkpoint_metadata.get( + "rollout_recovery_payload_sha256" + ) + is not None + ) + self._rollout_recovery_enabled = bool( + restoring_rollout_recovery + or ( + self._master_config.checkpointing["enabled"] + and self._master_config.checkpointing.get("save_data_plane") + and self._sampler.supports_buffer_checkpoint + ) + ) required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step) validate_sampler_buffer_capacity( self._async_cfg, @@ -269,6 +335,18 @@ 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() + 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() @@ -321,7 +399,7 @@ def __init__( self._async_cfg.max_buffered_rollouts ) - self._trainer_version: int = actor_args.save_state.current_step + self._trainer_version: int = restored_trainer_version 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] = { @@ -358,7 +436,10 @@ async def run(self) -> dict[str, Any]: await self._sync_weights() self._rollout_manager.set_weight_version(self._trainer_version) - await self._maybe_restore_replay_buffer() + restored_replay_groups = await self._maybe_restore_replay_buffer() + await self._maybe_restore_rollout_recovery( + restored_replay_groups=restored_replay_groups + ) await self._maybe_restore_replacement_reserve() # Start the rollout and train pumps, plus the watchdog @@ -423,50 +504,456 @@ async def ping(self) -> dict[str, Any]: # ── internal helpers ─────────────────────────────────────────────────── - async def _maybe_restore_replay_buffer(self) -> None: - """Restore replay-buffer groups from the previous run's checkpoint. + async def _maybe_restore_replay_buffer(self) -> int: + """Restore the local replay index for the native TQ checkpoint. - Skipped with a warning when the checkpoint was written under a - different sampler: restored groups carry the saving sampler's - weight/target-step stamps, which another policy may never select. + Recovery is authoritative only for samplers that explicitly support + buffered-group restoration. The native snapshot and replay metadata file + must both be present and agree on their manifest and group count. """ if self._last_checkpoint_path is None: - return - buffer_path = os.path.join(self._last_checkpoint_path, "replay_buffer.pt") - if not os.path.exists(buffer_path): - print( - f"⚠️ No replay buffer checkpoint found at {buffer_path}. " - "Starting with an empty replay buffer.", - flush=True, + return 0 + 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" ) - return - saved_sampler_name = self._save_state.sampler_name - current_sampler_name = self._async_cfg.sampler.name - if saved_sampler_name != current_sampler_name: + if not self._sampler.supports_buffer_checkpoint: + return 0 + 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"⚠️ Replay buffer checkpoint was saved with sampler " - f"{saved_sampler_name!r} but this run uses " - f"{current_sampler_name!r}; skipping the buffer restore.", + f"⚠️ No native replay metadata found at {metadata_path}. " + "Starting with an empty replay buffer.", flush=True, ) - return - print(f"📦 Restoring replay buffer from checkpoint: {buffer_path}") - # weights_only=False: groups hold pickled KVBatchMeta/TensorDicts, - # not plain tensors. The checkpoint is a trusted same-job artifact. + return 0 + print(f"📦 Restoring replay buffer metadata: {metadata_path}") + # weights_only=False: the replay metadata file contains pickled KVBatchMeta + # objects but no rollout tensor payloads. It is a trusted same-job artifact. buffer_state = await asyncio.to_thread( - torch.load, buffer_path, weights_only=False + 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._algo_cfg.num_generations_per_prompt, + expected_manifest_digest=expected_manifest_digest_value, ) - # Each buffered group holds one _buffer_capacity permit; the load - # truncation guarantees restored <= capacity, so this never blocks. + 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() + return restored + + async def _maybe_restore_rollout_recovery( + self, + *, + restored_replay_groups: int, + ) -> None: + """Restore unfinished ownership for prioritized rollout-pump redispatch.""" + if self._last_checkpoint_path is None: + return + recovery_path = Path( + self._last_checkpoint_path, + ROLLOUT_RECOVERY_STATE_FILENAME, + ) + metadata = self._data_plane_checkpoint_metadata or {} + expected_payload_sha256 = metadata.get("rollout_recovery_payload_sha256") + if expected_payload_sha256 is None: + if recovery_path.is_file(): + raise RuntimeError( + f"{ROLLOUT_RECOVERY_STATE_FILENAME} exists, but the matching " + "native TQ checkpoint does not advertise rollout recovery" + ) + return + if not isinstance(expected_payload_sha256, str): + raise TypeError( + "rollout_recovery_payload_sha256 must be a string in native " + "TQ checkpoint metadata" + ) + expected_schema_version = metadata.get("rollout_recovery_schema_version") + if ( + isinstance(expected_schema_version, bool) + or expected_schema_version != ROLLOUT_RECOVERY_SCHEMA_VERSION + ): + raise ValueError( + "native TQ checkpoint rollout recovery schema mismatch: " + f"checkpoint={expected_schema_version!r}, " + f"expected={ROLLOUT_RECOVERY_SCHEMA_VERSION}" + ) + expected_group_count = metadata.get("rollout_recovery_group_count") + if ( + isinstance(expected_group_count, bool) + or not isinstance(expected_group_count, int) + or expected_group_count < 0 + ): + raise TypeError( + "rollout_recovery_group_count must be an integer in native " + "TQ checkpoint metadata" + ) + if not recovery_path.is_file(): + raise FileNotFoundError( + "native TQ checkpoint advertises rollout recovery, but the " + f"sidecar is missing at {recovery_path}" + ) + + payload = await asyncio.to_thread(recovery_path.read_bytes) + actual_payload_sha256 = hashlib.sha256(payload).hexdigest() + if actual_payload_sha256 != expected_payload_sha256: + raise ValueError( + "rollout recovery sidecar checksum mismatch: " + f"checkpoint={expected_payload_sha256}, " + f"actual={actual_payload_sha256}" + ) + state = await asyncio.to_thread( + torch.load, + io.BytesIO(payload), + weights_only=True, + ) + parsed_state = parse_rollout_recovery_state(state) + if len(parsed_state.ledger_state["groups"]) != expected_group_count: + raise ValueError( + "rollout recovery sidecar group count does not match native " + "TQ checkpoint metadata" + ) + + recovery_ledger = self._rollout_manager.recovery_ledger + async with self._data_plane_checkpoint_barrier.mutation() as cut: + recovery_ledger.load_state_dict(cut, parsed_state.ledger_state) + self._batch_shortfall = parsed_state.batch_shortfall + canonical_state = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + canonical_group_ids = { + group["group_id"] for group in canonical_state["groups"] + } + recovery_ledger.discard_canonical_groups(cut, canonical_group_ids) + await self._rehydrate_rollout_recovery_prompts(cut) + self._sampler_stamps_target_steps = ( + parsed_state.sampler_stamps_target_steps + if parsed_state.sampler_stamps_target_steps is not None + else any( + group.target_step is not None for group in recovery_ledger.groups() + ) + or any( + group.get("target_step") is not None + for group in canonical_state["groups"] + ) + ) + + groups_to_recover = recovery_ledger.groups() + if groups_to_recover: + print( + f"📦 Loaded {len(groups_to_recover)} unfinished rollout " + f"group(s) next to {restored_replay_groups} canonical group(s); " + "the rollout pump will redispatch them before new dataloader work", + flush=True, + ) + + async def _rehydrate_rollout_recovery_prompts( + self, + cut: DataPlaneMutationCut, + ) -> None: + """Resolve positional prompt references against a stable map-style dataset. + + This assumes the dataset exposes integer ``__getitem__`` and retains the + same ordering across checkpoint and restart. + """ + recovery_ledger = self._rollout_manager.recovery_ledger + groups = recovery_ledger.groups() + if not groups: + return + + dataset = getattr(self._dataloader, "dataset", None) + if dataset is None: + raise RuntimeError( + "cannot restore unfinished rollouts because the dataloader does " + "not expose its source dataset" + ) + + resolved_prompts: dict[str, DatumSpec] = {} + for group in groups: + sample_id = group.prompt_ref.sample_id + try: + sample_index = int(sample_id) + except ValueError as error: + raise ValueError( + f"recovery group {group.group_id!r} has a non-integer " + f"dataset sample_id={sample_id!r}" + ) from error + if sample_index < 0 or str(sample_index) != sample_id: + raise ValueError( + f"recovery group {group.group_id!r} has a non-canonical " + f"dataset sample_id={sample_id!r}" + ) + + prompt = resolved_prompts.get(sample_id) + if prompt is None: + try: + dataset_prompt = await asyncio.to_thread( + dataset.__getitem__, sample_index + ) + except (IndexError, KeyError) as error: + raise RuntimeError( + f"cannot rehydrate recovery group {group.group_id!r}: " + f"dataset sample_id={sample_id!r} is unavailable" + ) from error + if not isinstance(dataset_prompt, dict): + raise TypeError( + f"dataset sample_id={sample_id!r} resolved to " + f"{type(dataset_prompt).__name__}, expected a DatumSpec " + "dictionary" + ) + + # Re-run one-row collation to reconstruct the tensor scalars, + # optional fields, and multimodal wrappers expected by RolloutManager. + collate_fn = getattr(self._dataloader, "collate_fn", None) + if collate_fn is None: + prompt = dataset_prompt + else: + prompt_batch = await asyncio.to_thread( + collate_fn, + [dataset_prompt], + ) + if isinstance(prompt_batch, BatchedDataDict): + if prompt_batch.size != 1: + raise ValueError( + "recovery collation must return exactly one prompt; " + f"sample_id={sample_id!r}, size={prompt_batch.size}" + ) + prompt = {key: value[0] for key, value in prompt_batch.items()} + elif isinstance(prompt_batch, dict): + # Identity-style collators used by lightweight/custom + # dataloaders may return the DatumSpec directly. + prompt = prompt_batch + else: + raise TypeError( + "recovery collation for " + f"sample_id={sample_id!r} returned " + f"{type(prompt_batch).__name__}, expected a mapping" + ) + resolved_prompts[sample_id] = cast(DatumSpec, prompt) + recovery_ledger.bind_runtime_prompt( + cut, + group.group_id, + cast(DatumSpec, prompt), + ) + + async def _admit_reserved_prompt_groups( + self, + group_ids: list[str], + ) -> tuple[Optional[int], list[str], int]: + """Commit one admission and atomically reconcile restored canonical groups. + + Returns: + The target-step stamp, IDs that still require rollout dispatch, and the + number of already-canonical groups that replaced reservations in this + admission. + """ + if not group_ids: + raise ValueError("sampler admission requires at least one prompt group") + + def _commit( + cut: DataPlaneMutationCut, + target_step: Optional[int], + ) -> tuple[Optional[int], list[str], int]: + if target_step is not None: + self._sampler_stamps_target_steps = True + for group_id in group_ids: + self._rollout_manager.mark_prompt_group_admitted( + cut, + group_id, + target_step=target_step, + ) + + buffered = 0 + dispatch_group_ids = group_ids + if target_step is not None: + buffered = self._buffer.count_for_target_step(target_step) + if buffered: + dispatch_count = max(0, len(group_ids) - buffered) + dispatch_group_ids = group_ids[:dispatch_count] + for group_id in group_ids[dispatch_count:]: + self._rollout_manager.discard_prompt_group(cut, group_id) + return target_step, dispatch_group_ids, buffered + + if isinstance(self._sampler, TransactionalAdmissionSampler): + await self._sampler.wait_until_admissible( + trainer_version_fn=lambda: self._trainer_version + ) + async with self._data_plane_checkpoint_barrier.mutation() as cut: + target_step = self._sampler.commit_admission(cut) + return _commit(cut, target_step) + + # Custom samplers retain their existing monolithic admission API. Hold + # the mutation cut across it for correctness. Contract: a custom admit() + # must not wait on anything beyond a single trainer_version increment -- + # a checkpoint drains mutation slots while blocking the train pump, so a + # longer wait deadlocks the run. Implement TransactionalAdmissionSampler + # to keep the gate wait outside the mutation cut entirely. + async with self._data_plane_checkpoint_barrier.mutation() as cut: + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version + ) + return _commit(cut, target_step) + + async def _redispatch_restored_rollouts( + self, + launch: Callable[[DatumSpec, Optional[int], str], Awaitable[None]], + ) -> None: + """Prioritize durable unfinished groups while the train pump drains TQ. + + Launching happens inside the ordinary rollout pump so restored groups use + the same in-flight and replay-capacity semaphores as new work. The train + pump runs concurrently and releases replay capacity as it consumes + canonical or newly recovered groups; therefore recovery cannot deadlock + merely because the checkpoint contained more unfinished ownership records + than free replay slots. + """ + recovery_ledger = self._rollout_manager.recovery_ledger + groups_to_recover = recovery_ledger.groups() + if not groups_to_recover: + return + + recognized_phases = ( + PromptGroupPhase.ADMITTED, + PromptGroupPhase.RESERVED, + ) + unhandled_groups = [ + group for group in groups_to_recover if group.phase not in recognized_phases + ] + if unhandled_groups: + details = ", ".join( + f"{group.group_id}={group.phase!r}" for group in unhandled_groups + ) + raise RuntimeError(f"unrecognized rollout recovery phase(s): {details}") + + # ADMITTED groups may be the only work capable of advancing the trainer and + # opening the sampler gate. Launch them before waiting to re-admit RESERVED + # groups, or restore can deadlock with the trainer waiting for recovered work + # that this method has not launched yet. + redispatched = 0 + for group in groups_to_recover: + if group.phase is PromptGroupPhase.ADMITTED: + await launch( + group.prompt_payload, + group.target_step, + group.group_id, + ) + redispatched += 1 + + # A checkpoint may land after dataloader ownership is recorded but before + # sampler admission commits. Re-admit each original dataloader batch once and + # launch it immediately; do not wait for every reserved batch to pass its gate. + reserved_admissions: dict[str, list[str]] = {} + for group in groups_to_recover: + if group.phase is PromptGroupPhase.RESERVED: + reserved_admissions.setdefault(group.admission_id, []).append( + group.group_id + ) + for group_ids in reserved_admissions.values(): + _, dispatch_group_ids, _ = await self._admit_reserved_prompt_groups( + group_ids + ) + for group_id in dispatch_group_ids: + group = recovery_ledger.get_group(group_id) + await launch( + group.prompt_payload, + group.target_step, + group.group_id, + ) + redispatched += 1 + + print( + f"📦 Redispatched {redispatched} unfinished rollout " + "group(s) before new dataloader work", + flush=True, + ) + + async def _validate_replay_inventory( + self, replay_metadata: TQReplayMetadataState + ) -> None: + """Require the canonical TQ keys to match the SC replay index exactly. + + Live checkpoint callers must hold the exclusive data-plane barrier so + commits and clears cannot race this inventory read. Restore calls are + also safe before the rollout and train pumps start any live writers. + """ + 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 " + f"{REPLAY_BUFFER_METADATA_FILENAME}: " + 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 _maybe_restore_replacement_reserve(self) -> None: """Restore spare prompts diverted before the previous run's checkpoint. @@ -485,7 +972,7 @@ async def _maybe_restore_replacement_reserve(self) -> None: if self._last_checkpoint_path is None: return reserve_path = os.path.join( - self._last_checkpoint_path, "replacement_reserve.pt" + self._last_checkpoint_path, REPLACEMENT_RESERVE_FILENAME ) # Absent for every run that never diverted a batch, which is every run that # does not use "replace" -- so silence here rather than the buffer restore's @@ -519,6 +1006,95 @@ async def _call_dp(self, method_name: str, **kwargs) -> Any: return await result return result + 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: PathLike, + replay_metadata: Optional[TQReplayMetadataState] = None, + rollout_recovery_payload_sha256: Optional[str] = None, + rollout_recovery_group_count: Optional[int] = 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 replay index 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, + ) + save_state = self._save_state + checkpoint_trainer_version = save_state.trainer_version + if checkpoint_trainer_version is None: + raise RuntimeError( + "Cannot save a data-plane checkpoint before trainer_version " + "is captured in the controller save state" + ) + metadata: DataPlaneCheckpointMetadata = { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": save_state.current_step, + "single_controller_trainer_version": checkpoint_trainer_version, + "single_controller_epoch": save_state.current_epoch, + "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["replay_metadata_schema_version"] = ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION + ) + metadata["replay_manifest_digest"] = replay_metadata["manifest_digest"] + metadata["replay_group_count"] = len(replay_metadata["groups"]) + if rollout_recovery_payload_sha256 is not None: + if rollout_recovery_group_count is None: + raise ValueError("rollout recovery payload hash requires a group count") + metadata["rollout_recovery_schema_version"] = ( + ROLLOUT_RECOVERY_SCHEMA_VERSION + ) + metadata["rollout_recovery_payload_sha256"] = ( + rollout_recovery_payload_sha256 + ) + metadata["rollout_recovery_group_count"] = rollout_recovery_group_count + elif rollout_recovery_group_count is not None: + raise ValueError("rollout recovery group count requires a payload hash") + 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 ─────────────────────── async def _rollout_pump(self) -> None: @@ -553,6 +1129,7 @@ async def _rollout_pump(self) -> None: async def _dispatch_one_prompt( prompt: DatumSpec, target_step: Optional[int], + lineage_group_id: Optional[str], task_started_event: asyncio.Event, ) -> None: task_started_event.set() @@ -567,11 +1144,19 @@ async def _dispatch_one_prompt( try: while True: try: - outcome = await self._rollout_manager.generate_and_push( - prompt, - target_step=target_step, - inflight_registry=self._inflight_by_group_id, - ) + if lineage_group_id is None: + outcome = await self._rollout_manager.generate_and_push( + prompt, + target_step=target_step, + inflight_registry=self._inflight_by_group_id, + ) + else: + outcome = await self._rollout_manager.generate_and_push( + prompt, + target_step=target_step, + inflight_registry=self._inflight_by_group_id, + lineage_group_id=lineage_group_id, + ) except BaseException: # On success ownership transfers to the train pump, which # releases this permit after consuming the committed group. @@ -581,12 +1166,42 @@ async def _dispatch_one_prompt( if outcome is not RolloutOutcome.SKIPPED: break - replacement = self._take_replacement(target_step, replacements) + if self._rollout_recovery_enabled: + assert lineage_group_id is not None + async with ( + self._data_plane_checkpoint_barrier.mutation() + ) as cut: + replacement = self._take_replacement( + target_step, replacements + ) + # A skipped tracked prompt remains ledger-owned until this + # controller transition. Dropping the old owner, reserving + # a replacement, or crediting the target step short must be + # one checkpoint-atomic decision. + self._rollout_manager.discard_prompt_group( + cut, lineage_group_id + ) + if replacement is not None: + lender_step = self._promote_into_step(target_step) + if lender_step is not None: + target_step = lender_step + lineage_group_id = ( + self._rollout_manager.reserve_prompt_group( + cut, + replacement, + target_step=target_step, + ) + ) + else: + self._credit_shortfall(target_step) + else: + replacement = self._take_replacement(target_step, replacements) if replacement is None: # Nothing was committed, so the train pump will never see this # group and never release its permit on our behalf. self._buffer_capacity.release() - self._credit_shortfall(target_step) + if not self._rollout_recovery_enabled: + self._credit_shortfall(target_step) return replacements += 1 @@ -601,9 +1216,10 @@ async def _dispatch_one_prompt( # Attempted only now that a spare is in hand, because the borrow is a # debt and the spare is what repays it. Borrowing without one would # leave the lender short instead: the same hole, one step later. - lender_step = self._promote_into_step(target_step) - if lender_step is not None: - target_step = lender_step + if not self._rollout_recovery_enabled: + lender_step = self._promote_into_step(target_step) + if lender_step is not None: + target_step = lender_step # A substitution is a fresh rollout, not a continuation of the one # that failed, so it observes the same pause a first dispatch does # instead of pushing new generation into a weight-sync window. @@ -638,7 +1254,16 @@ def _release_permits_if_task_not_started( self._buffer_capacity.release() sem.release() - async def _launch(prompt: DatumSpec, target_step: Optional[int]) -> None: + async def _launch( + prompt: DatumSpec, + target_step: Optional[int], + lineage_group_id: Optional[str], + ) -> None: + if self._rollout_recovery_enabled and lineage_group_id is None: + raise RuntimeError( + "recovery-enabled rollout dispatch requires a pre-reserved " + "prompt-group ID" + ) # check if buffer is full await self._buffer_capacity.acquire() # check if inflight rollouts is full @@ -649,7 +1274,12 @@ async def _launch(prompt: DatumSpec, target_step: Optional[int]) -> None: task_started_event = asyncio.Event() # dispatch rollout task = rollout_tasks.create_task( - _dispatch_one_prompt(prompt, target_step, task_started_event) + _dispatch_one_prompt( + prompt, + target_step, + lineage_group_id, + task_started_event, + ) ) self._dispatched_rollouts.add(task) task.add_done_callback(self._dispatched_rollouts.discard) @@ -662,36 +1292,88 @@ async def _launch(prompt: DatumSpec, target_step: Optional[int]) -> None: max_epochs = self._algo_cfg.max_num_epochs async with asyncio.TaskGroup() as rollout_tasks: + if self._rollout_recovery_enabled: + await self._redispatch_restored_rollouts(_launch) while max_epochs is None or self._current_epoch < max_epochs: - for prompt_batch in self._dataloader: - if self._divert_batch_to_reserve(prompt_batch): - continue + if not self._rollout_recovery_enabled: + for prompt_batch in self._dataloader: + if self._divert_batch_to_reserve(prompt_batch): + continue + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version + ) + if target_step is not None: + self._sampler_stamps_target_steps = True + num_prompts = prompt_batch.size + if target_step is not None: + buffered = self._buffer.count_for_target_step(target_step) + if buffered: + num_prompts = max(0, prompt_batch.size - buffered) + print( + f" target_step={target_step}: {buffered} group(s) " + f"already buffered; dispatching {num_prompts} of " + f"{prompt_batch.size} prompt(s), dropping the rest", + flush=True, + ) + for prompt_idx in range(num_prompts): + prompt: DatumSpec = { # type: ignore + k: v[prompt_idx] for k, v in prompt_batch.items() + } + await _launch(prompt, target_step, None) + self._current_epoch += 1 + continue - target_step = await self._sampler.admit( - trainer_version_fn=lambda: self._trainer_version + dataloader_iterator = iter(self._dataloader) + while True: + prompt_dispatches: list[tuple[DatumSpec, str]] = [] + async with self._data_plane_checkpoint_barrier.mutation() as cut: + try: + prompt_batch = next(dataloader_iterator) + except StopIteration: + self._current_epoch += 1 + break + if self._divert_batch_to_reserve(prompt_batch): + continue + admission_id = str(uuid.uuid4()) + for prompt_idx in range(prompt_batch.size): + prompt = { # type: ignore + k: v[prompt_idx] for k, v in prompt_batch.items() + } + group_id = self._rollout_manager.reserve_prompt_group( + cut, + prompt, + target_step=None, + admitted=False, + admission_id=admission_id, + ) + prompt_dispatches.append((prompt, group_id)) + + ( + target_step, + dispatch_group_ids, + buffered, + ) = await self._admit_reserved_prompt_groups( + [group_id for _, group_id in prompt_dispatches] ) - if target_step is not None: - self._sampler_stamps_target_steps = True - num_prompts = prompt_batch.size if target_step is not None: - buffered = self._buffer.count_for_target_step(target_step) if buffered: - num_prompts = max(0, prompt_batch.size - buffered) print( f" target_step={target_step}: {buffered} group(s) " - f"already buffered; dispatching {num_prompts} of " - f"{prompt_batch.size} prompt(s), dropping the rest", + f"already buffered; dispatching " + f"{len(dispatch_group_ids)} of " + f"{len(prompt_dispatches)} prompt(s), dropping the rest", flush=True, ) + dispatch_group_id_set = set(dispatch_group_ids) + prompt_dispatches = [ + (prompt, group_id) + for prompt, group_id in prompt_dispatches + if group_id in dispatch_group_id_set + ] - for prompt_idx in range(num_prompts): - prompt: DatumSpec = { # type: ignore - k: v[prompt_idx] for k, v in prompt_batch.items() - } - await _launch(prompt, target_step) - - self._current_epoch += 1 + for prompt, group_id in prompt_dispatches: + await _launch(prompt, target_step, group_id) # Only now that every dispatched rollout has settled is the pool genuinely # spare. Draining it inside the group above would race them for it, and a @@ -752,7 +1434,8 @@ def _divert_batch_to_reserve( return True async def _drain_reserve_into_steps( - self, launch: Callable[[DatumSpec, Optional[int]], Awaitable[None]] + self, + launch: Callable[[DatumSpec, Optional[int], Optional[str]], Awaitable[None]], ) -> None: """Train on the leftover spares once the dataloader has nothing more to give. @@ -776,6 +1459,46 @@ async def _drain_reserve_into_steps( """ num_prompts_per_step = self._algo_cfg.num_prompts_per_step while len(self._replacement_reserve) >= num_prompts_per_step: + if self._rollout_recovery_enabled: + prompt_dispatches: list[tuple[DatumSpec, str]] = [] + async with self._data_plane_checkpoint_barrier.mutation() as cut: + step_prompts = [ + self._replacement_reserve.popleft() + for _ in range(num_prompts_per_step) + ] + admission_id = str(uuid.uuid4()) + for prompt in step_prompts: + group_id = self._rollout_manager.reserve_prompt_group( + cut, + prompt, + target_step=None, + admitted=False, + admission_id=admission_id, + ) + prompt_dispatches.append((prompt, group_id)) + ( + target_step, + dispatch_group_ids, + buffered, + ) = await self._admit_reserved_prompt_groups( + [group_id for _, group_id in prompt_dispatches] + ) + dispatch_group_id_set = set(dispatch_group_ids) + prompt_dispatches = [ + (prompt, group_id) + for prompt, group_id in prompt_dispatches + if group_id in dispatch_group_id_set + ] + print( + f" dataloader exhausted; training on {len(prompt_dispatches)} " + f"pooled spare(s) as target_step={target_step}" + + (f" ({buffered} group(s) already buffered)" if buffered else ""), + flush=True, + ) + for prompt, group_id in prompt_dispatches: + await launch(prompt, target_step, group_id) + continue + # Take the step's prompts out before the first await. A drop resolving # concurrently draws from this same pool, and could otherwise claim one of # them and leave the step it is filling one group short. @@ -791,7 +1514,7 @@ async def _drain_reserve_into_steps( flush=True, ) for prompt in step_prompts: - await launch(prompt, target_step) + await launch(prompt, target_step, None) if self._replacement_reserve: print( @@ -1199,11 +1922,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 chunks_dispatched += 1 @@ -1935,20 +2654,39 @@ async def _check_env_health(self, timeout_s: float) -> list[str]: async def _abort_stale_inflight(self) -> int: """Abort in-flight rollouts that the sampler can no longer select.""" - stale_tasks = [ - task - for task, start_version in self._inflight_by_group_id.values() - if self._sampler.should_abort_inflight( - start_weight_version=start_version, - current_train_weight=self._trainer_version, - ) - ] - if not stale_tasks: - return 0 - for task in stale_tasks: - task.cancel() + def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: + stale_groups: list[tuple[str, asyncio.Task[None]]] = [] + for group_id, inflight in self._inflight_by_group_id.items(): + task, start_version = inflight + if self._sampler.should_abort_inflight( + start_weight_version=start_version, + current_train_weight=self._trainer_version, + ): + stale_groups.append((group_id, task)) + return stale_groups + + if self._rollout_recovery_enabled: + async with self._data_plane_checkpoint_barrier.mutation() as cut: + # Re-evaluate after acquiring the cut: a rollout may have completed + # while a checkpoint holder delayed this mutation. + stale_groups = _stale_groups() + for group_id, _ in stale_groups: + # This is an intentional live abort, not a process failure. Remove + # durable ownership before cancellation cleanup removes the unready + # TQ slot, so a concurrent checkpoint cannot resurrect the prompt. + self._rollout_manager.discard_prompt_group(cut, group_id) + for _, task in stale_groups: + task.cancel() + else: + stale_groups = _stale_groups() + for _, task in stale_groups: + task.cancel() + if not stale_groups: + return 0 + + stale_tasks = [task for _, task in stale_groups] results = await asyncio.gather(*stale_tasks, return_exceptions=True) failures = [ result @@ -1982,26 +2720,6 @@ async def _save_checkpoint( stepped. """ save_state = self._save_state - save_state.current_step = self._train_steps - save_state.total_steps = self._train_steps - save_state.current_epoch = self._current_epoch - save_state.consumed_samples = self._consumed_samples - save_state.total_valid_tokens = self._total_valid_tokens - # The restore skips the replay buffer when the resuming run uses a - # different sampler (its stamps may never be selectable there). - save_state.sampler_name = self._async_cfg.sampler.name - # Snapshot before any await so it can't interleave with - # _rollout_pump iterating this same dataloader. - dataloader_state = self._dataloader.state_dict() - # The spare pool has to be saved with that snapshot, not left out of the - # checkpoint: diverting a batch already advanced the iterator, so the state - # above records those prompts as consumed while they are still only in memory. - # Without this a resumed replace-mode run comes back with an empty pool and a - # dataloader positioned past the diverted batch, silently losing it -- and - # losing it for good, since _drain_reserve_into_steps only ever recovers spares - # held by the process that diverted them. Snapshotted here, in the same - # await-free window, so the pair cannot disagree. - reserve_state = list(self._replacement_reserve) # SC has no validation loop yet; drop the default sentinel instead of # persisting a bogus val_reward. if hasattr(save_state, "val_reward"): @@ -2031,12 +2749,81 @@ async def _save_checkpoint( await asyncio.to_thread(self._checkpointer.finalize_pending) print(f"Saving checkpoint for step {self._train_steps}...") - checkpoint_path: PathLike = await asyncio.to_thread( # pyrefly: ignore[bad-assignment] the PathLike alias resolves inconsistently under pyrefly's import-cycle breaking - self._checkpointer.init_tmp_checkpoint, - self._train_steps, - vars(save_state), - self._master_config, - ) + replay_metadata: Optional[TQReplayMetadataState] = None + rollout_recovery_state: Optional[RolloutRecoveryState] = None + rollout_recovery_payload: Optional[bytes] = None + rollout_recovery_payload_sha256: Optional[str] = None + + # Admission, dataloader movement, replay mutations, and canonical TQ writes + # all take the mutation side of this barrier. Capture every restart-facing + # controller artifact under the exclusive side so the checkpoint cannot + # contain a cursor without its prompt owner, or two durable owners for one + # canonical group. + async with self._data_plane_checkpoint_barrier.checkpoint(): + save_state.current_step = self._train_steps + save_state.total_steps = self._train_steps + save_state.trainer_version = self._trainer_version + save_state.current_epoch = self._current_epoch + save_state.consumed_samples = self._consumed_samples + save_state.total_valid_tokens = self._total_valid_tokens + save_state.sampler_name = self._async_cfg.sampler.name + save_state.sampler_dispatch_index = self._sampler.dispatch_index + dataloader_state = self._dataloader.state_dict() + # The spare pool and dataloader advance together under the same mutation + # cut in recovery-enabled dispatch, so preserve them in this cut too. + reserve_state = list(self._replacement_reserve) + + checkpoint_path: PathLike = await asyncio.to_thread( # pyrefly: ignore[bad-assignment] the PathLike alias resolves inconsistently under pyrefly's import-cycle breaking + self._checkpointer.init_tmp_checkpoint, + self._train_steps, + vars(save_state), + self._master_config, + ) + + if self._master_config.checkpointing.get("save_data_plane"): + if self._sampler.supports_buffer_checkpoint: + replay_metadata = self._buffer.metadata_state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + if replay_metadata is not None: + await self._validate_replay_inventory(replay_metadata) + + if self._rollout_recovery_enabled: + rollout_recovery_state = build_rollout_recovery_state( + self._rollout_manager.recovery_ledger, + batch_shortfall=self._batch_shortfall, + sampler_stamps_target_steps=(self._sampler_stamps_target_steps), + ) + if replay_metadata is not None: + canonical_group_ids = { + group["group_id"] for group in replay_metadata["groups"] + } + rollout_recovery_state["groups"] = [ + group + for group in rollout_recovery_state["groups"] + if group["group_id"] not in canonical_group_ids + ] + payload_buffer = io.BytesIO() + await asyncio.to_thread( + torch.save, + rollout_recovery_state, + payload_buffer, + ) + rollout_recovery_payload = payload_buffer.getvalue() + rollout_recovery_payload_sha256 = hashlib.sha256( + rollout_recovery_payload + ).hexdigest() + + await self._save_data_plane_checkpoint( + checkpoint_path, + replay_metadata=replay_metadata, + rollout_recovery_payload_sha256=(rollout_recovery_payload_sha256), + rollout_recovery_group_count=( + len(rollout_recovery_state["groups"]) + if rollout_recovery_state is not None + else None + ), + ) # Save value model if self._is_ppo: @@ -2088,17 +2875,22 @@ async def _save_checkpoint( await asyncio.to_thread( torch.save, reserve_state, - os.path.join(checkpoint_path, "replacement_reserve.pt"), + os.path.join(checkpoint_path, REPLACEMENT_RESERVE_FILENAME), + ) + if replay_metadata is not None: + await asyncio.to_thread( + torch.save, + replay_metadata, + os.path.join(checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME), + ) + if rollout_recovery_payload is not None: + await asyncio.to_thread( + Path( + checkpoint_path, + ROLLOUT_RECOVERY_STATE_FILENAME, + ).write_bytes, + rollout_recovery_payload, ) - buffer_state = await self._buffer.state_dict( - saved_capacity=self._async_cfg.max_buffered_rollouts - ) - await asyncio.to_thread( - torch.save, - buffer_state, - os.path.join(checkpoint_path, "replay_buffer.pt"), - ) - # Rename happens in the background once the async weight writes # finish; flushed at the next save or on exit. self._checkpointer.begin_finalization( @@ -2361,7 +3153,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: return meta, True 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, diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 25daea5f0e7..4f77ba4c802 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -22,6 +22,7 @@ from __future__ import annotations import time +import warnings from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from functools import partial @@ -35,7 +36,17 @@ from transformers.tokenization_utils_base import PreTrainedTokenizerBase from nemo_rl.algorithms import opd as opd_module -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointMetadata, + TQReplayBuffer, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + sampler_supports_buffer_checkpoint, +) from nemo_rl.algorithms.grpo import ( GRPOSaveState, _get_effort_config, @@ -59,7 +70,12 @@ from nemo_rl.algorithms.utils import set_seed from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.utils import load_dataloader_state, setup_response_data -from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client +from nemo_rl.data_plane import ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneClient, + build_data_plane_client, + data_plane_supports_checkpointing, +) from nemo_rl.data_plane.schema import ( SC_ROLLOUT_SCHEMA_FIELDS, fields_with_optional_routed_experts, @@ -132,13 +148,10 @@ class SingleControllerActorArgs: partition_id: str save_state: GRPOSaveState last_checkpoint_path: Optional[str] - # Defaulted fields must follow the required ones above, so these two stay last. - # None when async_rl.generation_fleet_health is disabled; the SingleController drives the - # probe loop when it is present. + data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None + # None when async_rl.generation_fleet_health is disabled. fleet_monitor: Optional[GenerationFleetHealth] = None - # None unless async_rl.generation_router is enabled; the SingleController pushes the - # serving backend set to it. Parameterized with the Impl class because the decorated - # GenerationRouterActor name is an ActorClass instance, not a type. + # None unless async_rl.generation_router is enabled. generation_router: Optional[ray.actor.ActorHandle[GenerationRouterImpl]] = None # Populated only for text MOPD. Aliases may outnumber worker groups when # multiple agents share one deduplicated teacher checkpoint. @@ -150,6 +163,102 @@ class SingleControllerActorArgs: value_loss_fn: Optional[LossFunction] = 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[DataPlaneCheckpointMetadata]: + """Load and validate an authoritative native TQ checkpoint when present. + + The replay metadata file 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." + ) + print( + f"⚠️ No {REPLAY_BUFFER_METADATA_FILENAME} found in checkpoint " + f"{checkpoint_path}. The matching TQ checkpoint will not be loaded, " + "and recovery will use an empty replay buffer. The dataloader cursor " + "is still restored, so any prompt groups buffered at checkpoint time " + "will be discarded.", + flush=True, + ) + 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) + raw_metadata = policy.load_data_plane_checkpoint(data_plane_path) + if not isinstance(raw_metadata, dict): + raise TypeError( + "Native TQ checkpoint load must return a metadata dictionary, " + f"got {type(raw_metadata).__name__}" + ) + metadata = cast(DataPlaneCheckpointMetadata, raw_metadata) + expected_values: DataPlaneCheckpointMetadata = { + "data_plane_checkpoint_schema_version": (DATA_PLANE_CHECKPOINT_SCHEMA_VERSION), + "single_controller_train_steps": save_state.current_step, + "single_controller_trainer_version": ( + save_state.trainer_version + if save_state.trainer_version is not None + else 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( + f"📦 Native TQ checkpoint restored and validated: groups={group_count}", + flush=True, + ) + return metadata + + def _non_colocated_teacher_node_count(master_config: MasterConfig) -> int: """Validate teacher GPU geometry and return its deduplicated node count.""" if not opd_module.is_non_colocated_teachers_enabled(master_config): @@ -809,6 +918,43 @@ def setup_single_controller( "master_config.data_plane.enabled=True. The async-RL " "SingleController path is built on the TransferQueue data plane." ) + data_plane_checkpointing_supported = data_plane_supports_checkpointing(dp_config) + if ( + master_config.checkpointing.get("save_data_plane") + and not data_plane_checkpointing_supported + ): + raise NotImplementedError( + "SingleController data-plane checkpointing is not supported for " + f"data_plane.backend={dp_config['backend']!r}." + ) + if master_config.checkpointing["enabled"]: + sampler_supports_replay_recovery = sampler_supports_buffer_checkpoint( + master_config.async_rl.sampler + ) + if sampler_supports_replay_recovery and not master_config.checkpointing.get( + "save_data_plane" + ): + error_message = ( + "SingleController checkpointing with a replay-checkpoint-capable " + "sampler requires checkpointing.save_data_plane=true so " + "completed, unconsumed rollouts are recoverable." + ) + if not data_plane_checkpointing_supported: + error_message += ( + f" The configured data_plane.backend={dp_config['backend']!r} " + "does not support data-plane checkpointing; use " + "data_plane.backend='simple' or set " + "checkpointing.enabled=false." + ) + raise ValueError(error_message) + if not sampler_supports_replay_recovery: + warnings.warn( + f"Sampler {master_config.async_rl.sampler.name!r} cannot recover " + "completed buffered rollouts. On resume, the dataloader cursor " + "is restored while buffered prompt groups are discarded.", + UserWarning, + stacklevel=2, + ) assert generation_config is not None, ( "single_controller_utils.setup requires policy.generation in master_config" @@ -1144,6 +1290,17 @@ def _build_generation_then_trainer( if "value_time" in time_metrics: setup_timing_metrics.value_init_time_s = time_metrics["value_time"] + # Native TQ restore must run through the trainer's bootstrap client before + # the normal SC data-plane client is created or any rollout/train data-plane + # operation starts. + data_plane_checkpoint_metadata = _maybe_restore_native_data_plane_checkpoint( + trainer, + last_checkpoint_path=last_checkpoint_path, + save_state=save_state, + partition_id=partition_id, + sampler_name=master_config.async_rl.sampler.name, + ) + if use_nemo_gym: # the two fields are only meaningful when use_nemo_gym enabled setup_timing_metrics.generation_init_reserve_time_s = gen_reserve_time @@ -1281,6 +1438,7 @@ def _build_generation_then_trainer( partition_id=partition_id, save_state=save_state, last_checkpoint_path=last_checkpoint_path, + data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, fleet_monitor=fleet_monitor, generation_router=generation_router, teacher_worker_groups=teacher_worker_groups, diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py index 56b19178a1c..c97346ed4d2 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -21,18 +21,22 @@ from nemo_rl.data_plane.codec import materialize from nemo_rl.data_plane.factory import build_data_plane_client from nemo_rl.data_plane.interfaces import ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, DataPlaneClient, DataPlaneConfig, KVBatchMeta, + data_plane_supports_checkpointing, ) from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event __all__ = [ + "DATA_PLANE_CHECKPOINT_SCHEMA_VERSION", "DataPlaneClient", "DataPlaneConfig", "KVBatchMeta", "MetricsDataPlaneClient", "build_data_plane_client", + "data_plane_supports_checkpointing", "log_event", "materialize", ] diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py index 1c5b00a5e44..6f98800e505 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,67 @@ 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") + if "partitions" not in state: + raise ValueError( + f"NoOp checkpoint at {checkpoint_file} has no 'partitions' key. " + "It was written by an incompatible version of this adapter, or the " + "write was interrupted. Delete it and re-run the test that produced " + "it; there is nothing to recover because this adapter holds no " + "training state." + ) + 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 b5e803f1f6d..0ad7b6b3dfc 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -25,6 +25,7 @@ import contextlib import glob import ipaddress +import json import os import resource import socket @@ -33,6 +34,7 @@ import warnings import weakref from importlib import resources +from pathlib import Path from queue import Empty, SimpleQueue from typing import Any, cast @@ -49,6 +51,7 @@ DataPlaneConfig, KVBatchMeta, backend_config, + data_plane_supports_checkpointing, ) from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS @@ -725,6 +728,8 @@ 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._supports_checkpointing = data_plane_supports_checkpointing(cfg) self._promote_1d = cfg["backend"] == "mooncake_cpu" if bootstrap: @@ -733,11 +738,36 @@ 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 # Fields whose schema this process has already warmed, per partition. - # See register_partition: the controller's field map is append-only, - # so a field only ever needs warming once. + # The controller's field map is append-only, so each field only needs + # warming once for the lifetime of this client. self._warmed_fields: dict[str, set[str]] = {} + def _require_checkpointing_support(self) -> None: + """Reject backends that cannot round-trip all data-plane state.""" + if not self._supports_checkpointing: + raise NotImplementedError( + "TQ checkpointing is not supported for " + f"data_plane.backend={self._backend!r}: the backend cannot " + "persist and restore all storage rows." + ) + + 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 ─────────────────────────────────────────────── def register_partition( @@ -774,6 +804,7 @@ def register_partition( # (``0@field`` at the Mooncake storage layer). Mooncake does not # support upsert, so repeated schema warmups can collide with # stale metadata from a previous registration. + self._mark_data_operation_started() schema_key = ( f"__schema__:{partition_id}:{os.getpid()}:{id(self)}:{time.time_ns()}" ) @@ -803,6 +834,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] = {} @@ -874,6 +906,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( @@ -915,6 +948,7 @@ def put_samples( 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), @@ -939,6 +973,7 @@ def get_samples( ) -> TensorDict: if not sample_ids: return TensorDict({}, batch_size=(0,)) + self._mark_data_operation_started() td = tq.kv_batch_get( keys=list(sample_ids), partition_id=partition_id, @@ -946,9 +981,16 @@ def get_samples( ) return _from_wire(td) + 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 @@ -968,11 +1010,47 @@ 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.""" + self._require_checkpointing_support() + _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. + """ + self._require_checkpointing_support() + 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 00000000000..aee3ef9b536 --- /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 8ee07141c27..fcff2e398f4 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -37,11 +37,14 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict from pydantic import BaseModel from tensordict import TensorDict +DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 + class SimpleStorageConfig(BaseModel, extra="allow"): """Sizing for ``backend="simple"``. Ignored by every other backend. @@ -128,6 +131,19 @@ class DataPlaneConfig(TypedDict): observability: NotRequired["ObservabilityConfig"] +_CHECKPOINTABLE_BACKENDS: frozenset[str] = frozenset({"simple"}) + + +def data_plane_supports_checkpointing(cfg: DataPlaneConfig) -> bool: + """Return whether the configured backend supports complete save/load. + + This is a static allow-list so an unrecognized future backend defaults to + unsupported until its storage payload and controller metadata are both + known to round-trip through a checkpoint. + """ + return cfg["backend"] in _CHECKPOINTABLE_BACKENDS + + _BACKEND_MODELS: dict[str, type[BaseModel]] = { "simple": SimpleStorageConfig, "mooncake_cpu": MooncakeCpuConfig, @@ -337,7 +353,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* — @@ -489,6 +506,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, @@ -520,6 +553,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 63e551dc209..d569740cf18 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: str) -> list[str]: + 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/experience/payload.py b/nemo_rl/experience/payload.py index 4d1ec62a5f5..e10a9026fa0 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -129,6 +129,7 @@ def pack_payload( *, weight_version: int, group_id: str, + prompt_idx: int, ) -> tuple[list[str], TensorDict, list[dict[str, Any]]]: """Pack a producer batch into (sample_ids, fields, tags) for put_samples. @@ -136,6 +137,7 @@ def pack_payload( train_batch: Mapping with at least input_lengths plus the tensor/object fields to send. weight_version: Trainer weight version stamped on every row's tag. group_id: Per-group identifier used as the sample_id prefix; the caller owns uniqueness. + prompt_idx: Stable dataset prompt index stamped on every row's tag. Returns: sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row @@ -154,5 +156,12 @@ def pack_payload( ) sample_ids = [f"{group_id}_g{i}" for i in range(n)] violations = train_batch.get(_VIOLATION_COUNTS_KEY, [{}] * n) - tags = [{"weight_version": weight_version, **violations[i]} for i in range(n)] + tags = [ + { + "weight_version": weight_version, + "prompt_idx": prompt_idx, + **violations[i], + } + for i in range(n) + ] return sample_ids, fields_td, tags diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 6717542d642..bcadd4f5271 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -25,6 +25,7 @@ from wandb import Table from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneMutationCut, PostWriteEnrichmentError, TQReplayBuffer, ) @@ -43,6 +44,10 @@ ) from nemo_rl.experience.interfaces import Completion, PromptGroupRecord from nemo_rl.experience.metric_utils import calculate_single_metric, pct +from nemo_rl.experience.rollout_recovery import ( + PromptGroupPhase, + RolloutRecoveryLedger, +) from nemo_rl.experience.rollouts import ( EffortLevelsConfig, _apply_effort_shaping, @@ -83,7 +88,8 @@ class RolloutOutcome(str, enum.Enum): # The prompt was given up on within a budget: its data-failure budget within # max_skipped_prompts, or its infrastructure budget within # max_consecutive_dropped_prompts. No group was committed, so the caller owns - # releasing its backpressure permit and crediting the step's shortfall. + # releasing its backpressure permit and atomically replacing the ledger owner or + # crediting the step's shortfall. SKIPPED = "skipped" @@ -1198,6 +1204,7 @@ def __init__( self._tokenizer = tokenizer self._num_generations_per_prompt = num_generations_per_prompt self._tq_buffer = tq_buffer + self._recovery_ledger = RolloutRecoveryLedger() self._weight_version: int = 0 # Run-wide, shared across concurrent generate_and_push calls. Safe as a plain # int: every caller runs on the SingleController's single event loop. @@ -1212,6 +1219,62 @@ def stats(self) -> RolloutStats: """Counters describing retry/skip activity so far.""" return self._stats + @property + def recovery_ledger(self) -> RolloutRecoveryLedger: + """Return the prompt-group ownership ledger shared with the controller.""" + return self._recovery_ledger + + def reserve_prompt_group( + self, + cut: DataPlaneMutationCut, + input_sample: DatumSpec, + *, + target_step: Optional[int], + admitted: bool = True, + admission_id: Optional[str] = None, + ) -> str: + """Own a prompt before controller dispatch can yield or checkpoint.""" + prompt_idx = input_sample.get("idx") + if isinstance(prompt_idx, bool) or not isinstance(prompt_idx, int): + raise ValueError( + "rollout recovery requires every dataloader sample to contain " + f"a stable integer idx, got {prompt_idx!r}" + ) + record = self._recovery_ledger.reserve_group( + cut, + prompt_id=str(prompt_idx), + prompt_payload=input_sample, + expected_generations=self._num_generations_per_prompt, + target_step=target_step, + start_weight_version=self._weight_version, + admitted=admitted, + admission_id=admission_id, + ) + return record.group_id + + def mark_prompt_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: Optional[int], + ) -> None: + """Attach sampler admission state to a pre-admission reservation.""" + self._recovery_ledger.mark_group_admitted( + cut, + group_id, + target_step=target_step, + start_weight_version=self._weight_version, + ) + + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + """Release a reservation that will intentionally never be dispatched.""" + self._recovery_ledger.discard_group(cut, group_id) + def set_weight_version(self, version: int) -> None: """Set the weight_version used for rollout tags. @@ -1229,6 +1292,7 @@ async def generate_and_push( *, target_step: Optional[int] = None, inflight_registry: Optional[dict[str, tuple[asyncio.Task[None], int]]] = None, + lineage_group_id: Optional[str] = None, ) -> RolloutOutcome: """Roll out one prompt and commit it, re-dispatching on infrastructure failure. @@ -1249,14 +1313,18 @@ async def generate_and_push( target_step: Training step this rollout targets; stamped on the buffer slot for StalenessSampler.force_in_order. inflight_registry: Optional controller-owned mapping from group ID to its dispatch task and start weight version. + lineage_group_id: Stable group minted by the rollout ledger before + dataloader dispatch. TQ records this same ID rather than minting one. + ``None`` preserves the ordinary non-checkpointed fresh-ID retry path. Returns: ``COMMITTED`` when the group reached the buffer, or ``SKIPPED`` when the prompt was given up on within a budget: its data budget within ``max_skipped_prompts``, or its infra budget within ``max_consecutive_dropped_prompts``. A ``SKIPPED`` prompt committed nothing, - so the caller owns both its backpressure permit and the shortfall for the - training step it was stamped for. + so the caller owns both its backpressure permit and the checkpoint-atomic + transition from its retained ledger record to either a replacement prompt + or the shortfall for the training step it was stamped for. Raises: RolloutRedispatchExhausted: The infra budget ran out and the fleet has not @@ -1266,6 +1334,20 @@ async def generate_and_push( assert self._tq_buffer is not None, ( "generate_and_push requires tq_buffer to be set at __init__" ) + if lineage_group_id is not None: + lineage_group = self._recovery_ledger.get_group(lineage_group_id) + if lineage_group.phase is not PromptGroupPhase.ADMITTED: + raise RuntimeError( + f"lineage group {lineage_group_id!r} must be admitted " + "before dispatch" + ) + if lineage_group.expected_generations != self._num_generations_per_prompt: + raise ValueError( + f"lineage group {lineage_group_id!r} expects " + f"{lineage_group.expected_generations} generation(s), but " + "the resumed configuration requests " + f"{self._num_generations_per_prompt}" + ) policy = self._retry_policy infra_attempts = 0 data_attempts = 0 @@ -1277,15 +1359,17 @@ async def generate_and_push( # about the prompt rather than about the fleet. while infra_attempts < policy.max_infra_attempts: start_version = self._weight_version - # Reserved inside the loop so each attempt owns a fresh group_id: rows a - # failed attempt may have written cannot then collide with the retry's. + # A lineage-tracked prompt reuses its durable logical ID only after the + # prior attempt's buffer slot was removed successfully. Ordinary callers + # retain the existing fresh-ID-per-attempt behavior. group_id = self._tq_buffer.reserve( - weight_version=start_version, target_step=target_step + weight_version=start_version, + target_step=target_step, + group_id=lineage_group_id, ) try: - # Registered per ATTEMPT, not per prompt: each retry reserves a fresh - # group_id, so the controller's registry must follow the attempt that - # actually owns the slot it might abort. + # Registered per active attempt so cancellation follows the slot that + # currently owns the stable recovery group ID. if inflight_registry is not None: current_task = asyncio.current_task() assert current_task is not None @@ -1307,13 +1391,25 @@ async def generate_and_push( # A failed rollout must not leave an unready slot that can block an # in-order sampler. commit() rolls back any DataPlane rows it wrote. # Cleanup failure must not mask the error that caused it. + cleanup_failed = False try: await self._tq_buffer.remove_group(group_id) except Exception as cleanup_exc: + cleanup_failed = True print( f" warn: remove_group({group_id}) cleanup failed: {cleanup_exc!r}", flush=True, ) + if cleanup_failed: + # Fail fast for every caller, not only lineage-tracked ones: + # the failed remove leaves an unready slot the retry cannot + # reclaim (capacity accounting drifts), and a post-write + # failure may have left TQ rows that no owner records -- the + # next data-plane checkpoint's inventory check would reject + # those later with a less useful error. A lineage-tracked + # retry additionally must not reuse its stable ID while the + # previous slot may still exist. Re-raise the rollout error. + raise # The rollout itself succeeded. Re-running generation cannot repair # a required downstream stage (for example MOPD teacher inference), # and would spend the rollout retry budget on the wrong subsystem. @@ -1379,6 +1475,11 @@ async def generate_and_push( # the success path rather than in the infra handler so that a prompt which # succeeded on a retry also counts -- the fleet recovered either way. self._consecutive_infra_drops = 0 + if lineage_group_id is not None: + async with ( + self._tq_buffer.data_plane_checkpoint_barrier.mutation() + ) as cut: + self._recovery_ledger.discard_group(cut, lineage_group_id) return RolloutOutcome.COMMITTED # The infrastructure budget ran out. The same failure followed the prompt across @@ -1405,7 +1506,8 @@ async def generate_and_push( # Under the budget: give up on this prompt and let the run continue. The caller # owns the backpressure permit for a SKIPPED outcome, and -- because the prompt # may have been stamped for a specific training step that will now never fill -- - # owns crediting the shortfall so the train pump can close that step short. + # owns atomically replacing its retained ledger entry or crediting the shortfall + # so the train pump can close that step short. self._stats.record_infra_drop(reason, self._consecutive_infra_drops) print( f"dropping prompt idx={input_sample['idx']} after {infra_attempts} " diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py new file mode 100644 index 00000000000..98afc450ecb --- /dev/null +++ b/nemo_rl/experience/rollout_recovery.py @@ -0,0 +1,595 @@ +# 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. + +"""Versioned ownership state for unfinished SingleController prompt groups.""" + +from __future__ import annotations + +import copy +import uuid +from dataclasses import dataclass +from enum import StrEnum +from typing import TYPE_CHECKING, Any, NotRequired, TypedDict + +if TYPE_CHECKING: + from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneMutationCut + from nemo_rl.data.interfaces import DatumSpec + +ROLLOUT_RECOVERY_SCHEMA_VERSION = 1 +ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" + + +class PromptGroupPhase(StrEnum): + """Durable admission phase for an unfinished prompt group.""" + + RESERVED = "reserved" + ADMITTED = "admitted" + + +class PromptRefState(TypedDict): + """Serializable locator for rebuilding one prompt from the dataset.""" + + sample_id: str + task_name: str | None + + +class PromptGroupRecoveryState(TypedDict): + """Serializable ownership state for one unfinished prompt group.""" + + group_id: str + admission_id: str + prompt_id: str + prompt_ref: PromptRefState + expected_generations: int + target_step: int | None + start_weight_version: int + phase: str + + +class RolloutRecoveryLedgerState(TypedDict): + """Versioned prompt-group ownership state managed by the ledger.""" + + schema_version: int + groups: list[PromptGroupRecoveryState] + + +class RolloutRecoveryState(RolloutRecoveryLedgerState): + """Complete checkpoint sidecar for unfinished rollout scheduling state.""" + + batch_shortfall: NotRequired[dict[int, int]] + sampler_stamps_target_steps: NotRequired[bool] + + +@dataclass(frozen=True) +class PromptRef: + """Stable dataset identity for rebuilding one prompt.""" + + sample_id: str + task_name: str | None + + +@dataclass(frozen=True) +class PromptGroupRecoveryRecord: + """In-memory ownership record for one prompt group.""" + + group_id: str + admission_id: str + prompt_id: str + prompt_ref: PromptRef + runtime_prompt_payload: DatumSpec | None + expected_generations: int + target_step: int | None + start_weight_version: int + phase: PromptGroupPhase + + @property + def prompt_payload(self) -> DatumSpec: + """Return the rehydrated prompt required for rollout redispatch.""" + if self.runtime_prompt_payload is None: + raise RuntimeError( + f"recovery group {self.group_id!r} has not rehydrated prompt " + f"sample_id={self.prompt_ref.sample_id!r}" + ) + return self.runtime_prompt_payload + + +@dataclass(frozen=True) +class ParsedRolloutRecoveryState: + """Validated controller and ledger state loaded from one checkpoint sidecar.""" + + ledger_state: RolloutRecoveryLedgerState + batch_shortfall: dict[int, int] + sampler_stamps_target_steps: bool | None + + +def _require_int(value: Any, *, field: str, minimum: int) -> int: + """Validate one integer field without accepting booleans.""" + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise ValueError(f"{field} must be an integer >= {minimum}, got {value!r}") + return value + + +def _prompt_task_name(prompt_payload: DatumSpec) -> str | None: + task_name = prompt_payload.get("task_name") + if task_name is not None and not isinstance(task_name, str): + raise TypeError( + "prompt_payload.task_name must be a string or None, got " + f"{type(task_name).__name__}" + ) + return task_name + + +def _validate_prompt_identity( + prompt_ref: PromptRef, + prompt_payload: DatumSpec, + *, + group_id: str, +) -> None: + sample_id = prompt_payload.get("idx") + if isinstance(sample_id, bool) or not isinstance(sample_id, int): + raise ValueError( + f"recovery group {group_id!r} prompt payload must contain an integer idx" + ) + if str(sample_id) != prompt_ref.sample_id: + raise ValueError( + f"recovery group {group_id!r} resolved sample_id={sample_id!r}; " + f"expected {prompt_ref.sample_id!r}" + ) + task_name = _prompt_task_name(prompt_payload) + if task_name != prompt_ref.task_name: + raise ValueError( + f"recovery group {group_id!r} resolved task_name={task_name!r}; " + f"expected {prompt_ref.task_name!r}" + ) + + +class RolloutRecoveryLedger: + """Own prompts after dataloader advance and before canonical TQ commit. + + Every mutating operation requires a live data-plane cut so ownership cannot + change outside the checkpoint barrier's consistent snapshot boundary. + """ + + def __init__(self) -> None: + self._groups: dict[str, PromptGroupRecoveryRecord] = {} + + def reserve_group( + self, + cut: DataPlaneMutationCut, + *, + prompt_id: str, + prompt_payload: DatumSpec, + expected_generations: int, + target_step: int | None, + start_weight_version: int, + admitted: bool, + group_id: str | None = None, + admission_id: str | None = None, + ) -> PromptGroupRecoveryRecord: + """Record ownership before the prompt can disappear from the dataloader. + + Args: + cut: Live capability yielded by the shared data-plane barrier. + prompt_id: Dataset-level prompt identity used for diagnostics. + prompt_payload: Runtime prompt used for whole-group regeneration. Only + its stable dataset reference is checkpointed. + expected_generations: Number of GRPO siblings in the prompt group. + target_step: Original gated training step, when the sampler stamps one. + start_weight_version: Policy version visible at reservation time. + admitted: Whether sampler admission already completed. This is explicit + because ``target_step=None`` is also valid for admitted ungated groups. + group_id: Stable logical and canonical TQ group ID. Generated when absent. + admission_id: Stable identity shared by every prompt in one sampler + admission. Defaults to ``group_id`` for single-prompt direct callers. + + Returns: + A defensive copy of the new record. + """ + cut.require_live() + if not prompt_id: + raise ValueError("prompt_id must not be empty") + sample_id = prompt_payload.get("idx") + if isinstance(sample_id, bool) or not isinstance(sample_id, int): + raise ValueError("prompt_payload must contain an integer idx") + if prompt_id != str(sample_id): + raise ValueError( + f"prompt_id={prompt_id!r} does not match prompt_payload idx={sample_id!r}" + ) + _require_int( + expected_generations, + field="expected_generations", + minimum=1, + ) + _require_int( + start_weight_version, + field="start_weight_version", + minimum=0, + ) + if target_step is not None: + _require_int(target_step, field="target_step", minimum=0) + group_id = group_id or str(uuid.uuid4()) + if not group_id: + raise ValueError("group_id must not be empty") + if group_id in self._groups: + raise ValueError(f"duplicate recovery group_id={group_id!r}") + admission_id = admission_id or group_id + if not admission_id: + raise ValueError("admission_id must not be empty") + + record = PromptGroupRecoveryRecord( + group_id=group_id, + admission_id=admission_id, + prompt_id=prompt_id, + # The rollout path treats the dataloader sample as immutable and builds + # mutable environment inputs from copies. Retaining that sample by + # reference avoids cloning a potentially very long prompt on every + # dispatch; state_dict() persists only its dataset locator. + prompt_ref=PromptRef( + sample_id=prompt_id, + task_name=_prompt_task_name(prompt_payload), + ), + runtime_prompt_payload=prompt_payload, + expected_generations=expected_generations, + target_step=target_step, + start_weight_version=start_weight_version, + phase=( + PromptGroupPhase.ADMITTED if admitted else PromptGroupPhase.RESERVED + ), + ) + self._groups[group_id] = record + return copy.copy(record) + + def mark_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: int | None, + start_weight_version: int, + ) -> None: + """Attach the sampler result to a previously reserved prompt group.""" + cut.require_live() + record = self._require_group(group_id) + if record.phase is not PromptGroupPhase.RESERVED: + raise ValueError( + f"recovery group {group_id!r} is already {record.phase.value}" + ) + if target_step is not None: + _require_int(target_step, field="target_step", minimum=0) + _require_int( + start_weight_version, + field="start_weight_version", + minimum=0, + ) + self._groups[group_id] = PromptGroupRecoveryRecord( + group_id=record.group_id, + admission_id=record.admission_id, + prompt_id=record.prompt_id, + prompt_ref=record.prompt_ref, + runtime_prompt_payload=record.runtime_prompt_payload, + expected_generations=record.expected_generations, + target_step=target_step, + start_weight_version=start_weight_version, + phase=PromptGroupPhase.ADMITTED, + ) + + def bind_runtime_prompt( + self, + cut: DataPlaneMutationCut, + group_id: str, + prompt_payload: DatumSpec, + ) -> None: + """Attach a dataset-rehydrated prompt after identity validation. + + The current reference is a positional index into a map-style dataset. + Recovery therefore requires dataset ordering to remain unchanged between + checkpoint and restart. + """ + cut.require_live() + record = self._require_group(group_id) + _validate_prompt_identity( + record.prompt_ref, + prompt_payload, + group_id=group_id, + ) + self._groups[group_id] = PromptGroupRecoveryRecord( + group_id=record.group_id, + admission_id=record.admission_id, + prompt_id=record.prompt_id, + prompt_ref=PromptRef( + sample_id=record.prompt_ref.sample_id, + task_name=record.prompt_ref.task_name, + ), + runtime_prompt_payload=prompt_payload, + expected_generations=record.expected_generations, + target_step=record.target_step, + start_weight_version=record.start_weight_version, + phase=record.phase, + ) + + def get_group(self, group_id: str) -> PromptGroupRecoveryRecord: + """Return a record copy while sharing its immutable runtime prompt.""" + return copy.copy(self._require_group(group_id)) + + def groups(self) -> list[PromptGroupRecoveryRecord]: + """Return record copies in reservation order without cloning prompts.""" + return [copy.copy(record) for record in self._groups.values()] + + def discard_group(self, cut: DataPlaneMutationCut, group_id: str) -> None: + """Release ownership after canonical commit or intentional discard.""" + cut.require_live() + self._require_group(group_id) + del self._groups[group_id] + + def discard_canonical_groups( + self, + cut: DataPlaneMutationCut, + group_ids: set[str], + ) -> int: + """Drop ledger copies already owned by canonical replay metadata.""" + cut.require_live() + discarded = 0 + for group_id in list(self._groups): + if group_id in group_ids: + del self._groups[group_id] + discarded += 1 + return discarded + + def state_dict(self) -> RolloutRecoveryLedgerState: + """Return versioned references without serializing full prompt payloads.""" + groups: list[PromptGroupRecoveryState] = [] + for record in self._groups.values(): + prompt_payload = record.runtime_prompt_payload + if prompt_payload is None: + raise RuntimeError( + f"cannot checkpoint recovery group {record.group_id!r} before " + "its prompt is rehydrated" + ) + _validate_prompt_identity( + record.prompt_ref, + prompt_payload, + group_id=record.group_id, + ) + # sample_id is currently a positional index into a map-style dataset, + # not a dataset-independent identity. The checkpoint is recoverable only + # when that dataset's ordering remains unchanged across the restart. + groups.append( + { + "group_id": record.group_id, + "admission_id": record.admission_id, + "prompt_id": record.prompt_id, + "prompt_ref": { + "sample_id": record.prompt_ref.sample_id, + "task_name": record.prompt_ref.task_name, + }, + "expected_generations": record.expected_generations, + "target_step": record.target_step, + "start_weight_version": record.start_weight_version, + "phase": record.phase.value, + } + ) + return { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": groups, + } + + def load_state_dict( + self, + cut: DataPlaneMutationCut, + state: RolloutRecoveryLedgerState, + ) -> None: + """Replace this empty ledger from a validated checkpoint payload.""" + cut.require_live() + if self._groups: + raise RuntimeError( + "cannot restore into a non-empty rollout recovery ledger" + ) + if not isinstance(state, dict): + raise TypeError( + "rollout recovery state must be a dictionary, got " + f"{type(state).__name__}" + ) + if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: + raise ValueError( + "unsupported rollout recovery schema_version=" + f"{state.get('schema_version')!r}; expected " + f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" + ) + groups = state.get("groups") + if not isinstance(groups, list): + raise TypeError("rollout recovery groups must be a list") + + restored: dict[str, PromptGroupRecoveryRecord] = {} + for index, raw_group in enumerate(groups): + if not isinstance(raw_group, dict): + raise TypeError( + f"rollout recovery groups[{index}] must be a dictionary" + ) + group_id = raw_group.get("group_id") + prompt_id = raw_group.get("prompt_id") + admission_id = raw_group.get("admission_id") + if not isinstance(group_id, str) or not group_id: + raise ValueError( + f"rollout recovery groups[{index}].group_id must be non-empty" + ) + if group_id in restored: + raise ValueError(f"duplicate recovery group_id={group_id!r}") + if not isinstance(admission_id, str) or not admission_id: + raise ValueError( + f"rollout recovery groups[{index}].admission_id must be non-empty" + ) + if not isinstance(prompt_id, str) or not prompt_id: + raise ValueError( + f"rollout recovery groups[{index}].prompt_id must be non-empty" + ) + expected_generations = _require_int( + raw_group.get("expected_generations"), + field=f"groups[{index}].expected_generations", + minimum=1, + ) + start_weight_version = _require_int( + raw_group.get("start_weight_version"), + field=f"groups[{index}].start_weight_version", + minimum=0, + ) + target_step = raw_group.get("target_step") + if target_step is not None: + target_step = _require_int( + target_step, + field=f"groups[{index}].target_step", + minimum=0, + ) + raw_phase = raw_group.get("phase") + if not isinstance(raw_phase, str): + raise ValueError( + f"rollout recovery groups[{index}].phase is invalid: {raw_phase!r}" + ) + try: + phase = PromptGroupPhase(raw_phase) + except ValueError as error: + raise ValueError( + f"rollout recovery groups[{index}].phase is invalid: {raw_phase!r}" + ) from error + raw_prompt_ref = raw_group.get("prompt_ref") + if not isinstance(raw_prompt_ref, dict): + raise TypeError( + f"rollout recovery groups[{index}].prompt_ref must be a dictionary" + ) + sample_id = raw_prompt_ref.get("sample_id") + task_name = raw_prompt_ref.get("task_name") + if not isinstance(sample_id, str) or not sample_id: + raise ValueError( + f"rollout recovery groups[{index}].prompt_ref.sample_id " + "must be non-empty" + ) + if sample_id != prompt_id: + raise ValueError( + f"rollout recovery groups[{index}] prompt_id and " + "prompt_ref.sample_id must match" + ) + if task_name is not None and not isinstance(task_name, str): + raise TypeError( + f"rollout recovery groups[{index}].prompt_ref.task_name " + "must be a string or None" + ) + restored[group_id] = PromptGroupRecoveryRecord( + group_id=group_id, + admission_id=admission_id, + prompt_id=prompt_id, + prompt_ref=PromptRef( + sample_id=sample_id, + task_name=task_name, + ), + runtime_prompt_payload=None, + expected_generations=expected_generations, + target_step=target_step, + start_weight_version=start_weight_version, + phase=phase, + ) + + admission_states: dict[str, tuple[PromptGroupPhase, int | None]] = {} + for record in restored.values(): + signature = (record.phase, record.target_step) + prior = admission_states.setdefault(record.admission_id, signature) + if prior != signature: + raise ValueError( + "rollout recovery groups sharing admission_id=" + f"{record.admission_id!r} disagree on phase or target_step" + ) + self._groups = restored + + def _require_group(self, group_id: str) -> PromptGroupRecoveryRecord: + try: + return self._groups[group_id] + except KeyError as error: + raise KeyError(f"unknown recovery group_id={group_id!r}") from error + + def __len__(self) -> int: + return len(self._groups) + + +def _validate_batch_shortfall(value: object) -> dict[int, int]: + """Return a defensive copy of per-step permanent rollout losses.""" + if not isinstance(value, dict): + raise TypeError("rollout recovery batch_shortfall must be a dictionary") + batch_shortfall: dict[int, int] = {} + for step, count in value.items(): + if ( + isinstance(step, bool) + or not isinstance(step, int) + or step < 0 + or isinstance(count, bool) + or not isinstance(count, int) + or count < 0 + ): + raise ValueError( + "rollout recovery batch_shortfall entries must contain " + f"non-negative integer steps and counts, got {step!r}: {count!r}" + ) + batch_shortfall[step] = count + return batch_shortfall + + +def build_rollout_recovery_state( + ledger: RolloutRecoveryLedger, + *, + batch_shortfall: dict[int, int], + sampler_stamps_target_steps: bool, +) -> RolloutRecoveryState: + """Build the complete versioned sidecar from ledger and controller state.""" + if not isinstance(sampler_stamps_target_steps, bool): + raise TypeError( + "rollout recovery sampler_stamps_target_steps must be a boolean" + ) + ledger_state = ledger.state_dict() + return { + "schema_version": ledger_state["schema_version"], + "groups": ledger_state["groups"], + "batch_shortfall": _validate_batch_shortfall(batch_shortfall), + "sampler_stamps_target_steps": sampler_stamps_target_steps, + } + + +def parse_rollout_recovery_state(state: object) -> ParsedRolloutRecoveryState: + """Validate and split a complete checkpoint sidecar by runtime owner.""" + if not isinstance(state, dict): + raise TypeError( + "rollout recovery sidecar must contain a dictionary, got " + f"{type(state).__name__}" + ) + if state.get("schema_version") != ROLLOUT_RECOVERY_SCHEMA_VERSION: + raise ValueError( + "unsupported rollout recovery schema_version=" + f"{state.get('schema_version')!r}; expected " + f"{ROLLOUT_RECOVERY_SCHEMA_VERSION}" + ) + groups = state.get("groups") + if not isinstance(groups, list): + raise TypeError("rollout recovery groups must be a list") + + raw_sampler_stamps = state.get("sampler_stamps_target_steps") + if raw_sampler_stamps is not None and not isinstance(raw_sampler_stamps, bool): + raise TypeError( + "rollout recovery sampler_stamps_target_steps must be a boolean" + ) + + ledger_state: RolloutRecoveryLedgerState = { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": groups, + } + return ParsedRolloutRecoveryState( + ledger_state=ledger_state, + batch_shortfall=_validate_batch_shortfall(state.get("batch_shortfall", {})), + sampler_stamps_target_steps=raw_sampler_stamps, + ) diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index bbc412cf1d0..9f68526a937 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -32,6 +32,7 @@ import warnings from collections import defaultdict from contextlib import nullcontext +from pathlib import Path from typing import Any, Optional import ray @@ -131,6 +132,10 @@ 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/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index 42938fa925f..6e4db00bde5 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -146,6 +146,9 @@ class CheckpointingConfig(TypedDict): model_repo_id (str): Repository ID for the model (for safetensors format). is_peft (bool): Whether the model uses PEFT. save_optimizer (bool): Whether to save optimizer state with checkpoints. + save_data_plane (bool): Whether SingleController checkpoints include the + native TQ snapshot and replay-buffer metadata. Currently supported only + with the simple data-plane backend. load_replay_buffer (bool): Whether async GRPO restores replay-buffer state when resuming from a checkpoint. Defaults to True. When False the buffer starts empty and a frontier-aligned resume regenerates the @@ -164,6 +167,7 @@ class CheckpointingConfig(TypedDict): checkpoint_must_save_by: NotRequired[str | None] pretrained_checkpoint: NotRequired[PretrainedCheckpointConfig] save_optimizer: NotRequired[bool] # Default: True + save_data_plane: NotRequired[bool] load_replay_buffer: NotRequired[bool] # Default: True (async GRPO only) # New nemo-automodel integration fields model_save_format: NotRequired[str | None] # Default: "safetensors" diff --git a/pyrefly.toml b/pyrefly.toml index 726e5e5ec23..5c93a2b0a46 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -118,6 +118,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/adapters/transfer_queue_env.py", "nemo_rl/data_plane/codec.py", "nemo_rl/data_plane/column_io.py", @@ -157,6 +158,7 @@ project-includes = [ "nemo_rl/experience/metric_utils.py", "nemo_rl/experience/payload.py", "nemo_rl/experience/rollout_manager.py", + "nemo_rl/experience/rollout_recovery.py", "nemo_rl/experience/rollouts.py", "nemo_rl/modelopt/__init__.py", "nemo_rl/modelopt/models/__init__.py", @@ -287,6 +289,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 751bf5a0ca0..21f15127f5a 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -173,6 +173,11 @@ run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/ # Checkpoint save/restore (upstream #3429). run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh +# Native TQ + metadata-only completed replay recovery (#3480). +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh +# Deterministic process restart with an admitted group held before canonical TQ +# commit, followed by exact-once redispatch at its stable group ID. +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_unfinished_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/_single_controller_rollout_recovery_hook.py b/tests/functional/_single_controller_rollout_recovery_hook.py new file mode 100644 index 00000000000..875333397f8 --- /dev/null +++ b/tests/functional/_single_controller_rollout_recovery_hook.py @@ -0,0 +1,137 @@ +# 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. + +"""Test-only SC entrypoint for deterministic unfinished-rollout recovery. + +The first process parks one selected rollout after controller admission. The +second process records its redispatch and successful canonical TQ commit. The +wrapper is injected driver-side into ``SingleControllerActorArgs`` so the +production request path has no environment-variable or timing hook. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from typing import Any, cast + +from examples import run_grpo_single_controller +from nemo_rl.experience.rollout_manager import RolloutManager, RolloutOutcome + + +class _InstrumentedRolloutManager: + """Delegate every operation except the deterministic recovery test cut.""" + + def __init__( + self, + delegate: Any, + *, + events_path: Path, + block_target_step: int | None, + ) -> None: + self._delegate = delegate + self._events_path = events_path + self._block_target_step = block_target_step + self._blocked = False + + def __getattr__(self, name: str) -> Any: + delegate = self.__dict__.get("_delegate") + if delegate is None: + raise AttributeError(name) + return getattr(delegate, name) + + @property + def _tq_buffer(self) -> Any: + return self._delegate._tq_buffer + + @_tq_buffer.setter + def _tq_buffer(self, value: Any) -> None: + self._delegate._tq_buffer = value + + def _append_event(self, event: str, **fields: Any) -> None: + self._events_path.parent.mkdir(parents=True, exist_ok=True) + payload = {"event": event, **fields} + with self._events_path.open("a", encoding="utf-8") as stream: + stream.write(json.dumps(payload, sort_keys=True) + "\n") + + async def generate_and_push( + self, + input_sample: Any, + *, + target_step: int | None = None, + inflight_registry: Any = None, + lineage_group_id: str | None = None, + ) -> RolloutOutcome: + fields = { + "group_id": lineage_group_id, + "prompt_idx": int(input_sample["idx"]), + "target_step": target_step, + } + self._append_event("dispatch", **fields) + + if ( + not self._blocked + and self._block_target_step is not None + and target_step == self._block_target_step + ): + if lineage_group_id is None: + raise RuntimeError("recovery test expected a lineage-tracked group") + self._blocked = True + self._append_event("blocked_before_tq_commit", **fields) + print( + "recovery functional hook: blocked admitted " + f"group_id={lineage_group_id} target_step={target_step}", + flush=True, + ) + # The phase-1 timeout checkpoint terminates the process and cancels + # this task. No polling or wall-clock race controls the checkpoint cut. + await asyncio.Event().wait() + + outcome = await self._delegate.generate_and_push( + input_sample, + target_step=target_step, + inflight_registry=inflight_registry, + lineage_group_id=lineage_group_id, + ) + if outcome is RolloutOutcome.COMMITTED: + self._append_event("canonical_tq_commit", **fields) + return outcome + + +_original_setup_single_controller = run_grpo_single_controller.setup_single_controller + + +def _setup_with_recovery_hook(*args: Any, **kwargs: Any) -> Any: + actor_args, timing_metrics = _original_setup_single_controller(*args, **kwargs) + events_path = Path(os.environ["SC_RECOVERY_TEST_EVENTS"]) + raw_target_step = os.environ.get("SC_RECOVERY_TEST_BLOCK_TARGET_STEP") + block_target_step = int(raw_target_step) if raw_target_step is not None else None + actor_args.rollout_manager = cast( + RolloutManager, + _InstrumentedRolloutManager( + actor_args.rollout_manager, + events_path=events_path, + block_target_step=block_target_step, + ), + ) + return actor_args, timing_metrics + + +run_grpo_single_controller.setup_single_controller = _setup_with_recovery_hook + + +if __name__ == "__main__": + run_grpo_single_controller.main() diff --git a/tests/functional/grpo_checkpoint_single_controller.sh b/tests/functional/grpo_checkpoint_single_controller.sh index fd41d4b272c..4528ae6982b 100755 --- a/tests/functional/grpo_checkpoint_single_controller.sh +++ b/tests/functional/grpo_checkpoint_single_controller.sh @@ -51,6 +51,7 @@ TRAIN_CMD=( checkpointing.checkpoint_dir=$CKPT_DIR checkpointing.save_period=2 checkpointing.metric_name=null + checkpointing.save_data_plane=true data_plane.enabled=true data_plane.impl=transfer_queue data_plane.backend=simple @@ -84,17 +85,22 @@ for artifact in \ "$STEP1/config.yaml" \ "$STEP1/policy/weights" \ "$STEP1/train_dataloader.pt" \ - "$STEP1/replay_buffer.pt"; do + "$STEP1/replay_buffer_metadata.pt" \ + "$STEP1/data_plane"; do if [[ ! -e "$artifact" ]]; then echo "FAIL: expected checkpoint artifact missing: $artifact" exit 1 fi done +if [[ -e "$STEP1/replay_buffer.pt" ]]; then + echo "FAIL: legacy tensor-bearing replay_buffer.pt should not be written" + exit 1 +fi if compgen -G "$CKPT_DIR/tmp_step_*" > /dev/null; then echo "FAIL: tmp_step_* leftovers — async finalization was not flushed" exit 1 fi -echo "✅ step_1 checkpoint complete (weights, dataloader, replay buffer), no tmp leftovers" +echo "✅ step_1 checkpoint complete (weights, dataloader, TQ replay), no tmp leftovers" if ! grep -q '"current_step": 1' "$STEP1/training_info.json"; then echo "FAIL: training_info.json does not record current_step=1" @@ -119,7 +125,7 @@ if ! grep -q "Restoring dataloader state from checkpoint" $EXP_DIR/run2.log; the echo "FAIL: dataloader restore log line not found in run2 output" exit 1 fi -if ! grep -q "Restoring replay buffer from checkpoint" $EXP_DIR/run2.log; then +if ! grep -q "Restoring replay buffer metadata" $EXP_DIR/run2.log; then echo "FAIL: replay buffer restore log line not found in run2 output" exit 1 fi @@ -134,6 +140,14 @@ if [[ ! -e "$STEP4/training_info.json" ]]; then echo "FAIL: run2 did not produce step_4 (resume did not reach step 4)" exit 1 fi +if [[ ! -f "$STEP4/replay_buffer_metadata.pt" || ! -d "$STEP4/data_plane" ]]; then + echo "FAIL: resumed run did not produce native TQ replay artifacts at step_4" + exit 1 +fi +if [[ -e "$STEP4/replay_buffer.pt" ]]; then + echo "FAIL: resumed run wrote legacy tensor-bearing replay_buffer.pt" + exit 1 +fi if ! grep -q '"current_step": 4' "$STEP4/training_info.json"; then echo "FAIL: step_4 training_info.json does not record current_step=4" cat "$STEP4/training_info.json" diff --git a/tests/functional/grpo_dp_single_controller.sh b/tests/functional/grpo_dp_single_controller.sh index 2dbbd7bae0e..7aaabbf3146 100755 --- a/tests/functional/grpo_dp_single_controller.sh +++ b/tests/functional/grpo_dp_single_controller.sh @@ -16,14 +16,15 @@ EXP_DIR=$SCRIPT_DIR/$EXP_NAME LOG_DIR=$EXP_DIR/logs JSON_METRICS=$EXP_DIR/metrics.json RUN_LOG=$EXP_DIR/run.log +SC_ENTRYPOINT=${SC_TEST_ENTRYPOINT:-$PROJECT_ROOT/examples/run_grpo_single_controller.py} export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-} 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 \ - $PROJECT_ROOT/examples/run_grpo_single_controller.py \ +uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ + $SC_ENTRYPOINT \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ grpo.num_generations_per_prompt=4 \ @@ -48,13 +49,15 @@ 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/num_masked_seqs_by_logprob_error"]) == 0' \ - 'max(data["train/max_seq_mult_prob_error"]) < 1000' \ - '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/num_masked_seqs_by_logprob_error"]) == 0' \ + 'max(data["train/max_seq_mult_prob_error"]) < 1000' \ + '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 00000000000..ea069d03b0c --- /dev/null +++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh @@ -0,0 +1,63 @@ +#!/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) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +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 +PHASE1_LOG=$TEST_DIR/phase1.log +PHASE2_LOG=$TEST_DIR/phase2.log + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +COMMON_OVERRIDES=( + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + checkpointing.save_data_plane=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 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" + +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" +REPLAY_GROUP_COUNT=$(uv run --directory "$PROJECT_ROOT" --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; print(metadata["replay_group_count"])' \ + "$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 +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -q "Native TQ checkpoint restored and validated: groups=${REPLAY_GROUP_COUNT}" "$PHASE2_LOG" +grep -q "Native TQ replay inventory validated" "$PHASE2_LOG" +grep -qF "Restored ${REPLAY_GROUP_COUNT} replay group(s) from checkpoint" "$PHASE2_LOG" +test -d "$CHECKPOINT_DIR/step_2/data_plane" +test -f "$CHECKPOINT_DIR/step_2/replay_buffer_metadata.pt" + +echo "=== Verify the standalone TQ checkpoint CLI round trip ===" +uv run --directory "$PROJECT_ROOT" --no-sync python \ + tools/verify_tq_data_plane_checkpoint.py \ + --checkpoint-dir "$TEST_DIR/verifier_bundle" + +echo "Native TQ recovery functional test passed." diff --git a/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh b/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh new file mode 100644 index 00000000000..d400af8db26 --- /dev/null +++ b/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh @@ -0,0 +1,85 @@ +#!/bin/bash +# Two-process functional test for one admitted, unfinished rollout group. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +PROJECT_ROOT=$(realpath "$SCRIPT_DIR/../..") +BASE_TEST=$SCRIPT_DIR/grpo_dp_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_dp_single_controller_unfinished_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +BASE_RUN_LOG=$SCRIPT_DIR/grpo_dp_single_controller/run.log +PHASE1_LOG=$TEST_DIR/phase1.log +PHASE2_LOG=$TEST_DIR/phase2.log +PHASE1_EVENTS=$TEST_DIR/phase1-events.jsonl +PHASE2_EVENTS=$TEST_DIR/phase2-events.jsonl +RECOVERY_HOOK=$SCRIPT_DIR/_single_controller_rollout_recovery_hook.py + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +COMMON_OVERRIDES=( + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + checkpointing.save_data_plane=true + async_rl.sampler.name=in_order + async_rl.sampler.max_lookahead_versions=1 + async_rl.max_inflight_prompts=4 + async_rl.max_buffered_rollouts=4 + # A recovery-ordering regression otherwise leaves zero rollouts in flight and + # only warns forever. Bound both slow generation and whole-run stalls in CI. + ++async_rl.rollout_failure.native.generation_timeout_s=60 + ++async_rl.stall_watchdog.interval_s=10 + ++async_rl.stall_watchdog.stall_timeout_s=180 + ++async_rl.stall_watchdog.stall_action=abort +) + +echo "=== Phase 1: checkpoint one admitted rollout before its TQ commit ===" +# The wrapper permanently parks one target-step-1 group. save_period=1 captures +# that ownership after train step 1; checkpoint_must_save_by only terminates the +# first process afterward. No sleep determines whether the group is unfinished. +SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ +SC_RECOVERY_TEST_EVENTS="$PHASE1_EVENTS" \ +SC_RECOVERY_TEST_BLOCK_TARGET_STEP=1 \ +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" \ + grpo.max_num_steps=2 \ + checkpointing.checkpoint_must_save_by=0:0:0:1 +cp "$BASE_RUN_LOG" "$PHASE1_LOG" + +STEP1=$CHECKPOINT_DIR/step_1 +test -d "$STEP1/data_plane" +test -f "$STEP1/replay_buffer_metadata.pt" +test -f "$STEP1/rollout_recovery.pt" +test ! -f "$STEP1/replay_buffer.pt" +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; metadata = json.load(open(sys.argv[1]))["user_metadata"]; assert metadata["mode"] == "authoritative", metadata; assert metadata["rollout_recovery_group_count"] > 0, metadata' \ + "$STEP1/data_plane/metadata.json" +BLOCKED_GROUP_ID=$(uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; events = [json.loads(line) for line in open(sys.argv[1])]; blocked = [event for event in events if event["event"] == "blocked_before_tq_commit"]; assert len(blocked) == 1, blocked; print(blocked[0]["group_id"])' \ + "$PHASE1_EVENTS") +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import sys, torch; state = torch.load(sys.argv[1], weights_only=True); group_id = sys.argv[2]; groups = [group for group in state["groups"] if group["group_id"] == group_id]; assert len(groups) == 1, state; assert groups[0]["phase"] == "admitted", groups[0]' \ + "$STEP1/rollout_recovery.pt" "$BLOCKED_GROUP_ID" + +echo "=== Phase 2: restore and canonically commit the same logical group once ===" +SC_TEST_ENTRYPOINT="$RECOVERY_HOOK" \ +SC_RECOVERY_TEST_EVENTS="$PHASE2_EVENTS" \ +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" grpo.max_num_steps=2 +cp "$BASE_RUN_LOG" "$PHASE2_LOG" + +grep -q "Native TQ checkpoint restored and validated" "$PHASE2_LOG" +grep -q "Loaded .* unfinished rollout group(s)" "$PHASE2_LOG" +test -d "$CHECKPOINT_DIR/step_2/data_plane" +test -f "$CHECKPOINT_DIR/step_2/replay_buffer_metadata.pt" +test -f "$CHECKPOINT_DIR/step_2/rollout_recovery.pt" +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import json, sys; events = [json.loads(line) for line in open(sys.argv[1])]; group_id = sys.argv[2]; dispatches = [event for event in events if event["event"] == "dispatch" and event["group_id"] == group_id]; commits = [event for event in events if event["event"] == "canonical_tq_commit" and event["group_id"] == group_id]; assert len(dispatches) == 1, dispatches; assert len(commits) == 1, commits' \ + "$PHASE2_EVENTS" "$BLOCKED_GROUP_ID" +uv run --directory "$PROJECT_ROOT" --no-sync python -c \ + 'import sys, torch; state = torch.load(sys.argv[1], weights_only=True); group_id = sys.argv[2]; assert group_id not in {group["group_id"] for group in state["groups"]}, state' \ + "$CHECKPOINT_DIR/step_2/rollout_recovery.pt" "$BLOCKED_GROUP_ID" + +echo "Unfinished rollout recovery functional test passed." diff --git a/tests/functional/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh index aa3cb1fd637..2861a038f60 100755 --- a/tests/functional/ppo_async_single_controller.sh +++ b/tests/functional/ppo_async_single_controller.sh @@ -65,6 +65,7 @@ TRAIN_CMD=( checkpointing.checkpoint_dir="${CKPT_DIR}" checkpointing.metric_name=null checkpointing.save_period=1 + +checkpointing.save_data_plane=true ) cd "${PROJECT_ROOT}" @@ -80,8 +81,12 @@ grep -q "weight_sync=CollectiveWeightSynchronizer" "${EXP_DIR}/run1.log" # policy_training_start_step=1, so step 0 trains the critic alone and step 1 is # where the policy joins. The banner fires exactly on that transition. test "$(grep -c "Critic warmup complete" "${EXP_DIR}/run1.log")" -eq 1 -test -f "${CKPT_DIR}/step_1/replay_buffer.pt" -test -f "${CKPT_DIR}/step_2/replay_buffer.pt" +test -f "${CKPT_DIR}/step_1/replay_buffer_metadata.pt" +test -d "${CKPT_DIR}/step_1/data_plane" +test ! -f "${CKPT_DIR}/step_1/replay_buffer.pt" +test -f "${CKPT_DIR}/step_2/replay_buffer_metadata.pt" +test -d "${CKPT_DIR}/step_2/data_plane" +test ! -f "${CKPT_DIR}/step_2/replay_buffer.pt" test -d "${CKPT_DIR}/step_1/value/weights" "${TRAIN_CMD[@]}" \ @@ -90,7 +95,7 @@ test -d "${CKPT_DIR}/step_1/value/weights" "$@" \ 2>&1 | tee "${EXP_DIR}/run2.log" -grep -q "Restoring replay buffer from checkpoint" "${EXP_DIR}/run2.log" +grep -q "Restoring replay buffer metadata" "${EXP_DIR}/run2.log" grep -qF "replay group(s) from checkpoint" "${EXP_DIR}/run2.log" # Warmup is behind us on the resumed run, so the policy trains every step and the # transition never happens again. diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index f7388ad6ac2..b3449aa7451 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -480,8 +480,10 @@ def test_get_grpo_save_state_handles_legacy_checkpoint_and_filters_metrics(): "total_steps": 13, "total_valid_tokens": 0, "val_reward": -99999999.0, - # SingleController-only field; None for every other algorithm. + # SingleController-only fields; None for every other algorithm. "sampler_name": None, + "trainer_version": None, + "sampler_dispatch_index": None, } assert "total_valid_tokens" not in loaded_state assert not hasattr(save_state, "val:accuracy") @@ -3599,8 +3601,8 @@ def test_async_grpo_colocated_save_defers_wake_until_after_checkpoint( policy_generation.finish_generation.side_effect = lambda *a, **k: events.append( ("finish_generation", k.get("release_gpu", True)) ) - policy_generation.prepare_for_generation.side_effect = ( - lambda *a, **k: events.append("wake_engine") + policy_generation.prepare_for_generation.side_effect = lambda *a, **k: ( + events.append("wake_engine") ) policy.offload_before_refit.side_effect = lambda *a, **k: events.append( "offload_before_refit" diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py index b0a9d95af99..ab961960098 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 00000000000..8fd8f4d1a3f --- /dev/null +++ b/tests/unit/data_plane/test_async_utils.py @@ -0,0 +1,75 @@ +# 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_backend_config.py b/tests/unit/data_plane/test_backend_config.py index 2054a2b2da9..b7bea76c7c8 100644 --- a/tests/unit/data_plane/test_backend_config.py +++ b/tests/unit/data_plane/test_backend_config.py @@ -30,6 +30,7 @@ MooncakeCpuConfig, SimpleStorageConfig, backend_config, + data_plane_supports_checkpointing, ) _BASE = { @@ -43,6 +44,20 @@ def _cfg(backend: str, **extra) -> dict: return {**_BASE, "backend": backend, **extra} +@pytest.mark.parametrize( + ("backend", "expected"), + [ + ("simple", True), + ("mooncake_cpu", False), + ("future_backend", False), + ], +) +def test_checkpointing_capability_defaults_to_unsupported( + backend: str, expected: bool +) -> None: + assert data_plane_supports_checkpointing(_cfg(backend)) is expected + + def test_nested_block_is_used() -> None: cfg = _cfg( "mooncake_cpu", diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py index 68752b5bb58..f2392701111 100644 --- a/tests/unit/data_plane/test_codec_mooncake.py +++ b/tests/unit/data_plane/test_codec_mooncake.py @@ -234,6 +234,7 @@ def fake_kv_batch_get( monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get) client = object.__new__(tq_adapter.TQDataPlaneClient) client._promote_1d = True + client._data_operations_started = False restored = client.get_samples( ["a", "b", "c"], "train", ["total_reward", "input_ids"] @@ -268,6 +269,7 @@ def fake_kv_batch_get( 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"]) diff --git a/tests/unit/data_plane/test_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py index 3426c3b5067..f414947c84f 100644 --- a/tests/unit/data_plane/test_interface_contract.py +++ b/tests/unit/data_plane/test_interface_contract.py @@ -20,6 +20,8 @@ from __future__ import annotations +import pickle + import pytest import torch from tensordict import TensorDict @@ -65,8 +67,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 +128,82 @@ 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() + + +def test_noop_checkpoint_load_explains_missing_partitions(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + checkpoint_dir.mkdir() + checkpoint_file = checkpoint_dir / "noop_state.pkl" + with checkpoint_file.open("wb") as state_file: + pickle.dump({"metadata": {}}, state_file) + + client = NoOpDataPlaneClient() + with pytest.raises(ValueError) as exc_info: + client.load_checkpoint(checkpoint_dir) + + message = str(exc_info.value) + assert str(checkpoint_file) in message + assert "has no 'partitions' key" in message + assert "Delete it and re-run the test" in message diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0d471bc2660..13a5f59a382 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 25505d10530..e373adf076f 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 4ab3b51abab..e51fe398741 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -22,6 +22,11 @@ from __future__ import annotations +import inspect +import json +from typing import Callable +from unittest.mock import MagicMock + import numpy as np import pytest import torch @@ -30,11 +35,73 @@ transfer_queue = pytest.importorskip("transfer_queue") # noqa: F841 from nemo_rl.data_plane.column_io import kv_first_write, read_columns -from nemo_rl.data_plane.interfaces import KVBatchMeta +from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta from nemo_rl.data_plane.schema import DP_TRAIN_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict +def _register_partition(client: DataPlaneClient) -> None: + client.register_partition( + partition_id="p", + fields=["x"], + num_samples=1, + consumer_tasks=["train"], + ) + + +def _claim_meta(client: DataPlaneClient) -> None: + client.claim_meta( + partition_id="p", + task_name="train", + required_fields=["x"], + batch_size=1, + ) + + +def _get_data(client: DataPlaneClient) -> None: + client.get_data( + KVBatchMeta( + partition_id="p", + task_name="train", + sample_ids=["sample-0"], + fields=["x"], + ) + ) + + +def _check_consumption_status(client: DataPlaneClient) -> None: + client.check_consumption_status("p", ["train"]) + + +def _put_samples(client: DataPlaneClient) -> None: + client.put_samples(["sample-0"], "p") + + +def _get_samples(client: DataPlaneClient) -> None: + client.get_samples(["sample-0"], "p", ["x"]) + + +def _list_sample_ids(client: DataPlaneClient) -> None: + client.list_sample_ids("p") + + +def _clear_samples(client: DataPlaneClient) -> None: + client.clear_samples(["sample-0"], "p") + + +_DATA_OPERATION_INVOKERS: dict[str, Callable[[DataPlaneClient], None]] = { + "register_partition": _register_partition, + "claim_meta": _claim_meta, + "get_data": _get_data, + "check_consumption_status": _check_consumption_status, + "put_samples": _put_samples, + "get_samples": _get_samples, + "list_sample_ids": _list_sample_ids, + "clear_samples": _clear_samples, +} +_LIFECYCLE_METHODS = {"save_checkpoint", "load_checkpoint", "close"} + + def test_register_partition_uses_unique_schema_warmup_key(monkeypatch) -> None: from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter @@ -104,6 +171,209 @@ def fake_clear(**kwargs): ] +def test_data_operation_guard_covers_the_full_interface() -> None: + public_abstract_methods = { + name + for name, member in inspect.getmembers(DataPlaneClient, inspect.isfunction) + if getattr(member, "__isabstractmethod__", False) + } + assert set(_DATA_OPERATION_INVOKERS) == public_abstract_methods - _LIFECYCLE_METHODS + + +@pytest.mark.parametrize("operation_name", _DATA_OPERATION_INVOKERS) +def test_each_public_data_operation_marks_the_client_dirty( + monkeypatch, + operation_name: str, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + tq_meta = MagicMock(size=1, global_indexes=[0], custom_meta=[{}]) + tq_client = MagicMock() + tq_client.get_meta.return_value = tq_meta + tq_client.kv_retrieve_keys.return_value = ["sample-0"] + tq_client.check_consumption_status.return_value = True + monkeypatch.setattr(tq_adapter.tq, "get_client", MagicMock(return_value=tq_client)) + monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", MagicMock()) + monkeypatch.setattr( + tq_adapter.tq, + "kv_batch_get", + MagicMock( + return_value=TensorDict( + {"x": torch.tensor([1])}, + batch_size=[1], + ) + ), + ) + monkeypatch.setattr( + tq_adapter.tq, + "kv_list", + MagicMock(return_value={"p": {"sample-0": {}}}), + ) + monkeypatch.setattr(tq_adapter.tq, "kv_clear", MagicMock()) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._data_operations_started = False + client._warmed_fields = {} + client._poll_interval_s = 0 + client._promote_1d = False + + _DATA_OPERATION_INVOKERS[operation_name](client) + + assert client._data_operations_started + + +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._supports_checkpointing = True + 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._supports_checkpointing = True + 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._supports_checkpointing = True + 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._supports_checkpointing = False + 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. diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 9d717dd5372..5029d3ebb63 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -124,6 +124,7 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None: train_batch, weight_version=3, group_id="group", + prompt_idx=17, ) assert sample_ids == ["group_g0", "group_g1"] assert "routed_experts" in fields @@ -138,8 +139,8 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None: "num_assistant_messages": 1, } assert tags == [ - {"weight_version": 3, **no_violations}, - {"weight_version": 3, **no_violations}, + {"weight_version": 3, "prompt_idx": 17, **no_violations}, + {"weight_version": 3, "prompt_idx": 17, **no_violations}, ] @@ -156,6 +157,7 @@ def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: train_batch, weight_version=3, group_id="group", + prompt_idx=17, ) assert "routed_experts" not in fields @@ -195,7 +197,12 @@ def test_record_to_train_batch_backfills_routes_for_failed_completion() -> None: # It is fully loss-masked either way. assert train_batch["token_mask"][1, :2].tolist() == [0, 0] - _, fields, _ = pack_payload(train_batch, weight_version=3, group_id="group") + _, fields, _ = pack_payload( + train_batch, + weight_version=3, + group_id="group", + prompt_idx=17, + ) assert "routed_experts" in fields assert list(fields["routed_experts"].unbind())[1].shape == (2, 2, 2) @@ -214,24 +221,32 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None: _record(completions), pad_value_dict={"token_ids": 0, "input_ids": 0}, ) - _, fields, tags = pack_payload(train_batch, weight_version=7, group_id="g") + _, fields, tags = pack_payload( + train_batch, + weight_version=7, + group_id="g", + prompt_idx=17, + ) assert "violation_counts" not in fields assert tags == [ { "weight_version": 7, + "prompt_idx": 17, "num_invalid_tool_calls": 1, "num_malformed_thinking": 0, "num_assistant_messages": 1, }, { "weight_version": 7, + "prompt_idx": 17, "num_invalid_tool_calls": 0, "num_malformed_thinking": 1, "num_assistant_messages": 1, }, { "weight_version": 7, + "prompt_idx": 17, "num_invalid_tool_calls": 0, "num_malformed_thinking": 0, "num_assistant_messages": 0, diff --git a/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 88d4ab6529f..820d7243779 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -33,7 +33,10 @@ import pytest import torch -from nemo_rl.algorithms.async_utils.replay_buffer import PostWriteEnrichmentError +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + PostWriteEnrichmentError, +) from nemo_rl.data.collate_fn import rl_collate_fn from nemo_rl.data.datasets.response_datasets import NemoGymDataset from nemo_rl.data.interfaces import DatumSpec @@ -43,9 +46,11 @@ from nemo_rl.experience.rollout_manager import ( AsyncNemoGymRolloutImpl, RolloutManager, + RolloutOutcome, RolloutRetryPolicy, RolloutStats, ) +from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, run_async_nemo_gym_rollout, @@ -73,10 +78,19 @@ def _run(coro): return asyncio.run(coro) +def _with_cut(buffer, callback): + async def apply(): + async with buffer.data_plane_checkpoint_barrier.mutation() as cut: + return callback(cut) + + return _run(apply()) + + class _FakeBuffer: """Minimal TQReplayBuffer stand-in that records reserve/commit calls.""" def __init__(self) -> None: + self.data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() self.reserve_calls: list[int] = [] # weight_versions passed to reserve self.commit_calls: list[tuple[str, object, int, int]] = [] self.remove_calls: list[str] = [] @@ -141,6 +155,7 @@ def _make_manager( mgr._tokenizer = None mgr._num_generations_per_prompt = 1 mgr._tq_buffer = buffer + mgr._recovery_ledger = RolloutRecoveryLedger() mgr._weight_version = 0 mgr._retry_policy = ( retry_policy @@ -335,6 +350,96 @@ async def _logged_commit(*args, **kwargs): assert record == "r0" assert start_v == 0 assert end_v == 0 + assert len(mgr.recovery_ledger) == 0 + + def test_ledger_hands_ownership_to_canonical_buffer_on_commit(self): + buf = _FakeBuffer() + + async def _assert_ledger_owns_inflight_prompt(_sample): + groups = mgr.recovery_ledger.groups() + assert len(groups) == 1 + assert groups[0].group_id in buf._slots + + mgr = _make_manager( + buf, + _FakeImpl(on_run=_assert_ledger_owns_inflight_prompt), + ) + prompt = {"idx": 0, "message_log": [], "prompt": "p"} + group_id = _with_cut( + buf, + lambda cut: mgr.reserve_prompt_group( + cut, + prompt, + target_step=None, + ), + ) + + _run( + mgr.generate_and_push( + prompt, + lineage_group_id=group_id, + ) + ) + + assert len(mgr.recovery_ledger) == 0 + assert buf._slots == [group_id] + assert buf.commit_calls[0][0] == group_id + + def test_skipped_tracked_prompt_remains_owned_for_controller_handoff(self): + async def _fail_rollout(_sample): + raise RuntimeError("bad prompt") + + buf = _FakeBuffer() + mgr = _make_manager( + buf, + _FakeImpl(on_run=_fail_rollout), + RolloutRetryPolicy.single_attempt(max_skipped_prompts=1), + ) + group_id = _with_cut( + buf, + lambda cut: mgr.reserve_prompt_group( + cut, + {"idx": 7, "message_log": []}, + target_step=7, + ), + ) + + outcome = _run( + mgr.generate_and_push( + {"idx": 7, "message_log": []}, + target_step=7, + lineage_group_id=group_id, + ) + ) + + assert outcome is RolloutOutcome.SKIPPED + assert mgr.recovery_ledger.get_group(group_id).target_step == 7 + + def test_tracked_dispatch_rejects_changed_generations_per_prompt(self): + buf = _FakeBuffer() + mgr = _make_manager(buf, _FakeImpl()) + _with_cut( + buf, + lambda cut: mgr.recovery_ledger.reserve_group( + cut, + group_id="g0", + prompt_id="0", + prompt_payload={"idx": 0, "message_log": []}, + expected_generations=2, + target_step=0, + start_weight_version=0, + admitted=True, + ), + ) + + with pytest.raises(ValueError, match="expects 2 generation"): + _run( + mgr.generate_and_push( + {"idx": 0, "message_log": []}, + target_step=0, + lineage_group_id="g0", + ) + ) def test_start_weight_version_pinned_at_reserve_time(self): """If set_weight_version is called mid-rollout, start != end.""" diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py new file mode 100644 index 00000000000..ef96bbf0b93 --- /dev/null +++ b/tests/unit/experience/test_rollout_recovery.py @@ -0,0 +1,442 @@ +# 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 __future__ import annotations + +import asyncio +from collections.abc import Callable +from typing import Any, TypeVar + +import pytest +import torch +from torchdata.stateful_dataloader import StatefulDataLoader + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + DataPlaneMutationCut, +) +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_SCHEMA_VERSION, + PromptGroupPhase, + RolloutRecoveryLedger, + build_rollout_recovery_state, + parse_rollout_recovery_state, +) + +_T = TypeVar("_T") + + +def _mutate(callback: Callable[[DataPlaneMutationCut], _T]) -> _T: + async def apply() -> _T: + async with DataPlaneCheckpointBarrier().mutation() as cut: + return callback(cut) + + return asyncio.run(apply()) + + +def _reserve(ledger: RolloutRecoveryLedger, **kwargs: Any): + return _mutate(lambda cut: ledger.reserve_group(cut, **kwargs)) + + +def _load(ledger: RolloutRecoveryLedger, state) -> None: + _mutate(lambda cut: ledger.load_state_dict(cut, state)) + + +def _mark(ledger: RolloutRecoveryLedger, group_id: str, **kwargs: Any) -> None: + _mutate(lambda cut: ledger.mark_group_admitted(cut, group_id, **kwargs)) + + +def _bind(ledger: RolloutRecoveryLedger, group_id: str, prompt: DatumSpec) -> None: + _mutate(lambda cut: ledger.bind_runtime_prompt(cut, group_id, prompt)) + + +def _prompt(idx: int = 7) -> DatumSpec: + return { + "idx": idx, + "message_log": [{"role": "user", "content": f"prompt {idx}"}], + "length": 1, + "extra_env_info": None, + "loss_multiplier": 1.0, + } + + +def _single_prompt_batch(batch: list[DatumSpec]) -> DatumSpec: + assert len(batch) == 1 + return batch[0] + + +def _shuffled_prompt_loader(seed: int = 123) -> StatefulDataLoader: + return StatefulDataLoader( + [_prompt(idx) for idx in range(12)], + batch_size=1, + shuffle=True, + generator=torch.Generator().manual_seed(seed), + collate_fn=_single_prompt_batch, + num_workers=0, + ) + + +def _group_state( + idx: int = 7, + *, + target_step: int | None = 7, + phase: str = "admitted", +) -> dict: + return { + "group_id": f"g{idx}", + "admission_id": "batch-7", + "prompt_id": str(idx), + "prompt_ref": { + "sample_id": str(idx), + "task_name": None, + }, + "expected_generations": 2, + "target_step": target_step, + "start_weight_version": 7, + "phase": phase, + } + + +def test_ledger_round_trip_preserves_group_ownership() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + + state = ledger.state_dict() + restored = RolloutRecoveryLedger() + _load(restored, state) + + with pytest.raises(RuntimeError, match="has not rehydrated prompt"): + _ = restored.get_group("g7").prompt_payload + _bind(restored, "g7", _prompt()) + + assert restored.state_dict() == state + assert restored.get_group("g7").phase is PromptGroupPhase.ADMITTED + + +def test_checkpoint_state_round_trip_preserves_controller_and_ledger_state() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + + state = build_rollout_recovery_state( + ledger, + batch_shortfall={7: 1}, + sampler_stamps_target_steps=True, + ) + parsed = parse_rollout_recovery_state(state) + restored = RolloutRecoveryLedger() + _load(restored, parsed.ledger_state) + + assert [group.group_id for group in restored.groups()] == ["g7"] + assert parsed.batch_shortfall == {7: 1} + assert parsed.sampler_stamps_target_steps is True + + +def test_checkpoint_parser_defaults_fields_absent_from_older_state() -> None: + parsed = parse_rollout_recovery_state(RolloutRecoveryLedger().state_dict()) + + assert parsed.batch_shortfall == {} + assert parsed.sampler_stamps_target_steps is None + + +@pytest.mark.parametrize( + ("field", "value", "error_type"), + [ + ("batch_shortfall", [], TypeError), + ("batch_shortfall", {True: 1}, ValueError), + ("batch_shortfall", {7: -1}, ValueError), + ("sampler_stamps_target_steps", "yes", TypeError), + ], +) +def test_checkpoint_parser_rejects_malformed_controller_state( + field: str, + value: object, + error_type: type[Exception], +) -> None: + state: dict[str, object] = dict(RolloutRecoveryLedger().state_dict()) + state[field] = value + + with pytest.raises(error_type): + parse_rollout_recovery_state(state) + + +def test_ledger_rejects_an_expired_mutation_cut() -> None: + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + ledger = RolloutRecoveryLedger() + async with barrier.mutation() as cut: + ledger.reserve_group( + cut, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + + with pytest.raises(RuntimeError, match="no longer active"): + ledger.discard_group(cut, "g7") + + asyncio.run(exercise()) + + +def test_checkpoint_cut_can_guard_a_ledger_mutation() -> None: + async def exercise() -> None: + ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().checkpoint() as cut: + ledger.reserve_group( + cut, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=6, + admitted=True, + ) + + assert [group.group_id for group in ledger.groups()] == ["g7"] + + asyncio.run(exercise()) + + +def test_target_step_none_does_not_mean_unadmitted() -> None: + ledger = RolloutRecoveryLedger() + record = _reserve( + ledger, + group_id="windowed", + admission_id="batch-windowed", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=None, + start_weight_version=6, + admitted=True, + ) + + assert record.phase is PromptGroupPhase.ADMITTED + assert record.target_step is None + + +def test_reserved_group_can_be_admitted_exactly_once() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=None, + start_weight_version=6, + admitted=False, + ) + + _mark( + ledger, + "g7", + target_step=7, + start_weight_version=7, + ) + + record = ledger.get_group("g7") + assert record.phase is PromptGroupPhase.ADMITTED + assert record.target_step == 7 + assert record.start_weight_version == 7 + with pytest.raises(ValueError, match="already admitted"): + _mark( + ledger, + "g7", + target_step=8, + start_weight_version=8, + ) + + +def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> None: + ledger = RolloutRecoveryLedger() + for idx, group_id in enumerate(("canonical", "unfinished"), start=7): + _reserve( + ledger, + group_id=group_id, + admission_id="batch-7", + prompt_id=str(idx), + prompt_payload=_prompt(idx), + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + + assert _mutate(lambda cut: ledger.discard_canonical_groups(cut, {"canonical"})) == 1 + assert [group.group_id for group in ledger.groups()] == ["unfinished"] + + +def test_state_dict_stores_a_prompt_ref_without_the_full_payload() -> None: + ledger = RolloutRecoveryLedger() + prompt = _prompt() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=prompt, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + + state = ledger.state_dict() + group_state = state["groups"][0] + assert "prompt_payload" not in group_state + assert group_state["prompt_ref"] == { + "sample_id": "7", + "task_name": None, + } + group_state["prompt_ref"]["sample_id"] = "100" + + assert ledger.get_group("g7").prompt_ref.sample_id == "7" + + +def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> None: + ledger = RolloutRecoveryLedger() + original = _prompt() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=original, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + restored = RolloutRecoveryLedger() + _load(restored, ledger.state_dict()) + + changed = _prompt() + changed["message_log"][0]["content"] = "different prompt" + _bind(restored, "g7", changed) + + assert restored.get_group("g7").prompt_payload == changed + + +def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="g7", + admission_id="batch-7", + prompt_id="7", + prompt_payload=_prompt(), + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + restored = RolloutRecoveryLedger() + _load(restored, ledger.state_dict()) + + with pytest.raises(ValueError, match="expected '7'"): + _bind(restored, "g7", _prompt(8)) + + +def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: + """Shuffle position and prompt identity are independent recovery assets.""" + + dataloader = _shuffled_prompt_loader() + iterator = iter(dataloader) + fetched = [next(iterator) for _ in range(3)] + owned_prompt = fetched[-1] + + ledger = RolloutRecoveryLedger() + _reserve( + ledger, + group_id="unfinished", + admission_id="shuffled-batch", + prompt_id=str(owned_prompt["idx"]), + prompt_payload=owned_prompt, + expected_generations=2, + target_step=1, + start_weight_version=0, + admitted=True, + ) + ledger_state = ledger.state_dict() + dataloader_state = dataloader.state_dict() + expected_next_prompt = next(iterator) + + restored_dataloader = _shuffled_prompt_loader() + restored_dataloader.load_state_dict(dataloader_state) + assert next(iter(restored_dataloader)) == expected_next_prompt + + restored_ledger = RolloutRecoveryLedger() + _load(restored_ledger, ledger_state) + restored_group = restored_ledger.get_group("unfinished") + dataset_prompt = restored_dataloader.dataset[ + int(restored_group.prompt_ref.sample_id) + ] + _bind(restored_ledger, "unfinished", dataset_prompt) + + assert restored_ledger.get_group("unfinished").prompt_payload == owned_prompt + + +@pytest.mark.parametrize( + "state", + [ + {"schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION + 1, "groups": []}, + {"schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, "groups": {}}, + { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": [_group_state(phase="unknown")], + }, + { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": [ + _group_state(idx, target_step=target_step, phase=phase) + for idx, target_step, phase in ( + (7, None, "reserved"), + (8, 7, "admitted"), + ) + ], + }, + ], +) +def test_restore_rejects_incompatible_or_malformed_state(state: dict) -> None: + with pytest.raises((TypeError, ValueError)): + _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 659f0191b77..4fc0ac00080 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -128,6 +128,7 @@ checkpointing: # the buffered window fresh instead of reusing completed groups (which skew # toward short rollouts at a save boundary). load_replay_buffer: true + save_data_plane: false model_save_format: "safetensors" save_consolidated: false save_optimizer: true diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py new file mode 100644 index 00000000000..5e657615746 --- /dev/null +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -0,0 +1,566 @@ +# 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. + +"""Scenario harness for the checkpoint no-data-loss property. + +The property under test, in one sentence: **pausing a run to checkpoint and +resuming it must train on exactly the same prompt groups as never pausing at +all.** Concretely, every prompt group the dataloader has already handed out and +the trainer has not yet consumed must come back after a restore -- whether it +finished generating or not. If it does not come back, that prompt is lost for +good, because the dataloader cursor is saved where it stands and never rewinds. + +Everything here runs on CPU. The real ``TQReplayBuffer``, the real samplers, the +real ``DataPlaneCheckpointBarrier`` and the real ``NoOpDataPlaneClient`` +save/load are exercised. Only two things are stubbed, and neither is on the path +under test: + +* ``record_to_train_batch`` -- the tensor converter, so a scenario can use empty + prompt records instead of building real rollouts. +* the trainer/generation side -- absent entirely; this harness is the buffer and + the samplers, which is where save/restore lives. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional + +import torch + +from nemo_rl.algorithms.async_utils import replay_buffer as _rb +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + TQReplayBuffer, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + ReadyFirstSamplerConfig, + WeightFifoSamplerConfig, + WindowedSamplerConfig, + create_sampler, +) +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.interfaces import PromptGroupRecord +from nemo_rl.experience.rollout_recovery import RolloutRecoveryLedger + +PARTITION = "rollout_data" +ROLLOUTS_PER_GROUP = 2 # rollouts_per_prompt_group +GROUPS_PER_STEP = 3 # prompt_groups per training step +CAPACITY = 64 # max_buffered_rollouts +_FIELDS = ["input_ids", "input_lengths", "total_reward"] + +SAMPLERS = ("windowed", "ready_first", "weight_fifo", "in_order") + + +def sampler_config(name: str, lag: int): + """Build the real discriminated sampler config for ``name``.""" + if name == "windowed": + return WindowedSamplerConfig(max_staleness_versions=lag) + if name == "ready_first": + return ReadyFirstSamplerConfig(max_staleness_versions=lag) + if name == "weight_fifo": + return WeightFifoSamplerConfig(max_staleness_versions=lag) + if name == "in_order": + return InOrderSamplerConfig(max_lookahead_versions=lag) + raise ValueError(f"unknown sampler {name!r}") + + +# ── scenario description ──────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class Group: + """One prompt group at checkpoint time. + + Args: + gid: Prompt-group number, matching the order the dataloader served it. + done: How many of its ``ROLLOUTS_PER_GROUP`` rollouts have finished. + ``ROLLOUTS_PER_GROUP`` means the group committed; anything less + means it is still in flight. + weight: Weight version the group was dispatched at. + target: ``target_step`` stamp, used by the gated samplers. + evicted: The sampler deliberately dropped it (too stale). An evicted + group is an intentional discard, not data loss, so it is excluded + from what a restore must return. + """ + + gid: int + done: int + weight: int = 0 + target: Optional[int] = None + evicted: bool = False + + +@dataclass(frozen=True) +class Scenario: + """A checkpoint taken mid-run. + + Args: + name: Short label used in test ids. + groups: Every prompt group the dataloader has handed out so far. + cursor: The next prompt-group number the dataloader would serve. Every + group below it has already been handed out and will never be + handed out again after a restore. + trained: Groups the trainer has already consumed. Not necessarily + contiguous -- some samplers train whatever is ready and leave an + earlier unfinished group behind. + lag: How far generation may run ahead, in steps. + """ + + name: str + groups: tuple[Group, ...] + cursor: int + trained: frozenset[int] = field(default_factory=frozenset) + lag: int = 1 + + def must_survive(self) -> set[str]: + """Groups a restore has to return, or data is lost. + + Handed out, not trained, not deliberately evicted -- and below the + cursor, so the dataloader will never produce them again. This is the + no-data-loss bar enforced by the recovery matrix. + """ + return { + _gid(g.gid) + for g in self.groups + if not g.evicted and g.gid not in self.trained and g.gid < self.cursor + } + + def committed_outstanding(self) -> set[str]: + """Completed, unconsumed groups covered by #3480's TQ recovery.""" + return { + _gid(g.gid) + for g in self.groups + if not g.evicted + and g.gid not in self.trained + and g.done == ROLLOUTS_PER_GROUP + } + + def expected_stamps(self) -> dict[str, tuple[int | None, int]]: + """Target-step and start-weight stamps every restored group must retain.""" + return { + _gid(group.gid): (group.target, group.weight) + for group in self.groups + if not group.evicted and group.gid not in self.trained + } + + +def _gid(n: int) -> str: + return f"g{n:02d}" + + +@dataclass(frozen=True) +class Case: + """One row of the test matrix: a scenario run under one sampler. + + Args: + scenario: The buffer state at checkpoint time. + sampler: Which sampler the run is configured with. + why: Optional diagnostic context for an expected behavior gap. + """ + + scenario: Scenario + sampler: str + why: str = "" + + @property + def id(self) -> str: + return f"{self.sampler}::{self.scenario.name}" + + +# ── the round trip ────────────────────────────────────────────────────────── + + +def _record() -> PromptGroupRecord: + return PromptGroupRecord( + prompt_idx=0, + prompt=[], + extra_env_info=None, + metadata={}, + completions=[], + rollout_metrics={}, + ) + + +def _stub_converter(record: PromptGroupRecord, *, pad_value_dict: Any): + del record, pad_value_dict + return BatchedDataDict[Any]( + { + "input_ids": torch.ones((ROLLOUTS_PER_GROUP, 3), dtype=torch.long), + "input_lengths": torch.full((ROLLOUTS_PER_GROUP,), 3, dtype=torch.long), + "total_reward": torch.zeros(ROLLOUTS_PER_GROUP, dtype=torch.float32), + } + ) + + +def patch_converter(monkeypatch) -> None: + """Swap the tensor converter so scenarios can use empty prompt records.""" + monkeypatch.setattr(_rb, "record_to_train_batch", _stub_converter) + + +def _fresh_client(register: bool) -> NoOpDataPlaneClient: + dp = NoOpDataPlaneClient() + if register: + dp.register_partition( + partition_id=PARTITION, + fields=list(_FIELDS), + num_samples=CAPACITY * ROLLOUTS_PER_GROUP, + consumer_tasks=["train"], + ) + return dp + + +def _new_buffer(dp: NoOpDataPlaneClient) -> TQReplayBuffer: + buf = TQReplayBuffer( + dp, + partition_id=PARTITION, + pad_value_dict={"input_ids": 0}, + require_routed_experts=False, + ) + buf.set_data_plane_checkpoint_barrier(DataPlaneCheckpointBarrier()) + return buf + + +async def _fill(buf: TQReplayBuffer, scenario: Scenario) -> None: + """Recreate the scenario through the buffer's real reserve/commit path. + + Only groups still held at checkpoint time are added. A trained group is gone + -- ``_finalize_selection`` removes it from the buffer as it hands it to the + trainer -- and an evicted one was dropped by the staleness rule. + """ + for g in scenario.groups: + if g.evicted or g.gid in scenario.trained: + continue + gid = buf.reserve( + weight_version=g.weight, target_step=g.target, group_id=_gid(g.gid) + ) + if g.done == ROLLOUTS_PER_GROUP: + await buf.commit( + gid, + _record(), + start_weight_version=g.weight, + end_weight_version=g.weight, + ) + # done < ROLLOUTS_PER_GROUP: still generating, so no commit yet. + + +@dataclass +class RoundTrip: + """What a save/restore cycle returned. + + ``recovered`` is deliberately *presence*, not readiness: a group counts as + recovered if the restored buffer knows about it at all. That keeps the + assertions independent of whether recovery regenerates the whole group or + later resumes only missing siblings. A group could come back already + committed, or as a reserved slot waiting to be finished -- either way the + run has not lost the prompt, and either way these tests notice. + ``ready`` and ``pending`` are reported separately for diagnosis only; + nothing asserts on them. ``stamps`` records each restored group's + ``target_step`` and start weight. ``selected`` and ``selected_count`` report + the optional restore-then-select result used to verify each sampler's + recovery key at multiple gate lags. + """ + + recovered: set[str] + ready: set[str] + pending: set[str] + saved_sidecar: bool + rows_before: set[str] + rows_after: set[str] + stamps: dict[str, tuple[int | None, int]] + selected: set[str] + selected_count: int + + +async def _round_trip( + scenario: Scenario, + sampler_name: str, + tmp_path: Path, + *, + select_current_train_weight: int | None = None, +) -> RoundTrip: + dp_a = _fresh_client(register=True) + buf_a = _new_buffer(dp_a) + await _fill(buf_a, scenario) + sampler_a = create_sampler(buf_a, sampler_config(sampler_name, scenario.lag)) + + # Mirrors SingleControllerActor._save_checkpoint: the sidecar is written + # only when the sampler says it can restore one. + sidecar = ( + buf_a.metadata_state_dict(saved_capacity=CAPACITY) + if sampler_a.supports_buffer_checkpoint + else None + ) + recovery_ledger_a = RolloutRecoveryLedger() + async with buf_a.data_plane_checkpoint_barrier.mutation() as cut: + for group in scenario.groups: + if ( + group.evicted + or group.gid in scenario.trained + or group.done == ROLLOUTS_PER_GROUP + ): + continue + recovery_ledger_a.reserve_group( + cut, + group_id=_gid(group.gid), + admission_id=f"batch-{group.target}", + prompt_id=str(group.gid), + prompt_payload={"idx": group.gid, "message_log": []}, + expected_generations=ROLLOUTS_PER_GROUP, + target_step=group.target, + start_weight_version=group.weight, + admitted=True, + ) + recovery_sidecar = recovery_ledger_a.state_dict() + rows_before = set(dp_a.list_sample_ids(PARTITION)) + dp_a.save_checkpoint(tmp_path / "data_plane") + + # ---- restart: brand new process, nothing carried over in memory ---- + dp_b = _fresh_client(register=False) # load_checkpoint demands a clean client + dp_b.load_checkpoint(tmp_path / "data_plane") + buf_b = _new_buffer(dp_b) + sampler_b = create_sampler(buf_b, sampler_config(sampler_name, scenario.lag)) + + # Mirrors SingleControllerActor._maybe_restore_replay_buffer. + if sidecar is not None and sampler_b.supports_buffer_checkpoint: + await buf_b.load_state_dict( + sidecar, + max_groups=CAPACITY, + expected_partition_id=PARTITION, + expected_group_size=ROLLOUTS_PER_GROUP, + expected_manifest_digest=sidecar["manifest_digest"], + ) + recovery_ledger_b = RolloutRecoveryLedger() + async with buf_b.data_plane_checkpoint_barrier.mutation() as cut: + recovery_ledger_b.load_state_dict(cut, recovery_sidecar) + recovery_ledger_b.discard_canonical_groups(cut, set(buf_b._group_ids)) + for group in recovery_ledger_b.groups(): + group_id = buf_b.reserve( + weight_version=group.start_weight_version, + target_step=group.target_step, + group_id=group.group_id, + ) + await buf_b.commit( + group_id, + _record(), + start_weight_version=group.start_weight_version, + end_weight_version=group.start_weight_version, + ) + async with buf_b.data_plane_checkpoint_barrier.mutation() as cut: + recovery_ledger_b.discard_group(cut, group_id) + + ready = { + gid for gid, is_ready in zip(buf_b._group_ids, buf_b.ready_list) if is_ready + } + recovered = set(buf_b._group_ids) + pending = recovered - ready + rows_after = set(dp_b.list_sample_ids(PARTITION)) + stamps = { + group_id: ( + buf_b.target_step_list[index], + buf_b.start_weight_list[index], + ) + for index, group_id in enumerate(buf_b._group_ids) + } + selected: set[str] = set() + selected_count = 0 + if select_current_train_weight is not None: + selected_meta, selected_count = await sampler_b.select( + current_train_weight=select_current_train_weight, + min_prompt_groups=GROUPS_PER_STEP, + max_prompt_groups=GROUPS_PER_STEP, + ) + if selected_meta is not None: + selected = { + sample_id.rpartition("_g")[0] for sample_id in selected_meta.sample_ids + } + return RoundTrip( + recovered=recovered, + ready=ready, + pending=pending, + saved_sidecar=sidecar is not None, + rows_before=rows_before, + rows_after=rows_after, + stamps=stamps, + selected=selected, + selected_count=selected_count, + ) + + +def round_trip( + scenario: Scenario, + sampler_name: str, + tmp_path: Path, + *, + select_current_train_weight: int | None = None, +) -> RoundTrip: + """Save the scenario, restore it into a fresh buffer, report what came back.""" + return asyncio.run( + _round_trip( + scenario, + sampler_name, + tmp_path, + select_current_train_weight=select_current_train_weight, + ) + ) + + +def assert_no_data_loss( + scenario: Scenario, sampler_name: str, tmp_path: Path +) -> RoundTrip: + """Fail if the restore dropped any group the run still needs. + + The full bar: everything handed out and not yet trained comes back. + """ + result = round_trip(scenario, sampler_name, tmp_path) + missing = sorted(scenario.must_survive() - result.recovered) + assert not missing, ( + f"{sampler_name}/{scenario.name}: restore dropped {missing}. " + f"These groups were handed out by the dataloader, never trained, and sit " + f"below the saved cursor ({scenario.cursor}), so nothing will produce them " + f"again. recovered={sorted(result.recovered)}" + ) + return result + + +def assert_completed_groups_survive( + scenario: Scenario, sampler_name: str, tmp_path: Path +) -> RoundTrip: + """Fail if #3480 loses a completed, unconsumed prompt group.""" + result = round_trip(scenario, sampler_name, tmp_path) + lost = sorted(scenario.committed_outstanding() - result.recovered) + assert not lost, ( + f"{sampler_name}/{scenario.name}: restore dropped completed groups {lost}. " + "Their rows and replay index should both be present in the #3480 " + f"checkpoint. recovered={sorted(result.recovered)}" + ) + return result + + +# ── the scenarios ─────────────────────────────────────────────────────────── +# Numbering follows the worked example: groups 09-11 are trained, 12+ are not. + +S_ALL_COMPLETE = Scenario( + name="lag1-next-step-complete", + groups=( + Group(9, 2, weight=0), + Group(10, 2, weight=0), + Group(11, 2, weight=0), + Group(12, 2, weight=1, target=5), + Group(13, 2, weight=1, target=5), + Group(14, 2, weight=1, target=5), + ), + cursor=15, + trained=frozenset({9, 10, 11}), + lag=1, +) + +S_ZERO_LAG_ALL_COMPLETE = Scenario( + name="lag0-current-step-complete", + groups=( + Group(9, 2, weight=4, target=4), + Group(10, 2, weight=4, target=4), + Group(11, 2, weight=4, target=4), + Group(12, 2, weight=5, target=5), + Group(13, 2, weight=5, target=5), + Group(14, 2, weight=5, target=5), + ), + cursor=15, + trained=frozenset({9, 10, 11}), + lag=0, +) + +S_PARTIAL = Scenario( + name="lag1-next-step-partly-generated", + groups=( + Group(9, 2, weight=0), + Group(10, 2, weight=0), + Group(11, 2, weight=0), + Group(12, 1, weight=1, target=5), # one rollout still running + Group(13, 2, weight=1, target=5), + Group(14, 0, weight=1, target=5), # not started + ), + cursor=15, + trained=frozenset({9, 10, 11}), + lag=1, +) + +S_LAG2 = Scenario( + name="lag2-two-batches-in-flight", + groups=( + Group(9, 2, weight=0), + Group(10, 2, weight=0), + Group(11, 2, weight=0), + Group(12, 1, weight=1, target=5), + Group(13, 2, weight=1, target=5), + Group(14, 0, weight=1, target=5), + Group(15, 2, weight=2, target=6), + Group(16, 1, weight=2, target=6), + Group(17, 0, weight=2, target=6), + ), + cursor=18, + trained=frozenset({9, 10, 11}), + lag=2, +) + +S_EVICTED = Scenario( + name="lag1-with-an-evicted-group", + groups=( + Group(9, 2, weight=0), + Group(10, 2, weight=0, evicted=True), # dropped on purpose: too stale + Group(11, 2, weight=0), + Group(12, 1, weight=1, target=5), + Group(13, 2, weight=1, target=5), + Group(14, 0, weight=1, target=5), + ), + cursor=15, + trained=frozenset({9, 11}), + lag=1, +) + +S_TRAINED_OUT_OF_ORDER = Scenario( + name="trained-what-was-ready-leaving-a-hole", + groups=( + Group(9, 2, weight=0), + Group(10, 2, weight=0), + Group(11, 2, weight=0), + Group(12, 1, weight=1, target=5), # skipped: still generating + Group(13, 2, weight=1, target=5), # trained ahead of 12 + Group(14, 0, weight=1, target=5), + ), + cursor=15, + trained=frozenset({9, 10, 11, 13}), + lag=1, +) + +S_STALE_ONLY = Scenario( + name="one-group-far-outside-the-staleness-window", + groups=( + Group(9, ROLLOUTS_PER_GROUP, weight=0, target=1), + Group(10, ROLLOUTS_PER_GROUP, weight=7, target=8), + ), + cursor=11, + trained=frozenset(), + lag=1, +) + +# Everything fully generated -- the case this PR set out to recover. +FULLY_GENERATED = (S_ZERO_LAG_ALL_COMPLETE, S_ALL_COMPLETE, S_STALE_ONLY) +# At least one group still generating when the snapshot was taken. +WITH_IN_FLIGHT = (S_PARTIAL, S_LAG2, S_EVICTED, S_TRAINED_OUT_OF_ORDER) +ALL_SCENARIOS = FULLY_GENERATED + WITH_IN_FLIGHT diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py index 1c5d9a82054..3d63f5915ef 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_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py new file mode 100644 index 00000000000..062b2595d62 --- /dev/null +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -0,0 +1,1385 @@ +# 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. + +"""Controller checkpoint cuts around in-order rollout admission. + +These tests isolate the liveness hole that a replay-buffer-only checkpoint +cannot close: + +* the dataloader has advanced past a batch; +* the sampler has admitted that batch and persisted dispatch_index=7; +* none of its prompt groups committed before the data-plane snapshot. + +Restoring only the cursor correctly makes the next *new* admission step 8. +Recovery must therefore replay the owned batch at its saved target step 7 +without admitting it a second time. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import threading +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any, TypeVar, cast + +import pytest +import torch + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_FILENAME, + DataPlaneCheckpointBarrier, + DataPlaneMutationCut, + TQReplayBuffer, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import InOrderSampler +from nemo_rl.algorithms.grpo import _initial_grpo_save_state +from nemo_rl.algorithms.metric_utils import SetupTimingMetrics +from nemo_rl.algorithms.single_controller import SingleControllerActor +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.rollout_manager import RolloutOutcome +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_SCHEMA_VERSION, + ROLLOUT_RECOVERY_STATE_FILENAME, + PromptGroupPhase, + RolloutRecoveryLedger, + build_rollout_recovery_state, +) +from tests.unit.single_controller._checkpoint_scenarios import ( + _record, + patch_converter, +) +from tests.unit.single_controller.test_checkpointing import ( + _actor_master_config, + _FakeDataloader, + _make_actor_args, +) + +_ASYNC_TEST_TIMEOUT_S = 10.0 +_T = TypeVar("_T") + + +def _with_mutation_cut(callback: Callable[[DataPlaneMutationCut], _T]) -> _T: + async def apply() -> _T: + async with DataPlaneCheckpointBarrier().mutation() as cut: + return callback(cut) + + return asyncio.run(apply()) + + +async def _wait_for_event_or_pump( + event: asyncio.Event, + pump: asyncio.Task[None], +) -> None: + """Wait for a test hook while surfacing an early rollout-pump failure.""" + event_waiter = asyncio.create_task(event.wait()) + try: + done, _ = await asyncio.wait( + {event_waiter, pump}, + timeout=_ASYNC_TEST_TIMEOUT_S, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + raise TimeoutError("rollout pump did not reach the expected test hook") + if pump in done: + await pump + raise AssertionError("rollout pump completed before the expected test hook") + await event_waiter + finally: + if not event_waiter.done(): + event_waiter.cancel() + await asyncio.gather(event_waiter, return_exceptions=True) + + +class _CountingInOrderSampler(InOrderSampler): + """Real in-order sampler with observable admission calls.""" + + def __init__(self) -> None: + super().__init__(None, max_lookahead_versions=1) + self.admit_calls = 0 + self.admission_commits = 0 + + async def admit(self, *, trainer_version_fn): + self.admit_calls += 1 + return await super().admit(trainer_version_fn=trainer_version_fn) + + def commit_admission(self, cut: DataPlaneMutationCut): + self.admission_commits += 1 + return super().commit_admission(cut) + + +class _BlockingBeforeAdmissionSampler(_CountingInOrderSampler): + """Pause after the dataloader advances but before admission mutates state.""" + + def __init__(self) -> None: + super().__init__() + self.admission_entered = asyncio.Event() + self.release_admission = asyncio.Event() + + async def wait_until_admissible(self, *, trainer_version_fn): + self.admission_entered.set() + await self.release_admission.wait() + await super().wait_until_admissible(trainer_version_fn=trainer_version_fn) + + +@dataclass(frozen=True) +class _PendingGroup: + group_id: str + target_step: int | None + prompt_payload: dict[str, Any] + + +class _PendingLedger: + """Small stand-in for the group-level recovery ledger contract.""" + + def __init__(self, group: _PendingGroup | None = None) -> None: + self._groups = [group] if group is not None else [] + self.prepare_calls = 0 + + def prepare_for_restart(self) -> None: + self.prepare_calls += 1 + + def groups(self) -> list[_PendingGroup]: + return list(self._groups) + + def expected_staging_keys(self) -> set[str]: + return set() + + def record(self, group: _PendingGroup) -> None: + self._groups.append(group) + + def assign_target_step(self, group_id: str, target_step: int) -> None: + self._groups = [ + _PendingGroup( + group_id=group.group_id, + target_step=target_step, + prompt_payload=group.prompt_payload, + ) + if group.group_id == group_id + else group + for group in self._groups + ] + + def state_dict(self) -> dict[str, Any]: + return { + "schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "groups": [ + { + "group_id": group.group_id, + "admission_id": group.group_id, + "prompt_id": str(group.prompt_payload.get("idx", "unknown")), + "target_step": group.target_step, + "prompt_ref": { + "sample_id": str(group.prompt_payload.get("idx", "unknown")), + "task_name": group.prompt_payload.get("task_name"), + }, + "expected_generations": 2, + "start_weight_version": 7, + "phase": ("reserved" if group.target_step is None else "admitted"), + } + for group in self._groups + ], + } + + def release(self, group_id: str) -> None: + self._groups = [group for group in self._groups if group.group_id != group_id] + + +class _RecoveryRolloutManager: + def __init__(self, ledger: RolloutRecoveryLedger) -> None: + self.recovery_ledger = ledger + self.recovered: list[tuple[str, int | None]] = [] + + async def complete_recovery( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + group = self.recovery_ledger.get_group(group_id) + self.recovered.append((group.group_id, group.target_step)) + self.recovery_ledger.discard_group(cut, group_id) + + def mark_prompt_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: int | None, + ) -> None: + self.recovery_ledger.mark_group_admitted( + cut, + group_id, + target_step=target_step, + start_weight_version=7, + ) + + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + self.recovery_ledger.discard_group(cut, group_id) + + +class _BlockingRolloutManager: + """Hold one admitted rollout unfinished while the checkpoint is written.""" + + def __init__(self, ledger: _PendingLedger) -> None: + self.recovery_ledger = ledger + self.started = asyncio.Event() + self.release = asyncio.Event() + self.weight_version = 0 + + def set_weight_version(self, version: int) -> None: + self.weight_version = version + + def reserve_prompt_group( + self, + cut: DataPlaneMutationCut | None, + prompt: DatumSpec, + *, + target_step: int | None = None, + admitted: bool = True, + admission_id: str | None = None, + ) -> str: + del admitted, admission_id + batch_label = "fetched" if target_step is None else str(target_step) + group_id = f"batch-{batch_label}-prompt-{prompt['idx']}" + if not self.recovery_ledger.groups(): + self.recovery_ledger.record( + _PendingGroup( + group_id=group_id, + target_step=target_step, + prompt_payload=dict(prompt), + ) + ) + return group_id + + def mark_prompt_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: int | None, + ) -> None: + del cut + if target_step is None: + return + self.recovery_ledger.assign_target_step(group_id, target_step) + + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + del cut + self.recovery_ledger.release(group_id) + + async def generate_and_push( + self, + prompt: DatumSpec, + *, + target_step: int | None = None, + inflight_registry: dict[str, Any] | None = None, + lineage_group_id: str | None = None, + ) -> RolloutOutcome: + del inflight_registry + if lineage_group_id is None: + lineage_group_id = self.reserve_prompt_group( + None, + prompt, + target_step=target_step, + ) + self.started.set() + await self.release.wait() + return RolloutOutcome.COMMITTED + + +class _BlockingNoOpDataPlaneClient(NoOpDataPlaneClient): + """Hold the native data-plane save while a commit tries to publish.""" + + def __init__(self) -> None: + super().__init__() + self.save_started = threading.Event() + self.release_save = threading.Event() + + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = 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 _LedgerFacade: + """Minimal RolloutManager ownership surface for reserve-pool cuts.""" + + def __init__(self) -> None: + self.recovery_ledger = RolloutRecoveryLedger() + + def reserve_prompt_group( + self, + cut: DataPlaneMutationCut, + prompt: DatumSpec, + *, + target_step: int | None = None, + admitted: bool = True, + admission_id: str | None = None, + ) -> str: + record = self.recovery_ledger.reserve_group( + cut, + prompt_id=str(prompt["idx"]), + prompt_payload=prompt, + expected_generations=2, + target_step=target_step, + start_weight_version=7, + admitted=admitted, + admission_id=admission_id, + ) + return record.group_id + + def mark_prompt_group_admitted( + self, + cut: DataPlaneMutationCut, + group_id: str, + *, + target_step: int | None, + ) -> None: + self.recovery_ledger.mark_group_admitted( + cut, + group_id, + target_step=target_step, + start_weight_version=7, + ) + + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + self.recovery_ledger.discard_group(cut, group_id) + + +def _reserve_prompt(idx: int) -> DatumSpec: + return { + "idx": idx, + "message_log": [{"role": "user", "content": f"prompt {idx}"}], + "length": 1, + "extra_env_info": None, + "loss_multiplier": 1.0, + } + + +def _identity_dict_collator(batch: list[DatumSpec]) -> DatumSpec: + """Return one directly usable prompt rather than a BatchedDataDict.""" + assert len(batch) == 1 + prompt = dict(batch[0]) + prompt["length"] = 99 + return cast(DatumSpec, prompt) + + +def _two_row_collator(_batch: list[DatumSpec]) -> BatchedDataDict: + """Return an invalid two-row recovery batch.""" + return BatchedDataDict({"idx": [7, 8]}) + + +def _non_mapping_collator(_batch: list[DatumSpec]) -> list[str]: + """Return an invalid collator result type.""" + return ["not-a-prompt"] + + +def _rehydration_controller( + collate_fn: Callable[[list[DatumSpec]], Any], +) -> tuple[Any, RolloutRecoveryLedger]: + """Build a restored ledger whose prompt must be resolved from the dataset.""" + dataset_prompt = _reserve_prompt(7) + saved_ledger = RolloutRecoveryLedger() + _with_mutation_cut( + lambda cut: saved_ledger.reserve_group( + cut, + group_id="rehydrate-7", + prompt_id="7", + prompt_payload=dataset_prompt, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + ) + restored_ledger = RolloutRecoveryLedger() + _with_mutation_cut( + lambda cut: restored_ledger.load_state_dict(cut, saved_ledger.state_dict()) + ) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._rollout_manager = SimpleNamespace(recovery_ledger=restored_ledger) + controller._dataloader = SimpleNamespace( + dataset={7: dataset_prompt}, + collate_fn=collate_fn, + ) + return controller, restored_ledger + + +def _run_rehydration(controller: Any) -> None: + async def rehydrate() -> None: + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + await controller._rehydrate_rollout_recovery_prompts(cut) + + asyncio.run(rehydrate()) + + +def test_recovery_rehydration_accepts_an_identity_dict_collator() -> None: + controller, ledger = _rehydration_controller(_identity_dict_collator) + + _run_rehydration(controller) + + assert ledger.get_group("rehydrate-7").prompt_payload["length"] == 99 + + +@pytest.mark.parametrize( + ("collate_fn", "expected_error", "match"), + [ + pytest.param( + _two_row_collator, + ValueError, + "must return exactly one prompt", + id="multiple-prompts", + ), + pytest.param( + _non_mapping_collator, + TypeError, + "expected a mapping", + id="non-mapping", + ), + ], +) +def test_recovery_rehydration_rejects_invalid_collator_results( + collate_fn: Callable[[list[DatumSpec]], Any], + expected_error: type[Exception], + match: str, +) -> None: + controller, _ = _rehydration_controller(collate_fn) + + with pytest.raises(expected_error, match=match): + _run_rehydration(controller) + + +def _reserve_controller() -> Any: + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._replacement_reserve = deque() + controller._sampler_stamps_target_steps = True + controller._rollout_recovery_enabled = True + controller._async_cfg = SimpleNamespace( + rollout_failure=SimpleNamespace( + on_dropped_prompt="replace", + replacement_reserve_prompts=2, + max_replacement_attempts=1, + ) + ) + controller._algo_cfg = SimpleNamespace(num_prompts_per_step=2) + controller._rollout_manager = _LedgerFacade() + return controller + + +def test_dispatch_cursor_alone_assigns_the_next_batch_to_step_8() -> None: + """The exact cursor is correct; it cannot recreate the missing step-7 batch.""" + + async def exercise() -> int | None: + sampler = _CountingInOrderSampler() + sampler.restore_dispatch_index(7) + return await sampler.admit(trainer_version_fn=lambda: 7) + + assert asyncio.run(exercise()) == 8 + + +def test_checkpoint_after_fetch_before_admit_owns_the_prompt(tmp_path) -> None: + """A checkpoint cut inside admit retains the fetched batch for recovery.""" + + async def exercise() -> None: + save_state = _initial_grpo_save_state() + save_state.current_step = 7 + save_state.total_steps = 7 + save_state.trainer_version = 7 + save_state.sampler_dispatch_index = 6 + + ledger = _PendingLedger() + rollout_manager = _BlockingRolloutManager(ledger) + dataloader = _FakeDataloader( + [ + BatchedDataDict( + { + "idx": [70], + "message_log": [[{"role": "user", "content": "batch 7"}]], + } + ) + ], + state={"next_batch": 8}, + ) + config = _actor_master_config( + tmp_path, + max_num_steps=8, + num_prompts_per_step=1, + max_num_epochs=1, + data_plane_checkpoint=True, + ) + actor_args = _make_actor_args( + save_state=save_state, + dataloader=dataloader, + ) + actor_args.rollout_manager = rollout_manager # type: ignore[assignment] + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = controller_cls(config, actor_args, SetupTimingMetrics()) + sampler = _BlockingBeforeAdmissionSampler() + sampler.restore_dispatch_index(6) + controller._sampler = sampler + pump = asyncio.create_task(controller._rollout_pump()) + await _wait_for_event_or_pump(sampler.admission_entered, pump) + assert controller._sampler.dispatch_index == 6 + + try: + await controller._save_checkpoint( + {"loss": 1.0}, + is_policy_training_step=True, + ) + finally: + sampler.release_admission.set() + await _wait_for_event_or_pump(rollout_manager.started, pump) + rollout_manager.release.set() + await asyncio.wait_for(pump, timeout=_ASYNC_TEST_TIMEOUT_S) + controller._checkpointer.shutdown() + + checkpoint = tmp_path / "checkpoints" / "step_7" + recovery_state = torch.load( + checkpoint / "rollout_recovery.pt", + weights_only=False, + ) + assert len(recovery_state["groups"]) == 1 + assert recovery_state["groups"][0]["target_step"] is None + assert recovery_state["groups"][0]["prompt_ref"]["sample_id"] == "70" + assert "prompt_payload" not in recovery_state["groups"][0] + assert torch.load( + checkpoint / "train_dataloader.pt", + weights_only=False, + ) == {"next_batch": 8} + + asyncio.run(exercise()) + + +def test_checkpoint_owns_batch_7_while_its_rollout_is_unfinished(tmp_path) -> None: + """A finalized checkpoint cannot contain a cursor hole for target step 7.""" + + async def exercise() -> None: + save_state = _initial_grpo_save_state() + save_state.current_step = 7 + save_state.total_steps = 7 + save_state.trainer_version = 7 + save_state.sampler_dispatch_index = 6 + + # Start empty: generate_and_push records the batch-7 prompt only after + # the sampler has admitted it. + ledger = _PendingLedger() + rollout_manager = _BlockingRolloutManager(ledger) + dataloader = _FakeDataloader( + [ + BatchedDataDict( + { + "idx": [70], + "message_log": [[{"role": "user", "content": "batch 7"}]], + } + ) + ], + state={"next_batch": 8}, + ) + config = _actor_master_config( + tmp_path, + max_num_steps=8, + num_prompts_per_step=1, + max_num_epochs=1, + data_plane_checkpoint=True, + ) + actor_args = _make_actor_args( + save_state=save_state, + dataloader=dataloader, + ) + actor_args.rollout_manager = rollout_manager # type: ignore[assignment] + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = controller_cls(config, actor_args, SetupTimingMetrics()) + + pump = asyncio.create_task(controller._rollout_pump()) + await _wait_for_event_or_pump(rollout_manager.started, pump) + assert controller._sampler.dispatch_index == 7 + + try: + await controller._save_checkpoint( + {"loss": 1.0}, + is_policy_training_step=True, + ) + finally: + rollout_manager.release.set() + await asyncio.wait_for(pump, timeout=_ASYNC_TEST_TIMEOUT_S) + controller._checkpointer.shutdown() + + checkpoint = tmp_path / "checkpoints" / "step_7" + recovery_path = checkpoint / "rollout_recovery.pt" + assert recovery_path.is_file() + recovery_state = torch.load(recovery_path, weights_only=False) + assert [group["target_step"] for group in recovery_state["groups"]] == [7] + assert torch.load( + checkpoint / "train_dataloader.pt", + weights_only=False, + ) == {"next_batch": 8} + + asyncio.run(exercise()) + + +def test_commit_contending_with_checkpoint_has_exactly_one_saved_owner( + tmp_path, + monkeypatch, +) -> None: + """The checkpoint records the group as canonical or pending, never neither.""" + patch_converter(monkeypatch) + + async def exercise() -> None: + dp_client = _BlockingNoOpDataPlaneClient() + dp_client.register_partition( + partition_id="rollout_data", + fields=["input_ids", "input_lengths", "total_reward"], + num_samples=8, + consumer_tasks=["train"], + grpo_group_size=2, + ) + buffer = TQReplayBuffer( + dp_client, + partition_id="rollout_data", + pad_value_dict={"input_ids": 0}, + require_routed_experts=False, + ) + group_id = buffer.reserve( + weight_version=7, + target_step=7, + group_id="batch-7-prompt-70", + ) + ledger = _PendingLedger( + _PendingGroup( + group_id=group_id, + target_step=7, + prompt_payload={"idx": 70, "message_log": []}, + ) + ) + rollout_manager = _BlockingRolloutManager(ledger) + + save_state = _initial_grpo_save_state() + save_state.current_step = 7 + save_state.total_steps = 7 + save_state.trainer_version = 7 + save_state.sampler_dispatch_index = 7 + config = _actor_master_config( + tmp_path, + max_num_steps=8, + num_prompts_per_step=1, + data_plane_checkpoint=True, + ) + actor_args = _make_actor_args( + save_state=save_state, + dataloader=_FakeDataloader(state={"next_batch": 8}), + tq_buffer=buffer, # type: ignore[arg-type] + dp_client=dp_client, # type: ignore[arg-type] + ) + actor_args.rollout_manager = rollout_manager # type: ignore[assignment] + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = controller_cls(config, actor_args, SetupTimingMetrics()) + save_task = asyncio.create_task( + controller._save_checkpoint( + {"loss": 1.0}, + is_policy_training_step=True, + ) + ) + save_started = await asyncio.to_thread(dp_client.save_started.wait, 5.0) + assert save_started + + commit_task = asyncio.create_task( + controller._buffer.commit( + group_id, + _record(), + start_weight_version=7, + end_weight_version=7, + ) + ) + await asyncio.sleep(0) + assert not commit_task.done() + assert controller._buffer.ready_list == [False] + + dp_client.release_save.set() + await asyncio.wait_for(save_task, timeout=5.0) + await asyncio.wait_for(commit_task, timeout=5.0) + controller._checkpointer.shutdown() + + checkpoint = tmp_path / "checkpoints" / "step_7" + replay_state = torch.load( + checkpoint / REPLAY_BUFFER_METADATA_FILENAME, + weights_only=False, + ) + recovery_state = torch.load( + checkpoint / "rollout_recovery.pt", + weights_only=False, + ) + canonical_ids = {group["group_id"] for group in replay_state["groups"]} + pending_ids = {group["group_id"] for group in recovery_state["groups"]} + + assert int(group_id in canonical_ids) + int(group_id in pending_ids) == 1 + assert group_id not in canonical_ids + assert group_id in pending_ids + assert controller._buffer.ready_list == [True] + + asyncio.run(exercise()) + + +def test_canonical_replay_wins_over_stale_ledger_entry( + tmp_path, + monkeypatch, +) -> None: + """A completed group appears exactly once when ledger cleanup loses the cut.""" + patch_converter(monkeypatch) + + async def exercise() -> None: + dp_client = NoOpDataPlaneClient() + dp_client.register_partition( + partition_id="rollout_data", + fields=["input_ids", "input_lengths", "total_reward"], + num_samples=8, + consumer_tasks=["train"], + grpo_group_size=2, + ) + buffer = TQReplayBuffer( + dp_client, + partition_id="rollout_data", + pad_value_dict={"input_ids": 0}, + require_routed_experts=False, + ) + group_id = buffer.reserve( + weight_version=7, + target_step=7, + group_id="batch-7-prompt-70", + ) + + # Model the narrow cut after the canonical commit but before the live + # ledger entry is released. The checkpoint must not persist both owners. + ledger = _PendingLedger( + _PendingGroup( + group_id=group_id, + target_step=7, + prompt_payload={"idx": 70, "message_log": []}, + ) + ) + rollout_manager = _BlockingRolloutManager(ledger) + + save_state = _initial_grpo_save_state() + save_state.current_step = 7 + save_state.total_steps = 7 + save_state.trainer_version = 7 + save_state.sampler_dispatch_index = 7 + config = _actor_master_config( + tmp_path, + max_num_steps=8, + num_prompts_per_step=1, + data_plane_checkpoint=True, + ) + actor_args = _make_actor_args( + save_state=save_state, + dataloader=_FakeDataloader(state={"next_batch": 8}), + tq_buffer=buffer, # type: ignore[arg-type] + dp_client=dp_client, # type: ignore[arg-type] + ) + actor_args.rollout_manager = rollout_manager # type: ignore[assignment] + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = controller_cls(config, actor_args, SetupTimingMetrics()) + await buffer.commit( + group_id, + _record(), + start_weight_version=7, + end_weight_version=7, + ) + assert buffer.ready_list == [True] + + try: + await controller._save_checkpoint( + {"loss": 1.0}, + is_policy_training_step=True, + ) + finally: + controller._checkpointer.shutdown() + + checkpoint = tmp_path / "checkpoints" / "step_7" + replay_state = torch.load( + checkpoint / REPLAY_BUFFER_METADATA_FILENAME, + weights_only=False, + ) + recovery_state = torch.load( + checkpoint / "rollout_recovery.pt", + weights_only=False, + ) + canonical_ids = {group["group_id"] for group in replay_state["groups"]} + pending_ids = {group["group_id"] for group in recovery_state["groups"]} + + assert group_id in canonical_ids + assert group_id not in pending_ids + assert int(group_id in canonical_ids) + int(group_id in pending_ids) == 1 + + asyncio.run(exercise()) + + +def test_recovery_replays_step_7_without_readmitting_the_batch(tmp_path) -> None: + """An admitted batch keeps target_step=7 across a process restart.""" + + async def exercise() -> None: + sampler = _CountingInOrderSampler() + sampler.restore_dispatch_index(7) + dataset_prompt: DatumSpec = { + "idx": 70, + "message_log": [], + "length": 1, + "extra_env_info": None, + "loss_multiplier": 1.0, + } + prompt_batch = rl_collate_fn([dataset_prompt]) + dispatched_prompt = cast( + DatumSpec, + {key: value[0] for key, value in prompt_batch.items()}, + ) + saved_ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + saved_ledger.reserve_group( + cut, + group_id="batch-7-prompt-0", + admission_id="batch-7", + prompt_id="70", + prompt_payload=dispatched_prompt, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + saved_state = build_rollout_recovery_state( + saved_ledger, + batch_shortfall={6: 1}, + sampler_stamps_target_steps=True, + ) + recovery_path = tmp_path / ROLLOUT_RECOVERY_STATE_FILENAME + torch.save(saved_state, recovery_path) + payload_sha256 = hashlib.sha256(recovery_path.read_bytes()).hexdigest() + + ledger = RolloutRecoveryLedger() + rollout_manager = _RecoveryRolloutManager(ledger) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._sampler = sampler + controller._rollout_manager = rollout_manager + controller._last_checkpoint_path = str(tmp_path) + controller._data_plane_checkpoint_metadata = { + "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "rollout_recovery_payload_sha256": payload_sha256, + "rollout_recovery_group_count": 1, + } + controller._async_cfg = SimpleNamespace( + max_buffered_rollouts=4, + max_inflight_prompts=2, + ) + controller._buffer_capacity = asyncio.Semaphore(4) + controller._trainer_version = 7 + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._dataloader = SimpleNamespace( + dataset={70: dataset_prompt}, + collate_fn=rl_collate_fn, + ) + controller._buffer = SimpleNamespace( + count_for_target_step=lambda _target_step: 0, + metadata_state_dict=lambda *, saved_capacity: { + "groups": [], + "saved_capacity": saved_capacity, + }, + ) + + async def _recover( + _prompt: dict[str, Any], + _target_step: int | None, + group_id: str, + ) -> None: + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + await rollout_manager.complete_recovery(cut, group_id) + + await controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + await controller._redispatch_restored_rollouts(_recover) + + assert rollout_manager.recovered == [("batch-7-prompt-0", 7)] + assert sampler.admit_calls == 0 + assert sampler.admission_commits == 0 + assert sampler.dispatch_index == 7 + assert controller._batch_shortfall == {6: 1} + assert controller._sampler_stamps_target_steps is True + + asyncio.run(exercise()) + + +def test_recovery_rejects_an_unhandled_phase_before_redispatch() -> None: + """A future phase must fail loudly instead of remaining owned forever.""" + + async def exercise() -> None: + recovery_ledger = SimpleNamespace( + groups=lambda: [ + SimpleNamespace(group_id="future-group", phase="future-phase") + ] + ) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._rollout_manager = SimpleNamespace(recovery_ledger=recovery_ledger) + launched = False + + async def _recover( + _prompt: dict[str, Any], + _target_step: int | None, + _group_id: str, + ) -> None: + nonlocal launched + launched = True + + with pytest.raises( + RuntimeError, + match=r"unrecognized rollout recovery phase.*future-group='future-phase'", + ): + await controller._redispatch_restored_rollouts(_recover) + + assert not launched + + asyncio.run(exercise()) + + +def test_recovery_readmits_one_reserved_batch_only_once(tmp_path) -> None: + """Two prompts fetched together consume one sampler admission on restart.""" + + async def exercise() -> None: + saved_ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + for prompt_idx in (70, 71): + saved_ledger.reserve_group( + cut, + group_id=f"batch-7-prompt-{prompt_idx}", + admission_id="batch-7", + prompt_id=str(prompt_idx), + prompt_payload={"idx": prompt_idx, "message_log": []}, + expected_generations=2, + target_step=None, + start_weight_version=7, + admitted=False, + ) + recovery_path = tmp_path / ROLLOUT_RECOVERY_STATE_FILENAME + torch.save(saved_ledger.state_dict(), recovery_path) + + sampler = _CountingInOrderSampler() + sampler.restore_dispatch_index(6) + rollout_manager = _RecoveryRolloutManager(RolloutRecoveryLedger()) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._sampler = sampler + controller._rollout_manager = rollout_manager + controller._last_checkpoint_path = str(tmp_path) + controller._data_plane_checkpoint_metadata = { + "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "rollout_recovery_payload_sha256": hashlib.sha256( + recovery_path.read_bytes() + ).hexdigest(), + "rollout_recovery_group_count": 2, + } + controller._async_cfg = SimpleNamespace( + max_buffered_rollouts=4, + max_inflight_prompts=2, + ) + controller._buffer_capacity = asyncio.Semaphore(4) + controller._trainer_version = 7 + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._dataloader = SimpleNamespace( + dataset={ + prompt_idx: {"idx": prompt_idx, "message_log": []} + for prompt_idx in (70, 71) + } + ) + controller._buffer = SimpleNamespace( + count_for_target_step=lambda _target_step: 0, + metadata_state_dict=lambda *, saved_capacity: { + "groups": [], + "saved_capacity": saved_capacity, + }, + ) + + async def _recover( + _prompt: dict[str, Any], + _target_step: int | None, + group_id: str, + ) -> None: + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + await rollout_manager.complete_recovery(cut, group_id) + + await controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + await controller._redispatch_restored_rollouts(_recover) + + assert sampler.admit_calls == 0 + assert sampler.admission_commits == 1 + assert sampler.dispatch_index == 7 + assert set(rollout_manager.recovered) == { + ("batch-7-prompt-70", 7), + ("batch-7-prompt-71", 7), + } + + asyncio.run(exercise()) + + +def test_recovery_launches_admitted_groups_before_waiting_to_readmit() -> None: + """Recovered work can open the gate that a reserved batch is waiting on.""" + + async def exercise() -> None: + sampler = _CountingInOrderSampler() + sampler.restore_dispatch_index(7) + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + ledger.reserve_group( + cut, + group_id="admitted-step-7", + admission_id="batch-7", + prompt_id="70", + prompt_payload={"idx": 70, "message_log": []}, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + ledger.reserve_group( + cut, + group_id="reserved-step-8", + admission_id="batch-8", + prompt_id="80", + prompt_payload={"idx": 80, "message_log": []}, + expected_generations=2, + target_step=None, + start_weight_version=7, + admitted=False, + ) + rollout_manager = _RecoveryRolloutManager(ledger) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._sampler = sampler + controller._rollout_manager = rollout_manager + controller._trainer_version = 6 + controller._data_plane_checkpoint_barrier = barrier + controller._buffer = SimpleNamespace( + count_for_target_step=lambda _target_step: 0, + ) + + launched: list[tuple[str, int | None]] = [] + + async def _recover( + _prompt: dict[str, Any], + target_step: int | None, + group_id: str, + ) -> None: + launched.append((group_id, target_step)) + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + await rollout_manager.complete_recovery(cut, group_id) + if group_id == "admitted-step-7": + # Model the concurrent train pump consuming recovered step 7. This + # opens the in-order gate so the reserved batch can become step 8. + controller._trainer_version = 7 + + await asyncio.wait_for( + controller._redispatch_restored_rollouts(_recover), + timeout=1.0, + ) + + assert launched == [ + ("admitted-step-7", 7), + ("reserved-step-8", 8), + ] + assert sampler.admission_commits == 1 + assert sampler.dispatch_index == 8 + assert len(ledger) == 0 + + asyncio.run(exercise()) + + +def test_recovery_load_does_not_require_every_unfinished_group_to_fit_at_once( + tmp_path, +) -> None: + """The train pump may free replay slots while recovery is redispatching.""" + + async def exercise() -> None: + saved_ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + for prompt_idx in (70, 71): + saved_ledger.reserve_group( + cut, + group_id=f"batch-7-prompt-{prompt_idx}", + admission_id="batch-7", + prompt_id=str(prompt_idx), + prompt_payload={"idx": prompt_idx, "message_log": []}, + expected_generations=2, + target_step=7, + start_weight_version=7, + admitted=True, + ) + recovery_path = tmp_path / ROLLOUT_RECOVERY_STATE_FILENAME + torch.save(saved_ledger.state_dict(), recovery_path) + + rollout_manager = _RecoveryRolloutManager(RolloutRecoveryLedger()) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._rollout_manager = rollout_manager + controller._last_checkpoint_path = str(tmp_path) + controller._data_plane_checkpoint_metadata = { + "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "rollout_recovery_payload_sha256": hashlib.sha256( + recovery_path.read_bytes() + ).hexdigest(), + "rollout_recovery_group_count": 2, + } + controller._async_cfg = SimpleNamespace(max_buffered_rollouts=4) + controller._dataloader = SimpleNamespace( + dataset={ + prompt_idx: {"idx": prompt_idx, "message_log": []} + for prompt_idx in (70, 71) + } + ) + controller._buffer = SimpleNamespace( + metadata_state_dict=lambda *, saved_capacity: { + "groups": [{"group_id": f"canonical-{idx}"} for idx in range(3)], + "saved_capacity": saved_capacity, + } + ) + + # Three canonical groups plus two unfinished groups exceed capacity four, + # but only the canonical groups occupy slots at restore time. Recovery is + # launched beside the train pump, which releases capacity as it consumes. + await controller._maybe_restore_rollout_recovery(restored_replay_groups=3) + + assert len(rollout_manager.recovery_ledger) == 2 + + asyncio.run(exercise()) + + +def test_checkpoint_waits_for_replacement_reserve_refill() -> None: + """The dataloader-owned batch is visible in the pool after the mutation cut.""" + + async def exercise() -> None: + controller = _reserve_controller() + mutation_applied = asyncio.Event() + release_mutation = asyncio.Event() + checkpoint_entered = asyncio.Event() + batch = BatchedDataDict( + { + "idx": [20, 21], + "message_log": [ + [{"role": "user", "content": "prompt 20"}], + [{"role": "user", "content": "prompt 21"}], + ], + "length": [1, 1], + "extra_env_info": [None, None], + "loss_multiplier": [1.0, 1.0], + } + ) + + async def refill() -> None: + async with controller._data_plane_checkpoint_barrier.mutation(): + assert controller._divert_batch_to_reserve(batch) + mutation_applied.set() + await release_mutation.wait() + + async def checkpoint_snapshot() -> list[int]: + async with controller._data_plane_checkpoint_barrier.checkpoint(): + checkpoint_entered.set() + return [prompt["idx"] for prompt in controller._replacement_reserve] + + refill_task = asyncio.create_task(refill()) + await mutation_applied.wait() + checkpoint_task = asyncio.create_task(checkpoint_snapshot()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + + release_mutation.set() + assert await checkpoint_task == [20, 21] + await refill_task + + asyncio.run(exercise()) + + +def test_checkpoint_waits_for_replacement_pop_and_reownership() -> None: + """A skipped owner becomes its replacement atomically at checkpoint time.""" + + async def exercise() -> None: + controller = _reserve_controller() + manager = controller._rollout_manager + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + old_group_id = manager.reserve_prompt_group( + cut, + _reserve_prompt(20), + target_step=7, + ) + controller._replacement_reserve.append(_reserve_prompt(21)) + mutation_applied = asyncio.Event() + release_mutation = asyncio.Event() + checkpoint_entered = asyncio.Event() + + async def replace() -> None: + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + replacement = controller._take_replacement(7, 0) + assert replacement is not None + manager.discard_prompt_group(cut, old_group_id) + manager.reserve_prompt_group(cut, replacement, target_step=7) + mutation_applied.set() + await release_mutation.wait() + + async def checkpoint_snapshot() -> tuple[list[int], list[str]]: + async with controller._data_plane_checkpoint_barrier.checkpoint(): + checkpoint_entered.set() + reserve_ids = [ + prompt["idx"] for prompt in controller._replacement_reserve + ] + ledger_prompt_ids = [ + group.prompt_id for group in manager.recovery_ledger.groups() + ] + return reserve_ids, ledger_prompt_ids + + replace_task = asyncio.create_task(replace()) + await mutation_applied.wait() + checkpoint_task = asyncio.create_task(checkpoint_snapshot()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + + release_mutation.set() + reserve_ids, ledger_prompt_ids = await checkpoint_task + await replace_task + + assert reserve_ids == [] + assert ledger_prompt_ids == ["21"] + + asyncio.run(exercise()) + + +def test_reserve_drain_is_recoverable_before_sampler_admission() -> None: + """After pool removal, RESERVED ledger records own the whole batch.""" + + async def exercise() -> None: + controller = _reserve_controller() + controller._replacement_reserve.extend( + [_reserve_prompt(20), _reserve_prompt(21)] + ) + admission_started = asyncio.Event() + release_admission = asyncio.Event() + launched: list[tuple[int, int | None, str | None]] = [] + + async def block_admission( + group_ids: list[str], + ) -> tuple[int, list[str], int]: + admission_started.set() + await release_admission.wait() + async with controller._data_plane_checkpoint_barrier.mutation() as cut: + for group_id in group_ids: + controller._rollout_manager.mark_prompt_group_admitted( + cut, + group_id, + target_step=7, + ) + return 7, group_ids, 0 + + async def launch( + prompt: DatumSpec, + target_step: int | None, + group_id: str | None, + ) -> None: + launched.append((prompt["idx"], target_step, group_id)) + + controller._admit_reserved_prompt_groups = block_admission + drain_task = asyncio.create_task(controller._drain_reserve_into_steps(launch)) + await admission_started.wait() + + async with controller._data_plane_checkpoint_barrier.checkpoint(): + assert list(controller._replacement_reserve) == [] + groups = controller._rollout_manager.recovery_ledger.groups() + assert [group.prompt_id for group in groups] == ["20", "21"] + assert all(group.phase is PromptGroupPhase.RESERVED for group in groups) + assert len({group.admission_id for group in groups}) == 1 + + release_admission.set() + await drain_task + assert {(idx, target_step) for idx, target_step, _ in launched} == { + (20, 7), + (21, 7), + } + + asyncio.run(exercise()) + + +def test_recovery_rejects_a_corrupt_ledger_sidecar(tmp_path) -> None: + recovery_path = tmp_path / ROLLOUT_RECOVERY_STATE_FILENAME + recovery_path.write_bytes(b"corrupt checkpoint payload") + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._last_checkpoint_path = str(tmp_path) + controller._data_plane_checkpoint_metadata = { + "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "rollout_recovery_payload_sha256": "0" * 64, + "rollout_recovery_group_count": 1, + } + + with pytest.raises(ValueError, match="checksum mismatch"): + asyncio.run( + controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + ) + + +def test_recovery_rejects_a_missing_advertised_ledger_sidecar(tmp_path) -> None: + """Do not combine an older/missing ledger with the restored trainer step.""" + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._last_checkpoint_path = str(tmp_path) + controller._data_plane_checkpoint_metadata = { + "rollout_recovery_schema_version": ROLLOUT_RECOVERY_SCHEMA_VERSION, + "rollout_recovery_payload_sha256": "0" * 64, + "rollout_recovery_group_count": 1, + } + + with pytest.raises(FileNotFoundError, match="sidecar is missing"): + asyncio.run( + controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + ) diff --git a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py new file mode 100644 index 00000000000..c99e9cb15a4 --- /dev/null +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -0,0 +1,139 @@ +# 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. + +"""Checkpoint recovery contract across the built-in async samplers. + +The six scenarios come from #3827. That PR covered windowed, weight_fifo, and +in_order; a zero-lag completed row and ready_first are included here to cover +the complete built-in recovery contract. + +Unfinished rows are regenerated as whole prompt groups from the group-level +ledger. Sibling-level continuation remains outside this recovery foundation. +""" + +from __future__ import annotations + +import pytest + +from tests.unit.single_controller._checkpoint_scenarios import ( + ALL_SCENARIOS, + FULLY_GENERATED, + GROUPS_PER_STEP, + S_ALL_COMPLETE, + S_LAG2, + S_ZERO_LAG_ALL_COMPLETE, + SAMPLERS, + WITH_IN_FLIGHT, + Case, + assert_completed_groups_survive, + assert_no_data_loss, + patch_converter, + round_trip, +) + +ALL_CASES = [ + Case(scenario, sampler) for sampler in SAMPLERS for scenario in ALL_SCENARIOS +] +COMPLETED_CASES = [ + Case(scenario, sampler) for sampler in SAMPLERS for scenario in FULLY_GENERATED +] +UNFINISHED_CASES = [ + Case(scenario, sampler) for sampler in SAMPLERS for scenario in WITH_IN_FLIGHT +] +SELECTABLE_CASES = [ + Case(scenario, sampler) + for sampler in SAMPLERS + for scenario in (S_ZERO_LAG_ALL_COMPLETE, S_ALL_COMPLETE, S_LAG2) +] + + +@pytest.fixture(autouse=True) +def _converter(monkeypatch): + patch_converter(monkeypatch) + + +@pytest.mark.parametrize("case", ALL_CASES, ids=lambda case: case.id) +def test_completed_groups_survive_the_round_trip(case, tmp_path): + """A pending sibling must not hide a different group that already committed.""" + result = assert_completed_groups_survive( + case.scenario, + case.sampler, + tmp_path, + ) + assert result.saved_sidecar + + +@pytest.mark.parametrize("case", COMPLETED_CASES, ids=lambda case: case.id) +def test_fully_generated_scenarios_have_no_data_loss(case, tmp_path): + result = assert_no_data_loss(case.scenario, case.sampler, tmp_path) + assert result.recovered == case.scenario.must_survive() + + +@pytest.mark.parametrize("case", UNFINISHED_CASES, ids=lambda case: case.id) +def test_unfinished_groups_are_owned_across_restart(case, tmp_path): + """Handed-out unfinished groups remain recoverable.""" + assert_no_data_loss(case.scenario, case.sampler, tmp_path) + + +@pytest.mark.parametrize("case", ALL_CASES, ids=lambda case: case.id) +def test_restore_preserves_sampler_stamps(case, tmp_path): + """Every restored group retains the keys its sampler uses for selection.""" + result = round_trip(case.scenario, case.sampler, tmp_path) + + assert result.stamps == case.scenario.expected_stamps() + + +@pytest.mark.parametrize("sampler", SAMPLERS) +def test_restore_reuses_the_same_tq_rows(sampler, tmp_path): + """The replay sidecar restores the index; it must not duplicate tensor rows.""" + scenario = FULLY_GENERATED[0] + result = round_trip(scenario, sampler, tmp_path) + + assert result.rows_before + assert result.rows_after == result.rows_before + + +@pytest.mark.parametrize("sampler", SAMPLERS) +def test_intentionally_evicted_group_is_not_resurrected(sampler, tmp_path): + """Recovery restores owned work, not work the sampler deliberately discarded.""" + scenario = next(s for s in WITH_IN_FLIGHT if "evicted" in s.name) + result = round_trip(scenario, sampler, tmp_path) + + assert "g10" not in result.recovered + + +@pytest.mark.parametrize("case", SELECTABLE_CASES, ids=lambda case: case.id) +def test_restored_groups_are_selectable(case, tmp_path): + """Each sampler can select the restored batch at gate lags zero, one, and two.""" + first_outstanding = next( + group + for group in case.scenario.groups + if group.gid not in case.scenario.trained and not group.evicted + ) + current_train_weight = ( + first_outstanding.target + if case.sampler == "in_order" + else first_outstanding.weight + ) + assert current_train_weight is not None + + result = round_trip( + case.scenario, + case.sampler, + tmp_path, + select_current_train_weight=current_train_weight, + ) + + assert result.selected_count == GROUPS_PER_STEP + assert result.selected == {"g12", "g13", "g14"} diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 1a2d6d5e263..69b5caada71 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -27,8 +27,7 @@ - dataloader state: train_dataloader.pt written at save, position round-trip through a real StatefulDataLoader, dataset-swap guard, setup restore wiring + missing-file corruption check; - - replay buffer persistence (restore skipped on a sampler_name mismatch, - restored permits released by a live train pump); + - 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). """ @@ -36,6 +35,7 @@ from __future__ import annotations import asyncio +import hashlib import json import os import threading @@ -50,7 +50,19 @@ import yaml from torchdata.stateful_dataloader import StatefulDataLoader -from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.async_utils.replay_buffer import ( + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, + DataPlaneCheckpointBarrier, + DataPlaneCheckpointMetadata, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + WindowedSamplerConfig, + sampler_supports_buffer_checkpoint, +) from nemo_rl.algorithms.grpo import ( GRPOConfig, GRPOSaveState, @@ -67,7 +79,12 @@ ) 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.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_SCHEMA_VERSION, + ROLLOUT_RECOVERY_STATE_FILENAME, + RolloutRecoveryLedger, +) from nemo_rl.utils.checkpoint import CheckpointManager # Reuse the factory patches from the setup tests (same cross-module fixture @@ -161,8 +178,10 @@ def finalize_async_save(self) -> None: class _FakeSampler: """PromptGroupSampler stand-in: always returns a full, fresh batch.""" - def __init__(self) -> None: + def __init__(self, supports_buffer_checkpoint: bool = True) -> None: + self._supports_buffer_checkpoint = supports_buffer_checkpoint self._step = 0 + self._dispatch_index = -1 async def admit(self, *, trainer_version_fn) -> Optional[int]: return None @@ -193,14 +212,25 @@ async def select( 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 def set_gate_window(self, gate_window: int) -> None: self.gate_window = gate_window - def set_dispatch_index(self, resume_from_step: int) -> None: - pass + @property + def dispatch_index(self) -> int: + return self._dispatch_index + + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: + self._dispatch_index = resume_from_trainer_version - 1 + + def restore_dispatch_index(self, dispatch_index: int) -> None: + self._dispatch_index = dispatch_index class _ExhaustingSampler(_FakeSampler): @@ -217,12 +247,98 @@ async def select(self, **kwargs) -> tuple[Optional[KVBatchMeta], int]: return await super().select(**kwargs) +class _RestoredGroupsSampler(_FakeSampler): + """Drain the exact groups represented by a restored replay metadata file.""" + + def __init__(self, groups: list[dict[str, Any]]) -> None: + super().__init__() + self._groups = list(groups) + + async def select( + self, + *, + current_train_weight: int, + min_prompt_groups: int, + max_prompt_groups: int, + ) -> tuple[Optional[KVBatchMeta], int]: + del current_train_weight + selected = self._groups[:max_prompt_groups] + if len(selected) < min_prompt_groups: + return None, 0 + del self._groups[: len(selected)] + + metas = [group["meta"] for group in selected] + return ( + KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=[sid for meta in metas for sid in meta.sample_ids], + sequence_lengths=[ + length for meta in metas for length in (meta.sequence_lengths or []) + ], + tags=[tag for meta in metas for tag in (meta.tags or [])], + ), + len(selected), + ) + + class _FakeDPClient: - def __init__(self) -> None: + 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)) + cleared = set(sample_ids) + self.sample_ids = [sid for sid in self.sample_ids if sid not in cleared] + + 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: @@ -246,6 +362,7 @@ class _FakeRolloutManager: def __init__(self) -> None: self.weight_versions: list[int] = [] self._tq_buffer = None + self.recovery_ledger = RolloutRecoveryLedger() def set_weight_version(self, version: int) -> None: self.weight_versions.append(version) @@ -256,19 +373,39 @@ class _FakeTQBuffer: def __init__( self, - state: Optional[dict[str, Any]] = None, + metadata_state: Optional[dict[str, Any]] = None, load_return: int = 0, ) -> None: # Empty like a drained buffer; the pump's exhaustion checks len() it. self._num_groups = 0 - self._state = state if state is not None else {"fake_buffer_envelope": 1} + 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.state_dict_calls: list[int] = [] + self.metadata_state_dict_calls: list[int] = [] self.load_calls: list[dict[str, Any]] = [] + self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None - async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]: - self.state_dict_calls.append(saved_capacity) - return dict(self._state) + 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) + + def count_for_target_step(self, target_step: int) -> int: + """Return the number of ready fake groups owned by one gated step.""" + return sum( + group["target_step"] == target_step + for group in self._metadata_state["groups"] + ) async def load_state_dict( self, @@ -277,6 +414,7 @@ async def load_state_dict( max_groups: int, expected_partition_id: str, expected_group_size: int, + expected_manifest_digest: str, ) -> int: self.load_calls.append( { @@ -284,6 +422,7 @@ async def load_state_dict( "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 @@ -327,13 +466,19 @@ def _actor_master_config( ft_save_period: Optional[int] = None, num_prompts_per_step: int = 2, max_num_epochs: int = 1, + buffer_checkpoint: bool = False, + data_plane_checkpoint: bool = True, ) -> 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. """ - sampler_cfg = WindowedSamplerConfig(max_staleness_versions=1) + 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. @@ -367,10 +512,15 @@ def _actor_master_config( "keep_top_k": None, "save_period": save_period, "save_optimizer": save_optimizer, + "save_data_plane": data_plane_checkpoint, "checkpoint_must_save_by": checkpoint_must_save_by, "ft_save_period": ft_save_period, }, - data_plane={"enabled": True, "impl": "transfer_queue"}, + data_plane={ + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + }, async_rl=AsyncRLConfig( sampler=sampler_cfg, min_groups_for_streaming_train=1, @@ -386,7 +536,9 @@ def _make_actor_args( save_state: Optional[GRPOSaveState] = 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[DataPlaneCheckpointMetadata] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=object(), @@ -394,7 +546,7 @@ def _make_actor_args( env_handles={}, train_cluster=None, # type: ignore[arg-type] inference_cluster=None, # type: ignore[arg-type] - dp_client=_FakeDPClient(), + 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, @@ -406,9 +558,36 @@ def _make_actor_args( save_state if save_state is not None else _initial_grpo_save_state() ), last_checkpoint_path=last_checkpoint_path, + data_plane_checkpoint_metadata=data_plane_checkpoint_metadata, ) +def _data_plane_checkpoint_metadata( + *, + step: int = 0, + trainer_version: Optional[int] = None, + epoch: int = 0, + sampler_name: str = "in_order", + manifest_digest: str = "digest-1", + group_count: int = 0, +) -> DataPlaneCheckpointMetadata: + """Build the authoritative SC envelope used by actor-level restore tests.""" + 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": _PARTITION_ID, + "sampler_name": sampler_name, + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "replay_manifest_digest": manifest_digest, + "replay_group_count": group_count, + } + + def _run_train_pump( mc: MasterConfig, actor_args: SingleControllerActorArgs, @@ -427,7 +606,11 @@ def _run_train_pump( async def _main(): actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) - actor._sampler = _FakeSampler() + actor._sampler = _FakeSampler( + supports_buffer_checkpoint=sampler_supports_buffer_checkpoint( + mc.async_rl.sampler + ) + ) if seed is not None: seed(actor) # In-process runs have no Ray runtime; the pump only reads the GPU @@ -475,7 +658,10 @@ async def _main(): def _run_restore_then_train_pump( - mc: MasterConfig, actor_args: SingleControllerActorArgs + mc: MasterConfig, + actor_args: SingleControllerActorArgs, + *, + restored_groups: list[dict[str, Any]], ): """Restore the replay buffer, then drive a live _train_pump. @@ -487,7 +673,7 @@ def _run_restore_then_train_pump( async def _main(): actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) await actor._maybe_restore_replay_buffer() - actor._sampler = _FakeSampler() + actor._sampler = _RestoredGroupsSampler(restored_groups) with patch("ray.cluster_resources", return_value={"GPU": 0}): await asyncio.wait_for(actor._train_pump(), timeout=60.0) actor._checkpointer.shutdown() @@ -530,11 +716,41 @@ def test_restore_from_step_n(self, tmp_path): 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._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 = _initial_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), + SetupTimingMetrics(), + ) + + assert actor._train_steps == 7 + assert actor._trainer_version == 11 + assert actor._sampler.dispatch_index == 10 + + def test_restores_exact_sampler_dispatch_index(self, tmp_path): + save_state = _initial_grpo_save_state() + save_state.current_step = 7 + save_state.trainer_version = 11 + save_state.sampler_dispatch_index = 13 + + actor = _ACTOR_CLS( + _actor_master_config(tmp_path), + _make_actor_args(save_state=save_state), + SetupTimingMetrics(), + ) + + assert actor._trainer_version == 11 + assert actor._sampler.dispatch_index == 13 + def test_fresh_start_defaults(self, tmp_path): actor = _ACTOR_CLS( _actor_master_config(tmp_path), _make_actor_args(), SetupTimingMetrics() @@ -542,7 +758,7 @@ def test_fresh_start_defaults(self, tmp_path): assert actor._train_steps == 0 assert actor._trainer_version == 0 - assert actor._sampler._dispatch_index == -1 + assert actor._sampler.dispatch_index == -1 assert actor._consumed_samples == 0 assert actor._current_epoch == 0 assert actor._total_valid_tokens == 0 @@ -566,7 +782,7 @@ def test_old_checkpoint_without_total_valid_tokens(self, tmp_path): ) assert actor._train_steps == 5 - assert actor._sampler._dispatch_index == 4 + assert actor._sampler.dispatch_index == 4 assert actor._total_valid_tokens == 0 def test_resumed_pump_continues_to_max_steps(self, tmp_path): @@ -607,6 +823,8 @@ def test_saves_on_period_boundary_and_last_step(self, tmp_path): info_2 = _training_info(ckpt_dir, 2) assert info_2["current_step"] == 2 + assert info_2["trainer_version"] == 2 + assert info_2["sampler_dispatch_index"] == -1 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. @@ -736,10 +954,294 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} +class TestDataPlaneCheckpoint: + def test_metadata_uses_pre_await_save_state_snapshot(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + data_plane_checkpoint=True, + ) + save_state = _initial_grpo_save_state() + save_state.current_step = 3 + save_state.trainer_version = 7 + save_state.current_epoch = 2 + dp_client = _FakeDPClient() + + async def _main() -> None: + actor = _ACTOR_CLS( + mc, + _make_actor_args(save_state=save_state, dp_client=dp_client), + SetupTimingMetrics(), + ) + # Simulate live fields diverging after _save_checkpoint captured + # save_state. The rollout pump can advance _current_epoch while + # checkpoint I/O awaits; both fields must come from one snapshot. + actor._trainer_version = 11 + actor._current_epoch = 5 + await actor._save_data_plane_checkpoint(str(tmp_path / "tmp_step_3")) + actor._checkpointer.shutdown() + + asyncio.run(_main()) + + metadata = dp_client.save_calls[0]["metadata"] + assert metadata["single_controller_train_steps"] == 3 + assert metadata["single_controller_trainer_version"] == 7 + assert metadata["single_controller_epoch"] == 2 + + 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" + ) + expected_metadata = _data_plane_checkpoint_metadata( + step=1, + trainer_version=1, + sampler_name="windowed", + group_count=1, + ) + assert { + key: save_call["metadata"][key] for key in expected_metadata + } == expected_metadata + assert save_call["metadata"]["rollout_recovery_schema_version"] == ( + ROLLOUT_RECOVERY_SCHEMA_VERSION + ) + assert save_call["metadata"]["rollout_recovery_group_count"] == 0 + 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() + recovery_path = step_dir / ROLLOUT_RECOVERY_STATE_FILENAME + assert recovery_path.is_file() + recovery_state = torch.load(recovery_path, weights_only=False) + assert recovery_state["batch_shortfall"] == {} + assert recovery_state["sampler_stamps_target_steps"] is False + assert ( + hashlib.sha256(recovery_path.read_bytes()).hexdigest() + == (save_call["metadata"]["rollout_recovery_payload_sha256"]) + ) + 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_writes_authoritative_tq_checkpoint(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"] == "authoritative" + step_dir = tmp_path / "checkpoints" / "step_1" + assert (step_dir / REPLAY_BUFFER_METADATA_FILENAME).exists() + assert not (step_dir / "replay_buffer.pt").exists() + assert buffer.metadata_state_dict_calls == [4] + + 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), SetupTimingMetrics() + ) + actor._train_steps = 1 + actor._trainer_version = 1 + save_task = asyncio.create_task( + actor._save_checkpoint({"loss": 1.0}, is_policy_training_step=True) + ) + 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), SetupTimingMetrics() + ) + 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_missing_sidecar_before_finalization_falls_back_to_previous_step( + self, tmp_path + ): + """A failed sidecar write leaves tmp_step_N invisible to resume lookup.""" + + mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=1) + checkpoint_dir = tmp_path / "checkpoints" + previous = checkpoint_dir / "step_1" + previous.mkdir(parents=True) + incomplete = checkpoint_dir / "tmp_step_2" + (incomplete / "data_plane").mkdir(parents=True) + # Model the cut after native TQ save but before rollout_recovery.pt is + # written and begin_finalization renames the bundle. + assert not (incomplete / ROLLOUT_RECOVERY_STATE_FILENAME).exists() + + checkpointer = CheckpointManager(mc.checkpointing) + try: + assert checkpointer.get_latest_checkpoint_path() == str(previous) + finally: + checkpointer.shutdown() + 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() @@ -832,8 +1334,10 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): checkpoint_path = tmp_path / "tmp_step_1" checkpoint_path.mkdir(parents=True, exist_ok=True) + actor._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() actor._save_state = SimpleNamespace() actor._train_steps = 1 + actor._trainer_version = 1 actor._current_epoch = 0 actor._consumed_samples = 0 actor._total_valid_tokens = 0 @@ -842,7 +1346,11 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): sampler=SimpleNamespace(name="in_order"), max_buffered_rollouts=4, ) - actor._master_config = SimpleNamespace(checkpointing={"metric_name": None}) + actor._sampler = _FakeSampler() + actor._master_config = SimpleNamespace( + checkpointing={"metric_name": None, "save_data_plane": False}, + data_plane={}, + ) actor._dataloader = SimpleNamespace(state_dict=lambda: {}) actor._buffer = SimpleNamespace(state_dict=AsyncMock(return_value={})) actor._checkpointer = MagicMock() @@ -1017,7 +1525,11 @@ def _setup_master_config(checkpoint_dir: str) -> MasterConfig: checkpointing block setup now reads. """ return MasterConfig.model_construct( - data_plane={"enabled": True, "impl": "transfer_queue"}, + data_plane={ + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + }, data={ "use_multiple_dataloader": False, "shuffle": False, @@ -1059,6 +1571,7 @@ def _setup_master_config(checkpoint_dir: str) -> MasterConfig: "keep_top_k": None, "save_period": 2, "save_optimizer": True, + "save_data_plane": True, "checkpoint_must_save_by": None, }, ) @@ -1259,44 +1772,107 @@ def test_setup_missing_dataloader_state_raises( # ── replay buffer persistence ──────────────────────────────────────────────── -def _matching_save_state() -> dict[str, Any]: - """save_state whose sampler_name matches _actor_master_config's sampler.""" +def _matching_save_state() -> GRPOSaveState: + """Return save state matching the default in-order actor config.""" save_state = _initial_grpo_save_state() - save_state.sampler_name = "windowed" + save_state.sampler_name = "in_order" return save_state class TestReplayBufferPersistence: - def test_save_writes_replay_buffer(self, tmp_path): - mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2) - envelope = {"groups": [], "sentinel": "abc"} - buffer = _FakeTQBuffer(state=envelope) + 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, + data_plane_checkpoint=False, + ) + + with pytest.raises( + ValueError, + match="replay-checkpoint-capable sampler requires", + ): + _ACTOR_CLS(mc, _make_actor_args(), SetupTimingMetrics()) + + def test_gated_sampler_writes_native_replay_metadata(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=2, + save_period=2, + buffer_checkpoint=False, + data_plane_checkpoint=True, + ) + buffer = _FakeTQBuffer() _run_train_pump(mc, _make_actor_args(tq_buffer=buffer)) ckpt_dir = tmp_path / "checkpoints" - buffer_path = ckpt_dir / "step_2" / "replay_buffer.pt" - assert buffer_path.exists() - assert torch.load(buffer_path, weights_only=False) == envelope - # state_dict is stamped with the capacity at save time; the sampler - # identity lands in training_info.json for the restore-side check. - assert buffer.state_dict_calls == [4] - assert _training_info(ckpt_dir, 2)["sampler_name"] == "windowed" - - def test_run_restores_replay_buffer_and_permits(self, tmp_path): + assert (ckpt_dir / "step_2" / "training_info.json").exists() + assert not (ckpt_dir / "step_2" / "replay_buffer.pt").exists() + assert (ckpt_dir / "step_2" / REPLAY_BUFFER_METADATA_FILENAME).exists() + assert buffer.metadata_state_dict_calls == [4] + + def test_run_rejects_legacy_replay_file(self, tmp_path): ckpt_dir = tmp_path / "resume_ckpt" ckpt_dir.mkdir() - envelope = {"groups": ["g0", "g1", "g2"]} - torch.save(envelope, ckpt_dir / "replay_buffer.pt") - mc = _actor_master_config(tmp_path, max_num_steps=0) - buffer = _FakeTQBuffer(load_return=3) + 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 = _data_plane_checkpoint_metadata(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), - save_state=_matching_save_state(), + data_plane_checkpoint_metadata=tq_metadata, ), ) @@ -1306,10 +1882,10 @@ def test_run_restores_replay_buffer_and_permits(self, tmp_path): "max_groups": 4, "expected_partition_id": _PARTITION_ID, "expected_group_size": 2, + "expected_manifest_digest": "digest-1", } ] - # Each restored group holds one _buffer_capacity permit. - assert actor._buffer_capacity._value == 4 - 3 + assert actor._buffer_capacity._value == 2 assert result["train_steps"] == 0 # run()'s finally must tear the synchronizer down exactly once. assert actor._weight_synchronizer.shutdown_count == 1 @@ -1323,17 +1899,43 @@ def test_restored_permits_are_released_by_a_live_pump(self, tmp_path): # acquisition shape. ckpt_dir = tmp_path / "resume_ckpt" ckpt_dir.mkdir() - torch.save({"groups": []}, ckpt_dir / "replay_buffer.pt") - mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2) + sample_ids = [f"g{i}-{j}" for i in range(4) for j in range(2)] + 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": None, + "group_id": f"g{i}", + } + for i in range(4) + ] + torch.save({"groups": groups}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + tq_metadata = _data_plane_checkpoint_metadata(group_count=4) + mc = _actor_master_config( + tmp_path, + max_num_steps=2, + save_period=2, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) buffer = _FakeTQBuffer(load_return=4) actor = _run_restore_then_train_pump( mc, _make_actor_args( tq_buffer=buffer, + dp_client=_FakeDPClient(sample_ids=sample_ids), last_checkpoint_path=str(ckpt_dir), - save_state=_matching_save_state(), + data_plane_checkpoint_metadata=tq_metadata, ), + restored_groups=groups, ) assert len(buffer.load_calls) == 1 @@ -1342,18 +1944,137 @@ def test_restored_permits_are_released_by_a_live_pump(self, tmp_path): # selected group (2 steps x 2 prompt groups), so all 4 came back. assert actor._buffer_capacity._value == 4 - def test_run_missing_replay_buffer_file_starts_empty(self, tmp_path, monkeypatch): - # Resuming from a checkpoint that predates replay-buffer persistence: - # no replay_buffer.pt. + @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() - mc = _actor_master_config(tmp_path, max_num_steps=0) + 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 = _data_plane_checkpoint_metadata(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_native_replay_metadata_rejects_group_count_mismatch(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, + ) buffer = _FakeTQBuffer() - printed: list[str] = [] - monkeypatch.setattr( - "builtins.print", - lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + + with pytest.raises(ValueError, match="group count does not match"): + _run_actor_run( + mc, + _make_actor_args( + tq_buffer=buffer, + last_checkpoint_path=str(ckpt_dir), + data_plane_checkpoint_metadata=_data_plane_checkpoint_metadata( + group_count=2 + ), + ), + ) + + assert buffer.load_calls == [] + + def test_native_replay_metadata_rejects_missing_manifest_digest(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, ) + buffer = _FakeTQBuffer() + tq_metadata = _data_plane_checkpoint_metadata() + del tq_metadata["replay_manifest_digest"] + + with pytest.raises(ValueError, match="missing a replay manifest digest"): + _run_actor_run( + mc, + _make_actor_args( + tq_buffer=buffer, + last_checkpoint_path=str(ckpt_dir), + data_plane_checkpoint_metadata=tq_metadata, + ), + ) + + assert buffer.load_calls == [] + + def test_run_missing_native_replay_metadata_starts_empty(self, tmp_path): + 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() actor, _ = _run_actor_run( mc, @@ -1362,37 +2083,39 @@ def test_run_missing_replay_buffer_file_starts_empty(self, tmp_path, monkeypatch assert buffer.load_calls == [] assert actor._buffer_capacity._value == 4 # zero permits consumed - assert any("No replay buffer checkpoint found" in line for line in printed) - def test_run_no_restore_on_sampler_mismatch(self, tmp_path, monkeypatch): - # File present but the checkpoint's training_info records a different - # sampler: warn and skip — the saved stamps may never be selectable - # under the current policy. + def test_run_restores_native_replay_state_with_in_order_sampler(self, tmp_path): ckpt_dir = tmp_path / "resume_ckpt" ckpt_dir.mkdir() - torch.save({"groups": []}, ckpt_dir / "replay_buffer.pt") - mc = _actor_master_config(tmp_path, max_num_steps=0) - save_state = _initial_grpo_save_state() - save_state.sampler_name = "in_order" # current run uses windowed - buffer = _FakeTQBuffer(load_return=2) - printed: list[str] = [] - monkeypatch.setattr( - "builtins.print", - lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + envelope = {"groups": []} + torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=False, + data_plane_checkpoint=True, ) + buffer = _FakeTQBuffer() - actor, _ = _run_actor_run( + _run_actor_run( mc, _make_actor_args( tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir), - save_state=save_state, + save_state=_matching_save_state(), + data_plane_checkpoint_metadata=_data_plane_checkpoint_metadata(), ), ) - assert buffer.load_calls == [] - assert actor._buffer_capacity._value == 4 - assert any("skipping the buffer restore" in line for line in printed) + assert buffer.load_calls == [ + { + "state": envelope, + "max_groups": 4, + "expected_partition_id": _PARTITION_ID, + "expected_group_size": 2, + "expected_manifest_digest": "digest-1", + } + ] # ── replacement reserve persistence ────────────────────────────────────────── @@ -1453,25 +2176,19 @@ def test_run_restores_the_pooled_spares(self, tmp_path): assert list(actor._replacement_reserve) == ["spare0", "spare1"] - def test_run_restores_the_pool_even_when_the_buffer_restore_is_skipped( - self, tmp_path - ): - """The sampler guard that protects the buffer must not cover the pool. + def test_run_restores_the_pool_when_replay_metadata_is_absent(self, tmp_path): + """An empty replay restore must not suppress the independent spare pool. Spares never reached `admit`, so they carry no target-step stamp and nothing - about them depends on which sampler wrote the checkpoint. Skipping them here - would strand the batch permanently for the one case -- a sampler change on - resume -- where the operator is least likely to look for it. + about them depends on replay metadata. Skipping them here would strand the + batch permanently even though starting with an empty replay index is valid. """ ckpt_dir = tmp_path / "resume_ckpt" ckpt_dir.mkdir() - torch.save({"groups": []}, ckpt_dir / "replay_buffer.pt") # One spare is less than num_prompts_per_step, so the full run() path here # holds it back rather than draining it, and the pool is still observable. torch.save(["spare0"], ckpt_dir / "replacement_reserve.pt") mc = _actor_master_config(tmp_path, max_num_steps=0) - save_state = _initial_grpo_save_state() - save_state.sampler_name = "in_order" # current run uses windowed buffer = _FakeTQBuffer(load_return=2) actor, _ = _run_actor_run( @@ -1479,7 +2196,7 @@ def test_run_restores_the_pool_even_when_the_buffer_restore_is_skipped( _make_actor_args( tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir), - save_state=save_state, + save_state=_matching_save_state(), ), ) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index 6c914abab0e..bcc3d791f38 100644 --- a/tests/unit/single_controller/test_ppo_setup.py +++ b/tests/unit/single_controller/test_ppo_setup.py @@ -103,7 +103,11 @@ def _make_master_config( block is active and the other one stays None. """ return MasterConfig.model_construct( - data_plane={"enabled": True, "impl": "transfer_queue"}, + data_plane={ + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + }, data={ "use_multiple_dataloader": False, "shuffle": False, diff --git a/tests/unit/single_controller/test_refit_recovery.py b/tests/unit/single_controller/test_refit_recovery.py index bc3ca02e03a..703adf9390f 100644 --- a/tests/unit/single_controller/test_refit_recovery.py +++ b/tests/unit/single_controller/test_refit_recovery.py @@ -156,6 +156,7 @@ def _make_controller( # is about, but both have to exist for it to reach the refit. ctrl._master_config = SimpleNamespace(env={}) ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False return ctrl, monitor, sync diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index ba36d1cf2dc..a57b0b9f13a 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -18,6 +18,8 @@ import asyncio from collections import deque +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from types import SimpleNamespace from typing import Any @@ -25,7 +27,11 @@ import ray import torch -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + DataPlaneMutationCut, + TQReplayBuffer, +) from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSampler, WeightFifoSampler, @@ -43,6 +49,10 @@ from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.rollout_manager import RolloutManager, RolloutOutcome +from nemo_rl.experience.rollout_recovery import ( + RolloutRecoveryLedger, + RolloutRecoveryLedgerState, +) # Reuse fixtures from the experience tests; same shape as test_async_rollout_manager. from tests.unit.experience.test_rollout_manager import ( @@ -91,6 +101,23 @@ def _init_pump_ledgers(ctrl: Any) -> None: ctrl._batch_replacements = {} ctrl._batch_promotions = {} ctrl._replacement_reserve = deque() + ctrl._rollout_recovery_enabled = False + + +class _PausingMutationBarrier(DataPlaneCheckpointBarrier): + """Hold a mutation after its body so a concurrent checkpoint can be observed.""" + + def __init__(self) -> None: + super().__init__() + self.mutation_applied = asyncio.Event() + self.release_mutation = asyncio.Event() + + @asynccontextmanager + async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: + async with super().mutation() as cut: + yield cut + self.mutation_applied.set() + await self.release_mutation.wait() class _RecordingBuffer: @@ -771,17 +798,42 @@ async def _main() -> None: stale = asyncio.create_task(asyncio.Event().wait()) await asyncio.sleep(0) + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + for group_id, prompt_idx, start_weight_version in ( + ("fresh", 50, 5), + ("stale", 10, 1), + ): + ledger.reserve_group( + cut, + group_id=group_id, + prompt_id=str(prompt_idx), + prompt_payload={"idx": prompt_idx, "message_log": []}, + expected_generations=2, + target_step=None, + start_weight_version=start_weight_version, + admitted=True, + ) + controller_cls = SingleControllerActor.__ray_metadata__.modified_class ctrl = object.__new__(controller_cls) ctrl._sampler = WindowedSampler(None, max_staleness_versions=2) ctrl._trainer_version = 5 ctrl._inflight_by_group_id = {"fresh": (fresh, 5), "stale": (stale, 1)} + ctrl._rollout_recovery_enabled = True + ctrl._data_plane_checkpoint_barrier = barrier + ctrl._rollout_manager = SimpleNamespace( + recovery_ledger=ledger, + discard_prompt_group=ledger.discard_group, + ) aborted = await ctrl._abort_stale_inflight() assert aborted == 1 assert stale.cancelled() assert not fresh.cancelled() + assert [group.group_id for group in ledger.groups()] == ["fresh"] fresh.cancel() with pytest.raises(asyncio.CancelledError): @@ -790,6 +842,110 @@ async def _main() -> None: asyncio.run(_main()) +def test_abort_stale_inflight_rechecks_registry_after_checkpoint_wait() -> None: + """A group completed while checkpoint-blocked is not subsequently aborted.""" + + async def _main() -> None: + completed = asyncio.create_task(asyncio.Event().wait()) + await asyncio.sleep(0) + ledger = RolloutRecoveryLedger() + barrier = DataPlaneCheckpointBarrier() + async with barrier.mutation() as cut: + ledger.reserve_group( + cut, + group_id="completed", + prompt_id="10", + prompt_payload={"idx": 10, "message_log": []}, + expected_generations=2, + target_step=None, + start_weight_version=1, + admitted=True, + ) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._sampler = WindowedSampler(None, max_staleness_versions=2) + ctrl._trainer_version = 5 + ctrl._inflight_by_group_id = {"completed": (completed, 1)} + ctrl._rollout_recovery_enabled = True + ctrl._data_plane_checkpoint_barrier = barrier + ctrl._rollout_manager = SimpleNamespace( + recovery_ledger=ledger, + discard_prompt_group=ledger.discard_group, + ) + + async with ctrl._data_plane_checkpoint_barrier.checkpoint() as cut: + abort_task = asyncio.create_task(ctrl._abort_stale_inflight()) + await asyncio.sleep(0) + assert not abort_task.done() + ctrl._inflight_by_group_id.pop("completed") + ledger.discard_group(cut, "completed") + + assert await asyncio.wait_for(abort_task, timeout=1.0) == 0 + assert not completed.cancelled() + + completed.cancel() + with pytest.raises(asyncio.CancelledError): + await completed + + asyncio.run(_main()) + + +def test_checkpoint_observes_stale_abort_ledger_discard() -> None: + """A checkpoint waiting on stale abort cannot persist its discarded owner.""" + + async def _main() -> None: + stale = asyncio.create_task(asyncio.Event().wait()) + await asyncio.sleep(0) + ledger = RolloutRecoveryLedger() + async with DataPlaneCheckpointBarrier().mutation() as cut: + ledger.reserve_group( + cut, + group_id="stale", + prompt_id="10", + prompt_payload={"idx": 10, "message_log": []}, + expected_generations=2, + target_step=None, + start_weight_version=1, + admitted=True, + ) + barrier = _PausingMutationBarrier() + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + ctrl = object.__new__(controller_cls) + ctrl._sampler = WindowedSampler(None, max_staleness_versions=2) + ctrl._trainer_version = 5 + ctrl._inflight_by_group_id = {"stale": (stale, 1)} + ctrl._rollout_recovery_enabled = True + ctrl._data_plane_checkpoint_barrier = barrier + ctrl._rollout_manager = SimpleNamespace( + recovery_ledger=ledger, + discard_prompt_group=ledger.discard_group, + ) + + abort_task = asyncio.create_task(ctrl._abort_stale_inflight()) + await asyncio.wait_for(barrier.mutation_applied.wait(), timeout=1.0) + + checkpoint_entered = asyncio.Event() + + async def checkpoint_snapshot() -> RolloutRecoveryLedgerState: + async with barrier.checkpoint(): + checkpoint_entered.set() + return ledger.state_dict() + + checkpoint_task = asyncio.create_task(checkpoint_snapshot()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + + barrier.release_mutation.set() + checkpoint_state = await asyncio.wait_for(checkpoint_task, timeout=1.0) + assert await asyncio.wait_for(abort_task, timeout=1.0) == 1 + assert checkpoint_state["groups"] == [] + assert stale.cancelled() + + asyncio.run(_main()) + + def test_abort_stale_inflight_aggregates_cleanup_failures() -> None: async def _main() -> None: async def _boom() -> None: @@ -806,6 +962,7 @@ async def _boom() -> None: ctrl._sampler = WindowedSampler(None, max_staleness_versions=0) ctrl._trainer_version = 5 ctrl._inflight_by_group_id = {"g": (task, 0)} + ctrl._rollout_recovery_enabled = False with pytest.raises(BaseExceptionGroup) as exc_info: await ctrl._abort_stale_inflight() @@ -1143,9 +1300,11 @@ def test_rollout_pump_writes_expected_tq_data( ) for tag in tags: assert tag["weight_version"] == 0 - # Tag schema: weight_version plus per-row violation counts. + assert tag["prompt_idx"] == input_sample["idx"] + # Tag schema: recovery identity plus per-row violation counts. assert set(tag) == { "weight_version", + "prompt_idx", "num_invalid_tool_calls", "num_malformed_thinking", "num_assistant_messages", diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 47d588f7174..a62c567d9ad 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -27,19 +27,23 @@ import pytest from pydantic import TypeAdapter, ValidationError +from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneCheckpointBarrier from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, InOrderSampler, InOrderSamplerConfig, PromptGroupSampler, ReadyFirstSampler, ReadyFirstSamplerConfig, SamplerConfig, + TransactionalAdmissionSampler, WeightFifoSampler, WeightFifoSamplerConfig, WindowedSampler, WindowedSamplerConfig, create_sampler, required_buffer_capacity_for_config, + sampler_supports_buffer_checkpoint, ) from nemo_rl.data_plane import KVBatchMeta @@ -103,9 +107,37 @@ class TestBuiltinsImplementInterface: ) def test_isinstance_protocol(self, sampler): assert isinstance(sampler, PromptGroupSampler) + assert isinstance(sampler, TransactionalAdmissionSampler) class TestAdmission: + def test_wait_does_not_advance_gated_dispatch_cursor(self): + sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) + + _run(sampler.wait_until_admissible(trainer_version_fn=lambda: 0)) + + assert sampler.dispatch_index == -1 + + async def commit() -> int | None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + return sampler.commit_admission(cut) + + assert _run(commit()) == 0 + assert sampler.dispatch_index == 0 + + def test_expired_cut_cannot_advance_gated_dispatch_cursor(self): + sampler = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) + + async def commit_after_cut_expires() -> None: + async with DataPlaneCheckpointBarrier().mutation() as cut: + pass + + with pytest.raises(RuntimeError, match="no longer active"): + sampler.commit_admission(cut) + + _run(commit_after_cut_expires()) + assert sampler.dispatch_index == -1 + def test_windowed_never_gates_and_never_stamps(self): s = WindowedSampler(FakeBuffer(), max_staleness_versions=2) # trainer stuck at 0, but over-sampled admission returns immediately. @@ -181,6 +213,23 @@ def test_unready_slot_is_never_evicted(self): class TestFactory: + @pytest.mark.parametrize( + ("config", "expected"), + [ + (WindowedSamplerConfig(), True), + (ReadyFirstSamplerConfig(), True), + (WeightFifoSamplerConfig(), True), + (InOrderSamplerConfig(), True), + ( + CustomSamplerConfig(target=f"{__name__}:EchoSampler"), + False, + ), + ], + ) + def test_capability_comes_from_sampler_class(self, config, expected): + assert sampler_supports_buffer_checkpoint(config) 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) @@ -200,6 +249,26 @@ def test_weight_fifo_config_builds_weight_fifo(self): assert isinstance(s, WeightFifoSampler) assert s.max_staleness_versions == 4 + def test_factory_rejects_dynamic_capability_before_construction(self): + PropertyCapabilitySampler.constructed = False + with pytest.raises(TypeError, match="boolean class attribute"): + create_sampler( + FakeBuffer(), + CustomSamplerConfig( + target=f"{__name__}:PropertyCapabilitySampler", + ), + ) + assert not PropertyCapabilitySampler.constructed + + def test_custom_checkpoint_capability_is_discoverable_without_construction(self): + CheckpointingEchoSampler.constructed = False + assert sampler_supports_buffer_checkpoint( + CustomSamplerConfig( + target=f"{__name__}:CheckpointingEchoSampler", + ) + ) + assert not CheckpointingEchoSampler.constructed + def test_ready_first_config_builds_ready_first_sampler(self): s = create_sampler( FakeBuffer(), @@ -324,12 +393,17 @@ def test_discriminated_union_parses_the_warmup_window(self): class TestCustomFqnSampler: + def test_custom_target_must_be_a_class(self): + with pytest.raises(TypeError, match="not a class"): + create_sampler( + FakeBuffer(), + CustomSamplerConfig( + target=f"{__name__}:NOT_A_SAMPLER_CLASS", + ), + ) + 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( @@ -487,20 +561,19 @@ def test_windowed_evict_skips_unready_stale(self): class TestDispatchCursorRestore: - """Checkpoint resume calls set_dispatch_index(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).""" + """Checkpoint resume restores the exact last admitted dispatch batch.""" - def test_resumed_in_order_stamps_from_trainer_version(self): + def test_resumed_in_order_stamps_after_exact_cursor(self): s = InOrderSampler(FakeBuffer(), max_lookahead_versions=1) - s.set_dispatch_index(7) + s.restore_dispatch_index(6) + assert s.dispatch_index == 6 assert _run(s.admit(trainer_version_fn=lambda: 7)) == 7 assert _run(s.admit(trainer_version_fn=lambda: 8)) == 8 + assert s.dispatch_index == 8 def test_resumed_gate_admits_window_then_blocks(self): s = WeightFifoSampler(FakeBuffer(), max_staleness_versions=0) - s.set_dispatch_index(7) + s.restore_dispatch_index(6) # 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 @@ -512,13 +585,13 @@ def test_fresh_start_is_a_noop_seed(self): s.set_dispatch_index(0) assert _run(s.admit(trainer_version_fn=lambda: 0)) == 0 - def test_negative_resume_step_rejected(self): - with pytest.raises(ValueError, match="resume_from_step"): - WindowedSampler(FakeBuffer(), max_staleness_versions=1).set_dispatch_index( - -1 - ) + def test_dispatch_index_below_initial_value_rejected(self): + with pytest.raises(ValueError, match="dispatch_index"): + WindowedSampler( + FakeBuffer(), max_staleness_versions=1 + ).restore_dispatch_index(-2) - def test_custom_fqn_sampler_supports_seeding(self): + def test_custom_fqn_sampler_supports_exact_restore(self): from nemo_rl.algorithms.async_utils.staleness_sampler import ( CustomSamplerConfig, ) @@ -529,7 +602,7 @@ def test_custom_fqn_sampler_supports_seeding(self): target=f"{__name__}:EchoSampler", max_lookahead_versions=1 ), ) - s.set_dispatch_index(6) + s.restore_dispatch_index(5) assert _run(s.admit(trainer_version_fn=lambda: 6)) == 6 @@ -565,3 +638,30 @@ def test_gated_samplers_never_abort_inflight(self, sampler): class EchoSampler(InOrderSampler): """Stand-in for a user-defined sampler loaded by FQN.""" + + +class CheckpointingEchoSampler(EchoSampler): + """Custom sampler with a static replay-checkpoint capability.""" + + supports_buffer_checkpoint = True + constructed = False + + def __init__(self, *args, **kwargs) -> None: + type(self).constructed = True + super().__init__(*args, **kwargs) + + +class PropertyCapabilitySampler: + """Invalid custom sampler whose capability requires construction.""" + + constructed = False + + def __init__(self, *args, **kwargs) -> None: + type(self).constructed = True + + @property + def supports_buffer_checkpoint(self) -> bool: + return True + + +NOT_A_SAMPLER_CLASS = object() diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 110b07f421d..f1ad9957a15 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -19,18 +19,34 @@ import contextlib import threading from pathlib import Path +from typing import Any, Optional from unittest.mock import MagicMock, patch import pytest +import torch from omegaconf import OmegaConf import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod from nemo_rl.algorithms.advantage_estimator import AdvEstimatorConfig +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointMetadata, +) from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, ReadyFirstSamplerConfig, SamplerConfig, + WindowedSampler, + WindowedSamplerConfig, +) +from nemo_rl.algorithms.grpo import ( + GRPOConfig, + GRPOSaveState, + _initial_grpo_save_state, ) -from nemo_rl.algorithms.grpo import GRPOConfig from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.opd import OnPolicyDistillationConfig from nemo_rl.algorithms.single_controller_utils import ( @@ -39,6 +55,7 @@ SingleControllerActorArgs, setup_single_controller, ) +from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.data_plane.schema import SC_ROLLOUT_SCHEMA_FIELDS from nemo_rl.experience.rollouts import EffortLevelsConfig from nemo_rl.models.generation.megatron.megatron_generation import MegatronGeneration @@ -48,6 +65,24 @@ _REAL_BUILD_GENERATION = sc_setup_mod._build_generation +class _CheckpointingCustomSampler(WindowedSampler): + """Custom sampler whose static capability must be validated during setup.""" + + supports_buffer_checkpoint = True + + def __init__(self, buffer: Any) -> None: + super().__init__(buffer, max_staleness_versions=1) + + +class _NonCheckpointingCustomSampler(WindowedSampler): + """Custom sampler that explicitly opts out of replay recovery.""" + + supports_buffer_checkpoint = False + + def __init__(self, buffer: Any) -> None: + super().__init__(buffer, max_staleness_versions=1) + + def _make_master_config( *, dp_enabled: bool = True, @@ -87,7 +122,11 @@ def _make_master_config( } policy_config["model_name"] = "test-model" return MasterConfig.model_construct( - data_plane={"enabled": dp_enabled, "impl": "transfer_queue"}, + data_plane={ + "enabled": dp_enabled, + "impl": "transfer_queue", + "backend": "simple", + }, data={ "use_multiple_dataloader": use_multiple_dataloader, "shuffle": False, @@ -129,6 +168,35 @@ def _make_master_config( ) +def _native_tq_metadata( + *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1 +) -> DataPlaneCheckpointMetadata: + 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, + } + + +def _save_state( + *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1 +) -> GRPOSaveState: + state = _initial_grpo_save_state() + state.current_step = step + state.current_epoch = epoch + state.trainer_version = trainer_version + return state + + @pytest.fixture def patched_factories(): """Patch every external factory setup calls. @@ -367,6 +435,78 @@ 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["backend"] = "mooncake_cpu" + mc.checkpointing["save_data_plane"] = True + with pytest.raises(NotImplementedError, match="backend='mooncake_cpu'"): + 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.checkpointing["save_data_plane"] = False + mc.async_rl.sampler = WindowedSamplerConfig(max_staleness_versions=1) + mc.data_plane["backend"] = "simple" + + with pytest.raises( + ValueError, + match=( + "replay-checkpoint-capable sampler requires " + "checkpointing.save_data_plane=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_checkpointing_error_explains_mooncake_incompatibility(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = False + mc.async_rl.sampler = WindowedSamplerConfig(max_staleness_versions=1) + mc.data_plane["backend"] = "mooncake_cpu" + + with pytest.raises( + ValueError, + match=( + "backend='mooncake_cpu'.*backend='simple'.*checkpointing.enabled=false" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_rejects_checkpointing_custom_sampler_without_native_tq(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = False + mc.async_rl.sampler = CustomSamplerConfig( + target=f"{__name__}:_CheckpointingCustomSampler" + ) + mc.data_plane["backend"] = "simple" + + with pytest.raises( + ValueError, + match=( + "replay-checkpoint-capable sampler requires " + "checkpointing.save_data_plane=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + + def test_warns_when_custom_sampler_cannot_recover_buffered_rollouts( + self, patched_factories + ): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.checkpointing["save_data_plane"] = False + mc.async_rl.sampler = CustomSamplerConfig( + target=f"{__name__}:_NonCheckpointingCustomSampler" + ) + mc.data_plane["backend"] = "simple" + + with pytest.warns( + UserWarning, match="cannot recover completed buffered rollouts" + ): + 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"): @@ -1240,3 +1380,147 @@ def test_megatron_fleet_health_rejected_with_clean_backend_error(self): match="not supported for the MegatronGeneration generation backend", ): sc_setup_mod._maybe_attach_fleet_health(generation, mc) + + +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() + torch.save({}, checkpoint_path / "train_dataloader.pt") + save_state = _save_state() + policy = patched_factories["fake_policy"] + 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 = vars(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_file_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 = _save_state() + + 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=_save_state(trainer_version=7), + 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=_save_state(), + 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, capsys + ): + 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=_save_state(), + partition_id="rollout_data", + sampler_name="in_order", + ) + + assert restored is None + policy.load_data_plane_checkpoint.assert_not_called() + output = capsys.readouterr().out + assert REPLAY_BUFFER_METADATA_FILENAME in output + assert "matching TQ checkpoint will not be loaded" in output + assert "dataloader cursor is still restored" in output + assert "buffered at checkpoint time will be discarded" in output + + def test_metadata_file_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=_save_state(), + 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=_save_state(), + partition_id="rollout_data", + sampler_name="in_order", + ) diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 926e0c92b55..d9c7d359b21 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -24,6 +24,7 @@ from tensordict import TensorDict import nemo_rl.algorithms.single_controller as single_controller +from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneCheckpointBarrier from nemo_rl.algorithms.async_utils.staleness_sampler import BaseSampler from nemo_rl.algorithms.grpo import GRPOConfig, _initial_grpo_save_state from nemo_rl.algorithms.loss import ClippedPGLossConfig @@ -47,6 +48,18 @@ class FakeWeightSynchronizer: pass +class _InitBuffer: + """Minimal non-optional TQ buffer contract for actor-init tests.""" + + def __init__(self) -> None: + self.checkpoint_barrier: DataPlaneCheckpointBarrier | None = None + + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + self.checkpoint_barrier = barrier + + def _checkpointing_config(tmp_path) -> dict: """Minimal checkpointing block for actors built through __init__.""" return { @@ -85,6 +98,7 @@ def _grpo_master_config(tmp_path) -> MasterConfig: def _actor_args_for_init(**overrides) -> SimpleNamespace: """Minimal actor args for a controller built through the real __init__.""" + tq_buffer = _InitBuffer() args = dict( partition_id="rollout_data", dp_client=None, @@ -94,8 +108,8 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: weight_synchronizer=FakeWeightSynchronizer(), advantage_estimator=None, loss_fn=None, - tq_buffer=None, - rollout_manager=SimpleNamespace(_tq_buffer=None), + tq_buffer=tq_buffer, + rollout_manager=SimpleNamespace(_tq_buffer=tq_buffer), env_handles={}, fleet_monitor=None, generation_router=None, @@ -103,6 +117,7 @@ def _actor_args_for_init(**overrides) -> SimpleNamespace: inference_cluster=None, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, + data_plane_checkpoint_metadata=None, ) args.update(overrides) return SimpleNamespace(**args) @@ -132,6 +147,7 @@ def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: logger={}, env={}, ) + tq_buffer = _InitBuffer() actor_args = SimpleNamespace( partition_id="rollout_data", dp_client=None, @@ -141,8 +157,8 @@ def test_rejects_multiple_optimizer_steps_per_rl_step(monkeypatch) -> None: weight_synchronizer=None, advantage_estimator=None, loss_fn=None, - tq_buffer=None, - rollout_manager=SimpleNamespace(_tq_buffer=None), + tq_buffer=tq_buffer, + rollout_manager=SimpleNamespace(_tq_buffer=tq_buffer), env_handles={}, fleet_monitor=None, generation_router=None, @@ -191,25 +207,7 @@ def test_logs_hyperparameters_and_concrete_weight_synchronizer( # __init__ builds a CheckpointManager + TimeoutChecker from this block. checkpointing=_checkpointing_config(tmp_path), ) - actor_args = SimpleNamespace( - partition_id="rollout_data", - dp_client=None, - gen_handle=None, - trainer_handle=None, - dataloader=None, - weight_synchronizer=FakeWeightSynchronizer(), - advantage_estimator=None, - loss_fn=None, - tq_buffer=None, - rollout_manager=SimpleNamespace(_tq_buffer=None), - env_handles={}, - fleet_monitor=None, - generation_router=None, - train_cluster=None, - inference_cluster=None, - save_state=_initial_grpo_save_state(), - last_checkpoint_path=None, - ) + actor_args = _actor_args_for_init() controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller_cls( @@ -320,26 +318,7 @@ def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: setup_metrics = SetupTimingMetrics( generation_init_time_s=1.5, policy_init_time_s=2.5 ) - actor_args = SimpleNamespace( - partition_id="rollout_data", - dp_client=None, - gen_handle=None, - trainer_handle=None, - dataloader=None, - weight_synchronizer=FakeWeightSynchronizer(), - advantage_estimator=None, - loss_fn=None, - tq_buffer=None, - rollout_manager=SimpleNamespace(_tq_buffer=None), - train_cluster=None, - inference_cluster=None, - # A real field of SingleControllerActorArgs. Read directly rather than via a - # getattr default, so omitting it breaks here instead of silently degrading - # watchdog.gym_subprocess_check into a no-op at runtime. - env_handles={}, - save_state=_initial_grpo_save_state(), - last_checkpoint_path=None, - ) + actor_args = _actor_args_for_init() controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller_cls( @@ -476,6 +455,7 @@ def test_sync_weights_honors_recompute_kv_cache_config( requires_kv_scale_sync=False, ) ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native # abort path (empty registry -> no-op) instead of the gym gate. ctrl._master_config = SimpleNamespace(env={}) @@ -505,6 +485,7 @@ def test_sync_weights_calibrates_and_forwards_fp8_kv_scales() -> None: calibrate_qkv_fp8_scales=MagicMock(return_value={"layers": {"layer.0": 0.5}}) ) ctrl._inflight_by_group_id = {} + ctrl._rollout_recovery_enabled = False # env={} -> should_use_nemo_gym is False, so _sync_weights takes the native # abort path (empty registry -> no-op) instead of the gym gate. ctrl._master_config = SimpleNamespace(env={}) @@ -1112,6 +1093,7 @@ def _train_pump_controller(*, sampler) -> object: ctrl._timer = Timer() ctrl._trainer_version = 0 ctrl._train_steps = 0 + ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() ctrl._batch_shortfall = {} ctrl._batch_replacements = {} ctrl._batch_promotions = {} diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 80081cbdd8c..edfa15a0b48 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -17,17 +17,20 @@ from __future__ import annotations import asyncio -import io +import threading from typing import Any import pytest import torch -from tensordict import TensorDict import nemo_rl.algorithms.async_utils.replay_buffer as _replay_buffer_module from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, + DataPlaneCheckpointBarrier, PostWriteEnrichmentError, TQReplayBuffer, + replay_manifest_digest, ) from nemo_rl.data_plane import KVBatchMeta from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -68,6 +71,7 @@ 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( @@ -99,11 +103,16 @@ 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], @@ -119,7 +128,7 @@ def get_samples( ), } ) - # Opaque per-group payload; load_state_dict must re-put it verbatim. + # Opaque payload used by tests that inspect direct DataPlane reads. return {"payload_for": list(sample_ids)} def depth(self) -> int: @@ -140,14 +149,22 @@ def put_samples( raise RuntimeError("injected put failure") +class FailAfterPutAndClearDataPlaneClient(FailAfterPutDataPlaneClient): + """Fail both the canonical write and its deterministic-ID rollback.""" + + def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None: + del sample_ids, partition_id + raise OSError("injected rollback failure") + + def _run(coro): return asyncio.run(coro) -def _make_record() -> PromptGroupRecord: +def _make_record(*, prompt_idx: int = 0) -> PromptGroupRecord: """Opaque PromptGroupRecord — converter is stubbed, so contents are unused.""" return PromptGroupRecord( - prompt_idx=0, + prompt_idx=prompt_idx, prompt=[], extra_env_info=None, metadata={}, @@ -160,13 +177,18 @@ def _make_buffer( dp: FakeDataPlaneClient, *, require_routed_experts: bool = False, + checkpoint_barrier: DataPlaneCheckpointBarrier | None = None, ) -> TQReplayBuffer: - return TQReplayBuffer( + buffer = TQReplayBuffer( dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0}, require_routed_experts=require_routed_experts, ) + buffer.set_data_plane_checkpoint_barrier( + checkpoint_barrier or DataPlaneCheckpointBarrier() + ) + return buffer def _add_group( @@ -188,7 +210,132 @@ def _add_group( ) +class TestDataPlaneCheckpointBarrier: + def test_mutation_and_checkpoint_cuts_expire_on_context_exit(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + + async with barrier.mutation() as mutation_cut: + mutation_cut.require_live() + with pytest.raises(RuntimeError, match="no longer active"): + mutation_cut.require_live() + + async with barrier.checkpoint() as checkpoint_cut: + checkpoint_cut.require_live() + with pytest.raises(RuntimeError, match="no longer active"): + checkpoint_cut.require_live() + + asyncio.run(exercise()) + + 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()) + + def test_two_checkpoints_serialize_without_deadlock(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + release = asyncio.Event() + entered: list[str] = [] + + async def checkpoint(tag: str) -> None: + async with barrier.checkpoint(): + entered.append(f"{tag}-enter") + await release.wait() + entered.append(f"{tag}-exit") + + first = asyncio.create_task(checkpoint("first")) + await asyncio.sleep(0) + second = asyncio.create_task(checkpoint("second")) + await asyncio.sleep(0) + assert entered == ["first-enter"] + + release.set() + await asyncio.wait_for(asyncio.gather(first, second), timeout=5.0) + assert entered == [ + "first-enter", + "first-exit", + "second-enter", + "second-exit", + ] + + 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_enriches_after_put_before_slot_becomes_ready(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) @@ -254,6 +401,29 @@ def test_commit_clears_rows_when_put_raises_after_writing(self): assert buf.ready_list == [False] assert buf.meta_list == [None] + def test_commit_reports_both_write_and_rollback_failures(self): + dp = FailAfterPutAndClearDataPlaneClient() + buf = _make_buffer(dp) + group_id = buf.reserve(weight_version=3) + + with pytest.raises(BaseExceptionGroup) as exc_info: + _run( + buf.commit( + group_id, + _make_record(), + start_weight_version=3, + end_weight_version=3, + ) + ) + + assert exc_info.value.subgroup(RuntimeError) is not None + assert exc_info.value.subgroup(OSError) is not None + # The failed rollback leaves uncertain external rows and an unready local + # slot. Both failures must remain visible so callers abort instead of retrying + # the same stable group ID over potentially orphaned data. + assert dp.depth() == _N_GENS + assert buf.ready_list == [False] + def test_reserve_appends_placeholder_unready(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) @@ -283,7 +453,7 @@ def test_commit_writes_tq_then_fills_meta(self, monkeypatch): meta = _run( buf.commit( group_id, - _make_record(), + _make_record(prompt_idx=418), start_weight_version=3, end_weight_version=4, ) @@ -300,8 +470,8 @@ def test_commit_writes_tq_then_fills_meta(self, monkeypatch): assert buf.end_weight_list == [4] assert buf.ready_list == [True] assert buf.meta_list[0].sample_ids == meta.sample_ids - # TQ tag uses start_weight_version (dispatch time). - assert meta.tags == [{"weight_version": 3}] * _N_GENS + # TQ tags preserve both dispatch-time weight and dataset identity. + assert meta.tags == [{"weight_version": 3, "prompt_idx": 418}] * _N_GENS assert len(dp.put_calls) == 1 assert len(trace_calls) == 1 assert trace_calls[0]["keys"] == meta.sample_ids @@ -392,6 +562,61 @@ def test_commit_appends_multiple_records_in_order(self): class TestTQReplayBufferRemove: + def test_remove_with_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.remove([0], remove_in_dp=True)) + + 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) + group_id = buf.reserve(weight_version=0) + await buf.commit( + group_id, + _make_record(), + start_weight_version=0, + end_weight_version=0, + ) + event_loop_thread_id = threading.get_ident() + await buf.remove([0], remove_in_dp=True) + 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) @@ -594,20 +819,23 @@ def _make_group_entry( "end_weight": weight, "target_step": target_step, "group_id": group_id, - "fields_data": {"payload_for": sids}, } -def _make_envelope( +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, - "groups": list(groups), + "manifest_digest": replay_manifest_digest(metadata_groups), + "groups": metadata_groups, } @@ -618,74 +846,82 @@ def _load( 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 TestReplayManifestDigest: + def test_rejects_non_json_metadata_with_field_path(self): + group = _make_group_entry("group-1", weight=1) + assert group["meta"].tags is not None + group["meta"].tags[0]["unsupported"] = torch.tensor(1) + + with pytest.raises( + TypeError, + match=r"groups\[0\]\.meta\.tags\[0\]\.unsupported", + ): + replay_manifest_digest([group]) + + def test_mapping_order_does_not_change_digest(self): + first = _make_group_entry("group-1", weight=1) + second = _make_group_entry("group-1", weight=1) + first["meta"].extra_info = {"a": 1, "b": [2, 3]} + second["meta"].extra_info = {"b": [2, 3], "a": 1} + + assert replay_manifest_digest([first]) == replay_manifest_digest([second]) + + class TestTQReplayBufferStateDict: - def test_state_dict_serializes_ready_and_skips_unready(self): + 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) # in-flight: must be excluded + buf.reserve(weight_version=3) - state = _run(buf.state_dict(saved_capacity=8)) + state = buf.metadata_state_dict(saved_capacity=8) - assert state["partition_id"] == "rollout_data" - assert state["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 [g["start_weight"] for g in state["groups"]] == [1, 2] - assert [g["end_weight"] for g in state["groups"]] == [1, 2] - assert [g["target_step"] for g in state["groups"]] == [None, None] - assert [g["group_id"] for g in state["groups"]] == [ - _group_id_of(metas[0]), - _group_id_of(metas[1]), - ] - # Payloads are fetched from the DataPlane rows of each group. - assert [c["sample_ids"] for c in dp.get_calls] == [ - list(metas[0].sample_ids), - list(metas[1].sample_ids), + 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 dp.get_calls[0]["select_fields"] == list(metas[0].fields) - assert state["groups"][0]["fields_data"] == { - "payload_for": list(metas[0].sample_ids) - } + assert state["manifest_digest"] == replay_manifest_digest(state["groups"]) + assert dp.get_calls == [] - def test_round_trip_restores_lists_and_rows(self): + 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 = _run(buf.state_dict(saved_capacity=8)) + state = buf.metadata_state_dict(saved_capacity=8) - dp2 = FakeDataPlaneClient() - buf2 = _make_buffer(dp2) - restored = _load(buf2, state) + restored_dp = FakeDataPlaneClient() + restored_buf = _make_buffer(restored_dp) + restored = _load( + restored_buf, + state, + expected_manifest_digest=state["manifest_digest"], + ) assert restored == 2 - assert buf2.size() == 2 - # Parallel lists rebuilt in order, all ready. - assert buf2.start_weight_list == [1, 2] - assert buf2.end_weight_list == [1, 2] - assert buf2.target_step_list == [None, None] - assert buf2.ready_list == [True, True] - assert buf2._group_ids == [_group_id_of(m) for m in metas] - assert [m.sample_ids for m in buf2.meta_list] == [ - list(metas[0].sample_ids), - list(metas[1].sample_ids), + 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 ] - # Rows re-put with identical sample_ids / fields payload / tags. - assert len(dp2.put_calls) == 2 - for put, meta in zip(dp2.put_calls, metas): - assert put["sample_ids"] == list(meta.sample_ids) - assert put["fields"] == {"payload_for": list(meta.sample_ids)} - assert put["tags"] == [dict(t) for t in meta.tags] + assert restored_dp.put_calls == [] def test_round_trip_preserves_end_weight_and_target_step(self): # start != end and a non-None target_step must survive the round-trip: @@ -695,7 +931,7 @@ def test_round_trip_preserves_end_weight_and_target_step(self): buf = _make_buffer(dp) _add_group(buf, weight=1, end_weight=2) _add_group(buf, weight=5, target_step=7) - state = _run(buf.state_dict(saved_capacity=8)) + state = buf.metadata_state_dict(saved_capacity=8) buf2 = _make_buffer(FakeDataPlaneClient()) assert _load(buf2, state) == 2 @@ -707,7 +943,7 @@ def test_round_trip_preserves_end_weight_and_target_step(self): def test_round_trip_empty_buffer(self): # Common resume shape: no group committed before the checkpoint. buf = _make_buffer(FakeDataPlaneClient()) - state = _run(buf.state_dict(saved_capacity=8)) + state = buf.metadata_state_dict(saved_capacity=8) assert state["groups"] == [] dp2 = FakeDataPlaneClient() @@ -725,7 +961,7 @@ def test_state_dict_skips_middle_unready(self): buf.reserve(weight_version=2) # in-flight, sandwiched third = _add_group(buf, weight=3) - state = _run(buf.state_dict(saved_capacity=8)) + state = buf.metadata_state_dict(saved_capacity=8) assert [g["start_weight"] for g in state["groups"]] == [1, 3] assert [g["group_id"] for g in state["groups"]] == [ @@ -737,36 +973,23 @@ def test_state_dict_skips_middle_unready(self): list(third.sample_ids), ] - def test_round_trip_tensordict_payload_through_torch_save(self): - # The production checkpoint file is torch.save(envelope) with - # TensorDict-valued fields_data; exercise that serialization for real - # (mixed dtypes + a non-contiguous view) instead of the opaque fake. - fields = TensorDict( - { - "input_ids": torch.arange(12, dtype=torch.long).reshape(2, 6), - "prev_logprobs": torch.randn(2, 12, dtype=torch.float32)[:, ::2], - "sample_mask": torch.ones(2, dtype=torch.long), - }, - batch_size=(2,), - ) - group = _make_group_entry("g0", weight=1) - group["fields_data"] = fields - state = _make_envelope([group]) - - buffer_bytes = io.BytesIO() - torch.save(state, buffer_bytes) - buffer_bytes.seek(0) - loaded_state = torch.load(buffer_bytes, weights_only=False) - + def test_state_dict_skips_older_long_tail_unready_group(self): + # Model a long-running rollout reserved on an older weight while newer + # rollouts finish. Completed-rollout recovery must checkpoint only the + # ready group; recovering the unfinished group belongs to partial- + # rollout checkpointing. dp = FakeDataPlaneClient() buf = _make_buffer(dp) - assert ( - _load(buf, loaded_state, expected_group_size=len(group["meta"].sample_ids)) - == 1 - ) - put_fields = dp.put_calls[0]["fields"] - for key in fields.keys(): - assert torch.equal(put_fields[key], fields[key]) + unfinished_group_id = buf.reserve(weight_version=1) + completed = _add_group(buf, weight=7) + + state = buf.metadata_state_dict(saved_capacity=8) + + assert [g["start_weight"] for g in state["groups"]] == [7] + assert [g["group_id"] for g in state["groups"]] == [_group_id_of(completed)] + assert unfinished_group_id not in { + group["group_id"] for group in state["groups"] + } class TestTQReplayBufferLoadPreflight: @@ -784,20 +1007,26 @@ def test_missing_envelope_keys(self): self._assert_rejected({"groups": []}, match="missing required keys") def test_partition_id_mismatch(self): - state = _make_envelope([], partition_id="other_partition") + 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) - del group["fields_data"] - self._assert_rejected(_make_envelope([group]), match="group missing keys") + 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_envelope([group]), match="misaligned") + self._assert_rejected(_make_metadata_envelope([group]), match="misaligned") def test_group_size_mismatch(self): - state = _make_envelope([_make_group_entry("g0", weight=1, n=2)]) + 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): @@ -805,75 +1034,38 @@ def test_duplicate_sample_ids_across_groups(self): g1 = _make_group_entry( "g1", weight=2, sample_ids=["g0_g0", "g1_g1"] ) # g0_g0 collides - self._assert_rejected(_make_envelope([g0, g1]), match="duplicate sample_id") - - -class TestTQReplayBufferLoadTruncation: - def test_capacity_change_truncates_to_freshest(self, monkeypatch): - state = _make_envelope( - [_make_group_entry(f"g{w}", weight=w) for w in (1, 2, 3)], - saved_capacity=8, - ) - dp = FakeDataPlaneClient() - buf = _make_buffer(dp) - printed: list[str] = [] - monkeypatch.setattr( - "builtins.print", - lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + self._assert_rejected( + _make_metadata_envelope([g0, g1]), match="duplicate sample_id" ) - restored = _load(buf, state, max_groups=2) + 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", + ) - assert restored == 2 - # The freshest max_groups groups survive, original order preserved. - assert buf.start_weight_list == [2, 3] - put_sample_ids = [sid for c in dp.put_calls for sid in c["sample_ids"]] - assert "g1_g0" not in put_sample_ids and "g1_g1" not in put_sample_ids - assert any("capacity changed" in line for line in printed) - - def test_over_capacity_target_stamped_groups_raise(self): - # start_weight and target_step are both monotonic for the InOrder - # family, so freshest-first would keep the far-future targets and drop - # the near-term ones. InOrderSampler.evict only drops - # target < current_train_weight, so those permits are never released: - # rollout blocks on capacity while the train pump waits for a target - # that can no longer be filled. Fail loudly instead of truncating. - state = _make_envelope( - [_make_group_entry(f"g{w}", weight=w, target_step=w) for w in (1, 2, 3)], - saved_capacity=8, + 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)] ) dp = FakeDataPlaneClient() buf = _make_buffer(dp) - with pytest.raises(ValueError, match="max_buffered_rollouts >= 3"): - _load(buf, state, max_groups=2) + with pytest.raises(ValueError) as exc_info: + _load( + buf, + state, + max_groups=2, + expected_manifest_digest=state["manifest_digest"], + ) - # Preflight semantics: nothing reached the DataPlane or the buffer. + message = str(exc_info.value) + assert "checkpoint=3, current=2" in message + assert "async_rl.max_buffered_rollouts >= 3" in message + assert "Deleting replay_buffer_metadata.pt" in message + assert "skips loading the matching TQ checkpoint" in message + assert "dataloader has already moved past them" in message assert dp.put_calls == [] assert buf.size() == 0 - - def test_over_capacity_mixed_stamps_raise(self): - # One target-stamped group is enough to make truncation unsafe. - state = _make_envelope( - [ - _make_group_entry("g1", weight=1), - _make_group_entry("g2", weight=2), - _make_group_entry("g3", weight=3, target_step=3), - ], - saved_capacity=8, - ) - buf = _make_buffer(FakeDataPlaneClient()) - - with pytest.raises(ValueError, match="target_step stamps"): - _load(buf, state, max_groups=2) - - def test_target_stamped_groups_within_capacity_load_fine(self): - # The guard is scoped to the over-capacity case only. - state = _make_envelope( - [_make_group_entry(f"g{w}", weight=w, target_step=w) for w in (1, 2)], - saved_capacity=8, - ) - buf = _make_buffer(FakeDataPlaneClient()) - - assert _load(buf, state, max_groups=2) == 2 - assert buf.target_step_list == [1, 2] 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 00000000000..406cc61df8c --- /dev/null +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -0,0 +1,244 @@ +# 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 pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest +from tensordict import TensorDict + +from nemo_rl.data_plane.interfaces import backend_config +from tools import verify_tq_data_plane_checkpoint as verifier + + +class _FakeDataPlaneClient: + def __init__( + self, + checkpoint_state: dict[str, Any], + *, + missing_samples: bool = False, + premature_consumption: bool = False, + ) -> None: + self._checkpoint_state = checkpoint_state + self._missing_samples = missing_samples + self._premature_consumption = premature_consumption + self._fields: TensorDict | None = None + self._consumed: set[str] = set() + + def register_partition(self, **kwargs: Any) -> None: + pass + + def put_samples(self, *, fields: TensorDict, **kwargs: Any) -> None: + self._fields = fields.clone() + + def claim_meta(self, *, batch_size: int, **kwargs: Any) -> MagicMock: + available = [ + sample_id + for sample_id in verifier.SAMPLE_IDS + if sample_id not in self._consumed + ] + sample_ids = available[:batch_size] + self._consumed.update(sample_ids) + return MagicMock(size=len(sample_ids), sample_ids=sample_ids) + + def save_checkpoint( + self, + checkpoint_dir: Path, + *, + metadata: dict[str, Any], + ) -> None: + del checkpoint_dir + self._checkpoint_state.update( + { + "fields": self._fields, + "consumed": set(self._consumed), + "metadata": dict(metadata), + } + ) + + def load_checkpoint(self, checkpoint_dir: Path) -> dict[str, Any]: + del checkpoint_dir + self._fields = self._checkpoint_state["fields"].clone() + self._consumed = set(self._checkpoint_state["consumed"]) + return dict(self._checkpoint_state["metadata"]) + + def get_samples(self, **kwargs: Any) -> TensorDict: + if self._missing_samples: + raise KeyError("injected missing sample") + assert self._fields is not None + return self._fields + + def check_consumption_status(self, *args: Any) -> bool: + if self._premature_consumption: + return True + return len(self._consumed) == len(verifier.SAMPLE_IDS) + + def close(self) -> None: + pass + + +def _checkpoint_state(*, schema_version: int | None = None) -> dict[str, Any]: + return { + "fields": verifier._expected_fields(), + "consumed": {verifier.SAMPLE_IDS[0]}, + "metadata": verifier._checkpoint_metadata( + [verifier.SAMPLE_IDS[0]], + schema_version=( + verifier.DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + if schema_version is None + else schema_version + ), + ), + } + + +def test_data_plane_config_uses_nested_simple_backend_config() -> None: + config = verifier._data_plane_config(num_storage_units=3) + + simple_config = backend_config(config) + + assert simple_config.num_storage_units == 3 + assert simple_config.storage_capacity == 1024 + + +def test_save_load_round_trip_exercises_payload_and_cursor_restore( + monkeypatch, tmp_path +) -> None: + checkpoint_state: dict[str, Any] = {} + clients = iter( + [ + _FakeDataPlaneClient(checkpoint_state), + _FakeDataPlaneClient(checkpoint_state), + ] + ) + monkeypatch.setattr( + verifier, + "build_data_plane_client", + lambda *args, **kwargs: next(clients), + ) + + verifier._save(tmp_path / "data_plane", num_storage_units=2) + verifier._load(tmp_path / "data_plane", num_storage_units=2) + + +@pytest.mark.parametrize( + ("client_kwargs", "error_type", "match"), + [ + ({"missing_samples": True}, KeyError, "injected missing sample"), + ( + {"premature_consumption": True}, + AssertionError, + "marked every row consumed", + ), + ], +) +def test_load_rejects_invalid_restored_state( + monkeypatch, + tmp_path, + client_kwargs, + error_type, + match, +) -> None: + client = _FakeDataPlaneClient(_checkpoint_state(), **client_kwargs) + monkeypatch.setattr( + verifier, + "build_data_plane_client", + lambda *args, **kwargs: client, + ) + + with pytest.raises(error_type, match=match): + verifier._load(tmp_path / "data_plane", num_storage_units=2) + + +def test_load_rejects_schema_mismatch(monkeypatch, tmp_path) -> None: + client = _FakeDataPlaneClient(_checkpoint_state(schema_version=-1)) + monkeypatch.setattr( + verifier, + "build_data_plane_client", + lambda *args, **kwargs: client, + ) + + with pytest.raises(AssertionError, match="Unexpected data-plane checkpoint schema"): + verifier._load(tmp_path / "data_plane", num_storage_units=2) + + +def test_run_child_uses_fresh_process(monkeypatch, tmp_path) -> None: + run = MagicMock() + monkeypatch.setattr(verifier.subprocess, "run", run) + + verifier._run_child("load", tmp_path / "step_1", num_storage_units=3) + + command = run.call_args.args[0] + assert command[0] == verifier.sys.executable + assert command[2:] == [ + "--phase", + "load", + "--checkpoint-dir", + str(tmp_path / "step_1"), + "--num-storage-units", + "3", + ] + assert run.call_args.kwargs == {"check": True} + + +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 00000000000..41ed44f0ddf --- /dev/null +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -0,0 +1,275 @@ +# 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 Any, cast + +import torch +from tensordict import TensorDict + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointMetadata, +) +from nemo_rl.data_plane import ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + 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 _checkpoint_metadata( + expected_consumed_ids: list[str], + *, + schema_version: int = DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, +) -> dict[str, Any]: + """Build the typed SC envelope plus smoke-test-only cursor metadata.""" + envelope: DataPlaneCheckpointMetadata = { + "data_plane_checkpoint_schema_version": schema_version, + "single_controller_train_steps": 0, + "single_controller_trainer_version": 0, + "single_controller_epoch": 0, + "partition_id": PARTITION_ID, + "sampler_name": "checkpoint_smoke", + "mode": "shadow", + } + return {**envelope, "expected_consumed_ids": expected_consumed_ids} + + +def _data_plane_config(num_storage_units: int) -> DataPlaneConfig: + return cast( + DataPlaneConfig, + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "claim_meta_poll_interval_s": 0.05, + "simple": { + "storage_capacity": 1024, + "num_storage_units": num_storage_units, + }, + }, + ) + + +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=_checkpoint_metadata(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) + checkpoint_metadata = cast(DataPlaneCheckpointMetadata, metadata) + + restored = dp_client.get_samples( + sample_ids=SAMPLE_IDS, + partition_id=PARTITION_ID, + select_fields=FIELDS, + ) + expected = _expected_fields() + for field in FIELDS: + restored_value = restored[field] + expected_value = expected[field] + assert isinstance(restored_value, torch.Tensor) + assert isinstance(expected_value, torch.Tensor) + if not torch.equal(restored_value, expected_value): + raise AssertionError(f"Restored field differs: {field}") + + if ( + checkpoint_metadata["data_plane_checkpoint_schema_version"] + != DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ): + 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()