From 1dec574eaae9d758c75859e610b3e2f8add6ccdb Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 2 Aug 2026 14:00:17 -0400 Subject: [PATCH 01/32] feat(sc): save native TQ state in checkpoints Signed-off-by: Anish Mahishi --- examples/configs/grpo_math_1B.yaml | 1 + ...po_math_1B_megatron_single_controller.yaml | 1 + .../algorithms/async_utils/replay_buffer.py | 29 ++- nemo_rl/algorithms/single_controller.py | 104 +++++++-- .../single_controller_utils/setup.py | 6 + nemo_rl/data_plane/adapters/noop.py | 56 +++++ nemo_rl/data_plane/adapters/transfer_queue.py | 36 +++ nemo_rl/data_plane/interfaces.py | 46 +++- nemo_rl/data_plane/observability.py | 23 ++ pyrefly.toml | 1 + .../test_architecture_invariants.py | 2 + .../data_plane/test_interface_contract.py | 62 ++++++ tests/unit/data_plane/test_observability.py | 36 +++ tests/unit/data_plane/test_tq_lifecycle.py | 73 ++++++ .../unit/reference_configs/grpo_math_1B.yaml | 1 + .../test_sc_checkpointing.py | 136 +++++++++++- .../test_single_controller_setup.py | 11 + .../test_tq_replay_buffer.py | 42 +++- tools/verify_tq_data_plane_checkpoint.py | 210 ++++++++++++++++++ 19 files changed, 852 insertions(+), 24 deletions(-) create mode 100644 tools/verify_tq_data_plane_checkpoint.py diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index d7cd6f528ef..93de9fe8772 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -520,6 +520,7 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + checkpointing_enabled: false # save TQ state inside algorithm checkpoints global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 6f1764718f8..2f94b10479a 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -58,6 +58,7 @@ logger: data_plane: enabled: true + checkpointing_enabled: true cluster: gpus_per_node: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 479038aa345..395e521ad3d 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -707,6 +707,13 @@ def __init__( self.target_step_list: list[Optional[int]] = [] self.ready_list: list[bool] = [] self._group_ids: list[str] = [] + self._data_plane_checkpoint_lock: Optional[asyncio.Lock] = None + + def set_data_plane_checkpoint_lock(self, lock: asyncio.Lock) -> None: + """Serialize replay-buffer-owned clears with SC checkpoints.""" + if self._data_plane_checkpoint_lock is not None: + raise RuntimeError("data-plane checkpoint lock is already configured") + self._data_plane_checkpoint_lock = lock def reserve( self, @@ -806,10 +813,8 @@ async def commit( # put_samples may have written rows before raising. Roll back by the # deterministic IDs known here; the caller removes the reserved slot. try: - await self._call_dp( - "clear_samples", + await self._clear_samples( sample_ids=list(sample_ids), - partition_id=self._partition_id, ) except BaseException as rollback_error: if isinstance(commit_error, asyncio.CancelledError): @@ -872,10 +877,8 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int: del self._group_ids[i] if remove_in_dp: - await self._call_dp( - "clear_samples", + await self._clear_samples( sample_ids=dropped_sample_ids, - partition_id=self._partition_id, ) return len(drop_idxs) @@ -1102,3 +1105,17 @@ async def _call_dp(self, method_name: str, **kwargs: Any) -> Any: if asyncio.iscoroutine(result): return await result return result + + async def _clear_samples(self, *, sample_ids: list[str]) -> None: + """Clear rows without overlapping a bound data-plane checkpoint.""" + if self._data_plane_checkpoint_lock is None: + raise RuntimeError( + "TQReplayBuffer must be bound to the controller data-plane " + "checkpoint lock before clearing samples" + ) + async with self._data_plane_checkpoint_lock: + await self._call_dp( + "clear_samples", + sample_ids=sample_ids, + partition_id=self._partition_id, + ) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 66d490a916c..891ad9f790c 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -73,6 +73,9 @@ Generation = Union[VllmGeneration, SGLangGeneration] +DATA_PLANE_CHECKPOINT_DIR = "data_plane" +DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 1 + @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: @@ -169,6 +172,16 @@ def __init__( ) # ── asyncio state ────────────────────────────────────────────────── + # TQ snapshots permit concurrent puts but not destructive clears. All + # clears currently owned by async SC use this lock, including rollback + # and eviction through TQReplayBuffer. A future staging/finalizer path + # must join the same barrier before native restore can be authoritative. + self._data_plane_checkpoint_lock: asyncio.Lock = asyncio.Lock() + if self._buffer is not None: + self._buffer.set_data_plane_checkpoint_lock( + self._data_plane_checkpoint_lock + ) + # Gate: cleared during _sync_weights, set when generation may proceed self._rollout_permitted: asyncio.Event = asyncio.Event() self._rollout_permitted.set() @@ -318,6 +331,61 @@ 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_lock: + await self._call_dp( + "clear_samples", + sample_ids=sample_ids, + partition_id=self._partition_id, + ) + + async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: + """Save a shadow TQ snapshot inside an SC checkpoint bundle.""" + checkpoint_dir = os.path.join( + checkpoint_path, + DATA_PLANE_CHECKPOINT_DIR, + ) + metadata = { + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), + "single_controller_train_steps": self._train_steps, + "single_controller_trainer_version": self._trainer_version, + "single_controller_epoch": self._current_epoch, + "partition_id": self._partition_id, + "mode": "shadow", + } + started = time.monotonic() + print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) + try: + method = getattr(self._dp_client, "save_checkpoint") + remote = getattr(method, "remote", None) + if remote is not None: + await self._ray_get( + remote(checkpoint_dir=checkpoint_dir, metadata=metadata) + ) + else: + result = await asyncio.to_thread( + method, + checkpoint_dir=checkpoint_dir, + metadata=metadata, + ) + if asyncio.iscoroutine(result): + await result + 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: @@ -590,11 +658,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 @@ -799,14 +863,28 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: dataloader_state, os.path.join(checkpoint_path, "train_dataloader.pt"), ) - 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"), - ) + buffer_state: Optional[dict[str, Any]] = None + if self._master_config.data_plane.get("checkpointing_enabled"): + # Capture the legacy replay payload and the native TQ snapshot + # under one clear barrier. Generation puts may continue, so TQ can + # contain a superset of the groups named by replay_buffer.pt. + async with self._data_plane_checkpoint_lock: + if self._sampler.supports_buffer_checkpoint: + buffer_state = await self._buffer.state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + await self._save_data_plane_checkpoint(checkpoint_path) + elif self._sampler.supports_buffer_checkpoint: + buffer_state = await self._buffer.state_dict( + saved_capacity=self._async_cfg.max_buffered_rollouts + ) + + if buffer_state is not None: + 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( diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index b2db77cb512..399ae0203bd 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -394,6 +394,12 @@ def setup_single_controller( "master_config.data_plane.enabled=True. The async-RL " "SingleController path is built on the TransferQueue data plane." ) + if dp_config.get("checkpointing_enabled") and dp_config["backend"] != "simple": + raise NotImplementedError( + "SingleController data-plane checkpointing currently requires " + "data_plane.backend='simple'; Mooncake storage cannot be restored " + "by TQ v0.1.9." + ) assert generation_config is not None, ( "single_controller_utils.setup requires policy.generation in master_config" diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py index 1c5b00a5e44..f01b40c0986 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 @@ -237,6 +240,59 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None for s in rec.consumed.values(): s.discard(sid) + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """Persist the trusted in-memory fixture for adapter contract tests. + + This test-only adapter uses pickle; callers must not load checkpoints + from untrusted paths. + """ + checkpoint_dir = Path(checkpoint_dir) + tmp_dir = checkpoint_dir.with_name(f"{checkpoint_dir.name}.tmp") + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + tmp_dir.mkdir(parents=True) + try: + with (tmp_dir / "noop_state.pkl").open("wb") as checkpoint_file: + pickle.dump( + { + "partitions": self._partitions, + "metadata": metadata or {}, + }, + checkpoint_file, + protocol=pickle.HIGHEST_PROTOCOL, + ) + if checkpoint_dir.exists(): + shutil.rmtree(checkpoint_dir) + tmp_dir.rename(checkpoint_dir) + except Exception: + if tmp_dir.exists(): + shutil.rmtree(tmp_dir) + raise + + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + """Restore the in-memory fixture into a clean client.""" + if self._partitions: + raise RuntimeError( + "load_checkpoint requires a clean data-plane client with no " + "registered partitions" + ) + checkpoint_file = Path(checkpoint_dir) / "noop_state.pkl" + if not checkpoint_file.is_file(): + raise FileNotFoundError(f"NoOp checkpoint not found: {checkpoint_file}") + with checkpoint_file.open("rb") as state_file: + state = pickle.load(state_file) + metadata = state.get("metadata", {}) + if not isinstance(metadata, dict): + raise ValueError("NoOp checkpoint metadata must be a dictionary") + self._partitions = state["partitions"] + self._closed = False + return dict(metadata) + def close(self) -> None: if self._closed: return diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 995cfa24c37..bf5ae6c31c7 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -23,12 +23,14 @@ from __future__ import annotations import ipaddress +import json import os import socket import subprocess import time import warnings from importlib import resources +from pathlib import Path from typing import Any, cast import torch @@ -469,6 +471,7 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: # is unaffected). Writer unsqueezes 1D → (N, 1) on put; reader # squeezes the trailing 1 back on get. Drop when upstream TQ # unifies the schema/data shapes for 1D fields. + self._backend = cfg["backend"] self._promote_1d = cfg["backend"] == "mooncake_cpu" if bootstrap: @@ -706,6 +709,39 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None # ── (C) lifecycle ────────────────────────────────────────────────── + def save_checkpoint( + self, + checkpoint_dir: str | Path, + *, + metadata: dict[str, Any] | None = None, + ) -> None: + """Save TQ controller metadata and storage data.""" + if self._backend == "mooncake_cpu": + raise NotImplementedError( + "TQ checkpointing is not supported for the mooncake_cpu " + "backend: MooncakeStorageManager cannot persist its in-memory " + "rows, so TQ would silently create a metadata-only checkpoint." + ) + _connect_existing() + tq.save_checkpoint(checkpoint_dir, metadata=metadata) + + def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: + """Restore TQ state after initialization and before data operations.""" + if self._backend == "mooncake_cpu": + raise NotImplementedError( + "TQ checkpoint restore is not supported for the mooncake_cpu " + "backend because its in-memory rows cannot be restored." + ) + _connect_existing() + tq.load_checkpoint(checkpoint_dir) + 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") + return dict(user_metadata) + def close(self) -> None: if self._closed: return diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 41a98f0c0ed..6f57bcade65 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -37,6 +37,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict from tensordict import TensorDict @@ -58,6 +59,11 @@ class DataPlaneConfig(TypedDict): ``backend == "mooncake_cpu"``; the simple backend ignores them. They are required (not NotRequired) so the YAML carries the full schema and there are no hidden Python defaults. + + ``checkpointing_enabled`` opts algorithms into saving native TQ state + inside their checkpoint bundle. It is optional during rollout because + existing configs predate data-plane checkpointing; exemplar configs carry + the recommended default explicitly. """ enabled: bool @@ -68,6 +74,7 @@ class DataPlaneConfig(TypedDict): claim_meta_poll_interval_s: float global_segment_size: int local_buffer_size: int + checkpointing_enabled: NotRequired[bool] controller_address: NotRequired[str] ack_timeout_ms: NotRequired[int] observability: NotRequired["ObservabilityConfig"] @@ -259,7 +266,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* — @@ -442,6 +450,42 @@ 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. + + 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..d7bcd88fca5 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 @@ -337,6 +338,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/pyrefly.toml b/pyrefly.toml index e0c35ebf542..f4c212c32f3 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -250,6 +250,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/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py index b0a9d95af99..5ff2d75fd28 100644 --- a/tests/unit/data_plane/test_architecture_invariants.py +++ b/tests/unit/data_plane/test_architecture_invariants.py @@ -80,6 +80,8 @@ def test_sync_trainer_rejects_message_level_advantage_penalties(): "get_samples", "clear_samples", "check_consumption_status", + "save_checkpoint", + "load_checkpoint", "close", ], ) diff --git a/tests/unit/data_plane/test_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py index 3426c3b5067..3a4009b164b 100644 --- a/tests/unit/data_plane/test_interface_contract.py +++ b/tests/unit/data_plane/test_interface_contract.py @@ -124,3 +124,65 @@ def test_kv_batch_put_rejects_non_tensor_leaves(client: DataPlaneClient): def test_close_is_idempotent(client: DataPlaneClient): client.close() client.close() + + +def test_checkpoint_round_trip_restores_data_and_consumption(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + source = NoOpDataPlaneClient() + source.register_partition( + partition_id="p", + fields=["x"], + num_samples=3, + consumer_tasks=["train"], + ) + source.put_samples( + sample_ids=["a", "b", "c"], + partition_id="p", + fields=TensorDict({"x": torch.tensor([10, 20, 30])}, batch_size=[3]), + ) + consumed = source.claim_meta( + partition_id="p", + task_name="train", + required_fields=["x"], + batch_size=1, + ) + source.save_checkpoint(checkpoint_dir, metadata={"step": 7}) + + restored = NoOpDataPlaneClient() + metadata = restored.load_checkpoint(checkpoint_dir) + assert metadata == {"step": 7} + data = restored.get_samples( + sample_ids=["a", "b", "c"], + partition_id="p", + select_fields=["x"], + ) + assert torch.equal(data["x"], torch.tensor([10, 20, 30])) + + remaining = restored.claim_meta( + partition_id="p", + task_name="train", + required_fields=["x"], + batch_size=3, + ) + assert consumed.sample_ids[0] not in remaining.sample_ids + assert set(consumed.sample_ids + remaining.sample_ids) == {"a", "b", "c"} + assert restored.check_consumption_status("p", ["train"]) + + source.close() + restored.close() + + +def test_checkpoint_load_requires_clean_client(tmp_path) -> None: + checkpoint_dir = tmp_path / "data-plane" + source = NoOpDataPlaneClient() + source.save_checkpoint(checkpoint_dir) + + source.register_partition( + partition_id="already-used", + fields=["x"], + num_samples=1, + consumer_tasks=["train"], + ) + with pytest.raises(RuntimeError, match="clean data-plane client"): + source.load_checkpoint(checkpoint_dir) + source.close() diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 0d471bc2660..2cdbdb1a626 100644 --- a/tests/unit/data_plane/test_observability.py +++ b/tests/unit/data_plane/test_observability.py @@ -131,6 +131,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_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index 3e79a50ff84..c6f1a98c067 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -22,6 +22,9 @@ from __future__ import annotations +import json +from unittest.mock import MagicMock + import numpy as np import pytest import torch @@ -82,6 +85,76 @@ def fake_clear(**kwargs): ] +def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect_calls = [] + save_calls = [] + load_calls = [] + monkeypatch.setattr( + tq_adapter, + "_connect_existing", + lambda: connect_calls.append(None), + ) + monkeypatch.setattr( + tq_adapter.tq, + "save_checkpoint", + lambda checkpoint_dir, *, metadata=None: save_calls.append( + (checkpoint_dir, metadata) + ), + ) + monkeypatch.setattr( + tq_adapter.tq, + "load_checkpoint", + lambda checkpoint_dir: load_calls.append(checkpoint_dir), + ) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + 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} + + +@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" + 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/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index e7f64db325c..9471acf49ee 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -512,6 +512,7 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence + checkpointing_enabled: false # save TQ state inside algorithm checkpoints global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index b9e0235590a..7829d7f5ef5 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -213,12 +213,49 @@ async def select(self, **kwargs) -> tuple[Optional[KVBatchMeta], int]: class _FakeDPClient: - def __init__(self) -> None: + def __init__(self, *, save_error: Optional[Exception] = None) -> None: self.clear_calls: list[tuple[list[str], str]] = [] + self.save_calls: list[dict[str, Any]] = [] + self.save_error = save_error def clear_samples(self, sample_ids: list[str], partition_id: str) -> None: self.clear_calls.append((list(sample_ids), partition_id)) + def save_checkpoint( + self, + checkpoint_dir: str, + *, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + self.save_calls.append( + { + "checkpoint_dir": checkpoint_dir, + "metadata": dict(metadata or {}), + } + ) + if self.save_error is not None: + raise self.save_error + os.makedirs(checkpoint_dir, exist_ok=True) + with open(os.path.join(checkpoint_dir, "metadata.json"), "w") as f: + json.dump({"user_metadata": metadata or {}}, f) + + +class _BlockingDPClient(_FakeDPClient): + def __init__(self) -> None: + super().__init__() + self.save_started = threading.Event() + self.release_save = threading.Event() + + def save_checkpoint( + self, + checkpoint_dir: str, + *, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + self.save_started.set() + assert self.release_save.wait(timeout=30.0), "test never released TQ save" + super().save_checkpoint(checkpoint_dir, metadata=metadata) + class _FakeWeightSynchronizer: def __init__(self) -> None: @@ -251,6 +288,10 @@ def __init__( self.load_return = load_return self.state_dict_calls: list[int] = [] self.load_calls: list[dict[str, Any]] = [] + self.checkpoint_lock: Optional[asyncio.Lock] = None + + def set_data_plane_checkpoint_lock(self, lock: asyncio.Lock) -> None: + self.checkpoint_lock = lock async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]: self.state_dict_calls.append(saved_capacity) @@ -313,6 +354,7 @@ def _actor_master_config( ft_save_period: Optional[int] = None, num_prompts_per_step: int = 2, max_num_epochs: int = 1, + data_plane_checkpoint: bool = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -355,7 +397,12 @@ def _actor_master_config( "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", + "checkpointing_enabled": data_plane_checkpoint, + }, async_rl=AsyncRLConfig( sampler=sampler_cfg, min_groups_for_streaming_train=1, @@ -371,6 +418,7 @@ 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, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( @@ -379,7 +427,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, @@ -699,6 +747,88 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} +class TestDataPlaneShadowCheckpoint: + def test_saves_tq_state_and_keeps_legacy_replay_payload(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + data_plane_checkpoint=True, + ) + dp_client = _FakeDPClient() + buffer = _FakeTQBuffer(state={"legacy_payload": "kept"}) + + _run_train_pump( + mc, + _make_actor_args(dp_client=dp_client, tq_buffer=buffer), + ) + + assert len(dp_client.save_calls) == 1 + save_call = dp_client.save_calls[0] + assert save_call["checkpoint_dir"] == str( + tmp_path / "checkpoints" / "tmp_step_1" / "data_plane" + ) + assert save_call["metadata"] == { + "data_plane_checkpoint_schema_version": 1, + "single_controller_train_steps": 1, + "single_controller_trainer_version": 1, + "single_controller_epoch": 0, + "partition_id": _PARTITION_ID, + "mode": "shadow", + } + step_dir = tmp_path / "checkpoints" / "step_1" + assert (step_dir / "data_plane" / "metadata.json").is_file() + assert torch.load(step_dir / "replay_buffer.pt", weights_only=False) == { + "legacy_payload": "kept" + } + assert buffer.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)) + actor._train_steps = 1 + actor._trainer_version = 1 + save_task = asyncio.create_task(actor._save_checkpoint({"loss": 1.0})) + started = await asyncio.to_thread(dp_client.save_started.wait, 30.0) + assert started + + clear_task = asyncio.create_task( + actor._clear_data_plane_samples(["sample-0"]) + ) + await asyncio.sleep(0) + assert dp_client.clear_calls == [] + + dp_client.release_save.set() + await save_task + await clear_task + actor._checkpointer.shutdown() + + asyncio.run(_main()) + assert dp_client.clear_calls == [(["sample-0"], _PARTITION_ID)] + + # ── async-save finalization ────────────────────────────────────────────────── diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index b66732dfb3a..83dc8d2cb63 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -224,6 +224,17 @@ def test_raises_when_data_plane_disabled(self): with pytest.raises(ValueError, match="data_plane.enabled=True"): setup_single_controller(mc, MagicMock()) + def test_rejects_mooncake_data_plane_checkpointing(self): + mc = _make_master_config() + mc.data_plane.update( + { + "backend": "mooncake_cpu", + "checkpointing_enabled": True, + } + ) + with pytest.raises(NotImplementedError, match="backend='simple'"): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_multiple_dataloader_not_supported(self): mc = _make_master_config(use_multiple_dataloader=True) with pytest.raises(NotImplementedError, match="use_multiple_dataloader"): diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 803f013f9ab..8847461e354 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -157,13 +157,16 @@ def _make_buffer( dp: FakeDataPlaneClient, *, require_routed_experts: bool = False, + checkpoint_lock: asyncio.Lock | 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_lock(checkpoint_lock or asyncio.Lock()) + return buffer def _add_group( @@ -347,6 +350,43 @@ def test_commit_appends_multiple_records_in_order(self): class TestTQReplayBufferRemove: + def test_dp_clear_fails_without_bound_checkpoint_lock(self): + dp = FakeDataPlaneClient() + buf = TQReplayBuffer( + dp, + partition_id="rollout_data", + pad_value_dict={"token_ids": 0}, + ) + + with pytest.raises(RuntimeError, match="must be bound"): + _run(buf._clear_samples(sample_ids=["sample-1"])) + + assert dp.clear_calls == [] + + def test_dp_clear_waits_for_bound_checkpoint_lock(self): + async def exercise() -> None: + dp = FakeDataPlaneClient() + checkpoint_lock = asyncio.Lock() + buf = _make_buffer(dp, checkpoint_lock=checkpoint_lock) + group_id = buf.reserve(weight_version=0) + await buf.commit( + group_id, + _make_record(), + start_weight_version=0, + end_weight_version=0, + ) + + await checkpoint_lock.acquire() + remove_task = asyncio.create_task(buf.remove([0], remove_in_dp=True)) + await asyncio.sleep(0) + assert dp.clear_calls == [] + + checkpoint_lock.release() + await remove_task + assert dp.clear_calls == [dp.put_calls[0]["sample_ids"]] + + asyncio.run(exercise()) + def test_remove_drops_indices_and_clears_dp_when_requested(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) diff --git a/tools/verify_tq_data_plane_checkpoint.py b/tools/verify_tq_data_plane_checkpoint.py new file mode 100644 index 00000000000..0f58c6d5d7c --- /dev/null +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -0,0 +1,210 @@ +# 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 restart. + +The save phase writes sample tensors, tags, and partial consumer progress to +TQ before checkpointing it. The load phase starts a fresh TQ instance, restores +the checkpoint 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 subprocess +import sys +from pathlib import Path +from typing import cast + +import torch +from tensordict import TensorDict + +from nemo_rl.data_plane import DataPlaneConfig, build_data_plane_client + +PARTITION_ID = "tq_checkpoint_smoke" +TASK_NAME = "train" +SAMPLE_IDS = [f"prompt-0:generation-{index}" for index in range(4)] +SEQ_LEN = 16 +FIELDS = ["token_ids", "token_mask", "generation_logprobs"] + + +def _data_plane_config(num_storage_units: int) -> DataPlaneConfig: + return cast( + DataPlaneConfig, + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "checkpointing_enabled": True, + "storage_capacity": 1024, + "num_storage_units": num_storage_units, + "claim_meta_poll_interval_s": 0.05, + "global_segment_size": 8 * 1024**3, + "local_buffer_size": 1024**3, + }, + ) + + +def _expected_fields() -> TensorDict: + token_ids = torch.arange(len(SAMPLE_IDS) * SEQ_LEN, dtype=torch.int64).reshape( + len(SAMPLE_IDS), + SEQ_LEN, + ) + return TensorDict( + { + "token_ids": token_ids, + "token_mask": torch.ones_like(token_ids), + "generation_logprobs": -token_ids.to(torch.float32) / 100.0, + }, + batch_size=[len(SAMPLE_IDS)], + ) + + +def _save(checkpoint_dir: Path, num_storage_units: int) -> None: + dp_client = build_data_plane_client( + _data_plane_config(num_storage_units), + bootstrap=True, + ) + try: + dp_client.register_partition( + partition_id=PARTITION_ID, + fields=FIELDS, + num_samples=len(SAMPLE_IDS), + consumer_tasks=[TASK_NAME], + ) + dp_client.put_samples( + sample_ids=SAMPLE_IDS, + partition_id=PARTITION_ID, + fields=_expected_fields(), + tags=[{"policy_version": 3, "prompt_id": "prompt-0"} for _ in SAMPLE_IDS], + ) + + consumed = dp_client.claim_meta( + partition_id=PARTITION_ID, + task_name=TASK_NAME, + required_fields=FIELDS, + batch_size=1, + timeout_s=30.0, + ) + if consumed.size != 1: + raise AssertionError(f"Expected one consumed row, got {consumed.size}") + + dp_client.save_checkpoint( + checkpoint_dir, + metadata={ + "data_plane_checkpoint_schema_version": 1, + "expected_consumed_ids": consumed.sample_ids, + }, + ) + finally: + dp_client.close() + + +def _load(checkpoint_dir: Path, num_storage_units: int) -> None: + dp_client = build_data_plane_client( + _data_plane_config(num_storage_units), + bootstrap=True, + ) + try: + metadata = dp_client.load_checkpoint(checkpoint_dir) + + restored = dp_client.get_samples( + sample_ids=SAMPLE_IDS, + partition_id=PARTITION_ID, + select_fields=FIELDS, + ) + expected = _expected_fields() + for field in FIELDS: + if not torch.equal(restored[field], expected[field]): + raise AssertionError(f"Restored field differs: {field}") + + if metadata["data_plane_checkpoint_schema_version"] != 1: + raise AssertionError("Unexpected data-plane checkpoint schema") + consumed_ids = set(metadata["expected_consumed_ids"]) + expected_remaining_ids = set(SAMPLE_IDS) - consumed_ids + + if dp_client.check_consumption_status(PARTITION_ID, [TASK_NAME]): + raise AssertionError( + "Restored consumer cursor marked every row consumed before " + "the expected remaining rows were claimed" + ) + remaining = dp_client.claim_meta( + partition_id=PARTITION_ID, + task_name=TASK_NAME, + required_fields=FIELDS, + batch_size=len(expected_remaining_ids), + timeout_s=30.0, + ) + if consumed_ids.intersection(remaining.sample_ids): + raise AssertionError("A previously consumed row was claimed after restore") + if set(remaining.sample_ids) != expected_remaining_ids: + raise AssertionError("Restored consumption state lost or added rows") + if not dp_client.check_consumption_status(PARTITION_ID, [TASK_NAME]): + raise AssertionError("Restored consumer cursor did not reach completion") + finally: + dp_client.close() + + +def _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) + 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(checkpoint_dir, args.num_storage_units) + return + if args.phase == "load": + _load(checkpoint_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 data-plane checkpoint survived a fresh process", flush=True) + + +if __name__ == "__main__": + main() From 03fba290c1f29a32169868409364f11450c60c29 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 2 Aug 2026 14:55:31 -0400 Subject: [PATCH 02/32] fix(sc): harden TQ checkpoint lifecycle Signed-off-by: Anish Mahishi --- examples/configs/grpo_math_1B.yaml | 2 +- ...po_math_1B_megatron_single_controller.yaml | 1 + .../algorithms/async_utils/replay_buffer.py | 7 ++- nemo_rl/algorithms/single_controller.py | 61 ++++++++++++------ nemo_rl/data_plane/adapters/transfer_queue.py | 40 +++++++++++- nemo_rl/data_plane/interfaces.py | 13 ++-- tests/unit/data_plane/test_codec_mooncake.py | 2 + tests/unit/data_plane/test_tq_lifecycle.py | 62 +++++++++++++++++++ .../unit/reference_configs/grpo_math_1B.yaml | 2 +- .../test_verify_tq_data_plane_checkpoint.py | 50 +++++++++++++++ tools/verify_tq_data_plane_checkpoint.py | 49 ++++++++++++--- 11 files changed, 250 insertions(+), 39 deletions(-) create mode 100644 tests/unit/tools/test_verify_tq_data_plane_checkpoint.py diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 93de9fe8772..fab79d3105d 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -520,7 +520,7 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - checkpointing_enabled: false # save TQ state inside algorithm checkpoints + checkpointing_enabled: false # SingleController only: save required shadow TQ state global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index 2f94b10479a..9f7f556c34f 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -58,6 +58,7 @@ logger: data_plane: enabled: true + # Required shadow snapshot: a save failure aborts checkpoint finalization. checkpointing_enabled: true cluster: diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 395e521ad3d..567e3fb6a1a 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -710,7 +710,12 @@ def __init__( self._data_plane_checkpoint_lock: Optional[asyncio.Lock] = None def set_data_plane_checkpoint_lock(self, lock: asyncio.Lock) -> None: - """Serialize replay-buffer-owned clears with SC checkpoints.""" + """Bind the controller's shared checkpoint/clear barrier exactly once. + + A private fallback lock would not coordinate with controller-owned + saves and clears, so destructive operations fail loudly until the SC + actor supplies its lock. + """ if self._data_plane_checkpoint_lock is not None: raise RuntimeError("data-plane checkpoint lock is already configured") self._data_plane_checkpoint_lock = lock diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 891ad9f790c..7e8eb11e3c3 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -174,8 +174,11 @@ def __init__( # ── asyncio state ────────────────────────────────────────────────── # TQ snapshots permit concurrent puts but not destructive clears. All # clears currently owned by async SC use this lock, including rollback - # and eviction through TQReplayBuffer. A future staging/finalizer path - # must join the same barrier before native restore can be authoritative. + # and eviction through TQReplayBuffer. Clear-dependent eviction waits + # during a save; _buffer_capacity bounds new rollout groups 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_lock: asyncio.Lock = asyncio.Lock() if self._buffer is not None: self._buffer.set_data_plane_checkpoint_lock( @@ -320,13 +323,33 @@ async def _ray_get(self, obj_ref: Any) -> Any: """Await a Ray ObjectRef without blocking the asyncio event loop.""" return await obj_ref - async def _call_dp(self, method_name: str, **kwargs) -> Any: - """Call a DataPlaneClient method or a Ray actor exposing that method.""" + async def _call_dp( + self, + method_name: str, + *, + offload_sync: bool = False, + **kwargs: Any, + ) -> Any: + """Call a local DataPlaneClient or a Ray actor exposing its methods. + + Args: + method_name: DataPlaneClient method to invoke. + offload_sync: Run a synchronous local implementation in a worker + thread. Use for blocking filesystem or RPC operations; Ray + methods are already asynchronous and ignore this setting. + **kwargs: Keyword arguments forwarded to the data-plane method. + + Returns: + The method result after awaiting Ray or coroutine results. + """ method = getattr(self._dp_client, method_name) remote = getattr(method, "remote", None) if remote is not None: return await self._ray_get(remote(**kwargs)) - result = method(**kwargs) + if offload_sync: + result = await asyncio.to_thread(method, **kwargs) + else: + result = method(**kwargs) if asyncio.iscoroutine(result): return await result return result @@ -341,7 +364,13 @@ async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: ) async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: - """Save a shadow TQ snapshot inside an SC checkpoint bundle.""" + """Save a required shadow TQ snapshot inside an SC checkpoint bundle. + + Although native TQ restore is not wired into SC yet, opting into this + shadow snapshot is intentionally fail-closed: any failure propagates so + a finalized bundle never silently omits the advertised data-plane + component. + """ checkpoint_dir = os.path.join( checkpoint_path, DATA_PLANE_CHECKPOINT_DIR, @@ -359,20 +388,12 @@ async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: started = time.monotonic() print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) try: - method = getattr(self._dp_client, "save_checkpoint") - remote = getattr(method, "remote", None) - if remote is not None: - await self._ray_get( - remote(checkpoint_dir=checkpoint_dir, metadata=metadata) - ) - else: - result = await asyncio.to_thread( - method, - checkpoint_dir=checkpoint_dir, - metadata=metadata, - ) - if asyncio.iscoroutine(result): - await result + await self._call_dp( + "save_checkpoint", + offload_sync=True, + checkpoint_dir=checkpoint_dir, + metadata=metadata, + ) except Exception as error: print( "data-plane checkpoint save failed: " diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index bf5ae6c31c7..266fb298591 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -480,6 +480,22 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: _connect_existing() self._poll_interval_s = cfg["claim_meta_poll_interval_s"] self._closed = False + # TQ restore is non-transactional and requires a globally clean system. + # This process-local guard catches incorrect ordering through this + # adapter; setup must still ensure no other client has touched TQ. + self._data_operations_started = False + + def _mark_data_operation_started(self) -> None: + """Make a later checkpoint load fail instead of mixing TQ states.""" + self._data_operations_started = True + + def _require_clean_for_load(self) -> None: + """Reject restore after this client has performed a data operation.""" + if self._data_operations_started: + raise RuntimeError( + "load_checkpoint requires a clean TQ client before any " + "register, claim, get, put, clear, or consumption operation" + ) # ── (A) task-mediated ─────────────────────────────────────────────── @@ -512,6 +528,7 @@ def register_partition( # stale metadata from a previous registration. if not fields: return + self._mark_data_operation_started() schema_key = ( f"__schema__:{partition_id}:{os.getpid()}:{id(self)}:{time.time_ns()}" ) @@ -537,6 +554,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] = {} @@ -608,6 +626,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( @@ -649,6 +668,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), @@ -673,6 +693,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, @@ -683,6 +704,7 @@ def get_samples( 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 @@ -704,6 +726,7 @@ 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) @@ -726,20 +749,31 @@ def save_checkpoint( 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.""" + """Restore TQ state after initialization and before data operations. + + The local lifecycle guard cannot observe operations issued by another + TQ client, so the recovery coordinator must also guarantee globally + clean setup ordering. + """ if self._backend == "mooncake_cpu": raise NotImplementedError( "TQ checkpoint restore is not supported for the mooncake_cpu " "backend because its in-memory rows cannot be restored." ) - _connect_existing() - tq.load_checkpoint(checkpoint_dir) + 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: diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 6f57bcade65..6bb23d1cbd5 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -60,10 +60,11 @@ class DataPlaneConfig(TypedDict): They are required (not NotRequired) so the YAML carries the full schema and there are no hidden Python defaults. - ``checkpointing_enabled`` opts algorithms into saving native TQ state - inside their checkpoint bundle. It is optional during rollout because - existing configs predate data-plane checkpointing; exemplar configs carry - the recommended default explicitly. + ``checkpointing_enabled`` opts SingleController into saving required + shadow TQ state inside its checkpoint bundle. Other algorithm entrypoints + do not consume this field. It is optional because existing configs predate + data-plane checkpointing; exemplar configs carry the recommended default + explicitly. """ enabled: bool @@ -474,7 +475,9 @@ 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. + 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 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_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index c6f1a98c067..0767d6cf6c0 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -54,6 +54,7 @@ def fake_clear(**kwargs): monkeypatch.setattr(tq_adapter.tq, "kv_clear", fake_clear) client = object.__new__(tq_adapter.TQDataPlaneClient) + client._data_operations_started = False client.register_partition( partition_id="obj-backend", fields=["msg_log"], @@ -111,6 +112,7 @@ def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: client = object.__new__(tq_adapter.TQDataPlaneClient) client._backend = "simple" + client._data_operations_started = False checkpoint_dir = tmp_path / "step-7" checkpoint_dir.mkdir() (checkpoint_dir / "metadata.json").write_text( @@ -123,6 +125,65 @@ def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: assert save_calls == [(checkpoint_dir, {"step": 7})] assert load_calls == [checkpoint_dir] assert metadata == {"step": 7} + assert client._data_operations_started + + +def test_checkpoint_load_rejects_client_after_data_operation( + monkeypatch, + tmp_path, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect = MagicMock() + load = MagicMock() + monkeypatch.setattr(tq_adapter, "_connect_existing", connect) + monkeypatch.setattr(tq_adapter.tq, "load_checkpoint", load) + monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", MagicMock()) + + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + client._promote_1d = False + client._data_operations_started = False + client.put_samples( + sample_ids=["sample-0"], + partition_id="rollout_data", + fields=TensorDict({"x": torch.tensor([1])}, batch_size=[1]), + ) + + with pytest.raises(RuntimeError, match="requires a clean TQ client"): + client.load_checkpoint(tmp_path / "data-plane") + + connect.assert_not_called() + load.assert_not_called() + + +def test_failed_checkpoint_load_leaves_client_in_dirty_state( + monkeypatch, + tmp_path, +) -> None: + from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + connect = MagicMock() + load = MagicMock(side_effect=RuntimeError("injected partial restore")) + monkeypatch.setattr(tq_adapter, "_connect_existing", connect) + monkeypatch.setattr(tq_adapter.tq, "load_checkpoint", load) + + checkpoint_dir = tmp_path / "data-plane" + checkpoint_dir.mkdir() + (checkpoint_dir / "metadata.json").write_text( + json.dumps({"storage_saved": True, "user_metadata": {"step": 7}}) + ) + client = object.__new__(tq_adapter.TQDataPlaneClient) + client._backend = "simple" + client._data_operations_started = False + + with pytest.raises(RuntimeError, match="injected partial restore"): + client.load_checkpoint(checkpoint_dir) + with pytest.raises(RuntimeError, match="requires a clean TQ client"): + client.load_checkpoint(checkpoint_dir) + + connect.assert_called_once_with() + load.assert_called_once_with(checkpoint_dir) @pytest.mark.parametrize("operation", ["save", "load"]) @@ -142,6 +203,7 @@ def test_mooncake_checkpoint_lifecycle_fails_loudly( client = object.__new__(tq_adapter.TQDataPlaneClient) client._backend = "mooncake_cpu" + client._data_operations_started = False checkpoint_dir = tmp_path / "step-7" with pytest.raises(NotImplementedError, match="mooncake_cpu"): diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 9471acf49ee..5fe4eaca44b 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -512,7 +512,7 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - checkpointing_enabled: false # save TQ state inside algorithm checkpoints + checkpointing_enabled: false # SingleController only: save required shadow TQ state global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/tests/unit/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..25c2455fe7b --- /dev/null +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest + +from tools import verify_tq_data_plane_checkpoint as verifier + + +def test_save_finalizes_by_renaming_parent_bundle(monkeypatch, tmp_path) -> None: + final_bundle = tmp_path / "step_7" + expected_staging_bundle = tmp_path / "tmp_step_7" + save_calls = [] + + def fake_save(checkpoint_dir, num_storage_units) -> None: + save_calls.append((checkpoint_dir, num_storage_units)) + checkpoint_dir.mkdir(parents=True) + (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() diff --git a/tools/verify_tq_data_plane_checkpoint.py b/tools/verify_tq_data_plane_checkpoint.py index 0f58c6d5d7c..3525424fff3 100644 --- a/tools/verify_tq_data_plane_checkpoint.py +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Verify TQ data-plane save/load across a fresh process restart. +"""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 before checkpointing it. The load phase starts a fresh TQ instance, restores -the checkpoint before any partition operations, and verifies both the tensors -and the consumer cursor. +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 \ @@ -27,6 +28,7 @@ from __future__ import annotations import argparse +import shutil import subprocess import sys from pathlib import Path @@ -42,6 +44,7 @@ SAMPLE_IDS = [f"prompt-0:generation-{index}" for index in range(4)] SEQ_LEN = 16 FIELDS = ["token_ids", "token_mask", "generation_logprobs"] +DATA_PLANE_DIR = "data_plane" def _data_plane_config(num_storage_units: int) -> DataPlaneConfig: @@ -161,6 +164,28 @@ def _load(checkpoint_dir: Path, num_storage_units: int) -> None: 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}" + ) + + 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, @@ -189,21 +214,29 @@ def main() -> None: default="round-trip", help=argparse.SUPPRESS, ) - parser.add_argument("--checkpoint-dir", type=Path, required=True) + 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(checkpoint_dir, args.num_storage_units) + _save_and_finalize_bundle(checkpoint_dir, args.num_storage_units) return if args.phase == "load": - _load(checkpoint_dir, args.num_storage_units) + _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 data-plane checkpoint survived a fresh process", flush=True) + print( + "PASS: TQ checkpoint survived a parent rename and fresh process", + flush=True, + ) if __name__ == "__main__": From 1477a15316cf0377eb0a1f007b940c2d468f8545 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Sun, 2 Aug 2026 21:01:48 -0400 Subject: [PATCH 03/32] fix(sc): harden TQ checkpoint I/O Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 25 +++--- nemo_rl/algorithms/single_controller.py | 49 +++--------- nemo_rl/data_plane/async_utils.py | 56 ++++++++++++++ tests/unit/data_plane/test_async_utils.py | 77 +++++++++++++++++++ .../test_sc_checkpointing.py | 17 ++++ .../test_tq_replay_buffer.py | 15 ++++ .../test_verify_tq_data_plane_checkpoint.py | 22 +++++- tools/verify_tq_data_plane_checkpoint.py | 3 + 8 files changed, 209 insertions(+), 55 deletions(-) create mode 100644 nemo_rl/data_plane/async_utils.py create mode 100644 tests/unit/data_plane/test_async_utils.py diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 567e3fb6a1a..c782550d504 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -26,6 +26,7 @@ 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, @@ -790,7 +791,8 @@ async def commit( ) trace_rollout_payload(keys=sample_ids, data=train_batch) try: - await self._call_dp( + await call_data_plane( + self._dp_client, "put_samples", sample_ids=sample_ids, partition_id=self._partition_id, @@ -926,7 +928,8 @@ async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]: groups: list[dict[str, Any]] = [] for meta, start_weight, end_weight, target_step, group_id in snapshot: - fields_data = await self._call_dp( + fields_data = await call_data_plane( + self._dp_client, "get_samples", sample_ids=meta.sample_ids, partition_id=self._partition_id, @@ -1069,7 +1072,8 @@ async def load_state_dict( for group in groups: meta = group["meta"] - await self._call_dp( + await call_data_plane( + self._dp_client, "put_samples", sample_ids=list(meta.sample_ids), partition_id=self._partition_id, @@ -1100,17 +1104,6 @@ def size(self) -> int: def __len__(self) -> int: return len(self.meta_list) - async def _call_dp(self, method_name: str, **kwargs: Any) -> Any: - """Call a DataPlaneClient method, awaiting Ray remotes if needed.""" - method = getattr(self._dp_client, method_name) - remote = getattr(method, "remote", None) - if remote is not None: - return await remote(**kwargs) - result = method(**kwargs) - if asyncio.iscoroutine(result): - return await result - return result - async def _clear_samples(self, *, sample_ids: list[str]) -> None: """Clear rows without overlapping a bound data-plane checkpoint.""" if self._data_plane_checkpoint_lock is None: @@ -1119,8 +1112,10 @@ async def _clear_samples(self, *, sample_ids: list[str]) -> None: "checkpoint lock before clearing samples" ) async with self._data_plane_checkpoint_lock: - await self._call_dp( + 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/single_controller.py b/nemo_rl/algorithms/single_controller.py index 7e8eb11e3c3..6fbe06a0d0e 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -62,6 +62,7 @@ ) from nemo_rl.data.interfaces import DatumSpec from nemo_rl.data_plane import KVBatchMeta +from nemo_rl.data_plane.async_utils import call_data_plane from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration @@ -319,46 +320,13 @@ async def _maybe_restore_replay_buffer(self) -> None: for _ in range(restored): await self._buffer_capacity.acquire() - async def _ray_get(self, obj_ref: Any) -> Any: - """Await a Ray ObjectRef without blocking the asyncio event loop.""" - return await obj_ref - - async def _call_dp( - self, - method_name: str, - *, - offload_sync: bool = False, - **kwargs: Any, - ) -> Any: - """Call a local DataPlaneClient or a Ray actor exposing its methods. - - Args: - method_name: DataPlaneClient method to invoke. - offload_sync: Run a synchronous local implementation in a worker - thread. Use for blocking filesystem or RPC operations; Ray - methods are already asynchronous and ignore this setting. - **kwargs: Keyword arguments forwarded to the data-plane method. - - Returns: - The method result after awaiting Ray or coroutine results. - """ - method = getattr(self._dp_client, method_name) - remote = getattr(method, "remote", None) - if remote is not None: - return await self._ray_get(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 - 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_lock: - await self._call_dp( + await call_data_plane( + self._dp_client, "clear_samples", + offload_sync=True, sample_ids=sample_ids, partition_id=self._partition_id, ) @@ -388,7 +356,8 @@ async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: started = time.monotonic() print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) try: - await self._call_dp( + await call_data_plane( + self._dp_client, "save_checkpoint", offload_sync=True, checkpoint_dir=checkpoint_dir, @@ -999,7 +968,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: return meta adv_cfg = self._advantage_cfg - data = await self._call_dp( + data = await call_data_plane( + self._dp_client, "get_samples", sample_ids=meta.sample_ids, partition_id=meta.partition_id, @@ -1049,7 +1019,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> KVBatchMeta: response_advantages.detach().cpu() ) - await self._call_dp( + await call_data_plane( + self._dp_client, "put_samples", sample_ids=meta.sample_ids, partition_id=meta.partition_id, diff --git a/nemo_rl/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/tests/unit/data_plane/test_async_utils.py b/tests/unit/data_plane/test_async_utils.py new file mode 100644 index 00000000000..25d29b5d711 --- /dev/null +++ b/tests/unit/data_plane/test_async_utils.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for local and Ray-style async data-plane dispatch.""" + +import asyncio +import threading + +from nemo_rl.data_plane.async_utils import call_data_plane + + +class _LocalClient: + def thread_id(self) -> int: + return threading.get_ident() + + async def async_value(self, *, value: int) -> int: + return value + + +class _RemoteMethod: + def __init__(self) -> None: + self.calls: list[int] = [] + + async def remote(self, *, value: int) -> int: + self.calls.append(value) + return value + + +class _RemoteClient: + def __init__(self) -> None: + self.value = _RemoteMethod() + + +def test_sync_call_stays_inline_by_default() -> None: + caller_thread_id = threading.get_ident() + + result = asyncio.run(call_data_plane(_LocalClient(), "thread_id")) + + assert result == caller_thread_id + + +def test_sync_call_can_be_offloaded() -> None: + caller_thread_id = threading.get_ident() + + result = asyncio.run( + call_data_plane(_LocalClient(), "thread_id", offload_sync=True) + ) + + assert result != caller_thread_id + + +def test_local_coroutine_result_is_awaited() -> None: + result = asyncio.run( + call_data_plane(_LocalClient(), "async_value", value=7) + ) + + assert result == 7 + + +def test_ray_style_remote_result_is_awaited() -> None: + client = _RemoteClient() + + result = asyncio.run(call_data_plane(client, "value", value=11)) + + assert result == 11 + assert client.value.calls == [11] diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 7829d7f5ef5..3475a313414 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -215,10 +215,12 @@ async def select(self, **kwargs) -> tuple[Optional[KVBatchMeta], int]: class _FakeDPClient: def __init__(self, *, save_error: Optional[Exception] = 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 def clear_samples(self, sample_ids: list[str], partition_id: str) -> None: + self.clear_thread_ids.append(threading.get_ident()) self.clear_calls.append((list(sample_ids), partition_id)) def save_checkpoint( @@ -828,6 +830,21 @@ async def _main() -> None: asyncio.run(_main()) assert dp_client.clear_calls == [(["sample-0"], _PARTITION_ID)] + def test_consumed_clear_does_not_block_actor_event_loop(self, tmp_path): + mc = _actor_master_config(tmp_path, max_num_steps=1, save_period=1) + dp_client = _FakeDPClient() + + async def _main() -> int: + actor = _ACTOR_CLS(mc, _make_actor_args(dp_client=dp_client)) + event_loop_thread_id = threading.get_ident() + await actor._clear_data_plane_samples(["sample-0"]) + actor._checkpointer.shutdown() + return event_loop_thread_id + + event_loop_thread_id = asyncio.run(_main()) + assert dp_client.clear_thread_ids + assert dp_client.clear_thread_ids[0] != event_loop_thread_id + # ── async-save finalization ────────────────────────────────────────────────── diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 8847461e354..d3d6f813a01 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -18,6 +18,7 @@ import asyncio import io +import threading from typing import Any import pytest @@ -65,6 +66,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( @@ -96,6 +98,7 @@ 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: @@ -387,6 +390,18 @@ async def exercise() -> None: asyncio.run(exercise()) + def test_dp_clear_does_not_block_actor_event_loop(self): + async def exercise() -> tuple[FakeDataPlaneClient, int]: + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + event_loop_thread_id = threading.get_ident() + await buf._clear_samples(sample_ids=["sample-1"]) + return dp, event_loop_thread_id + + dp, event_loop_thread_id = asyncio.run(exercise()) + assert dp.clear_thread_ids + assert dp.clear_thread_ids[0] != event_loop_thread_id + def test_remove_drops_indices_and_clears_dp_when_requested(self): dp = FakeDataPlaneClient() buf = _make_buffer(dp) diff --git a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py index 25c2455fe7b..6359fc672f7 100644 --- a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -26,7 +26,8 @@ def test_save_finalizes_by_renaming_parent_bundle(monkeypatch, tmp_path) -> None def fake_save(checkpoint_dir, num_storage_units) -> None: save_calls.append((checkpoint_dir, num_storage_units)) - checkpoint_dir.mkdir(parents=True) + assert checkpoint_dir.parent.is_dir() + checkpoint_dir.mkdir() (checkpoint_dir / "marker").write_text("saved") monkeypatch.setattr(verifier, "_save", fake_save) @@ -48,3 +49,22 @@ def test_save_refuses_to_replace_final_bundle(monkeypatch, tmp_path) -> None: 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 index 3525424fff3..33c9f50adaf 100644 --- a/tools/verify_tq_data_plane_checkpoint.py +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -177,6 +177,9 @@ def _save_and_finalize_bundle( 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) From e31f42c44586ecc22338b4a761a3516bf654df20 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Mon, 3 Aug 2026 15:28:53 -0400 Subject: [PATCH 04/32] feat(sc): recover replay buffer from native TQ checkpoints Signed-off-by: Anish Mahishi --- examples/configs/grpo_math_1B.yaml | 5 +- .../algorithms/async_utils/replay_buffer.py | 445 +++++++++----- .../async_utils/staleness_sampler.py | 62 +- nemo_rl/algorithms/grpo.py | 4 + nemo_rl/algorithms/single_controller.py | 234 ++++++-- .../single_controller_utils/setup.py | 119 +++- nemo_rl/data_plane/adapters/noop.py | 5 + nemo_rl/data_plane/adapters/transfer_queue.py | 8 +- nemo_rl/data_plane/interfaces.py | 27 +- nemo_rl/data_plane/observability.py | 7 + nemo_rl/models/policy/tq_policy.py | 5 + pyrefly.toml | 1 + .../L1_Functional_Tests_SingleController.sh | 1 + tests/functional/grpo_dp_single_controller.sh | 18 +- .../grpo_dp_single_controller_tq_recovery.sh | 53 ++ .../test_architecture_invariants.py | 1 + tests/unit/data_plane/test_async_utils.py | 4 +- .../data_plane/test_interface_contract.py | 2 + tests/unit/data_plane/test_observability.py | 16 + tests/unit/data_plane/test_smoke.py | 1 + tests/unit/data_plane/test_tq_lifecycle.py | 17 + .../unit/reference_configs/grpo_math_1B.yaml | 4 +- tests/unit/single_controller/_dp_fakes.py | 3 + .../test_sampler_interface.py | 37 +- .../test_sc_checkpointing.py | 543 +++++++++++++++--- .../test_single_controller_setup.py | 201 ++++++- .../test_tq_replay_buffer.py | 332 +++++------ 27 files changed, 1645 insertions(+), 510 deletions(-) create mode 100755 tests/functional/grpo_dp_single_controller_tq_recovery.sh diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index fab79d3105d..92abd99443f 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -520,7 +520,10 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - checkpointing_enabled: false # SingleController only: save required shadow TQ state + # SingleController only: save native TQ state. Required when trainer + # checkpointing is enabled with a replay-checkpoint-capable sampler; + # supported samplers restore from metadata-only replay indexes. + checkpointing_enabled: false global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index c782550d504..6e62fecedd7 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -14,12 +14,15 @@ import asyncio import gc +import hashlib +import json import statistics import threading as _threading import uuid from collections import Counter -from collections.abc import Mapping -from typing import Any, Iterable, Optional +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager +from typing import Any, Iterable, Literal, Optional, TypedDict import ray import torch @@ -37,6 +40,111 @@ from nemo_rl.utils.r3_trace import trace_rollout_payload +DATA_PLANE_CHECKPOINT_DIR = "data_plane" +DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 +REPLAY_BUFFER_METADATA_FILENAME = "replay_buffer_metadata.pt" +LEGACY_REPLAY_BUFFER_FILENAME = "replay_buffer.pt" +REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1 +REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" + + +class TQReplayGroupMetadata(TypedDict): + """Controller-local index for one training-ready group stored in TQ.""" + + meta: KVBatchMeta + start_weight: int + end_weight: int + target_step: Optional[int] + group_id: str + + +class TQReplayMetadataState(TypedDict): + """Versioned metadata-only replay sidecar paired with a TQ snapshot.""" + + schema_version: int + storage: Literal["tq_checkpoint"] + partition_id: str + saved_capacity: int + manifest_digest: str + groups: list[TQReplayGroupMetadata] + + +def replay_manifest_digest(groups: list[TQReplayGroupMetadata]) -> str: + """Return a stable digest binding replay metadata to a TQ checkpoint.""" + digest_input = [ + { + "group_id": group["group_id"], + "start_weight": group["start_weight"], + "end_weight": group["end_weight"], + "target_step": group["target_step"], + "meta": { + "partition_id": group["meta"].partition_id, + "task_name": group["meta"].task_name, + "sample_ids": list(group["meta"].sample_ids), + "fields": ( + list(group["meta"].fields) + if group["meta"].fields is not None + else None + ), + "sequence_lengths": ( + list(group["meta"].sequence_lengths) + if group["meta"].sequence_lengths is not None + else None + ), + "tags": group["meta"].tags, + "extra_info": group["meta"].extra_info, + }, + } + for group in groups + ] + encoded = json.dumps(digest_input, sort_keys=True, separators=(",", ":")).encode( + "utf-8" + ) + return hashlib.sha256(encoded).hexdigest() + + +class DataPlaneCheckpointBarrier: + """Allow concurrent mutations while giving checkpoints exclusive access.""" + + def __init__(self) -> None: + self._condition = asyncio.Condition() + self._checkpoint_active = False + self._active_mutations = 0 + + @asynccontextmanager + async def mutation(self) -> AsyncIterator[None]: + """Enter a commit/clear section, waiting only for an active checkpoint.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._checkpoint_active) + self._active_mutations += 1 + try: + yield + finally: + async with self._condition: + self._active_mutations -= 1 + if self._active_mutations == 0: + self._condition.notify_all() + + @asynccontextmanager + async def checkpoint(self) -> AsyncIterator[None]: + """Block new mutations and wait for active ones before snapshotting.""" + async with self._condition: + await self._condition.wait_for(lambda: not self._checkpoint_active) + self._checkpoint_active = True + try: + await self._condition.wait_for(lambda: self._active_mutations == 0) + except BaseException: + self._checkpoint_active = False + self._condition.notify_all() + raise + try: + yield + finally: + async with self._condition: + self._checkpoint_active = False + self._condition.notify_all() + + # Classes with @ray.remote can't be inherited from, so we split the implementation out. class ReplayBufferImpl(ReplayBufferProtocol): """Replay buffer storing per-prompt groups. @@ -708,18 +816,20 @@ def __init__( self.target_step_list: list[Optional[int]] = [] self.ready_list: list[bool] = [] self._group_ids: list[str] = [] - self._data_plane_checkpoint_lock: Optional[asyncio.Lock] = None + self._data_plane_checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None - def set_data_plane_checkpoint_lock(self, lock: asyncio.Lock) -> None: - """Bind the controller's shared checkpoint/clear barrier exactly once. + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + """Bind the controller's shared checkpoint/mutation barrier once. - A private fallback lock would not coordinate with controller-owned + A private fallback barrier would not coordinate with controller-owned saves and clears, so destructive operations fail loudly until the SC - actor supplies its lock. + actor supplies its barrier. """ - if self._data_plane_checkpoint_lock is not None: - raise RuntimeError("data-plane checkpoint lock is already configured") - self._data_plane_checkpoint_lock = lock + if self._data_plane_checkpoint_barrier is not None: + raise RuntimeError("data-plane checkpoint barrier is already configured") + self._data_plane_checkpoint_barrier = barrier def reserve( self, @@ -778,6 +888,11 @@ 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 @@ -790,47 +905,48 @@ async def commit( "the async message-log flattening path." ) trace_rollout_payload(keys=sample_ids, data=train_batch) - try: - await call_data_plane( - self._dp_client, - "put_samples", - sample_ids=sample_ids, - partition_id=self._partition_id, - fields=fields, - tags=tags, - ) - - # mirrors kv_first_write - lengths = train_batch["input_lengths"] - meta = KVBatchMeta( - partition_id=self._partition_id, - task_name="train", - sample_ids=list(sample_ids), - fields=list(fields.keys()), - sequence_lengths=[int(s) for s in lengths.tolist()], - tags=[dict(t) for t in tags], - ) - - idx = self._group_ids.index(group_id) - self.meta_list[idx] = meta - self.end_weight_list[idx] = end_weight_version - self.ready_list[idx] = True - return meta - except BaseException as commit_error: - # put_samples may have written rows before raising. Roll back by the - # deterministic IDs known here; the caller removes the reserved slot. + async with self._data_plane_checkpoint_barrier.mutation(): try: - await self._clear_samples( - sample_ids=list(sample_ids), + await call_data_plane( + self._dp_client, + "put_samples", + sample_ids=sample_ids, + partition_id=self._partition_id, + fields=fields, + tags=tags, ) - except BaseException as rollback_error: - if isinstance(commit_error, asyncio.CancelledError): - raise commit_error from rollback_error - raise BaseExceptionGroup( - f"commit and rollback both failed for group_id={group_id!r}", - [commit_error, rollback_error], + + # mirrors kv_first_write + lengths = train_batch["input_lengths"] + meta = KVBatchMeta( + partition_id=self._partition_id, + task_name="train", + sample_ids=list(sample_ids), + fields=list(fields.keys()), + sequence_lengths=[int(s) for s in lengths.tolist()], + tags=[dict(t) for t in tags], ) - raise + + idx = self._group_ids.index(group_id) + self.meta_list[idx] = meta + self.end_weight_list[idx] = end_weight_version + self.ready_list[idx] = True + return meta + except BaseException as commit_error: + # put_samples may have written rows before raising. Roll back by the + # deterministic IDs while retaining the barrier mutation slot. + try: + await self._clear_samples_unlocked( + sample_ids=list(sample_ids), + ) + except BaseException as rollback_error: + if isinstance(commit_error, asyncio.CancelledError): + raise commit_error from rollback_error + raise BaseExceptionGroup( + f"commit and rollback both failed for group_id={group_id!r}", + [commit_error, rollback_error], + ) + raise async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> int: """Remove the live slot identified by ``group_id``. @@ -845,11 +961,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. @@ -863,14 +985,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] @@ -883,71 +1015,48 @@ 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._clear_samples( + if clear_data_plane: + await self._clear_samples_unlocked( sample_ids=dropped_sample_ids, ) 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 sidecar + and native snapshot describe one exact set of training-ready groups. + Every operation that mutates the canonical rollout partition or its + controller-local replay membership must participate in that barrier + across the complete publish/index or clear/remove transition. This + includes future finalizer paths; canonical writes are not required to + originate specifically from :meth:`commit`. + In-flight reservations are intentionally omitted. """ - 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 call_data_plane( - self._dp_client, - "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, } @@ -958,16 +1067,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 sidecar never contains tensor payloads and this method never writes + to the DataPlane. TQ must be restored first; the caller binds the two + artifacts by passing the manifest digest returned by TQ checkpoint + loading. Staleness is intentionally NOT handled here — load only loads. The train pump's first ``sampler.evict`` drops any restored group that is @@ -975,7 +1082,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 @@ -983,6 +1090,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 metadata sidecar. Returns: Number of groups restored into the buffer. @@ -990,14 +1099,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: " @@ -1012,16 +1142,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 @@ -1042,44 +1181,30 @@ 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}" ) - indices_to_keep = sorted(prioritized[num_truncated:]) - groups = [groups[i] for i in indices_to_keep] for group in groups: meta = group["meta"] - await call_data_plane( - self._dp_client, - "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"]) @@ -1087,10 +1212,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: @@ -1106,16 +1231,20 @@ def __len__(self) -> int: async def _clear_samples(self, *, sample_ids: list[str]) -> None: """Clear rows without overlapping a bound data-plane checkpoint.""" - if self._data_plane_checkpoint_lock is None: + if self._data_plane_checkpoint_barrier is None: raise RuntimeError( "TQReplayBuffer must be bound to the controller data-plane " - "checkpoint lock before clearing samples" - ) - async with self._data_plane_checkpoint_lock: - await call_data_plane( - self._dp_client, - "clear_samples", - offload_sync=True, - sample_ids=sample_ids, - partition_id=self._partition_id, + "checkpoint barrier before clearing samples" ) + async with self._data_plane_checkpoint_barrier.mutation(): + await self._clear_samples_unlocked(sample_ids=sample_ids) + + async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None: + """Clear rows while the caller holds a barrier mutation slot.""" + await call_data_plane( + self._dp_client, + "clear_samples", + offload_sync=True, + sample_ids=sample_ids, + partition_id=self._partition_id, + ) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index c178edaafaf..fa234e886bf 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, @@ -109,6 +110,11 @@ def is_on_policy(self) -> bool: """True when the policy admits zero staleness (sync mode).""" ... + @property + def supports_buffer_checkpoint(self) -> 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.""" ... @@ -196,6 +202,10 @@ def should_abort_inflight( def is_on_policy(self) -> bool: return self._eviction_window() == 0 + @property + def supports_buffer_checkpoint(self) -> bool: + return False + def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: return None @@ -266,6 +276,11 @@ def __init__( def _eviction_window(self) -> int: return self.max_staleness_versions + @property + def supports_buffer_checkpoint(self) -> bool: + # Ungated restored groups are ordinary in-window candidates. + return True + def should_abort_inflight( self, *, @@ -443,6 +458,7 @@ async def evict(self, *, current_train_weight: int) -> int: class WindowedSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = True name: Literal["windowed"] = "windowed" # Max weight-version gap a selected group may have from the trainer. max_staleness_versions: NonNegativeInt = 1 @@ -451,18 +467,22 @@ class WindowedSamplerConfig(BaseModel, extra="allow"): class WeightFifoSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["weight_fifo"] = "weight_fifo" # Lookahead + selectable weight window, in trainer versions. max_staleness_versions: NonNegativeInt = 1 class InOrderSamplerConfig(BaseModel, extra="allow"): + supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["in_order"] = "in_order" # How far generation may run ahead of the trainer, in dispatch batches. max_lookahead_versions: NonNegativeInt = 1 class CustomSamplerConfig(BaseModel, extra="allow"): + # A custom implementation's capability is known only after construction. + supports_buffer_checkpoint: ClassVar[Optional[bool]] = None name: Literal["custom"] = "custom" # "module:ClassName" of a PromptGroupSampler defined outside this repo. # Extra keys are forwarded to the constructor (after ``buffer``). @@ -505,20 +525,30 @@ 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: PromptGroupSampler if isinstance(cfg, WindowedSamplerConfig): - return WindowedSampler( + sampler = WindowedSampler( buffer, max_staleness_versions=cfg.max_staleness_versions, sample_freshest_first=cfg.sample_freshest_first, ) - if isinstance(cfg, WeightFifoSamplerConfig): - return WeightFifoSampler( - buffer, max_staleness_versions=cfg.max_staleness_versions + elif isinstance(cfg, WeightFifoSamplerConfig): + sampler = WeightFifoSampler( + buffer, + max_staleness_versions=cfg.max_staleness_versions, ) - if isinstance(cfg, InOrderSamplerConfig): - return InOrderSampler(buffer, max_lookahead_versions=cfg.max_lookahead_versions) - if isinstance(cfg, CustomSamplerConfig): + elif isinstance(cfg, InOrderSamplerConfig): + sampler = InOrderSampler( + buffer, + max_lookahead_versions=cfg.max_lookahead_versions, + ) + elif isinstance(cfg, CustomSamplerConfig): module_name, sep, class_name = cfg.target.partition(":") if not sep: raise ValueError( @@ -532,5 +562,17 @@ def create_sampler( f"interface (needs admit/select/evict/should_abort_inflight, " f"set_dispatch_index, is_on_policy, required_buffer_capacity)" ) - return sampler - raise ValueError(f"unknown sampler config {type(cfg).__name__}") + else: + raise ValueError(f"unknown sampler config {type(cfg).__name__}") + + expected_capability = cfg.supports_buffer_checkpoint + if ( + expected_capability is not None + and sampler.supports_buffer_checkpoint != expected_capability + ): + raise RuntimeError( + f"{type(cfg).__name__}.supports_buffer_checkpoint=" + f"{expected_capability} disagrees with {type(sampler).__name__}." + f"supports_buffer_checkpoint={sampler.supports_buffer_checkpoint}" + ) + return sampler diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index cdc73ab971f..9bf15717cd4 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -315,6 +315,9 @@ 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. @@ -329,6 +332,7 @@ def _initial_grpo_save_state() -> GRPOSaveState: total_steps=0, total_valid_tokens=0, val_reward=-99999999.0, + trainer_version=None, sampler_name=None, ) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 6fbe06a0d0e..29414f93cc8 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -43,6 +43,15 @@ import ray import torch +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + TQReplayMetadataState, +) from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler from nemo_rl.algorithms.grpo import GRPOSaveState, _write_latest_checkpoint_status from nemo_rl.algorithms.metric_utils import SetupTimingMetrics @@ -74,9 +83,6 @@ Generation = Union[VllmGeneration, SGLangGeneration] -DATA_PLANE_CHECKPOINT_DIR = "data_plane" -DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 1 - @ray.remote(num_cpus=1, num_gpus=0) # pragma: no cover class SingleControllerActor: @@ -155,6 +161,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[dict[str, Any]] = ( + actor_args.data_plane_checkpoint_metadata + ) self._consumed_samples: int = actor_args.save_state.consumed_samples self._total_valid_tokens: int = actor_args.save_state.total_valid_tokens @@ -162,9 +171,24 @@ 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._master_config.grpo.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) + self._sampler.set_dispatch_index(restored_trainer_version) + if ( + self._master_config.checkpointing["enabled"] + and self._sampler.supports_buffer_checkpoint + and not self._master_config.data_plane.get("checkpointing_enabled") + ): + raise ValueError( + "SingleController checkpointing with a replay-checkpoint-capable " + "sampler requires data_plane.checkpointing_enabled=true so " + "completed, unconsumed rollouts are recoverable." + ) required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step) validate_sampler_buffer_capacity( self._async_cfg, @@ -173,17 +197,17 @@ def __init__( ) # ── asyncio state ────────────────────────────────────────────────── - # TQ snapshots permit concurrent puts but not destructive clears. All - # clears currently owned by async SC use this lock, including rollback - # and eviction through TQReplayBuffer. Clear-dependent eviction waits - # during a save; _buffer_capacity bounds new rollout groups and - # eventually stalls dispatch instead of allowing unbounded TQ growth. + # 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_lock: asyncio.Lock = asyncio.Lock() + self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() if self._buffer is not None: - self._buffer.set_data_plane_checkpoint_lock( - self._data_plane_checkpoint_lock + self._buffer.set_data_plane_checkpoint_barrier( + self._data_plane_checkpoint_barrier ) # Gate: cleared during _sync_weights, set when generation may proceed @@ -210,7 +234,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] = { @@ -276,53 +300,128 @@ 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. + """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 metadata sidecar + 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, + metadata_path = os.path.join( + self._last_checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME + ) + if ( + os.path.exists(metadata_path) + and not self._sampler.supports_buffer_checkpoint + ): + raise RuntimeError( + "The checkpoint contains native replay state, but the configured " + f"sampler {self._async_cfg.sampler.name!r} does not support " + "replay-buffer recovery" ) + if not self._sampler.supports_buffer_checkpoint: return - 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 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. + print(f"📦 Restoring replay buffer metadata: {metadata_path}") + # weights_only=False: the metadata sidecar contains pickled KVBatchMeta + # objects but no rollout tensor payloads. It is a trusted same-job artifact. buffer_state = await asyncio.to_thread( - torch.load, 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._master_config.grpo.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() + async def _validate_replay_inventory( + self, replay_metadata: TQReplayMetadataState + ) -> None: + """Require the canonical TQ keys to match the SC replay index exactly.""" + expected_sample_ids = { + sample_id + for group in replay_metadata["groups"] + for sample_id in group["meta"].sample_ids + } + actual_sample_ids = set( + await call_data_plane( + self._dp_client, + "list_sample_ids", + offload_sync=True, + partition_id=self._partition_id, + ) + ) + missing_sample_ids = sorted(expected_sample_ids - actual_sample_ids) + unexpected_sample_ids = sorted(actual_sample_ids - expected_sample_ids) + if missing_sample_ids or unexpected_sample_ids: + raise RuntimeError( + "Native TQ checkpoint inventory does not match the replay " + "metadata sidecar: " + f"missing={missing_sample_ids[:10]!r} " + f"(total={len(missing_sample_ids)}), " + f"unexpected={unexpected_sample_ids[:10]!r} " + f"(total={len(unexpected_sample_ids)})" + ) + print( + "📦 Native TQ replay inventory validated: " + f"samples={len(actual_sample_ids)}", + flush=True, + ) + async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: """Clear consumed rows without overlapping a data-plane checkpoint.""" - async with self._data_plane_checkpoint_lock: + async with self._data_plane_checkpoint_barrier.mutation(): await call_data_plane( self._dp_client, "clear_samples", @@ -331,13 +430,18 @@ async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: partition_id=self._partition_id, ) - async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: - """Save a required shadow TQ snapshot inside an SC checkpoint bundle. + async def _save_data_plane_checkpoint( + self, + checkpoint_path: str, + replay_metadata: Optional[TQReplayMetadataState] = None, + ) -> None: + """Save a required TQ snapshot inside an SC checkpoint bundle. - Although native TQ restore is not wired into SC yet, opting into this - shadow snapshot is intentionally fail-closed: any failure propagates so - a finalized bundle never silently omits the advertised data-plane - component. + A sampler with replay-buffer recovery writes an authoritative native + TQ snapshot bound to its metadata-only sidecar by a digest. Other + samplers retain shadow-mode snapshots until their recovery contract is + defined. Failures propagate so a finalized bundle never silently omits + the advertised data-plane component. """ checkpoint_dir = os.path.join( checkpoint_path, @@ -351,8 +455,19 @@ async def _save_data_plane_checkpoint(self, checkpoint_path: str) -> None: "single_controller_trainer_version": self._trainer_version, "single_controller_epoch": self._current_epoch, "partition_id": self._partition_id, - "mode": "shadow", + "sampler_name": self._async_cfg.sampler.name, + "mode": "authoritative" if replay_metadata is not None else "shadow", } + if replay_metadata is not None: + metadata.update( + { + "replay_metadata_schema_version": ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION + ), + "replay_manifest_digest": replay_metadata["manifest_digest"], + "replay_group_count": len(replay_metadata["groups"]), + } + ) started = time.monotonic() print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) try: @@ -803,11 +918,10 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: save_state = self._save_state save_state.current_step = self._train_steps save_state.total_steps = self._train_steps + save_state.trainer_version = self._trainer_version save_state.current_epoch = self._current_epoch save_state.consumed_samples = self._consumed_samples save_state.total_valid_tokens = self._total_valid_tokens - # 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. @@ -853,27 +967,27 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: dataloader_state, os.path.join(checkpoint_path, "train_dataloader.pt"), ) - buffer_state: Optional[dict[str, Any]] = None + replay_metadata: Optional[TQReplayMetadataState] = None if self._master_config.data_plane.get("checkpointing_enabled"): - # Capture the legacy replay payload and the native TQ snapshot - # under one clear barrier. Generation puts may continue, so TQ can - # contain a superset of the groups named by replay_buffer.pt. - async with self._data_plane_checkpoint_lock: + # Commits and destructive clears take the same barrier. Generation + # may continue while a snapshot is written, but completed groups + # wait at commit, so TQ and the metadata sidecar describe exactly + # the same set of training-ready groups. + async with self._data_plane_checkpoint_barrier.checkpoint(): if self._sampler.supports_buffer_checkpoint: - buffer_state = await self._buffer.state_dict( + replay_metadata = self._buffer.metadata_state_dict( saved_capacity=self._async_cfg.max_buffered_rollouts ) - await self._save_data_plane_checkpoint(checkpoint_path) - elif self._sampler.supports_buffer_checkpoint: - buffer_state = await self._buffer.state_dict( - saved_capacity=self._async_cfg.max_buffered_rollouts - ) - - if buffer_state is not None: + await self._save_data_plane_checkpoint( + checkpoint_path, replay_metadata=replay_metadata + ) + if replay_metadata is not None: + await self._validate_replay_inventory(replay_metadata) + if replay_metadata is not None: await asyncio.to_thread( torch.save, - buffer_state, - os.path.join(checkpoint_path, "replay_buffer.pt"), + replay_metadata, + os.path.join(checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME), ) # Rename happens in the background once the async weight writes # finish; flushed at the next save or on exit. diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 399ae0203bd..460984d2395 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -32,7 +32,14 @@ from transformers import AutoProcessor from transformers.tokenization_utils_base import PreTrainedTokenizerBase -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + TQReplayBuffer, +) from nemo_rl.algorithms.grpo import ( GRPOSaveState, _create_advantage_estimator, @@ -98,6 +105,94 @@ class SingleControllerActorArgs: partition_id: str save_state: GRPOSaveState last_checkpoint_path: Optional[str] + data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None + + +def _maybe_restore_native_data_plane_checkpoint( + policy: TQPolicy, + *, + last_checkpoint_path: Optional[str], + save_state: GRPOSaveState, + partition_id: str, + sampler_name: str, +) -> Optional[dict[str, Any]]: + """Load and validate an authoritative native TQ checkpoint when present. + + The metadata-only replay sidecar is the format marker. Checkpoints without + any replay artifact resume trainer state with an empty replay buffer; + legacy tensor-bearing replay files are rejected rather than silently + ignored. Rollout tensors are never serialized into a controller-side + replay checkpoint. + """ + if last_checkpoint_path is None: + return None + checkpoint_path = Path(last_checkpoint_path) + replay_metadata_path = checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME + if not replay_metadata_path.is_file(): + legacy_replay_path = checkpoint_path / LEGACY_REPLAY_BUFFER_FILENAME + if legacy_replay_path.is_file(): + raise RuntimeError( + "Checkpoint contains legacy replay_buffer.pt state, which " + "predates authoritative native TQ replay recovery. Resume it " + "with the older implementation or explicitly start without " + "restoring buffered rollouts." + ) + return None + + data_plane_path = checkpoint_path / DATA_PLANE_CHECKPOINT_DIR + if not data_plane_path.is_dir(): + raise FileNotFoundError( + "Metadata-only replay checkpoint requires a matching native TQ " + f"checkpoint at {data_plane_path}" + ) + + print(f"📦 Restoring native TQ checkpoint: {data_plane_path}", flush=True) + metadata = policy.load_data_plane_checkpoint(data_plane_path) + if not isinstance(metadata, dict): + raise TypeError( + "Native TQ checkpoint load must return a metadata dictionary, " + f"got {type(metadata).__name__}" + ) + expected_values: dict[str, Any] = { + "data_plane_checkpoint_schema_version": (DATA_PLANE_CHECKPOINT_SCHEMA_VERSION), + "single_controller_train_steps": save_state.current_step, + "single_controller_trainer_version": ( + save_state.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 _build_clusters( @@ -400,6 +495,16 @@ def setup_single_controller( "data_plane.backend='simple'; Mooncake storage cannot be restored " "by TQ v0.1.9." ) + if ( + master_config.checkpointing["enabled"] + and master_config.async_rl.sampler.supports_buffer_checkpoint is True + and not dp_config.get("checkpointing_enabled") + ): + raise ValueError( + "SingleController checkpointing with a replay-checkpoint-capable " + "sampler requires data_plane.checkpointing_enabled=true so " + "completed, unconsumed rollouts are recoverable." + ) assert generation_config is not None, ( "single_controller_utils.setup requires policy.generation in master_config" @@ -580,6 +685,17 @@ def _build_generation_then_trainer( setup_timing_metrics.policy_init_time_s = trainer_time setup_timing_metrics.generation_init_time_s = gen_reserve_time + gen_load_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: env_handles["nemo_gym"], gym_time = results["nemo_gym"] setup_timing_metrics.nemo_gym_init_time_s = gym_time @@ -660,5 +776,6 @@ 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, ) return actor_args, setup_timing_metrics diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py index f01b40c0986..34d2f76c7de 100644 --- a/nemo_rl/data_plane/adapters/noop.py +++ b/nemo_rl/data_plane/adapters/noop.py @@ -223,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: diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 266fb298591..47b13d59528 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -494,7 +494,7 @@ def _require_clean_for_load(self) -> None: if self._data_operations_started: raise RuntimeError( "load_checkpoint requires a clean TQ client before any " - "register, claim, get, put, clear, or consumption operation" + "register, claim, get, list, put, clear, or consumption operation" ) # ── (A) task-mediated ─────────────────────────────────────────────── @@ -701,6 +701,12 @@ 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: diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 6bb23d1cbd5..465f42b706b 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -60,11 +60,12 @@ class DataPlaneConfig(TypedDict): They are required (not NotRequired) so the YAML carries the full schema and there are no hidden Python defaults. - ``checkpointing_enabled`` opts SingleController into saving required - shadow TQ state inside its checkpoint bundle. Other algorithm entrypoints - do not consume this field. It is optional because existing configs predate - data-plane checkpointing; exemplar configs carry the recommended default - explicitly. + ``checkpointing_enabled`` opts SingleController into saving required TQ + state inside its checkpoint bundle. Samplers that support replay-buffer + recovery pair the native snapshot with a metadata-only local index; other + samplers save it in shadow mode. Other algorithm entrypoints do not consume + this field. It is optional because existing configs predate data-plane + checkpointing; exemplar configs carry the recommended default explicitly. """ enabled: bool @@ -420,6 +421,22 @@ def get_samples( ``TensorDict`` keyed by field name, batched along ``sample_ids``. """ + @abstractmethod + def list_sample_ids(self, partition_id: str) -> list[str]: + """List the sample IDs currently stored in a partition. + + This metadata-only operation is intended for recovery validation and + reconciliation. It must not fetch tensor payloads or advance consumer + cursors. + + Args: + partition_id: Partition whose stored keys should be listed. + + Returns: + Stable, sorted sample IDs. An unknown or empty partition returns + an empty list. + """ + @abstractmethod def clear_samples( self, diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index d7bcd88fca5..0e66d4c1d90 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -323,6 +323,13 @@ def get_samples(self, sample_ids, partition_id, select_fields): n_keys=len(sample_ids), ) + def list_sample_ids(self, partition_id): + return self._run( + "list_sample_ids", + partition_id, + lambda: self._inner.list_sample_ids(partition_id), + ) + def clear_samples(self, sample_ids, partition_id): sample_ids_list = ( sample_ids diff --git a/nemo_rl/models/policy/tq_policy.py b/nemo_rl/models/policy/tq_policy.py index 13f04d7d77c..d11275906e5 100644 --- a/nemo_rl/models/policy/tq_policy.py +++ b/nemo_rl/models/policy/tq_policy.py @@ -33,6 +33,7 @@ from collections import defaultdict from contextlib import nullcontext from dataclasses import replace +from pathlib import Path from typing import Any, Optional import ray @@ -134,6 +135,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/pyrefly.toml b/pyrefly.toml index f4c212c32f3..0bc96c95838 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -116,6 +116,7 @@ project-includes = [ "nemo_rl/data_plane/adapters/__init__.py", "nemo_rl/data_plane/adapters/noop.py", "nemo_rl/data_plane/adapters/transfer_queue.py", + "nemo_rl/data_plane/async_utils.py", "nemo_rl/data_plane/codec.py", "nemo_rl/data_plane/column_io.py", "nemo_rl/data_plane/factory.py", diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 7cb20f7f289..82cabe67478 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -37,6 +37,7 @@ run_test() { run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh run_test uv run --no-sync bash ./tests/functional/grpo_checkpoint_single_controller.sh +run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_dp_single_controller.sh b/tests/functional/grpo_dp_single_controller.sh index 855760445d7..f62ff6aac0f 100755 --- a/tests/functional/grpo_dp_single_controller.sh +++ b/tests/functional/grpo_dp_single_controller.sh @@ -22,7 +22,7 @@ rm -rf $EXP_DIR $LOG_DIR mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT -uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ +uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ $PROJECT_ROOT/examples/run_grpo_single_controller.py \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ @@ -47,11 +47,13 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE $@ \ 2>&1 | tee $RUN_LOG -uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS +if [[ "${RUN_CONVERGENCE_CHECKS:-1}" == "1" ]]; then + uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS -uv run tests/check_metrics.py $JSON_METRICS \ - 'max(data["train/gen_kl_error"]) < 0.002' \ - 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ - 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ - 'max(data["train/probs_ratio_clamped_max"]) < 1.21' + uv run tests/check_metrics.py $JSON_METRICS \ + 'max(data["train/gen_kl_error"]) < 0.002' \ + 'min(data["train/probs_ratio_clamped_min"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_min"]) < 1.21' \ + 'min(data["train/probs_ratio_clamped_max"]) > 0.79' \ + 'max(data["train/probs_ratio_clamped_max"]) < 1.21' +fi diff --git a/tests/functional/grpo_dp_single_controller_tq_recovery.sh b/tests/functional/grpo_dp_single_controller_tq_recovery.sh new file mode 100755 index 00000000000..671b67dd4ec --- /dev/null +++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# Two-process functional test for native TQ + metadata-only replay recovery. + +set -eou pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd) +BASE_TEST=$SCRIPT_DIR/grpo_dp_single_controller.sh +TEST_DIR=$SCRIPT_DIR/grpo_dp_single_controller_tq_recovery +CHECKPOINT_DIR=$TEST_DIR/checkpoints +BASE_RUN_LOG=$SCRIPT_DIR/grpo_dp_single_controller/run.log + +rm -rf "$TEST_DIR" +mkdir -p "$TEST_DIR" + +COMMON_OVERRIDES=( + checkpointing.enabled=true + checkpointing.checkpoint_dir="$CHECKPOINT_DIR" + checkpointing.save_period=1 + data_plane.checkpointing_enabled=true + async_rl.sampler.name=windowed + '~async_rl.sampler.max_lookahead_versions' + '+async_rl.sampler.max_staleness_versions=1' + async_rl.max_inflight_prompts=8 + async_rl.max_buffered_rollouts=8 +) + +echo "=== Phase 1: save an authoritative native TQ checkpoint ===" +# Keep the two-step training horizon identical across both processes so the +# Megatron optimizer scheduler can be restored. The timeout makes phase 1 save +# after its first completed step and exit early, simulating an interrupted job. +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" \ + grpo.max_num_steps=2 \ + checkpointing.checkpoint_must_save_by=0:0:0:1 + +test -d "$CHECKPOINT_DIR/step_1/data_plane" +test -f "$CHECKPOINT_DIR/step_1/replay_buffer_metadata.pt" +test ! -f "$CHECKPOINT_DIR/step_1/replay_buffer.pt" +uv run --no-sync python -c \ + 'import json, sys; metadata = json.load(open(sys.argv[1]))["user_metadata"]; assert metadata["mode"] == "authoritative"; assert metadata["replay_group_count"] > 0, metadata' \ + "$CHECKPOINT_DIR/step_1/data_plane/metadata.json" + +echo "=== Phase 2: start a fresh process, restore TQ, and train one more step ===" +RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \ + "${COMMON_OVERRIDES[@]}" grpo.max_num_steps=2 + +grep -q "Native TQ checkpoint restored and validated" "$BASE_RUN_LOG" +grep -q "Native TQ replay inventory validated" "$BASE_RUN_LOG" +grep -Eq "Restored [1-9][0-9]* replay group" "$BASE_RUN_LOG" +test -d "$CHECKPOINT_DIR/step_2/data_plane" +test -f "$CHECKPOINT_DIR/step_2/replay_buffer_metadata.pt" + +echo "Native TQ recovery functional test passed." diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py index 5ff2d75fd28..ab961960098 100644 --- a/tests/unit/data_plane/test_architecture_invariants.py +++ b/tests/unit/data_plane/test_architecture_invariants.py @@ -78,6 +78,7 @@ 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", diff --git a/tests/unit/data_plane/test_async_utils.py b/tests/unit/data_plane/test_async_utils.py index 25d29b5d711..8fd8f4d1a3f 100644 --- a/tests/unit/data_plane/test_async_utils.py +++ b/tests/unit/data_plane/test_async_utils.py @@ -61,9 +61,7 @@ def test_sync_call_can_be_offloaded() -> None: def test_local_coroutine_result_is_awaited() -> None: - result = asyncio.run( - call_data_plane(_LocalClient(), "async_value", value=7) - ) + result = asyncio.run(call_data_plane(_LocalClient(), "async_value", value=7)) assert result == 7 diff --git a/tests/unit/data_plane/test_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py index 3a4009b164b..7c9c9002599 100644 --- a/tests/unit/data_plane/test_interface_contract.py +++ b/tests/unit/data_plane/test_interface_contract.py @@ -65,8 +65,10 @@ def test_register_put_get_clear(client: DataPlaneClient): out = client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) assert torch.equal(out["x"], torch.arange(4)) + assert client.list_sample_ids("p") == keys client.clear_samples(sample_ids=None, partition_id="p") + assert client.list_sample_ids("p") == [] with pytest.raises(KeyError): client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"]) diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py index 2cdbdb1a626..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 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 0767d6cf6c0..0f7c79594e4 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -128,6 +128,23 @@ def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: 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, diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 5fe4eaca44b..7838a759a33 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -512,7 +512,9 @@ data_plane: storage_capacity: 1000000 # max samples retained per partition num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - checkpointing_enabled: false # SingleController only: save required shadow TQ state + # SingleController only: save native TQ state. Supported samplers restore + # from metadata-only replay indexes. + checkpointing_enabled: false global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" # observability: # NotRequired diff --git a/tests/unit/single_controller/_dp_fakes.py b/tests/unit/single_controller/_dp_fakes.py index 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_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index e6e648f5507..7168f4a88b6 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -27,10 +27,12 @@ import pytest from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, InOrderSampler, InOrderSamplerConfig, PromptGroupSampler, WeightFifoSampler, + WeightFifoSamplerConfig, WindowedSampler, WindowedSamplerConfig, create_sampler, @@ -151,6 +153,24 @@ def test_unready_slot_is_never_evicted(self): class TestFactory: + @pytest.mark.parametrize( + ("config", "expected"), + [ + (WindowedSamplerConfig(), True), + (WeightFifoSamplerConfig(), False), + (InOrderSamplerConfig(), False), + ( + CustomSamplerConfig(target=f"{__name__}:EchoSampler"), + None, + ), + ], + ) + def test_config_declares_checkpoint_capability_without_serializing_it( + self, config, expected + ): + assert config.supports_buffer_checkpoint is expected + assert "supports_buffer_checkpoint" not in config.model_dump() + def test_windowed_config_builds_windowed(self): s = create_sampler( FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3) @@ -164,24 +184,25 @@ def test_in_order_config_builds_in_order(self): assert s.max_lookahead_versions == 2 def test_weight_fifo_config_builds_weight_fifo(self): - from nemo_rl.algorithms.async_utils.staleness_sampler import ( - WeightFifoSamplerConfig, - ) - s = create_sampler( FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4) ) assert isinstance(s, WeightFifoSampler) assert s.max_staleness_versions == 4 + def test_factory_rejects_config_implementation_capability_drift(self, monkeypatch): + monkeypatch.setattr( + WindowedSampler, + "supports_buffer_checkpoint", + property(lambda _self: False), + ) + with pytest.raises(RuntimeError, match="disagrees with"): + create_sampler(FakeBuffer(), WindowedSamplerConfig()) + class TestCustomFqnSampler: def test_custom_target_loads_out_of_repo_sampler(self): # A user sampler defined anywhere importable; here, this test module. - from nemo_rl.algorithms.async_utils.staleness_sampler import ( - CustomSamplerConfig, - ) - s = create_sampler( FakeBuffer(), CustomSamplerConfig( diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 3475a313414..c744c8d7d94 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_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). """ @@ -48,7 +47,18 @@ 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 ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + WindowedSamplerConfig, +) from nemo_rl.algorithms.grpo import ( GRPOConfig, GRPOSaveState, @@ -159,7 +169,8 @@ 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 async def admit(self, *, trainer_version_fn) -> Optional[int]: @@ -191,6 +202,10 @@ 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 @@ -212,16 +227,63 @@ 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 metadata sidecar.""" + + 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, *, save_error: Optional[Exception] = None) -> 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, @@ -281,23 +343,32 @@ 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_lock: Optional[asyncio.Lock] = None + self.checkpoint_barrier: Optional[DataPlaneCheckpointBarrier] = None - def set_data_plane_checkpoint_lock(self, lock: asyncio.Lock) -> None: - self.checkpoint_lock = lock + def set_data_plane_checkpoint_barrier( + self, barrier: DataPlaneCheckpointBarrier + ) -> None: + self.checkpoint_barrier = barrier - async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]: - self.state_dict_calls.append(saved_capacity) - return dict(self._state) + def metadata_state_dict(self, *, saved_capacity: int) -> dict[str, Any]: + self.metadata_state_dict_calls.append(saved_capacity) + return dict(self._metadata_state) async def load_state_dict( self, @@ -306,6 +377,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( { @@ -313,6 +385,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 @@ -356,6 +429,7 @@ 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 = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -363,7 +437,11 @@ def _actor_master_config( 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. @@ -422,6 +500,7 @@ def _make_actor_args( tq_buffer: Optional[_FakeTQBuffer] = None, dp_client: Optional[_FakeDPClient] = None, last_checkpoint_path: Optional[str] = None, + data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=object(), @@ -441,6 +520,7 @@ 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, ) @@ -458,7 +538,9 @@ def _run_train_pump( async def _main(): actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) - actor._sampler = _FakeSampler() + actor._sampler = _FakeSampler( + supports_buffer_checkpoint=(mc.async_rl.sampler.name == "windowed") + ) # In-process runs have no Ray runtime; the pump only reads the GPU # count for a throughput metric. with patch("ray.cluster_resources", return_value={"GPU": 0}): @@ -488,7 +570,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. @@ -500,7 +585,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() @@ -548,6 +633,21 @@ def test_restore_from_step_n(self, tmp_path): 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_fresh_start_defaults(self, tmp_path): actor = _ACTOR_CLS( _actor_master_config(tmp_path), _make_actor_args(), SetupTimingMetrics() @@ -620,6 +720,7 @@ 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["total_steps"] == 2 assert info_2["consumed_samples"] == 4 # 2 prompts/step * 2 steps # No validation ran, so the default val_reward is dropped. @@ -749,16 +850,46 @@ def test_ft_save_period_triggers_saves(self, tmp_path): assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"} -class TestDataPlaneShadowCheckpoint: - def test_saves_tq_state_and_keeps_legacy_replay_payload(self, tmp_path): +class TestDataPlaneCheckpoint: + def test_saves_authoritative_tq_state_and_metadata_only_replay_index( + self, tmp_path + ): mc = _actor_master_config( tmp_path, max_num_steps=1, save_period=1, + buffer_checkpoint=True, data_plane_checkpoint=True, ) - dp_client = _FakeDPClient() - buffer = _FakeTQBuffer(state={"legacy_payload": "kept"}) + 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, @@ -771,19 +902,108 @@ def test_saves_tq_state_and_keeps_legacy_replay_payload(self, tmp_path): tmp_path / "checkpoints" / "tmp_step_1" / "data_plane" ) assert save_call["metadata"] == { - "data_plane_checkpoint_schema_version": 1, + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), "single_controller_train_steps": 1, "single_controller_trainer_version": 1, "single_controller_epoch": 0, "partition_id": _PARTITION_ID, - "mode": "shadow", + "sampler_name": "windowed", + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "replay_manifest_digest": "digest-1", + "replay_group_count": 1, } step_dir = tmp_path / "checkpoints" / "step_1" assert (step_dir / "data_plane" / "metadata.json").is_file() - assert torch.load(step_dir / "replay_buffer.pt", weights_only=False) == { - "legacy_payload": "kept" + assert ( + torch.load(step_dir / REPLAY_BUFFER_METADATA_FILENAME, weights_only=False) + == replay_metadata + ) + assert not (step_dir / "replay_buffer.pt").exists() + assert buffer.metadata_state_dict_calls == [4] + + @pytest.mark.parametrize( + ("actual_sample_ids", "error_fragment"), + [ + (["g0-0"], r"missing=\['g0-1'\]"), + ( + ["g0-0", "g0-1", "orphan-0"], + r"unexpected=\['orphan-0'\]", + ), + ], + ) + def test_tq_save_rejects_inventory_mismatch( + self, tmp_path, actual_sample_ids, error_fragment + ): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + sample_ids = ["g0-0", "g0-1"] + replay_metadata = { + "schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "storage": REPLAY_BUFFER_METADATA_STORAGE, + "partition_id": _PARTITION_ID, + "saved_capacity": 4, + "manifest_digest": "digest-1", + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name="train", + sample_ids=sample_ids, + fields=["input_ids"], + sequence_lengths=[16, 16], + tags=[ + {"weight_version": 0}, + {"weight_version": 0}, + ], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": None, + "group_id": "g0", + } + ], } - assert buffer.state_dict_calls == [4] + + with pytest.raises(RuntimeError, match=error_fragment): + _run_train_pump( + mc, + _make_actor_args( + dp_client=_FakeDPClient(sample_ids=actual_sample_ids), + tq_buffer=_FakeTQBuffer(metadata_state=replay_metadata), + ), + ) + + assert not (tmp_path / "checkpoints" / "step_1").exists() + + def test_gated_sampler_keeps_tq_checkpoint_in_shadow_mode(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=1, + save_period=1, + buffer_checkpoint=False, + data_plane_checkpoint=True, + ) + dp_client = _FakeDPClient() + buffer = _FakeTQBuffer() + + _run_train_pump( + mc, + _make_actor_args(dp_client=dp_client, tq_buffer=buffer), + ) + + assert dp_client.save_calls[0]["metadata"]["mode"] == "shadow" + step_dir = tmp_path / "checkpoints" / "step_1" + assert not (step_dir / REPLAY_BUFFER_METADATA_FILENAME).exists() + assert not (step_dir / "replay_buffer.pt").exists() + assert buffer.metadata_state_dict_calls == [] def test_tq_save_failure_aborts_checkpoint(self, tmp_path): mc = _actor_master_config( @@ -809,7 +1029,9 @@ def test_consumed_clear_waits_for_tq_save(self, tmp_path): dp_client = _BlockingDPClient() async def _main() -> None: - actor = _ACTOR_CLS(mc, _make_actor_args(dp_client=dp_client)) + 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})) @@ -835,7 +1057,9 @@ def test_consumed_clear_does_not_block_actor_event_loop(self, tmp_path): dp_client = _FakeDPClient() async def _main() -> int: - actor = _ACTOR_CLS(mc, _make_actor_args(dp_client=dp_client)) + 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() @@ -1227,44 +1451,97 @@ 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.""" - save_state = _initial_grpo_save_state() - save_state.sampler_name = "windowed" - return save_state +class TestReplayBufferPersistence: + def test_checkpoint_capable_sampler_without_native_tq_is_rejected(self, tmp_path): + mc = _actor_master_config( + tmp_path, + max_num_steps=2, + save_period=2, + buffer_checkpoint=True, + ) + with pytest.raises( + ValueError, + match="replay-checkpoint-capable sampler requires", + ): + _ACTOR_CLS(mc, _make_actor_args(), SetupTimingMetrics()) -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_no_replay_buffer_with_gated_sampler(self, tmp_path): + mc = _actor_master_config( + tmp_path, max_num_steps=2, save_period=2, buffer_checkpoint=False + ) + buffer = _FakeTQBuffer() _run_train_pump(mc, _make_actor_args(tq_buffer=buffer)) ckpt_dir = tmp_path / "checkpoints" - 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 not (ckpt_dir / "step_2" / REPLAY_BUFFER_METADATA_FILENAME).exists() + + def test_run_rejects_legacy_replay_file(self, tmp_path): ckpt_dir = tmp_path / "resume_ckpt" ckpt_dir.mkdir() - 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 = { + "replay_manifest_digest": "digest-1", + "replay_group_count": 2, + } + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, + ) + buffer = _FakeTQBuffer(load_return=2) actor, result = _run_actor_run( mc, _make_actor_args( tq_buffer=buffer, + dp_client=_FakeDPClient(sample_ids=sample_ids), last_checkpoint_path=str(ckpt_dir), - save_state=_matching_save_state(), + data_plane_checkpoint_metadata=tq_metadata, ), ) @@ -1274,10 +1551,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 def test_restored_permits_are_released_by_a_live_pump(self, tmp_path): @@ -1289,17 +1566,46 @@ 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 = { + "replay_manifest_digest": "digest-1", + "replay_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 @@ -1308,39 +1614,90 @@ 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) - buffer = _FakeTQBuffer() - printed: list[str] = [] - monkeypatch.setattr( - "builtins.print", - lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), + envelope = { + "groups": [ + { + "meta": KVBatchMeta( + partition_id=_PARTITION_ID, + task_name=None, + sample_ids=["g0-0", "g0-1"], + sequence_lengths=[16, 16], + tags=[{"weight_version": 0}, {"weight_version": 0}], + ), + "start_weight": 0, + "end_weight": 0, + "target_step": 0, + "group_id": "g0", + } + ] + } + torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + tq_metadata = { + "replay_manifest_digest": "digest-1", + "replay_group_count": 1, + } + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + buffer_checkpoint=True, + data_plane_checkpoint=True, ) - actor, _ = _run_actor_run( - mc, - _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)), + 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, ) - 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) + with pytest.raises(RuntimeError, match="native TQ checkpoint was not restored"): + _run_actor_run( + mc, + _make_actor_args(last_checkpoint_path=str(ckpt_dir)), + ) - def test_run_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_missing_native_replay_metadata_starts_empty( + self, tmp_path, monkeypatch + ): 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) + 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", @@ -1349,13 +1706,21 @@ def test_run_no_restore_on_sampler_mismatch(self, tmp_path, monkeypatch): actor, _ = _run_actor_run( mc, - _make_actor_args( - tq_buffer=buffer, - last_checkpoint_path=str(ckpt_dir), - save_state=save_state, - ), + _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)), ) assert buffer.load_calls == [] - assert actor._buffer_capacity._value == 4 - assert any("skipping the buffer restore" in line for line in printed) + assert actor._buffer_capacity._value == 4 # zero permits consumed + assert any("No native replay metadata found" in line for line in printed) + + def test_run_rejects_native_replay_state_with_gated_sampler(self, tmp_path): + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + torch.save({"groups": []}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) + mc = _actor_master_config(tmp_path, max_num_steps=0, buffer_checkpoint=False) + + with pytest.raises(RuntimeError, match="does not support replay-buffer"): + _run_actor_run( + mc, + _make_actor_args(last_checkpoint_path=str(ckpt_dir)), + ) diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 83dc8d2cb63..9e011a311e8 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -16,12 +16,26 @@ from __future__ import annotations +from typing import Any, Optional from unittest.mock import MagicMock, patch import pytest +import torch import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod -from nemo_rl.algorithms.grpo import GRPOConfig +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DATA_PLANE_CHECKPOINT_DIR, + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + LEGACY_REPLAY_BUFFER_FILENAME, + REPLAY_BUFFER_METADATA_FILENAME, + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, +) +from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.grpo import ( + GRPOConfig, + GRPOSaveState, + _initial_grpo_save_state, +) from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.single_controller_utils import ( AsyncRLConfig, @@ -99,6 +113,35 @@ def _make_master_config( ) +def _native_tq_metadata( + *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1 +) -> dict[str, Any]: + return { + "data_plane_checkpoint_schema_version": (DATA_PLANE_CHECKPOINT_SCHEMA_VERSION), + "single_controller_train_steps": step, + "single_controller_trainer_version": ( + step if trainer_version is None else trainer_version + ), + "single_controller_epoch": epoch, + "partition_id": "rollout_data", + "sampler_name": "in_order", + "mode": "authoritative", + "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + "replay_manifest_digest": "digest-1", + "replay_group_count": 2, + } + + +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. @@ -235,6 +278,26 @@ def test_rejects_mooncake_data_plane_checkpointing(self): with pytest.raises(NotImplementedError, match="backend='simple'"): setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_rejects_windowed_checkpointing_without_native_tq(self): + mc = _make_master_config() + mc.checkpointing["enabled"] = True + mc.async_rl.sampler = WindowedSamplerConfig(max_staleness_versions=1) + mc.data_plane.update( + { + "backend": "simple", + "checkpointing_enabled": False, + } + ) + + with pytest.raises( + ValueError, + match=( + "replay-checkpoint-capable sampler requires " + "data_plane.checkpointing_enabled=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_multiple_dataloader_not_supported(self): mc = _make_master_config(use_multiple_dataloader=True) with pytest.raises(NotImplementedError, match="use_multiple_dataloader"): @@ -641,3 +704,139 @@ def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend): ): setup_single_controller(mc, MagicMock(pad_token_id=0)) mock_spinup.assert_not_called() + + +class TestNativeTQRecoverySetup: + def test_setup_loads_tq_before_creating_single_controller_client( + self, tmp_path, patched_factories + ): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + 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_sidecar_exists( + self, tmp_path + ): + checkpoint_path = tmp_path / "step_3" + (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True) + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + policy = MagicMock() + metadata = _native_tq_metadata() + policy.load_data_plane_checkpoint.return_value = metadata + save_state = _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): + 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() + + def test_metadata_sidecar_requires_matching_tq_directory(self, tmp_path): + checkpoint_path = tmp_path / "step_3" + checkpoint_path.mkdir() + (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch() + + with pytest.raises(FileNotFoundError, match="matching native TQ checkpoint"): + sc_setup_mod._maybe_restore_native_data_plane_checkpoint( + MagicMock(), + last_checkpoint_path=str(checkpoint_path), + save_state=_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_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index d3d6f813a01..3e700c7eb16 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -17,16 +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 TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLAY_BUFFER_METADATA_STORAGE, + DataPlaneCheckpointBarrier, + TQReplayBuffer, + replay_manifest_digest, +) from nemo_rl.data_plane import KVBatchMeta from nemo_rl.distributed.batched_data_dict import BatchedDataDict from nemo_rl.experience.interfaces import PromptGroupRecord @@ -104,6 +108,10 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None 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 +127,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: @@ -160,7 +168,7 @@ def _make_buffer( dp: FakeDataPlaneClient, *, require_routed_experts: bool = False, - checkpoint_lock: asyncio.Lock | None = None, + checkpoint_barrier: DataPlaneCheckpointBarrier | None = None, ) -> TQReplayBuffer: buffer = TQReplayBuffer( dp, @@ -168,7 +176,9 @@ def _make_buffer( pad_value_dict={"token_ids": 0}, require_routed_experts=require_routed_experts, ) - buffer.set_data_plane_checkpoint_lock(checkpoint_lock or asyncio.Lock()) + buffer.set_data_plane_checkpoint_barrier( + checkpoint_barrier or DataPlaneCheckpointBarrier() + ) return buffer @@ -191,7 +201,87 @@ def _add_group( ) +class TestDataPlaneCheckpointBarrier: + def test_mutations_run_concurrently_without_checkpoint(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + both_entered = asyncio.Event() + release = asyncio.Event() + active = 0 + + async def mutate() -> None: + nonlocal active + async with barrier.mutation(): + active += 1 + if active == 2: + both_entered.set() + await release.wait() + active -= 1 + + tasks = [asyncio.create_task(mutate()) for _ in range(2)] + await asyncio.wait_for(both_entered.wait(), timeout=5.0) + assert active == 2 + release.set() + await asyncio.gather(*tasks) + + asyncio.run(exercise()) + + def test_checkpoint_waits_for_active_mutation(self): + async def exercise() -> None: + barrier = DataPlaneCheckpointBarrier() + mutation_entered = asyncio.Event() + release_mutation = asyncio.Event() + checkpoint_entered = asyncio.Event() + + async def mutate() -> None: + async with barrier.mutation(): + mutation_entered.set() + await release_mutation.wait() + + async def checkpoint() -> None: + async with barrier.checkpoint(): + checkpoint_entered.set() + + mutation_task = asyncio.create_task(mutate()) + await mutation_entered.wait() + checkpoint_task = asyncio.create_task(checkpoint()) + await asyncio.sleep(0) + assert not checkpoint_entered.is_set() + + release_mutation.set() + await asyncio.gather(mutation_task, checkpoint_task) + assert checkpoint_entered.is_set() + + asyncio.run(exercise()) + + class TestTQReplayBufferReserveCommit: + def test_commit_waits_for_active_checkpoint(self): + async def exercise() -> None: + dp = FakeDataPlaneClient() + checkpoint_barrier = DataPlaneCheckpointBarrier() + buf = _make_buffer(dp, checkpoint_barrier=checkpoint_barrier) + group_id = buf.reserve(weight_version=3) + + async with checkpoint_barrier.checkpoint(): + commit_task = asyncio.create_task( + buf.commit( + group_id, + _make_record(), + start_weight_version=3, + end_weight_version=3, + ) + ) + await asyncio.sleep(0) + assert dp.put_calls == [] + assert buf.ready_list == [False] + + await commit_task + assert len(dp.put_calls) == 1 + assert buf.ready_list == [True] + + asyncio.run(exercise()) + def test_commit_clears_rows_when_put_raises_after_writing(self): dp = FailAfterPutDataPlaneClient() buf = _make_buffer(dp) @@ -353,7 +443,7 @@ def test_commit_appends_multiple_records_in_order(self): class TestTQReplayBufferRemove: - def test_dp_clear_fails_without_bound_checkpoint_lock(self): + def test_dp_clear_fails_without_bound_checkpoint_barrier(self): dp = FakeDataPlaneClient() buf = TQReplayBuffer( dp, @@ -366,11 +456,11 @@ def test_dp_clear_fails_without_bound_checkpoint_lock(self): assert dp.clear_calls == [] - def test_dp_clear_waits_for_bound_checkpoint_lock(self): + def test_dp_clear_waits_for_active_checkpoint(self): async def exercise() -> None: dp = FakeDataPlaneClient() - checkpoint_lock = asyncio.Lock() - buf = _make_buffer(dp, checkpoint_lock=checkpoint_lock) + checkpoint_barrier = DataPlaneCheckpointBarrier() + buf = _make_buffer(dp, checkpoint_barrier=checkpoint_barrier) group_id = buf.reserve(weight_version=0) await buf.commit( group_id, @@ -379,12 +469,11 @@ async def exercise() -> None: end_weight_version=0, ) - await checkpoint_lock.acquire() - remove_task = asyncio.create_task(buf.remove([0], remove_in_dp=True)) - await asyncio.sleep(0) - assert dp.clear_calls == [] + 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 == [] - checkpoint_lock.release() await remove_task assert dp.clear_calls == [dp.put_calls[0]["sample_ids"]] @@ -535,20 +624,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, } @@ -559,74 +651,61 @@ 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 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: @@ -636,7 +715,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 @@ -648,7 +727,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() @@ -666,7 +745,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"]] == [ @@ -678,37 +757,6 @@ 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) - - 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]) - class TestTQReplayBufferLoadPreflight: """Malformed envelopes raise ValueError before any DataPlane write.""" @@ -725,20 +773,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): @@ -746,75 +800,25 @@ 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) - - 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_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", ) - dp = FakeDataPlaneClient() - buf = _make_buffer(dp) - - with pytest.raises(ValueError, match="max_buffered_rollouts >= 3"): - _load(buf, state, max_groups=2) - - # Preflight semantics: nothing reached the DataPlane or the buffer. - 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, + 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)] ) - 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, + self._assert_rejected( + state, + match="more replay groups than the current buffer capacity", + max_groups=2, + expected_manifest_digest=state["manifest_digest"], ) - buf = _make_buffer(FakeDataPlaneClient()) - - assert _load(buf, state, max_groups=2) == 2 - assert buf.target_step_list == [1, 2] From 9eb253afe567936c96bd7923dcb67492fcb25fdd Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 13 Aug 2026 23:43:32 -0400 Subject: [PATCH 05/32] fix(sc): address TQ recovery review feedback Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 73 ++++++-- .../async_utils/staleness_sampler.py | 114 +++++++----- nemo_rl/algorithms/single_controller.py | 16 +- .../single_controller_utils/setup.py | 12 +- nemo_rl/data_plane/__init__.py | 2 + nemo_rl/data_plane/interfaces.py | 3 + nemo_rl/data_plane/observability.py | 2 +- .../test_sampler_interface.py | 72 ++++++-- .../test_sc_checkpointing.py | 14 +- .../test_single_controller_setup.py | 35 +++- .../test_tq_replay_buffer.py | 50 ++++++ .../test_verify_tq_data_plane_checkpoint.py | 164 ++++++++++++++++++ tools/verify_tq_data_plane_checkpoint.py | 21 ++- 13 files changed, 481 insertions(+), 97 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 6e62fecedd7..b02f6f8db03 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -16,19 +16,21 @@ import gc import hashlib import json +import math import statistics import threading as _threading import uuid from collections import Counter from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager +from numbers import Integral, Real from typing import Any, Iterable, Literal, 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 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 ROUTED_EXPERTS_FIELD from nemo_rl.experience.interfaces import ( @@ -39,14 +41,18 @@ 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" -DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 REPLAY_BUFFER_METADATA_FILENAME = "replay_buffer_metadata.pt" LEGACY_REPLAY_BUFFER_FILENAME = "replay_buffer.pt" REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1 REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint" +# 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.""" @@ -69,6 +75,38 @@ class TQReplayMetadataState(TypedDict): groups: list[TQReplayGroupMetadata] +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 = [ @@ -91,20 +129,35 @@ def replay_manifest_digest(groups: list[TQReplayGroupMetadata]) -> str: if group["meta"].sequence_lengths is not None else None ), - "tags": group["meta"].tags, - "extra_info": group["meta"].extra_info, + "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 in groups + for group_index, group in enumerate(groups) ] - encoded = json.dumps(digest_input, sort_keys=True, separators=(",", ":")).encode( - "utf-8" - ) + encoded = json.dumps( + digest_input, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") return hashlib.sha256(encoded).hexdigest() class DataPlaneCheckpointBarrier: - """Allow concurrent mutations while giving checkpoints exclusive access.""" + """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() diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index fa234e886bf..649100bce86 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -110,16 +110,14 @@ def is_on_policy(self) -> bool: """True when the policy admits zero staleness (sync mode).""" ... - @property - def supports_buffer_checkpoint(self) -> bool: - """Whether completed buffered groups can be restored safely.""" - ... + 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: + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: """Seed the dispatch cursor when resuming from a checkpoint.""" ... @@ -132,28 +130,30 @@ 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: + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: """Seed the dispatch cursor when resuming from a checkpoint. 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 weight version this run starts + from — 0 for a fresh run, the restored trainer version when + resuming. Sets the cursor to one before that version so gated + ``admit`` and ``InOrderSampler`` target-step stamps line up with + the restored trainer version. Call before the first ``admit``. """ - 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 # ── rollout-pump side ──────────────────────────────────────────────── @abc.abstractmethod @@ -202,10 +202,6 @@ def should_abort_inflight( def is_on_policy(self) -> bool: return self._eviction_window() == 0 - @property - def supports_buffer_checkpoint(self) -> bool: - return False - def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: return None @@ -257,6 +253,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, @@ -276,11 +275,6 @@ def __init__( def _eviction_window(self) -> int: return self.max_staleness_versions - @property - def supports_buffer_checkpoint(self) -> bool: - # Ungated restored groups are ordinary in-window candidates. - return True - def should_abort_inflight( self, *, @@ -458,7 +452,6 @@ async def evict(self, *, current_train_weight: int) -> int: class WindowedSamplerConfig(BaseModel, extra="allow"): - supports_buffer_checkpoint: ClassVar[bool] = True name: Literal["windowed"] = "windowed" # Max weight-version gap a selected group may have from the trainer. max_staleness_versions: NonNegativeInt = 1 @@ -467,25 +460,24 @@ class WindowedSamplerConfig(BaseModel, extra="allow"): class WeightFifoSamplerConfig(BaseModel, extra="allow"): - supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["weight_fifo"] = "weight_fifo" # Lookahead + selectable weight window, in trainer versions. max_staleness_versions: NonNegativeInt = 1 class InOrderSamplerConfig(BaseModel, extra="allow"): - supports_buffer_checkpoint: ClassVar[bool] = False name: Literal["in_order"] = "in_order" # How far generation may run ahead of the trainer, in dispatch batches. max_lookahead_versions: NonNegativeInt = 1 class CustomSamplerConfig(BaseModel, extra="allow"): - # A custom implementation's capability is known only after construction. - supports_buffer_checkpoint: ClassVar[Optional[bool]] = None name: Literal["custom"] = "custom" # "module:ClassName" of a PromptGroupSampler defined outside this repo. - # Extra keys are forwarded to the constructor (after ``buffer``). + # 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 @@ -521,6 +513,46 @@ 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_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: type + if isinstance(cfg, WindowedSamplerConfig): + sampler_cls = WindowedSampler + elif isinstance(cfg, WeightFifoSamplerConfig): + sampler_cls = WeightFifoSampler + elif isinstance(cfg, InOrderSamplerConfig): + sampler_cls = InOrderSampler + elif isinstance(cfg, CustomSamplerConfig): + sampler_cls = _custom_sampler_class(cfg) + else: + raise ValueError(f"unknown sampler config {type(cfg).__name__}") + + 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, @@ -549,30 +581,16 @@ def create_sampler( max_lookahead_versions=cfg.max_lookahead_versions, ) elif isinstance(cfg, CustomSamplerConfig): - module_name, sep, class_name = cfg.target.partition(":") - if not sep: - raise ValueError( - f"custom sampler target must be 'module:ClassName', got {cfg.target!r}" - ) - sampler_cls = getattr(importlib.import_module(module_name), class_name) + sampler_cls = _custom_sampler_class(cfg) + 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"set_dispatch_index, is_on_policy, supports_buffer_checkpoint, " + f"required_buffer_capacity)" ) else: raise ValueError(f"unknown sampler config {type(cfg).__name__}") - - expected_capability = cfg.supports_buffer_checkpoint - if ( - expected_capability is not None - and sampler.supports_buffer_checkpoint != expected_capability - ): - raise RuntimeError( - f"{type(cfg).__name__}.supports_buffer_checkpoint=" - f"{expected_capability} disagrees with {type(sampler).__name__}." - f"supports_buffer_checkpoint={sampler.supports_buffer_checkpoint}" - ) return sampler diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 29414f93cc8..ab35fa5a8a3 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -45,11 +45,10 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( DATA_PLANE_CHECKPOINT_DIR, - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, - DataPlaneCheckpointBarrier, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointBarrier, TQReplayMetadataState, ) from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler @@ -70,7 +69,7 @@ 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 @@ -388,7 +387,12 @@ async def _maybe_restore_replay_buffer(self) -> None: async def _validate_replay_inventory( self, replay_metadata: TQReplayMetadataState ) -> None: - """Require the canonical TQ keys to match the SC replay index exactly.""" + """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"] @@ -978,11 +982,11 @@ async def _save_checkpoint(self, step_metrics: dict[str, Any]) -> None: 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) await self._save_data_plane_checkpoint( checkpoint_path, replay_metadata=replay_metadata ) - if replay_metadata is not None: - await self._validate_replay_inventory(replay_metadata) if replay_metadata is not None: await asyncio.to_thread( torch.save, diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 460984d2395..510ec80727c 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -34,12 +34,14 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( DATA_PLANE_CHECKPOINT_DIR, - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, TQReplayBuffer, ) +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + sampler_supports_buffer_checkpoint, +) from nemo_rl.algorithms.grpo import ( GRPOSaveState, _create_advantage_estimator, @@ -60,7 +62,11 @@ 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, +) from nemo_rl.distributed.virtual_cluster import RayVirtualCluster from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor @@ -497,7 +503,7 @@ def setup_single_controller( ) if ( master_config.checkpointing["enabled"] - and master_config.async_rl.sampler.supports_buffer_checkpoint is True + and sampler_supports_buffer_checkpoint(master_config.async_rl.sampler) and not dp_config.get("checkpointing_enabled") ): raise ValueError( diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py index 56b19178a1c..36574328e88 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -21,6 +21,7 @@ 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, @@ -28,6 +29,7 @@ from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event __all__ = [ + "DATA_PLANE_CHECKPOINT_SCHEMA_VERSION", "DataPlaneClient", "DataPlaneConfig", "KVBatchMeta", diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 465f42b706b..4b354ba397a 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -43,6 +43,9 @@ from tensordict import TensorDict +DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 + + class DataPlaneConfig(TypedDict): """Feature-gated config; defaults to disabled. diff --git a/nemo_rl/data_plane/observability.py b/nemo_rl/data_plane/observability.py index 0e66d4c1d90..d569740cf18 100644 --- a/nemo_rl/data_plane/observability.py +++ b/nemo_rl/data_plane/observability.py @@ -323,7 +323,7 @@ def get_samples(self, sample_ids, partition_id, select_fields): n_keys=len(sample_ids), ) - def list_sample_ids(self, partition_id): + def list_sample_ids(self, partition_id: str) -> list[str]: return self._run( "list_sample_ids", partition_id, diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 7168f4a88b6..b89625af734 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -36,6 +36,7 @@ WindowedSampler, WindowedSamplerConfig, create_sampler, + sampler_supports_buffer_checkpoint, ) from nemo_rl.data_plane import KVBatchMeta @@ -161,14 +162,12 @@ class TestFactory: (InOrderSamplerConfig(), False), ( CustomSamplerConfig(target=f"{__name__}:EchoSampler"), - None, + False, ), ], ) - def test_config_declares_checkpoint_capability_without_serializing_it( - self, config, expected - ): - assert config.supports_buffer_checkpoint is expected + 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): @@ -190,17 +189,37 @@ def test_weight_fifo_config_builds_weight_fifo(self): assert isinstance(s, WeightFifoSampler) assert s.max_staleness_versions == 4 - def test_factory_rejects_config_implementation_capability_drift(self, monkeypatch): - monkeypatch.setattr( - WindowedSampler, - "supports_buffer_checkpoint", - property(lambda _self: False), + 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", + ) ) - with pytest.raises(RuntimeError, match="disagrees with"): - create_sampler(FakeBuffer(), WindowedSamplerConfig()) + assert not CheckpointingEchoSampler.constructed 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. s = create_sampler( @@ -352,7 +371,7 @@ def test_fresh_start_is_a_noop_seed(self): 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"): + with pytest.raises(ValueError, match="resume_from_trainer_version"): WindowedSampler(FakeBuffer(), max_staleness_versions=1).set_dispatch_index( -1 ) @@ -404,3 +423,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_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index c744c8d7d94..82a69ac990e 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -49,11 +49,11 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, - DataPlaneCheckpointBarrier, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, + DataPlaneCheckpointBarrier, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, @@ -209,7 +209,7 @@ def supports_buffer_checkpoint(self) -> bool: def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: return None - def set_dispatch_index(self, resume_from_step: int) -> None: + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: pass @@ -1686,9 +1686,7 @@ def test_native_replay_metadata_requires_setup_side_tq_restore(self, tmp_path): _make_actor_args(last_checkpoint_path=str(ckpt_dir)), ) - def test_run_missing_native_replay_metadata_starts_empty( - self, tmp_path, monkeypatch - ): + 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( @@ -1698,11 +1696,6 @@ def test_run_missing_native_replay_metadata_starts_empty( data_plane_checkpoint=True, ) buffer = _FakeTQBuffer() - printed: list[str] = [] - monkeypatch.setattr( - "builtins.print", - lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)), - ) actor, _ = _run_actor_run( mc, @@ -1711,7 +1704,6 @@ def test_run_missing_native_replay_metadata_starts_empty( assert buffer.load_calls == [] assert actor._buffer_capacity._value == 4 # zero permits consumed - assert any("No native replay metadata found" in line for line in printed) def test_run_rejects_native_replay_state_with_gated_sampler(self, tmp_path): ckpt_dir = tmp_path / "resume_ckpt" diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 9e011a311e8..6b40612a0f4 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -30,7 +30,11 @@ REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, ) -from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + CustomSamplerConfig, + WindowedSampler, + WindowedSamplerConfig, +) from nemo_rl.algorithms.grpo import ( GRPOConfig, GRPOSaveState, @@ -45,6 +49,13 @@ ) +class _CheckpointingCustomSampler(WindowedSampler): + """Custom sampler whose static capability must be validated during setup.""" + + def __init__(self, buffer: Any) -> None: + super().__init__(buffer, max_staleness_versions=1) + + def _make_master_config( *, dp_enabled: bool = True, @@ -298,6 +309,28 @@ def test_rejects_windowed_checkpointing_without_native_tq(self): ): 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.async_rl.sampler = CustomSamplerConfig( + target=f"{__name__}:_CheckpointingCustomSampler" + ) + mc.data_plane.update( + { + "backend": "simple", + "checkpointing_enabled": False, + } + ) + + with pytest.raises( + ValueError, + match=( + "replay-checkpoint-capable sampler requires " + "data_plane.checkpointing_enabled=true" + ), + ): + setup_single_controller(mc, MagicMock(pad_token_id=0)) + def test_multiple_dataloader_not_supported(self): mc = _make_master_config(use_multiple_dataloader=True) with pytest.raises(NotImplementedError, match="use_multiple_dataloader"): diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 3e700c7eb16..77ae79da7ac 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -254,6 +254,35 @@ async def checkpoint() -> None: 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): @@ -666,6 +695,27 @@ def _load( ) +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_metadata_state_dict_omits_tensors_and_data_plane_reads(self): dp = FakeDataPlaneClient() diff --git a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py index 6359fc672f7..8f2f44c3b74 100644 --- a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -12,13 +12,177 @@ # 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 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": { + "data_plane_checkpoint_schema_version": ( + verifier.DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + if schema_version is None + else schema_version + ), + "expected_consumed_ids": [verifier.SAMPLE_IDS[0]], + }, + } + + +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" diff --git a/tools/verify_tq_data_plane_checkpoint.py b/tools/verify_tq_data_plane_checkpoint.py index 33c9f50adaf..c0bf41676f5 100644 --- a/tools/verify_tq_data_plane_checkpoint.py +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -37,7 +37,11 @@ import torch from tensordict import TensorDict -from nemo_rl.data_plane import DataPlaneConfig, build_data_plane_client +from nemo_rl.data_plane import ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, + DataPlaneConfig, + build_data_plane_client, +) PARTITION_ID = "tq_checkpoint_smoke" TASK_NAME = "train" @@ -111,7 +115,9 @@ def _save(checkpoint_dir: Path, num_storage_units: int) -> None: dp_client.save_checkpoint( checkpoint_dir, metadata={ - "data_plane_checkpoint_schema_version": 1, + "data_plane_checkpoint_schema_version": ( + DATA_PLANE_CHECKPOINT_SCHEMA_VERSION + ), "expected_consumed_ids": consumed.sample_ids, }, ) @@ -134,10 +140,17 @@ def _load(checkpoint_dir: Path, num_storage_units: int) -> None: ) expected = _expected_fields() for field in FIELDS: - if not torch.equal(restored[field], expected[field]): + 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 metadata["data_plane_checkpoint_schema_version"] != 1: + if ( + 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 From c781d74f163806992337374dbdf281e9204608bf Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 13 Aug 2026 21:18:54 -0700 Subject: [PATCH 06/32] fix: lint issues Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/async_utils/replay_buffer.py | 2 +- nemo_rl/data_plane/interfaces.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index b02f6f8db03..7ed7fe7d406 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -30,7 +30,7 @@ import torch from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol -from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta +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 ( diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 4b354ba397a..62cf30aa72a 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -42,7 +42,6 @@ from tensordict import TensorDict - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 From 83eafa34c96883022a5ae9a1e2716315a0582fcd Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 14 Aug 2026 08:48:32 -0700 Subject: [PATCH 07/32] fix(tests): import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from data_plane Ruff dropped the unused re-export from replay_buffer.py; the tests still needed the constant. Import it from its canonical location instead. Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_sc_checkpointing.py | 2 +- tests/unit/single_controller/test_single_controller_setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 82a69ac990e..963b542b655 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -48,13 +48,13 @@ from torchdata.stateful_dataloader import StatefulDataLoader from nemo_rl.algorithms.async_utils.replay_buffer import ( - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, DataPlaneCheckpointBarrier, ) +from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, WindowedSamplerConfig, diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 6b40612a0f4..496876cdec3 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -25,11 +25,11 @@ import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod from nemo_rl.algorithms.async_utils.replay_buffer import ( DATA_PLANE_CHECKPOINT_DIR, - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, ) +from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.algorithms.async_utils.staleness_sampler import ( CustomSamplerConfig, WindowedSampler, From da8b0a91b782ce14215d06bce77b932a92e76c9a Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 14 Aug 2026 09:03:57 -0700 Subject: [PATCH 08/32] fix: lint issues Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_sc_checkpointing.py | 3 +-- tests/unit/single_controller/test_single_controller_setup.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py index 963b542b655..0241d76e067 100644 --- a/tests/unit/single_controller/test_sc_checkpointing.py +++ b/tests/unit/single_controller/test_sc_checkpointing.py @@ -54,7 +54,6 @@ REPLAY_BUFFER_METADATA_STORAGE, DataPlaneCheckpointBarrier, ) -from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, WindowedSamplerConfig, @@ -75,7 +74,7 @@ ) 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.utils.checkpoint import CheckpointManager # Reuse the factory patches from the setup tests (same cross-module fixture diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py index 496876cdec3..3cc4110ab97 100644 --- a/tests/unit/single_controller/test_single_controller_setup.py +++ b/tests/unit/single_controller/test_single_controller_setup.py @@ -29,7 +29,6 @@ REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, ) -from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION from nemo_rl.algorithms.async_utils.staleness_sampler import ( CustomSamplerConfig, WindowedSampler, @@ -47,6 +46,7 @@ SingleControllerActorArgs, setup_single_controller, ) +from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION class _CheckpointingCustomSampler(WindowedSampler): From 0b4ceccbfcad6c3ec9d347cd09faab33ecca62bd Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 14 Aug 2026 12:34:52 -0700 Subject: [PATCH 09/32] fix(tests): update expected save_state and actor_args for new SC fields GRPOSaveState now has trainer_version; SingleControllerActorArgs now has data_plane_checkpoint_metadata. Two unit tests were still asserting the old shapes. Signed-off-by: Anish Mahishi --- tests/unit/algorithms/test_grpo.py | 3 ++- tests/unit/single_controller/test_single_controller.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index c3e4de00aee..e7109d43a86 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -392,8 +392,9 @@ 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, } assert "total_valid_tokens" not in loaded_state assert not hasattr(save_state, "val:accuracy") diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index 7dfcf331c32..f54179df065 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -132,6 +132,7 @@ def test_logs_hyperparameters_and_concrete_weight_synchronizer( inference_cluster=None, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, + data_plane_checkpoint_metadata=None, ) controller_cls = SingleControllerActor.__ray_metadata__.modified_class From 608e9bde376d23604b7c6472e1e1e19a43bf8345 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 14 Aug 2026 20:56:05 -0700 Subject: [PATCH 10/32] fix(tests): wire DataPlaneCheckpointBarrier into _train_pump_controller The train-pump epilogue now enters the data-plane barrier before clearing consumed samples. Two SC tests built controllers via object.__new__ and had no _data_plane_checkpoint_barrier attribute; the second test hit it once dispatch actually happened, raising AttributeError. Also add data_plane_checkpoint_metadata=None to the SetupTimingMetrics test's actor_args, matching the field _init_ now reads. Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_single_controller.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/single_controller/test_single_controller.py b/tests/unit/single_controller/test_single_controller.py index f54179df065..57266a06290 100644 --- a/tests/unit/single_controller/test_single_controller.py +++ b/tests/unit/single_controller/test_single_controller.py @@ -22,6 +22,7 @@ import torch import nemo_rl.algorithms.single_controller as single_controller +from nemo_rl.algorithms.async_utils.replay_buffer import DataPlaneCheckpointBarrier from nemo_rl.algorithms.grpo import GRPOConfig, _initial_grpo_save_state from nemo_rl.algorithms.loss import ClippedPGLossConfig from nemo_rl.algorithms.metric_utils import SetupTimingMetrics @@ -185,6 +186,7 @@ def test_logs_setup_timing_metrics(monkeypatch, tmp_path) -> None: inference_cluster=None, save_state=_initial_grpo_save_state(), last_checkpoint_path=None, + data_plane_checkpoint_metadata=None, ) controller_cls = SingleControllerActor.__ray_metadata__.modified_class @@ -366,6 +368,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._step_log_dict = { "rewards": [], "masked_advantages": [], From 9eda7b285694e494b571133272ddb48ae65572b5 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Tue, 25 Aug 2026 23:31:51 -0400 Subject: [PATCH 11/32] fix(sc): restore gated sampler checkpoint state Signed-off-by: Anish Mahishi --- .../async_utils/staleness_sampler.py | 58 ++++++++++-- nemo_rl/algorithms/grpo.py | 4 + nemo_rl/algorithms/single_controller.py | 11 ++- .../single_controller/test_checkpointing.py | 93 +++++++++++++++---- .../test_sampler_interface.py | 33 ++++--- 5 files changed, 153 insertions(+), 46 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 8bbcce2dc0e..1e8032c967f 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -117,8 +117,17 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: """Buffer-capacity the policy needs, or ``None`` if unconstrained.""" ... + @property + def dispatch_index(self) -> int: + """Last admitted dispatch batch index.""" + ... + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: - """Seed the dispatch cursor when resuming from a checkpoint.""" + """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.""" ... @@ -138,15 +147,17 @@ def __init__(self, buffer: TQReplayBuffer) -> None: # batch through a zero-staleness gate. self._dispatch_index: int = -1 + @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 dispatch cursor when resuming from a checkpoint. + """Seed the cursor for checkpoints that predate exact sampler state. Args: - resume_from_trainer_version: Trainer weight version this run starts - from — 0 for a fresh run, the restored trainer version when - resuming. Sets the cursor to one before that version so gated - ``admit`` and ``InOrderSampler`` target-step stamps line up with - the restored trainer version. 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_trainer_version < 0: raise ValueError( @@ -155,6 +166,19 @@ 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: + """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 async def admit( @@ -375,6 +399,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, @@ -413,6 +440,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 @@ -460,6 +490,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, @@ -647,7 +681,12 @@ def sampler_supports_buffer_checkpoint(cfg: SamplerConfig) -> bool: else: raise ValueError(f"unknown sampler config {type(cfg).__name__}") - capability = getattr(sampler_cls, "supports_buffer_checkpoint", None) + 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 " @@ -697,7 +736,8 @@ def create_sampler( 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, supports_buffer_checkpoint, " + f"dispatch_index, set_dispatch_index, restore_dispatch_index, " + f"is_on_policy, supports_buffer_checkpoint, " f"required_buffer_capacity)" ) else: diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index fa127b5b995..a1e2066e3d9 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -381,6 +381,9 @@ class GRPOSaveState: # 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: @@ -393,6 +396,7 @@ def _initial_grpo_save_state() -> GRPOSaveState: 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 868aabc672f..18486c9364d 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -227,7 +227,13 @@ def __init__( ) 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(restored_trainer_version) + 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 @@ -1860,8 +1866,9 @@ async def _save_checkpoint( 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 # Snapshot before any await so it can't interleave with - # _rollout_pump iterating this same dataloader. + # _rollout_pump advancing either the sampler cursor or this 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 diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index becd64eb336..81847d0e9f4 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -59,6 +59,7 @@ from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, WindowedSamplerConfig, + sampler_supports_buffer_checkpoint, ) from nemo_rl.algorithms.grpo import ( GRPOConfig, @@ -173,6 +174,7 @@ class _FakeSampler: 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 @@ -213,8 +215,15 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]: def set_gate_window(self, gate_window: int) -> None: self.gate_window = gate_window + @property + def dispatch_index(self) -> int: + return self._dispatch_index + def set_dispatch_index(self, resume_from_trainer_version: int) -> None: - pass + 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): @@ -438,7 +447,7 @@ def _actor_master_config( num_prompts_per_step: int = 2, max_num_epochs: int = 1, buffer_checkpoint: bool = False, - data_plane_checkpoint: bool = False, + data_plane_checkpoint: bool = True, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. @@ -552,7 +561,9 @@ def _run_train_pump( async def _main(): actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) actor._sampler = _FakeSampler( - supports_buffer_checkpoint=(mc.async_rl.sampler.name == "windowed") + supports_buffer_checkpoint=sampler_supports_buffer_checkpoint( + mc.async_rl.sampler + ) ) if seed is not None: seed(actor) @@ -659,7 +670,7 @@ 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 @@ -677,7 +688,22 @@ def test_restores_trainer_version_independently_from_train_step(self, tmp_path): assert actor._train_steps == 7 assert actor._trainer_version == 11 - assert actor._sampler._dispatch_index == 10 + 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( @@ -686,7 +712,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 @@ -710,7 +736,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): @@ -752,6 +778,7 @@ 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. @@ -1638,6 +1665,7 @@ def test_checkpoint_capable_sampler_without_native_tq_is_rejected(self, tmp_path max_num_steps=2, save_period=2, buffer_checkpoint=True, + data_plane_checkpoint=False, ) with pytest.raises( @@ -1646,9 +1674,13 @@ def test_checkpoint_capable_sampler_without_native_tq_is_rejected(self, tmp_path ): _ACTOR_CLS(mc, _make_actor_args(), SetupTimingMetrics()) - def test_no_replay_buffer_with_gated_sampler(self, tmp_path): + 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 + tmp_path, + max_num_steps=2, + save_period=2, + buffer_checkpoint=False, + data_plane_checkpoint=True, ) buffer = _FakeTQBuffer() @@ -1657,7 +1689,8 @@ def test_no_replay_buffer_with_gated_sampler(self, tmp_path): ckpt_dir = tmp_path / "checkpoints" assert (ckpt_dir / "step_2" / "training_info.json").exists() assert not (ckpt_dir / "step_2" / "replay_buffer.pt").exists() - assert not (ckpt_dir / "step_2" / REPLAY_BUFFER_METADATA_FILENAME).exists() + 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" @@ -1887,17 +1920,41 @@ def test_run_missing_native_replay_metadata_starts_empty(self, tmp_path): assert buffer.load_calls == [] assert actor._buffer_capacity._value == 4 # zero permits consumed - def test_run_rejects_native_replay_state_with_gated_sampler(self, tmp_path): + 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_METADATA_FILENAME) - mc = _actor_master_config(tmp_path, max_num_steps=0, buffer_checkpoint=False) + 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() - with pytest.raises(RuntimeError, match="does not support replay-buffer"): - _run_actor_run( - mc, - _make_actor_args(last_checkpoint_path=str(ckpt_dir)), - ) + _run_actor_run( + mc, + _make_actor_args( + tq_buffer=buffer, + last_checkpoint_path=str(ckpt_dir), + save_state=_matching_save_state(), + data_plane_checkpoint_metadata={ + "replay_manifest_digest": "digest-1", + "replay_group_count": 0, + }, + ), + ) + + 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 ────────────────────────────────────────── diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 1edfb442a59..7a586909a6a 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -187,9 +187,9 @@ class TestFactory: ("config", "expected"), [ (WindowedSamplerConfig(), True), - (ReadyFirstSamplerConfig(), False), - (WeightFifoSamplerConfig(), False), - (InOrderSamplerConfig(), False), + (ReadyFirstSamplerConfig(), True), + (WeightFifoSamplerConfig(), True), + (InOrderSamplerConfig(), True), ( CustomSamplerConfig(target=f"{__name__}:EchoSampler"), False, @@ -531,20 +531,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 @@ -556,13 +555,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_trainer_version"): - 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, ) @@ -573,7 +572,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 From 9ff964ba66028c85ab6aa0f66a6511159cf88b7c Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 26 Aug 2026 00:01:13 -0400 Subject: [PATCH 12/32] test(sc): align checkpoint tests with TQ recovery Signed-off-by: Anish Mahishi --- .../single_controller/test_checkpointing.py | 42 +++++++++++-------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 81847d0e9f4..cd2ad22fa66 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -1041,7 +1041,7 @@ def test_tq_save_rejects_inventory_mismatch( assert not (tmp_path / "checkpoints" / "step_1").exists() - def test_gated_sampler_keeps_tq_checkpoint_in_shadow_mode(self, tmp_path): + def test_gated_sampler_writes_authoritative_tq_checkpoint(self, tmp_path): mc = _actor_master_config( tmp_path, max_num_steps=1, @@ -1057,11 +1057,11 @@ def test_gated_sampler_keeps_tq_checkpoint_in_shadow_mode(self, tmp_path): _make_actor_args(dp_client=dp_client, tq_buffer=buffer), ) - assert dp_client.save_calls[0]["metadata"]["mode"] == "shadow" + assert dp_client.save_calls[0]["metadata"]["mode"] == "authoritative" step_dir = tmp_path / "checkpoints" / "step_1" - assert not (step_dir / REPLAY_BUFFER_METADATA_FILENAME).exists() + assert (step_dir / REPLAY_BUFFER_METADATA_FILENAME).exists() assert not (step_dir / "replay_buffer.pt").exists() - assert buffer.metadata_state_dict_calls == [] + assert buffer.metadata_state_dict_calls == [4] def test_tq_save_failure_aborts_checkpoint(self, tmp_path): mc = _actor_master_config( @@ -1092,7 +1092,9 @@ async def _main() -> None: ) actor._train_steps = 1 actor._trainer_version = 1 - save_task = asyncio.create_task(actor._save_checkpoint({"loss": 1.0})) + 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 @@ -1226,6 +1228,7 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): 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 @@ -1234,7 +1237,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}, + data_plane={"checkpointing_enabled": False}, + ) actor._dataloader = SimpleNamespace(state_dict=lambda: {}) actor._buffer = SimpleNamespace(state_dict=AsyncMock(return_value={})) actor._checkpointer = MagicMock() @@ -1409,7 +1416,12 @@ 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", + "checkpointing_enabled": True, + }, data={ "use_multiple_dataloader": False, "shuffle": False, @@ -2015,25 +2027,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( @@ -2041,7 +2047,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(), ), ) From 4e5147ca623b629c956b8c3c6c79bb78d7d8bd0f Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 26 Aug 2026 00:13:21 -0400 Subject: [PATCH 13/32] test(sc): declare custom sampler checkpoint capability Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_setup.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 39e1a0c46ee..9799d290df8 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -55,6 +55,8 @@ 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) From a5e965124a096f8f27a44e2d15d26691edb50786 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 26 Aug 2026 14:55:09 -0400 Subject: [PATCH 14/32] fix(sc): address TQ recovery review feedback Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 48 ++++- examples/configs/grpo_math_1B.yaml | 8 +- ...po_math_1B_megatron_single_controller.yaml | 5 +- .../algorithms/async_utils/replay_buffer.py | 70 +++++--- .../async_utils/staleness_sampler.py | 39 ++-- nemo_rl/algorithms/single_controller.py | 59 ++++--- .../single_controller_utils/setup.py | 74 +++++--- nemo_rl/data_plane/__init__.py | 2 + nemo_rl/data_plane/adapters/noop.py | 8 + nemo_rl/data_plane/adapters/transfer_queue.py | 33 ++-- nemo_rl/data_plane/interfaces.py | 21 ++- nemo_rl/experience/payload.py | 6 +- nemo_rl/utils/checkpoint.py | 4 + .../L1_Functional_Tests_SingleController.sh | 2 +- .../grpo_dp_single_controller_tq_recovery.sh | 24 ++- tests/unit/data_plane/test_backend_config.py | 15 ++ .../data_plane/test_interface_contract.py | 19 ++ tests/unit/data_plane/test_tq_lifecycle.py | 121 ++++++++++++- tests/unit/experience/test_payload.py | 14 +- .../unit/reference_configs/grpo_math_1B.yaml | 4 +- .../single_controller/test_checkpointing.py | 166 ++++++++++++++---- .../unit/single_controller/test_ppo_setup.py | 6 +- tests/unit/single_controller/test_setup.py | 93 +++++++--- .../test_single_controller_actor.py | 65 +++---- .../test_tq_replay_buffer.py | 68 +++++-- .../test_verify_tq_data_plane_checkpoint.py | 18 +- tools/verify_tq_data_plane_checkpoint.py | 42 +++-- 27 files changed, 755 insertions(+), 279 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index ace04ec811b..7ece80f3b68 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -54,7 +54,43 @@ 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. +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. + +## 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. +- 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. + +:::{warning} +This checkpointing path recovers completed groups that have been committed to TQ. It does not recover generations that were still in flight at the checkpoint boundary. +::: + +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 @@ -62,7 +98,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) @@ -73,13 +109,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 | @@ -87,6 +124,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) | @@ -127,7 +165,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` @@ -152,7 +190,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 164f9ddfb09..3cb4abd6c2d 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" @@ -543,10 +547,6 @@ data_plane: impl: transfer_queue backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - # SingleController only: save native TQ state. Required when trainer - # checkpointing is enabled with a replay-checkpoint-capable sampler; - # supported samplers restore from metadata-only replay indexes. - checkpointing_enabled: false # Backend-specific blocks: only the one named by `backend` is read, and an # absent block means that backend's defaults (see SimpleStorageConfig / # MooncakeCpuConfig in nemo_rl/data_plane/interfaces.py). diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index f9e7a13ff64..b50b60b17b1 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -121,6 +121,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: @@ -151,8 +154,6 @@ logger: data_plane: enabled: true - # Required shadow snapshot: a save failure aborts checkpoint finalization. - checkpointing_enabled: true cluster: gpus_per_node: 2 diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 1584b369e36..3edbd0bc46a 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -24,7 +24,7 @@ from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager from numbers import Integral, Real -from typing import Any, Iterable, Literal, Optional, TypedDict +from typing import Any, Iterable, Literal, NotRequired, Optional, TypedDict import ray import torch @@ -66,7 +66,7 @@ class TQReplayGroupMetadata(TypedDict): class TQReplayMetadataState(TypedDict): - """Versioned metadata-only replay sidecar paired with a TQ snapshot.""" + """Versioned metadata-only replay index paired with a TQ snapshot.""" schema_version: int storage: Literal["tq_checkpoint"] @@ -76,6 +76,25 @@ class TQReplayMetadataState(TypedDict): 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] + + 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)): @@ -1015,7 +1034,10 @@ async def commit( ) 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( @@ -1147,13 +1169,17 @@ def metadata_state_dict(self, *, saved_capacity: int) -> TQReplayMetadataState: The caller must hold the exclusive side of the shared data-plane checkpoint barrier through this capture and the matching TQ save. - Commits and destructive clears use shared mutation slots, so the sidecar - and native snapshot describe one exact set of training-ready groups. + 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 participate in that barrier - across the complete publish/index or clear/remove transition. This - includes future finalizer paths; canonical writes are not required to - originate specifically from :meth:`commit`. + 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. """ groups: list[TQReplayGroupMetadata] = [] @@ -1191,9 +1217,9 @@ async def load_state_dict( ) -> int: """Restore the local replay index for an already-restored TQ snapshot. - The sidecar never contains tensor payloads and this method never writes - to the DataPlane. TQ must be restored first; the caller binds the two - artifacts by passing the manifest digest returned by TQ checkpoint + 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 @@ -1211,7 +1237,7 @@ async def load_state_dict( 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 metadata sidecar. + TQ checkpoint load. It must match the replay metadata file. Returns: Number of groups restored into the buffer. @@ -1320,7 +1346,13 @@ async def load_state_dict( if len(groups) > max_groups: raise ValueError( "Native TQ checkpoint contains more replay groups than the current " - f"buffer capacity: checkpoint={len(groups)}, current={max_groups}" + 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." ) for group in groups: @@ -1394,16 +1426,6 @@ def size(self) -> int: def __len__(self) -> int: return len(self.meta_list) - async def _clear_samples(self, *, sample_ids: list[str]) -> None: - """Clear rows without overlapping a bound data-plane checkpoint.""" - if self._data_plane_checkpoint_barrier is None: - raise RuntimeError( - "TQReplayBuffer must be bound to the controller data-plane " - "checkpoint barrier before clearing samples" - ) - async with self._data_plane_checkpoint_barrier.mutation(): - await self._clear_samples_unlocked(sample_ids=sample_ids) - async def _clear_samples_unlocked(self, *, sample_ids: list[str]) -> None: """Clear rows while the caller holds a barrier mutation slot.""" await call_data_plane( diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 1e8032c967f..cb32575a151 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -661,25 +661,28 @@ def _custom_sampler_class(cfg: CustomSamplerConfig) -> type: 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: type - if isinstance(cfg, WindowedSamplerConfig): - sampler_cls = WindowedSampler - elif isinstance(cfg, ReadyFirstSamplerConfig): - sampler_cls = ReadyFirstSampler - elif isinstance(cfg, WeightFifoSamplerConfig): - sampler_cls = WeightFifoSampler - elif isinstance(cfg, InOrderSamplerConfig): - sampler_cls = InOrderSampler - elif isinstance(cfg, CustomSamplerConfig): - sampler_cls = _custom_sampler_class(cfg) - else: - raise ValueError(f"unknown sampler config {type(cfg).__name__}") + sampler_cls = _sampler_class_for_config(cfg) if isinstance(cfg, CustomSamplerConfig): # A custom subclass must opt in explicitly instead of inheriting a @@ -705,31 +708,31 @@ def create_sampler( 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): - sampler = WindowedSampler( + sampler = sampler_cls( buffer, max_staleness_versions=cfg.max_staleness_versions, sample_freshest_first=cfg.sample_freshest_first, ) elif isinstance(cfg, ReadyFirstSamplerConfig): - sampler = ReadyFirstSampler( + sampler = sampler_cls( buffer, max_staleness_versions=cfg.max_staleness_versions, ) elif isinstance(cfg, WeightFifoSamplerConfig): - sampler = WeightFifoSampler( + sampler = sampler_cls( buffer, max_staleness_versions=cfg.max_staleness_versions, ) elif isinstance(cfg, InOrderSamplerConfig): - sampler = InOrderSampler( + sampler = sampler_cls( buffer, max_lookahead_versions=cfg.max_lookahead_versions, warmup_lookahead_versions=cfg.warmup_lookahead_versions, ) elif isinstance(cfg, CustomSamplerConfig): - sampler_cls = _custom_sampler_class(cfg) sampler_supports_buffer_checkpoint(cfg) sampler = sampler_cls(buffer, **(cfg.model_extra or {})) if not isinstance(sampler, PromptGroupSampler): diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 18486c9364d..9760f1e55c3 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -57,6 +57,7 @@ REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, DataPlaneCheckpointBarrier, + DataPlaneCheckpointMetadata, TQReplayMetadataState, ) from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler @@ -210,9 +211,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[dict[str, Any]] = ( - actor_args.data_plane_checkpoint_metadata - ) + 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 @@ -237,11 +238,11 @@ def __init__( if ( self._master_config.checkpointing["enabled"] and self._sampler.supports_buffer_checkpoint - and not self._master_config.data_plane.get("checkpointing_enabled") + and not self._master_config.checkpointing.get("save_data_plane") ): raise ValueError( "SingleController checkpointing with a replay-checkpoint-capable " - "sampler requires data_plane.checkpointing_enabled=true so " + "sampler requires checkpointing.save_data_plane=true so " "completed, unconsumed rollouts are recoverable." ) required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step) @@ -260,10 +261,9 @@ def __init__( # A future staging/finalizer path must join the same barrier before # native restore can be authoritative. self._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() - if self._buffer is not None: - self._buffer.set_data_plane_checkpoint_barrier( - self._data_plane_checkpoint_barrier - ) + 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() @@ -418,7 +418,7 @@ async def _maybe_restore_replay_buffer(self) -> None: """Restore the local replay index for the native TQ checkpoint. Recovery is authoritative only for samplers that explicitly support - buffered-group restoration. The native snapshot and metadata sidecar + 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: @@ -455,7 +455,7 @@ async def _maybe_restore_replay_buffer(self) -> None: ) return print(f"📦 Restoring replay buffer metadata: {metadata_path}") - # weights_only=False: the metadata sidecar contains pickled KVBatchMeta + # 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, metadata_path, weights_only=False @@ -526,8 +526,8 @@ async def _validate_replay_inventory( unexpected_sample_ids = sorted(actual_sample_ids - expected_sample_ids) if missing_sample_ids or unexpected_sample_ids: raise RuntimeError( - "Native TQ checkpoint inventory does not match the replay " - "metadata sidecar: " + "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} " @@ -609,7 +609,7 @@ async def _save_data_plane_checkpoint( """Save a required TQ snapshot inside an SC checkpoint bundle. A sampler with replay-buffer recovery writes an authoritative native - TQ snapshot bound to its metadata-only sidecar by a digest. Other + 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. @@ -618,27 +618,30 @@ async def _save_data_plane_checkpoint( checkpoint_path, DATA_PLANE_CHECKPOINT_DIR, ) - metadata = { + 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": self._train_steps, - "single_controller_trainer_version": self._trainer_version, - "single_controller_epoch": self._current_epoch, + "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.update( - { - "replay_metadata_schema_version": ( - REPLAY_BUFFER_METADATA_SCHEMA_VERSION - ), - "replay_manifest_digest": replay_metadata["manifest_digest"], - "replay_group_count": len(replay_metadata["groups"]), - } + 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"]) started = time.monotonic() print(f"data-plane checkpoint save started: {checkpoint_dir}", flush=True) try: @@ -1968,10 +1971,10 @@ async def _save_checkpoint( os.path.join(checkpoint_path, "replacement_reserve.pt"), ) replay_metadata: Optional[TQReplayMetadataState] = None - if self._master_config.data_plane.get("checkpointing_enabled"): + if self._master_config.checkpointing.get("save_data_plane"): # Commits and destructive clears take the same barrier. Generation # may continue while a snapshot is written, but completed groups - # wait at commit, so TQ and the metadata sidecar describe exactly + # wait at commit, so TQ and the replay metadata file describe exactly # the same set of training-ready groups. async with self._data_plane_checkpoint_barrier.checkpoint(): if self._sampler.supports_buffer_checkpoint: diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index f28f4a9b791..051daa8385c 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 @@ -39,6 +40,7 @@ LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointMetadata, TQReplayBuffer, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( @@ -71,6 +73,7 @@ DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, DataPlaneClient, build_data_plane_client, + data_plane_supports_checkpointing, ) from nemo_rl.distributed.virtual_cluster import ( RayVirtualCluster, @@ -134,7 +137,7 @@ class SingleControllerActorArgs: partition_id: str save_state: GRPOSaveState last_checkpoint_path: Optional[str] - data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None + 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. @@ -151,10 +154,10 @@ def _maybe_restore_native_data_plane_checkpoint( save_state: GRPOSaveState, partition_id: str, sampler_name: str, -) -> Optional[dict[str, Any]]: +) -> Optional[DataPlaneCheckpointMetadata]: """Load and validate an authoritative native TQ checkpoint when present. - The metadata-only replay sidecar is the format marker. Checkpoints without + 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 @@ -173,6 +176,14 @@ def _maybe_restore_native_data_plane_checkpoint( "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 @@ -183,13 +194,14 @@ def _maybe_restore_native_data_plane_checkpoint( ) print(f"📦 Restoring native TQ checkpoint: {data_plane_path}", flush=True) - metadata = policy.load_data_plane_checkpoint(data_plane_path) - if not isinstance(metadata, dict): + 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(metadata).__name__}" + f"got {type(raw_metadata).__name__}" ) - expected_values: dict[str, Any] = { + 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": ( @@ -718,22 +730,44 @@ def setup_single_controller( "master_config.data_plane.enabled=True. The async-RL " "SingleController path is built on the TransferQueue data plane." ) - if dp_config.get("checkpointing_enabled") and dp_config["backend"] != "simple": - raise NotImplementedError( - "SingleController data-plane checkpointing currently requires " - "data_plane.backend='simple'; Mooncake storage cannot be restored " - "by TQ v0.1.9." - ) + data_plane_checkpointing_supported = data_plane_supports_checkpointing(dp_config) if ( - master_config.checkpointing["enabled"] - and sampler_supports_buffer_checkpoint(master_config.async_rl.sampler) - and not dp_config.get("checkpointing_enabled") + master_config.checkpointing.get("save_data_plane") + and not data_plane_checkpointing_supported ): - raise ValueError( - "SingleController checkpointing with a replay-checkpoint-capable " - "sampler requires data_plane.checkpointing_enabled=true so " - "completed, unconsumed rollouts are recoverable." + 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" diff --git a/nemo_rl/data_plane/__init__.py b/nemo_rl/data_plane/__init__.py index 36574328e88..c97346ed4d2 100644 --- a/nemo_rl/data_plane/__init__.py +++ b/nemo_rl/data_plane/__init__.py @@ -25,6 +25,7 @@ DataPlaneClient, DataPlaneConfig, KVBatchMeta, + data_plane_supports_checkpointing, ) from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event @@ -35,6 +36,7 @@ "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 34d2f76c7de..6f98800e505 100644 --- a/nemo_rl/data_plane/adapters/noop.py +++ b/nemo_rl/data_plane/adapters/noop.py @@ -294,6 +294,14 @@ def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: 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) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index ea29707ac87..0ad7b6b3dfc 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -51,6 +51,7 @@ DataPlaneConfig, KVBatchMeta, backend_config, + data_plane_supports_checkpointing, ) from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS @@ -728,6 +729,7 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: # 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: @@ -740,6 +742,19 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: # 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. + # 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.""" @@ -753,11 +768,6 @@ def _require_clean_for_load(self) -> None: "register, claim, get, list, put, clear, or consumption operation" ) - # 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. - self._warmed_fields: dict[str, set[str]] = {} - # ── (A) task-mediated ─────────────────────────────────────────────── def register_partition( @@ -1013,12 +1023,7 @@ def save_checkpoint( metadata: dict[str, Any] | None = None, ) -> None: """Save TQ controller metadata and storage data.""" - if self._backend == "mooncake_cpu": - raise NotImplementedError( - "TQ checkpointing is not supported for the mooncake_cpu " - "backend: MooncakeStorageManager cannot persist its in-memory " - "rows, so TQ would silently create a metadata-only checkpoint." - ) + self._require_checkpointing_support() _connect_existing() tq.save_checkpoint(checkpoint_dir, metadata=metadata) @@ -1029,11 +1034,7 @@ def load_checkpoint(self, checkpoint_dir: str | Path) -> dict[str, Any]: TQ client, so the recovery coordinator must also guarantee globally clean setup ordering. """ - if self._backend == "mooncake_cpu": - raise NotImplementedError( - "TQ checkpoint restore is not supported for the mooncake_cpu " - "backend because its in-memory rows cannot be restored." - ) + self._require_checkpointing_support() self._require_clean_for_load() # Validate the adapter-owned metadata before starting TQ's # non-transactional storage/controller restore. diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 4aa1cb1dcab..fcff2e398f4 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -110,13 +110,6 @@ class DataPlaneConfig(TypedDict): has no static default, since no single value is right across cluster sizes, so a ``simple`` run without the block fails validation. - ``checkpointing_enabled`` opts SingleController into saving required TQ - state inside its checkpoint bundle. Samplers that support replay-buffer - recovery pair the native snapshot with a metadata-only local index; other - samplers save it in shadow mode. Other algorithm entrypoints do not consume - this field. It is optional because existing configs predate data-plane - checkpointing; exemplar configs carry the recommended default explicitly. - Required keys (always set in the exemplar YAML): ``enabled``, ``impl``, ``backend``, ``claim_meta_poll_interval_s``. @@ -131,7 +124,6 @@ class DataPlaneConfig(TypedDict): impl: Literal["transfer_queue"] backend: Literal["simple", "mooncake_cpu"] claim_meta_poll_interval_s: float - checkpointing_enabled: NotRequired[bool] simple: NotRequired[SimpleStorageConfig] mooncake_cpu: NotRequired[MooncakeCpuConfig] controller_address: NotRequired[str] @@ -139,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, diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index 0728401d728..32f4054de82 100644 --- a/nemo_rl/experience/payload.py +++ b/nemo_rl/experience/payload.py @@ -101,6 +101,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. @@ -108,6 +109,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 tags. @@ -124,5 +126,7 @@ def pack_payload( tensor_fields, lengths=lengths, token_aligned_fields=TOKEN_ALIGNED_FIELDS ) sample_ids = [f"{group_id}_g{i}" for i in range(n)] - tags = [{"weight_version": weight_version} for _ in range(n)] + tags = [ + {"weight_version": weight_version, "prompt_idx": prompt_idx} for _ in range(n) + ] return sample_ids, fields_td, tags diff --git a/nemo_rl/utils/checkpoint.py b/nemo_rl/utils/checkpoint.py index ab0e3705cd9..66a7b7b43c2 100644 --- a/nemo_rl/utils/checkpoint.py +++ b/nemo_rl/utils/checkpoint.py @@ -116,6 +116,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 @@ -134,6 +137,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/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 38913cfd6cc..c06ebf75b5a 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -82,7 +82,7 @@ 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 replay checkpoint recovery (#3480). -run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh +run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh cd ${PROJECT_ROOT}/tests if compgen -G ".coverage*" > /dev/null; then diff --git a/tests/functional/grpo_dp_single_controller_tq_recovery.sh b/tests/functional/grpo_dp_single_controller_tq_recovery.sh index 671b67dd4ec..ea069d03b0c 100755 --- a/tests/functional/grpo_dp_single_controller_tq_recovery.sh +++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh @@ -4,10 +4,13 @@ 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" @@ -16,7 +19,7 @@ COMMON_OVERRIDES=( checkpointing.enabled=true checkpointing.checkpoint_dir="$CHECKPOINT_DIR" checkpointing.save_period=1 - data_plane.checkpointing_enabled=true + checkpointing.save_data_plane=true async_rl.sampler.name=windowed '~async_rl.sampler.max_lookahead_versions' '+async_rl.sampler.max_staleness_versions=1' @@ -32,22 +35,29 @@ 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" -uv run --no-sync python -c \ - 'import json, sys; metadata = json.load(open(sys.argv[1]))["user_metadata"]; assert metadata["mode"] == "authoritative"; assert metadata["replay_group_count"] > 0, metadata' \ - "$CHECKPOINT_DIR/step_1/data_plane/metadata.json" +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" "$BASE_RUN_LOG" -grep -q "Native TQ replay inventory validated" "$BASE_RUN_LOG" -grep -Eq "Restored [1-9][0-9]* replay group" "$BASE_RUN_LOG" +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/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_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py index 7c9c9002599..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 @@ -188,3 +190,20 @@ def test_checkpoint_load_requires_clean_client(tmp_path) -> None: 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_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index 86442c32a8b..e51fe398741 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -22,7 +22,9 @@ from __future__ import annotations +import inspect import json +from typing import Callable from unittest.mock import MagicMock import numpy as np @@ -33,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 @@ -107,6 +171,57 @@ 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 @@ -133,6 +248,7 @@ def test_checkpoint_lifecycle_forwards_to_tq(monkeypatch, tmp_path) -> None: 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() @@ -180,6 +296,7 @@ def test_checkpoint_load_rejects_client_after_data_operation( 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( @@ -213,6 +330,7 @@ def test_failed_checkpoint_load_leaves_client_in_dirty_state( ) 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"): @@ -241,6 +359,7 @@ def test_mooncake_checkpoint_lifecycle_fails_loudly( 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" diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index 7b240676c5b..029b0c47860 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 @@ -132,7 +133,10 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None: packed_rows = list(packed_routes.unbind()) assert torch.equal(packed_rows[0], expected_routes[0]) assert torch.equal(packed_rows[1], expected_routes[1]) - assert tags == [{"weight_version": 3}, {"weight_version": 3}] + assert tags == [ + {"weight_version": 3, "prompt_idx": 17}, + {"weight_version": 3, "prompt_idx": 17}, + ] def test_record_to_train_batch_omits_routed_experts_when_absent() -> None: @@ -148,6 +152,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 @@ -187,6 +192,11 @@ 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) diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index c670b090f64..b8f35fc9fe6 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 @@ -524,9 +525,6 @@ data_plane: impl: transfer_queue backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - # SingleController only: save native TQ state. Supported samplers restore - # from metadata-only replay indexes. - checkpointing_enabled: false simple: storage_capacity: 1000000 # max samples retained per partition num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index cd2ad22fa66..e2598a1e041 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -55,6 +55,7 @@ REPLAY_BUFFER_METADATA_SCHEMA_VERSION, REPLAY_BUFFER_METADATA_STORAGE, DataPlaneCheckpointBarrier, + DataPlaneCheckpointMetadata, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSamplerConfig, @@ -241,7 +242,7 @@ async def select(self, **kwargs) -> tuple[Optional[KVBatchMeta], int]: class _RestoredGroupsSampler(_FakeSampler): - """Drain the exact groups represented by a restored metadata sidecar.""" + """Drain the exact groups represented by a restored replay metadata file.""" def __init__(self, groups: list[dict[str, Any]]) -> None: super().__init__() @@ -492,6 +493,7 @@ 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, }, @@ -499,7 +501,6 @@ def _actor_master_config( "enabled": True, "impl": "transfer_queue", "backend": "simple", - "checkpointing_enabled": data_plane_checkpoint, }, async_rl=AsyncRLConfig( sampler=sampler_cfg, @@ -518,7 +519,7 @@ def _make_actor_args( tq_buffer: Optional[_FakeTQBuffer] = None, dp_client: Optional[_FakeDPClient] = None, last_checkpoint_path: Optional[str] = None, - data_plane_checkpoint_metadata: Optional[dict[str, Any]] = None, + data_plane_checkpoint_metadata: Optional[DataPlaneCheckpointMetadata] = None, ) -> SingleControllerActorArgs: return SingleControllerActorArgs( gen_handle=object(), @@ -542,6 +543,34 @@ def _make_actor_args( ) +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, @@ -909,6 +938,39 @@ def test_ft_save_period_triggers_saves(self, tmp_path): 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 ): @@ -959,20 +1021,12 @@ def test_saves_authoritative_tq_state_and_metadata_only_replay_index( assert save_call["checkpoint_dir"] == str( tmp_path / "checkpoints" / "tmp_step_1" / "data_plane" ) - assert save_call["metadata"] == { - "data_plane_checkpoint_schema_version": ( - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION - ), - "single_controller_train_steps": 1, - "single_controller_trainer_version": 1, - "single_controller_epoch": 0, - "partition_id": _PARTITION_ID, - "sampler_name": "windowed", - "mode": "authoritative", - "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION, - "replay_manifest_digest": "digest-1", - "replay_group_count": 1, - } + assert save_call["metadata"] == _data_plane_checkpoint_metadata( + step=1, + trainer_version=1, + sampler_name="windowed", + group_count=1, + ) step_dir = tmp_path / "checkpoints" / "step_1" assert (step_dir / "data_plane" / "metadata.json").is_file() assert ( @@ -1239,8 +1293,8 @@ def _ppo_save_actor(tmp_path: Path, calls: list[str]): ) actor._sampler = _FakeSampler() actor._master_config = SimpleNamespace( - checkpointing={"metric_name": None}, - data_plane={"checkpointing_enabled": False}, + checkpointing={"metric_name": None, "save_data_plane": False}, + data_plane={}, ) actor._dataloader = SimpleNamespace(state_dict=lambda: {}) actor._buffer = SimpleNamespace(state_dict=AsyncMock(return_value={})) @@ -1420,7 +1474,6 @@ def _setup_master_config(checkpoint_dir: str) -> MasterConfig: "enabled": True, "impl": "transfer_queue", "backend": "simple", - "checkpointing_enabled": True, }, data={ "use_multiple_dataloader": False, @@ -1463,6 +1516,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, }, ) @@ -1748,10 +1802,7 @@ def test_run_restores_native_tq_replay_metadata_without_payload_reput( ] envelope = {"groups": groups} torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) - tq_metadata = { - "replay_manifest_digest": "digest-1", - "replay_group_count": 2, - } + tq_metadata = _data_plane_checkpoint_metadata(group_count=2) mc = _actor_master_config( tmp_path, max_num_steps=0, @@ -1811,10 +1862,7 @@ def test_restored_permits_are_released_by_a_live_pump(self, tmp_path): for i in range(4) ] torch.save({"groups": groups}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) - tq_metadata = { - "replay_manifest_digest": "digest-1", - "replay_group_count": 4, - } + tq_metadata = _data_plane_checkpoint_metadata(group_count=4) mc = _actor_master_config( tmp_path, max_num_steps=2, @@ -1874,10 +1922,7 @@ def test_native_restore_rejects_tq_inventory_mismatch( ] } torch.save(envelope, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME) - tq_metadata = { - "replay_manifest_digest": "digest-1", - "replay_group_count": 1, - } + tq_metadata = _data_plane_checkpoint_metadata(group_count=1) mc = _actor_master_config( tmp_path, max_num_steps=0, @@ -1913,6 +1958,58 @@ def test_native_replay_metadata_requires_setup_side_tq_restore(self, tmp_path): _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() + + 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() @@ -1951,10 +2048,7 @@ def test_run_restores_native_replay_state_with_in_order_sampler(self, tmp_path): tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir), save_state=_matching_save_state(), - data_plane_checkpoint_metadata={ - "replay_manifest_digest": "digest-1", - "replay_group_count": 0, - }, + data_plane_checkpoint_metadata=_data_plane_checkpoint_metadata(), ), ) diff --git a/tests/unit/single_controller/test_ppo_setup.py b/tests/unit/single_controller/test_ppo_setup.py index d63de950e1f..18e66248a35 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_setup.py b/tests/unit/single_controller/test_setup.py index 9799d290df8..6d0bd768748 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -28,6 +28,7 @@ LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + DataPlaneCheckpointMetadata, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( CustomSamplerConfig, @@ -61,6 +62,15 @@ 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, @@ -82,7 +92,11 @@ def _make_master_config( only the dict-shaped fields setup reads. """ 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, @@ -134,7 +148,7 @@ def _make_master_config( def _native_tq_metadata( *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1 -) -> dict[str, Any]: +) -> DataPlaneCheckpointMetadata: return { "data_plane_checkpoint_schema_version": (DATA_PLANE_CHECKPOINT_SCHEMA_VERSION), "single_controller_train_steps": step, @@ -287,31 +301,39 @@ def test_raises_when_data_plane_disabled(self): def test_rejects_mooncake_data_plane_checkpointing(self): mc = _make_master_config() - mc.data_plane.update( - { - "backend": "mooncake_cpu", - "checkpointing_enabled": True, - } - ) - with pytest.raises(NotImplementedError, match="backend='simple'"): + 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.update( - { - "backend": "simple", - "checkpointing_enabled": False, - } - ) + mc.data_plane["backend"] = "simple" with pytest.raises( ValueError, match=( "replay-checkpoint-capable sampler requires " - "data_plane.checkpointing_enabled=true" + "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)) @@ -319,25 +341,37 @@ def test_rejects_windowed_checkpointing_without_native_tq(self): 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.update( - { - "backend": "simple", - "checkpointing_enabled": False, - } - ) + mc.data_plane["backend"] = "simple" with pytest.raises( ValueError, match=( "replay-checkpoint-capable sampler requires " - "data_plane.checkpointing_enabled=true" + "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"): @@ -843,7 +877,7 @@ def test_setup_loads_tq_before_creating_single_controller_client( assert events == ["load", "build"] assert actor_args.data_plane_checkpoint_metadata == _native_tq_metadata() - def test_loads_authoritative_tq_checkpoint_when_metadata_sidecar_exists( + def test_loads_authoritative_tq_checkpoint_when_metadata_file_exists( self, tmp_path ): checkpoint_path = tmp_path / "step_3" @@ -902,7 +936,9 @@ def test_legacy_replay_checkpoint_is_rejected(self, tmp_path): policy.load_data_plane_checkpoint.assert_not_called() - def test_checkpoint_without_replay_artifacts_does_not_load_tq(self, tmp_path): + 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() @@ -917,8 +953,13 @@ def test_checkpoint_without_replay_artifacts_does_not_load_tq(self, tmp_path): 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_sidecar_requires_matching_tq_directory(self, tmp_path): + 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() diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 4cbae13b94b..1cc95e5da04 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -45,6 +45,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 { @@ -83,6 +95,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, @@ -92,8 +105,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, @@ -131,6 +144,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, @@ -140,8 +154,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, @@ -190,26 +204,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, - data_plane_checkpoint_metadata=None, - ) + actor_args = _actor_args_for_init() controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller_cls( @@ -320,27 +315,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, - data_plane_checkpoint_metadata=None, - ) + actor_args = _actor_args_for_init() controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller_cls( diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index cd63cf29c93..a4f1d19ef8d 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -152,10 +152,10 @@ 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={}, @@ -363,7 +363,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, ) @@ -380,8 +380,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 @@ -472,7 +472,7 @@ def test_commit_appends_multiple_records_in_order(self): class TestTQReplayBufferRemove: - def test_dp_clear_fails_without_bound_checkpoint_barrier(self): + def test_remove_with_dp_clear_fails_without_bound_checkpoint_barrier(self): dp = FakeDataPlaneClient() buf = TQReplayBuffer( dp, @@ -481,7 +481,7 @@ def test_dp_clear_fails_without_bound_checkpoint_barrier(self): ) with pytest.raises(RuntimeError, match="must be bound"): - _run(buf._clear_samples(sample_ids=["sample-1"])) + _run(buf.remove([0], remove_in_dp=True)) assert dp.clear_calls == [] @@ -512,8 +512,15 @@ 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._clear_samples(sample_ids=["sample-1"]) + await buf.remove([0], remove_in_dp=True) return dp, event_loop_thread_id dp, event_loop_thread_id = asyncio.run(exercise()) @@ -876,6 +883,26 @@ def test_state_dict_skips_middle_unready(self): list(third.sample_ids), ] + 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) + 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: """Malformed envelopes raise ValueError before any DataPlane write.""" @@ -935,9 +962,22 @@ def test_metadata_only_restore_rejects_capacity_truncation(self): state = _make_metadata_envelope( [_make_group_entry(f"g{w}", weight=w) for w in (1, 2, 3)] ) - self._assert_rejected( - state, - match="more replay groups than the current buffer capacity", - max_groups=2, - expected_manifest_digest=state["manifest_digest"], - ) + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + + with pytest.raises(ValueError) as exc_info: + _load( + buf, + state, + max_groups=2, + expected_manifest_digest=state["manifest_digest"], + ) + + 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 diff --git a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py index 8f2f44c3b74..406cc61df8c 100644 --- a/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py +++ b/tests/unit/tools/test_verify_tq_data_plane_checkpoint.py @@ -19,6 +19,7 @@ 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 @@ -92,17 +93,26 @@ def _checkpoint_state(*, schema_version: int | None = None) -> dict[str, Any]: return { "fields": verifier._expected_fields(), "consumed": {verifier.SAMPLE_IDS[0]}, - "metadata": { - "data_plane_checkpoint_schema_version": ( + "metadata": verifier._checkpoint_metadata( + [verifier.SAMPLE_IDS[0]], + schema_version=( verifier.DATA_PLANE_CHECKPOINT_SCHEMA_VERSION if schema_version is None else schema_version ), - "expected_consumed_ids": [verifier.SAMPLE_IDS[0]], - }, + ), } +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: diff --git a/tools/verify_tq_data_plane_checkpoint.py b/tools/verify_tq_data_plane_checkpoint.py index c0bf41676f5..41ed44f0ddf 100644 --- a/tools/verify_tq_data_plane_checkpoint.py +++ b/tools/verify_tq_data_plane_checkpoint.py @@ -32,11 +32,14 @@ import subprocess import sys from pathlib import Path -from typing import cast +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, @@ -51,6 +54,24 @@ 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, @@ -58,12 +79,11 @@ def _data_plane_config(num_storage_units: int) -> DataPlaneConfig: "enabled": True, "impl": "transfer_queue", "backend": "simple", - "checkpointing_enabled": True, - "storage_capacity": 1024, - "num_storage_units": num_storage_units, "claim_meta_poll_interval_s": 0.05, - "global_segment_size": 8 * 1024**3, - "local_buffer_size": 1024**3, + "simple": { + "storage_capacity": 1024, + "num_storage_units": num_storage_units, + }, }, ) @@ -114,12 +134,7 @@ def _save(checkpoint_dir: Path, num_storage_units: int) -> None: dp_client.save_checkpoint( checkpoint_dir, - metadata={ - "data_plane_checkpoint_schema_version": ( - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION - ), - "expected_consumed_ids": consumed.sample_ids, - }, + metadata=_checkpoint_metadata(consumed.sample_ids), ) finally: dp_client.close() @@ -132,6 +147,7 @@ def _load(checkpoint_dir: Path, num_storage_units: int) -> None: ) try: metadata = dp_client.load_checkpoint(checkpoint_dir) + checkpoint_metadata = cast(DataPlaneCheckpointMetadata, metadata) restored = dp_client.get_samples( sample_ids=SAMPLE_IDS, @@ -148,7 +164,7 @@ def _load(checkpoint_dir: Path, num_storage_units: int) -> None: raise AssertionError(f"Restored field differs: {field}") if ( - metadata["data_plane_checkpoint_schema_version"] + checkpoint_metadata["data_plane_checkpoint_schema_version"] != DATA_PLANE_CHECKPOINT_SCHEMA_VERSION ): raise AssertionError("Unexpected data-plane checkpoint schema") From 1c33a9a13608b5cf7520979cacdd6e7293f24589 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 26 Aug 2026 13:09:33 -0700 Subject: [PATCH 15/32] fix: lint issues Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 8 ++++---- nemo_rl/algorithms/single_controller_utils/setup.py | 5 ++--- tests/unit/single_controller/test_checkpointing.py | 4 +--- tests/unit/single_controller/test_sampler_interface.py | 2 +- tests/unit/single_controller/test_setup.py | 3 +-- tests/unit/single_controller/test_tq_replay_buffer.py | 4 +--- 6 files changed, 10 insertions(+), 16 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 9760f1e55c3..6e1a08c635a 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -211,9 +211,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._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 @@ -603,7 +603,7 @@ async def _clear_data_plane_samples(self, sample_ids: list[str]) -> None: async def _save_data_plane_checkpoint( self, - checkpoint_path: str, + checkpoint_path: PathLike, replay_metadata: Optional[TQReplayMetadataState] = None, ) -> None: """Save a required TQ snapshot inside an SC checkpoint bundle. diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py index 051daa8385c..b54d94d8d5d 100644 --- a/nemo_rl/algorithms/single_controller_utils/setup.py +++ b/nemo_rl/algorithms/single_controller_utils/setup.py @@ -743,9 +743,8 @@ def setup_single_controller( 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") + if sampler_supports_replay_recovery and not master_config.checkpointing.get( + "save_data_plane" ): error_message = ( "SingleController checkpointing with a replay-checkpoint-capable " diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index e2598a1e041..e3b8d98b330 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -554,9 +554,7 @@ def _data_plane_checkpoint_metadata( ) -> DataPlaneCheckpointMetadata: """Build the authoritative SC envelope used by actor-level restore tests.""" return { - "data_plane_checkpoint_schema_version": ( - DATA_PLANE_CHECKPOINT_SCHEMA_VERSION - ), + "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 diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 7a586909a6a..7d47453ce5a 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -40,8 +40,8 @@ WindowedSampler, WindowedSamplerConfig, create_sampler, - sampler_supports_buffer_checkpoint, required_buffer_capacity_for_config, + sampler_supports_buffer_checkpoint, ) from nemo_rl.data_plane import KVBatchMeta diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index 6d0bd768748..70e5d203897 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -332,8 +332,7 @@ def test_checkpointing_error_explains_mooncake_incompatibility(self): with pytest.raises( ValueError, match=( - "backend='mooncake_cpu'.*backend='simple'.*" - "checkpointing.enabled=false" + "backend='mooncake_cpu'.*backend='simple'.*checkpointing.enabled=false" ), ): setup_single_controller(mc, MagicMock(pad_token_id=0)) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index a4f1d19ef8d..8f84825bfea 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -896,9 +896,7 @@ def test_state_dict_skips_older_long_tail_unready_group(self): 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 [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"] } From d98004ead21ef9a81d6b884ed9b6870d294e6e01 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Wed, 26 Aug 2026 22:37:33 -0400 Subject: [PATCH 16/32] fix(sc): align recovery tests with native TQ checkpoints Signed-off-by: Anish Mahishi --- .../grpo_checkpoint_single_controller.sh | 20 ++++++++++++++++--- .../functional/ppo_async_single_controller.sh | 11 +++++++--- tests/unit/algorithms/test_grpo.py | 1 + .../single_controller/test_rollout_pump.py | 4 ++-- 4 files changed, 28 insertions(+), 8 deletions(-) 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/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh index aa3cb1fd637..b3df65ab72c 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 82228539b5f..a279b357ab2 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -481,6 +481,7 @@ def test_get_grpo_save_state_handles_legacy_checkpoint_and_filters_metrics(): # 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") diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 97f25574e34..8ce81627b4e 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -1143,5 +1143,5 @@ def test_rollout_pump_writes_expected_tq_data( ) for tag in tags: assert tag["weight_version"] == 0 - # Slim tag schema: weight_version is the only field producers stamp. - assert set(tag) == {"weight_version"} + assert tag["prompt_idx"] == input_sample["idx"] + assert set(tag) == {"weight_version", "prompt_idx"} From fb53e36adf86de0ffe343534d06680bd811d03c8 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 12:02:45 -0400 Subject: [PATCH 17/32] test(sc): add checkpoint recovery race coverage Signed-off-by: Anish Mahishi --- .../_checkpoint_scenarios.py | 455 ++++++++++++ .../test_checkpoint_dispatch_races.py | 679 ++++++++++++++++++ .../test_checkpoint_recovery_matrix.py | 118 +++ 3 files changed, 1252 insertions(+) create mode 100644 tests/unit/single_controller/_checkpoint_scenarios.py create mode 100644 tests/unit/single_controller/test_checkpoint_dispatch_races.py create mode 100644 tests/unit/single_controller/test_checkpoint_recovery_matrix.py diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py new file mode 100644 index 00000000000..c4691a1c378 --- /dev/null +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -0,0 +1,455 @@ +# 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 + +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 + bar the feature should eventually meet, not the bar it meets today. + """ + 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 _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: For a case that fails today, the line that drops the data. + """ + + 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 how a future partial-group restore is built. A + group could come back already committed (its missing rollouts regenerated + before the save), 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. + """ + + recovered: set[str] + ready: set[str] + pending: set[str] + saved_sidecar: bool + rows_before: set[str] + rows_after: set[str] + + +async def _round_trip( + scenario: Scenario, sampler_name: str, tmp_path: Path +) -> 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 + ) + 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"], + ) + + ready = { + gid for gid, is_ready in zip(buf_b._group_ids, buf_b.ready_list) if is_ready + } + return RoundTrip( + recovered=set(buf_b._group_ids), + ready=ready, + pending=set(buf_b._group_ids) - ready, + saved_sidecar=sidecar is not None, + rows_before=rows_before, + rows_after=set(dp_b.list_sample_ids(PARTITION)), + ) + + +def round_trip(scenario: Scenario, sampler_name: str, tmp_path: Path) -> RoundTrip: + """Save the scenario, restore it into a fresh buffer, report what came back.""" + return asyncio.run(_round_trip(scenario, sampler_name, tmp_path)) + + +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_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_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/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py new file mode 100644 index 00000000000..72aa50f25e1 --- /dev/null +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -0,0 +1,679 @@ +# 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 threading +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from nemo_rl.algorithms.async_utils.replay_buffer import ( + REPLAY_BUFFER_METADATA_FILENAME, + 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_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.experience.rollout_manager import RolloutOutcome +from tests.unit.single_controller._checkpoint_scenarios import ( + _record, + patch_converter, +) +from tests.unit.single_controller.test_checkpointing import ( + _FakeDataloader, + _actor_master_config, + _make_actor_args, +) + + +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 + + async def admit(self, *, trainer_version_fn): + self.admit_calls += 1 + return await super().admit(trainer_version_fn=trainer_version_fn) + + +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 admit(self, *, trainer_version_fn): + self.admission_entered.set() + await self.release_admission.wait() + return await super().admit(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": 1, + "groups": [ + { + "group_id": group.group_id, + "target_step": group.target_step, + "prompt_payload": group.prompt_payload, + } + 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: _PendingLedger) -> None: + self.recovery_ledger = ledger + self.recovered: list[tuple[str, int | None]] = [] + + async def recover_group(self, group_id: str) -> bool: + group = next( + group + for group in self.recovery_ledger.groups() + if group.group_id == group_id + ) + self.recovered.append((group.group_id, group.target_step)) + self.recovery_ledger.release(group_id) + return True + + +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, prompt: dict[str, Any], *, target_step: int | None = None + ) -> str: + 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, group_id: str, *, target_step: int) -> None: + self.recovery_ledger.assign_target_step(group_id, target_step) + + async def generate_and_push( + self, + prompt: dict[str, Any], + *, + target_step: int | None = None, + inflight_registry: dict[str, Any] | None = None, + recovery_group_id: str | None = None, + ) -> RolloutOutcome: + del inflight_registry + if recovery_group_id is None: + recovery_group_id = self.reserve_prompt_group( + 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) + + +def _enable_recovery_checkpoint_capture(controller: Any) -> None: + """Install the narrow recovery hooks expected by the future foundation.""" + + async def _inventory_is_valid(**_: Any) -> None: + return None + + controller._validate_rollout_recovery_inventory = _inventory_is_valid + controller._master_config.__dict__["token_capture"] = SimpleNamespace( + enabled=True, + staging_partition="rollout_staging", + ) + + +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 + + +@pytest.mark.xfail( + strict=True, + reason=( + "The dataloader cursor can advance before unfinished prompt ownership " + "is recorded in the checkpoint bundle." + ), +) +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 + _enable_recovery_checkpoint_capture(controller) + + pump = asyncio.create_task(controller._rollout_pump()) + await asyncio.wait_for(sampler.admission_entered.wait(), timeout=1.0) + 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 asyncio.wait_for(rollout_manager.started.wait(), timeout=1.0) + rollout_manager.release.set() + await asyncio.wait_for(pump, timeout=1.0) + 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_payload"]["idx"] == 70 + assert torch.load( + checkpoint / "train_dataloader.pt", + weights_only=False, + ) == {"next_batch": 8} + + asyncio.run(exercise()) + + +@pytest.mark.xfail( + strict=True, + reason=( + "An unfinished admitted batch is not yet persisted alongside the native " + "TQ checkpoint." + ), +) +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()) + + _enable_recovery_checkpoint_capture(controller) + + pump = asyncio.create_task(controller._rollout_pump()) + await asyncio.wait_for(rollout_manager.started.wait(), timeout=1.0) + 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=1.0) + 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()) + + +@pytest.mark.xfail( + strict=True, + reason=( + "A commit that loses the checkpoint-barrier race is not yet retained in " + "a durable unfinished-group ledger." + ), +) +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()) + _enable_recovery_checkpoint_capture(controller) + + 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()) + + +@pytest.mark.xfail( + strict=True, + reason=( + "A canonical replay group is not yet filtered out of the checkpointed " + "unfinished-group ledger." + ), +) +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()) + _enable_recovery_checkpoint_capture(controller) + + 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()) + + +@pytest.mark.xfail( + strict=True, + reason=( + "The controller does not yet persist and replay unfinished prompt-group " + "ownership alongside the TQ checkpoint." + ), +) +def test_recovery_replays_step_7_without_readmitting_the_batch() -> None: + """An admitted batch keeps target_step=7 across a process restart.""" + + async def exercise() -> None: + sampler = _CountingInOrderSampler() + sampler.restore_dispatch_index(7) + ledger = _PendingLedger( + _PendingGroup( + group_id="batch-7-prompt-0", + target_step=7, + prompt_payload={"idx": 70, "message_log": []}, + ) + ) + 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._data_plane_checkpoint_metadata = { + "rollout_recovery_payload_sha256": "checkpoint-cut-digest" + } + controller._async_cfg = SimpleNamespace(max_buffered_rollouts=4) + controller._buffer_capacity = asyncio.Semaphore(4) + + async def _inventory_is_valid(*, clear_unreferenced: bool) -> None: + assert clear_unreferenced + + controller._validate_rollout_recovery_inventory = _inventory_is_valid + + await controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + + assert ledger.prepare_calls == 1 + assert rollout_manager.recovered == [("batch-7-prompt-0", 7)] + assert sampler.admit_calls == 0 + assert sampler.dispatch_index == 7 + + asyncio.run(exercise()) 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..fe26792c733 --- /dev/null +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -0,0 +1,118 @@ +# 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; ready_first is included here because it now advertises the same +completed-buffer recovery capability. + +Normal test runs keep the unfinished-group rows as strict xfails so #3480 can +land without claiming partial-rollout recovery. During development of the +dispatch ledger, run this file with --runxfail: those rows become the red TDD +contract and must all pass before the xfail marks are removed. +""" + +from __future__ import annotations + +import pytest + +from tests.unit.single_controller._checkpoint_scenarios import ( + ALL_SCENARIOS, + FULLY_GENERATED, + 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, + "The dataloader has advanced, but metadata_state_dict() omits every " + "ready=False reservation. A dispatch ledger and recovery pump must " + "redispatch the prompt group after restart.", + ) + for sampler in SAMPLERS + for scenario in WITH_IN_FLIGHT +] + + +def _known_gap(case: Case): + return pytest.param( + case, + id=case.id, + marks=pytest.mark.xfail(strict=True, reason=case.why), + ) + + +@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", [_known_gap(case) for case in UNFINISHED_CASES]) +def test_unfinished_groups_are_owned_across_restart(case, tmp_path): + """Desired end state: handed-out unfinished groups remain recoverable.""" + assert_no_data_loss(case.scenario, case.sampler, tmp_path) + + +@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 From e34bb7f7a7640cb94af2d4b3a391f9e7ff8baaca Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 17:36:12 -0400 Subject: [PATCH 18/32] feat(sc): recover unfinished rollouts from checkpoints Signed-off-by: Anish Mahishi --- .../algorithms/async_utils/replay_buffer.py | 7 +- .../async_utils/staleness_sampler.py | 41 +- nemo_rl/algorithms/single_controller.py | 668 +++++++++++++++--- nemo_rl/experience/rollout_manager.py | 104 ++- nemo_rl/experience/rollout_recovery.py | 533 ++++++++++++++ pyrefly.toml | 1 + .../L1_Functional_Tests_SingleController.sh | 5 +- ...single_controller_rollout_recovery_hook.py | 137 ++++ tests/functional/grpo_dp_single_controller.sh | 3 +- ...p_single_controller_unfinished_recovery.sh | 79 +++ tests/unit/experience/test_rollout_manager.py | 79 +++ .../unit/experience/test_rollout_recovery.py | 285 ++++++++ .../_checkpoint_scenarios.py | 48 +- .../test_checkpoint_dispatch_races.py | 595 +++++++++++++--- .../test_checkpoint_recovery_matrix.py | 36 +- .../single_controller/test_checkpointing.py | 46 +- .../single_controller/test_rollout_pump.py | 1 + .../test_sampler_interface.py | 11 + .../test_tq_replay_buffer.py | 31 + 19 files changed, 2479 insertions(+), 231 deletions(-) create mode 100644 nemo_rl/experience/rollout_recovery.py create mode 100644 tests/functional/_single_controller_rollout_recovery_hook.py create mode 100644 tests/functional/grpo_dp_single_controller_unfinished_recovery.sh create mode 100644 tests/unit/experience/test_rollout_recovery.py diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 3edbd0bc46a..fc8152893a6 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -93,6 +93,9 @@ class DataPlaneCheckpointMetadata(TypedDict): 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: @@ -982,7 +985,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. diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index cb32575a151..90874620a19 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -131,6 +131,21 @@ def restore_dispatch_index(self, dispatch_index: int) -> None: ... +@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) -> Optional[int]: + """Advance the admission cursor and return the batch target-step stamp.""" + ... + + class BaseSampler(abc.ABC): """Shared machinery for the built-in policies. @@ -311,10 +326,20 @@ def should_abort_inflight( ) return start_weight_version < min_valid_version - async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]: - # Over-sampled: dispatch is bounded by buffer capacity, not by 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) -> Optional[int]: + """Return the unstamped admission result without changing a cursor.""" return None + 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() + async def select( self, *, @@ -378,12 +403,22 @@ 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) -> Optional[int]: + """Advance the cursor after the controller enters its mutation cut.""" 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 diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 6e1a08c635a..8dffa71b2d4 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -39,14 +39,18 @@ from __future__ import annotations import asyncio +import hashlib +import io import logging import math import os import time +import uuid import warnings from collections import deque 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 @@ -60,7 +64,10 @@ DataPlaneCheckpointMetadata, TQReplayMetadataState, ) -from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + TransactionalAdmissionSampler, + create_sampler, +) from nemo_rl.algorithms.grpo import ( GRPOSaveState, _write_latest_checkpoint_status, @@ -92,6 +99,12 @@ from nemo_rl.environments.nemo_gym import should_use_nemo_gym from nemo_rl.experience.failures import RolloutStall 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, +) from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmGeneration from nemo_rl.models.policy.tq_policy import TQPolicy @@ -245,6 +258,21 @@ def __init__( "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, @@ -349,7 +377,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 @@ -414,7 +445,7 @@ async def ping(self) -> dict[str, Any]: # ── internal helpers ─────────────────────────────────────────────────── - async def _maybe_restore_replay_buffer(self) -> None: + async def _maybe_restore_replay_buffer(self) -> int: """Restore the local replay index for the native TQ checkpoint. Recovery is authoritative only for samplers that explicitly support @@ -422,7 +453,7 @@ async def _maybe_restore_replay_buffer(self) -> None: must both be present and agree on their manifest and group count. """ if self._last_checkpoint_path is None: - return + return 0 metadata_path = os.path.join( self._last_checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME ) @@ -436,7 +467,7 @@ async def _maybe_restore_replay_buffer(self) -> None: "replay-buffer recovery" ) if not self._sampler.supports_buffer_checkpoint: - return + return 0 if not os.path.exists(metadata_path): legacy_path = os.path.join( self._last_checkpoint_path, LEGACY_REPLAY_BUFFER_FILENAME @@ -453,7 +484,7 @@ async def _maybe_restore_replay_buffer(self) -> None: "Starting with an empty replay buffer.", flush=True, ) - return + 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. @@ -499,6 +530,280 @@ async def _maybe_restore_replay_buffer(self) -> None: 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, + ) + if not isinstance(state, dict): + raise TypeError( + "rollout recovery sidecar must contain a dictionary, got " + f"{type(state).__name__}" + ) + groups = state.get("groups") + if not isinstance(groups, list) or len(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 + recovery_ledger.load_state_dict(cast(RolloutRecoveryState, state)) + raw_batch_shortfall = state.get("batch_shortfall", {}) + if not isinstance(raw_batch_shortfall, dict): + raise TypeError("rollout recovery batch_shortfall must be a dictionary") + restored_batch_shortfall: dict[int, int] = {} + for step, count in raw_batch_shortfall.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}" + ) + restored_batch_shortfall[step] = count + 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" + ) + self._batch_shortfall = restored_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(canonical_group_ids) + await self._rehydrate_rollout_recovery_prompts() + self._sampler_stamps_target_steps = ( + raw_sampler_stamps + if raw_sampler_stamps 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) -> None: + """Resolve durable prompt references against the restored dataset.""" + 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: + 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(prompt, dict): + raise TypeError( + f"dataset sample_id={sample_id!r} resolved to " + f"{type(prompt).__name__}, expected a DatumSpec dictionary" + ) + resolved_prompts[sample_id] = cast(DatumSpec, prompt) + recovery_ledger.bind_runtime_prompt(group.group_id, 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(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( + 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(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(): + target_step = self._sampler.commit_admission() + return _commit(target_step) + + # Custom samplers retain their existing monolithic admission API. Hold + # the mutation cut across it for correctness; custom implementations can + # opt into TransactionalAdmissionSampler to avoid delaying checkpoints + # while their gate waits. + async with self._data_plane_checkpoint_barrier.mutation(): + target_step = await self._sampler.admit( + trainer_version_fn=lambda: self._trainer_version + ) + return _commit(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 + + # A checkpoint may land after dataloader ownership is recorded but before + # sampler admission commits. Re-admit each original dataloader batch once; + # ADMITTED groups retain their original target step and cursor position. + 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(): + await self._admit_reserved_prompt_groups(group_ids) + + refreshed_groups = recovery_ledger.groups() + for group in refreshed_groups: + await launch( + group.prompt_payload, + group.target_step, + group.group_id, + ) + + print( + f"📦 Redispatched {len(refreshed_groups)} unfinished rollout " + "group(s) before new dataloader work", + flush=True, + ) async def _validate_replay_inventory( self, replay_metadata: TQReplayMetadataState @@ -605,6 +910,8 @@ 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. @@ -642,6 +949,18 @@ async def _save_data_plane_checkpoint( ) 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: @@ -699,6 +1018,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() @@ -713,11 +1033,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. @@ -727,12 +1055,37 @@ 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(): + 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(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( + 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 @@ -747,9 +1100,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. @@ -784,7 +1138,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 @@ -795,7 +1158,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) @@ -808,36 +1176,87 @@ 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(): + 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( + 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 @@ -898,7 +1317,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. @@ -922,6 +1342,45 @@ 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(): + 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( + 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. @@ -937,7 +1396,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( @@ -1862,26 +2321,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.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 - # Snapshot before any await so it can't interleave with - # _rollout_pump advancing either the sampler cursor or this 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"): @@ -1911,12 +2350,85 @@ 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 = ( + self._rollout_manager.recovery_ledger.state_dict() + ) + rollout_recovery_state["batch_shortfall"] = ( + self._batch_shortfall.copy() + ) + rollout_recovery_state["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: @@ -1970,28 +2482,20 @@ async def _save_checkpoint( reserve_state, os.path.join(checkpoint_path, "replacement_reserve.pt"), ) - replay_metadata: Optional[TQReplayMetadataState] = None - if self._master_config.checkpointing.get("save_data_plane"): - # Commits and destructive clears take the same barrier. Generation - # may continue while a snapshot is written, but completed groups - # wait at commit, so TQ and the replay metadata file describe exactly - # the same set of training-ready groups. - async with self._data_plane_checkpoint_barrier.checkpoint(): - if self._sampler.supports_buffer_checkpoint: - replay_metadata = self._buffer.metadata_state_dict( - saved_capacity=self._async_cfg.max_buffered_rollouts - ) - if replay_metadata is not None: - await self._validate_replay_inventory(replay_metadata) - await self._save_data_plane_checkpoint( - checkpoint_path, replay_metadata=replay_metadata - ) if replay_metadata is not None: await asyncio.to_thread( torch.save, replay_metadata, os.path.join(checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME), ) + if rollout_recovery_payload is not None: + await asyncio.to_thread( + Path( + checkpoint_path, + ROLLOUT_RECOVERY_STATE_FILENAME, + ).write_bytes, + rollout_recovery_payload, + ) # Rename happens in the background once the async weight writes # finish; flushed at the next save or on exit. self._checkpointer.begin_finalization( diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 784ae5a7042..e64212b086b 100644 --- a/nemo_rl/experience/rollout_manager.py +++ b/nemo_rl/experience/rollout_manager.py @@ -40,6 +40,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, @@ -69,7 +73,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" @@ -1184,6 +1189,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. @@ -1198,6 +1204,54 @@ 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, + 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( + 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, + group_id: str, + *, + target_step: Optional[int], + ) -> None: + """Attach sampler admission state to a pre-admission reservation.""" + self._recovery_ledger.mark_group_admitted( + group_id, + target_step=target_step, + start_weight_version=self._weight_version, + ) + + def discard_prompt_group(self, group_id: str) -> None: + """Release a reservation that will intentionally never be dispatched.""" + self._recovery_ledger.discard_group(group_id) + def set_weight_version(self, version: int) -> None: """Set the weight_version used for rollout tags. @@ -1215,6 +1269,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. @@ -1235,14 +1290,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 @@ -1252,6 +1311,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 @@ -1263,15 +1336,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 @@ -1293,13 +1368,19 @@ 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: + # A retry cannot safely reuse the stable ID while the previous + # slot may still exist. Re-raise the original rollout failure. + raise reason = type(error).__name__ if classify_rollout_failure(error) is FailureClass.INFRA: @@ -1360,6 +1441,8 @@ 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: + self._recovery_ledger.discard_group(lineage_group_id) return RolloutOutcome.COMMITTED # The infrastructure budget ran out. The same failure followed the prompt across @@ -1386,7 +1469,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..08c3aaafb19 --- /dev/null +++ b/nemo_rl/experience/rollout_recovery.py @@ -0,0 +1,533 @@ +# 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 hashlib +import pickle +import uuid +from dataclasses import dataclass, replace +from enum import StrEnum +from typing import TYPE_CHECKING, Any, NotRequired, TypedDict + +if TYPE_CHECKING: + from nemo_rl.data.interfaces import DatumSpec + +ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 +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 + payload_sha256: str + + +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 RolloutRecoveryState(TypedDict): + """Versioned checkpoint sidecar for unfinished prompt groups.""" + + schema_version: int + groups: list[PromptGroupRecoveryState] + batch_shortfall: NotRequired[dict[int, int]] + sampler_stamps_target_steps: NotRequired[bool] + + +@dataclass(frozen=True) +class PromptRef: + """Stable dataset identity and integrity check for one prompt.""" + + sample_id: str + task_name: str | None + payload_sha256: str | None = 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 + + +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_payload_sha256(prompt_payload: object) -> str: + """Fingerprint a prompt so dataset rehydration cannot silently change it.""" + return hashlib.sha256( + pickle.dumps(prompt_payload, protocol=pickle.HIGHEST_PROTOCOL) + ).hexdigest() + + +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}" + ) + + +def _validate_prompt_ref( + prompt_ref: PromptRef, + prompt_payload: DatumSpec, + *, + group_id: str, +) -> str: + _validate_prompt_identity(prompt_ref, prompt_payload, group_id=group_id) + payload_sha256 = prompt_payload_sha256(prompt_payload) + if ( + prompt_ref.payload_sha256 is not None + and payload_sha256 != prompt_ref.payload_sha256 + ): + raise ValueError( + f"recovery group {group_id!r} prompt fingerprint mismatch for " + f"sample_id={prompt_ref.sample_id!r}" + ) + return payload_sha256 + + +class RolloutRecoveryLedger: + """Own prompts after dataloader advance and before canonical TQ commit.""" + + def __init__(self) -> None: + self._groups: dict[str, PromptGroupRecoveryRecord] = {} + + def reserve_group( + self, + *, + 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: + prompt_id: Dataset-level prompt identity used for diagnostics. + prompt_payload: Runtime prompt used for whole-group regeneration. Only + its stable dataset reference and fingerprint are 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. + """ + 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 locator and fingerprint. + 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, + group_id: str, + *, + target_step: int | None, + start_weight_version: int, + ) -> None: + """Attach the sampler result to a previously reserved prompt group.""" + 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, + group_id: str, + prompt_payload: DatumSpec, + ) -> None: + """Attach and verify a dataset-rehydrated prompt after checkpoint load.""" + record = self._require_group(group_id) + payload_sha256 = _validate_prompt_ref( + 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, + payload_sha256=payload_sha256, + ), + 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, group_id: str) -> None: + """Release ownership after canonical commit or intentional discard.""" + self._require_group(group_id) + del self._groups[group_id] + + def discard_canonical_groups(self, group_ids: set[str]) -> int: + """Drop ledger copies already owned by canonical replay metadata.""" + 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) -> RolloutRecoveryState: + """Return versioned references without serializing full prompt payloads.""" + groups: list[PromptGroupRecoveryState] = [] + for group_id, record in list(self._groups.items()): + 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, + ) + payload_sha256 = record.prompt_ref.payload_sha256 + if payload_sha256 is None: + payload_sha256 = prompt_payload_sha256(prompt_payload) + record = replace( + record, + prompt_ref=replace( + record.prompt_ref, + payload_sha256=payload_sha256, + ), + ) + # Prompts are immutable after dataloader processing. Cache the + # first durable fingerprint so repeated checkpoints do not + # serialize the same long prompt merely to hash it again. + self._groups[group_id] = record + 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, + "payload_sha256": payload_sha256, + }, + "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, state: RolloutRecoveryState) -> None: + """Replace this empty ledger from a validated checkpoint payload.""" + 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") + payload_sha256 = raw_prompt_ref.get("payload_sha256") + 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" + ) + if ( + not isinstance(payload_sha256, str) + or len(payload_sha256) != 64 + or any( + character not in "0123456789abcdef" for character in payload_sha256 + ) + ): + raise ValueError( + f"rollout recovery groups[{index}].prompt_ref.payload_sha256 " + "must be a lowercase SHA-256 digest" + ) + 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, + payload_sha256=payload_sha256, + ), + 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) diff --git a/pyrefly.toml b/pyrefly.toml index 2a203aa00be..cff3ea008be 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -156,6 +156,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", diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index c06ebf75b5a..55f33c16eae 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -81,8 +81,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 replay checkpoint recovery (#3480). +# 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 +# Full mode: deterministic process restart with an admitted group held before +# canonical TQ commit, followed by exact-once redispatch at its stable group ID. +run_test 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_dp_single_controller.sh b/tests/functional/grpo_dp_single_controller.sh index 003f762347e..7aaabbf3146 100755 --- a/tests/functional/grpo_dp_single_controller.sh +++ b/tests/functional/grpo_dp_single_controller.sh @@ -16,6 +16,7 @@ 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 @@ -23,7 +24,7 @@ mkdir -p $EXP_DIR $LOG_DIR cd $PROJECT_ROOT uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \ - $PROJECT_ROOT/examples/run_grpo_single_controller.py \ + $SC_ENTRYPOINT \ policy.model_name=Qwen/Qwen3-0.6B \ grpo.num_prompts_per_step=2 \ grpo.num_generations_per_prompt=4 \ 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..b52f716f440 --- /dev/null +++ b/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh @@ -0,0 +1,79 @@ +#!/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 +) + +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/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 0a73393e7db..9589e9e3281 100644 --- a/tests/unit/experience/test_rollout_manager.py +++ b/tests/unit/experience/test_rollout_manager.py @@ -42,9 +42,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, @@ -140,6 +142,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 @@ -244,6 +247,82 @@ 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 = mgr.reserve_prompt_group( + 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") + + mgr = _make_manager( + _FakeBuffer(), + _FakeImpl(on_run=_fail_rollout), + RolloutRetryPolicy.single_attempt(max_skipped_prompts=1), + ) + group_id = mgr.reserve_prompt_group( + {"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): + mgr = _make_manager(_FakeBuffer(), _FakeImpl()) + mgr.recovery_ledger.reserve_group( + 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..eef76f3d97d --- /dev/null +++ b/tests/unit/experience/test_rollout_recovery.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import pytest +import torch +from torchdata.stateful_dataloader import StatefulDataLoader + +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.experience.rollout_recovery import ( + ROLLOUT_RECOVERY_SCHEMA_VERSION, + PromptGroupPhase, + RolloutRecoveryLedger, + prompt_payload_sha256, +) + + +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: + prompt = _prompt(idx) + return { + "group_id": f"g{idx}", + "admission_id": "batch-7", + "prompt_id": str(idx), + "prompt_ref": { + "sample_id": str(idx), + "task_name": None, + "payload_sha256": prompt_payload_sha256(prompt), + }, + "expected_generations": 2, + "target_step": target_step, + "start_weight_version": 7, + "phase": phase, + } + + +def test_ledger_round_trip_preserves_group_ownership() -> None: + ledger = RolloutRecoveryLedger() + ledger.reserve_group( + 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() + restored.load_state_dict(state) + + with pytest.raises(RuntimeError, match="has not rehydrated prompt"): + _ = restored.get_group("g7").prompt_payload + restored.bind_runtime_prompt("g7", _prompt()) + + assert restored.state_dict() == state + assert restored.get_group("g7").phase is PromptGroupPhase.ADMITTED + + +def test_target_step_none_does_not_mean_unadmitted() -> None: + ledger = RolloutRecoveryLedger() + record = ledger.reserve_group( + 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() + ledger.reserve_group( + 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, + ) + + ledger.mark_group_admitted( + "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"): + ledger.mark_group_admitted( + "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): + ledger.reserve_group( + 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 ledger.discard_canonical_groups({"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() + ledger.reserve_group( + 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, + "payload_sha256": prompt_payload_sha256(prompt), + } + group_state["prompt_ref"]["sample_id"] = "100" + + assert ledger.get_group("g7").prompt_ref.sample_id == "7" + + +def test_bind_runtime_prompt_rejects_changed_dataset_content() -> None: + ledger = RolloutRecoveryLedger() + original = _prompt() + ledger.reserve_group( + 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() + restored.load_state_dict(ledger.state_dict()) + + changed = _prompt() + changed["message_log"][0]["content"] = "different prompt" + with pytest.raises(ValueError, match="fingerprint mismatch"): + restored.bind_runtime_prompt("g7", changed) + + +def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: + ledger = RolloutRecoveryLedger() + ledger.reserve_group( + 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() + restored.load_state_dict(ledger.state_dict()) + + with pytest.raises(ValueError, match="expected '7'"): + restored.bind_runtime_prompt("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() + ledger.reserve_group( + 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() + restored_ledger.load_state_dict(ledger_state) + restored_group = restored_ledger.get_group("unfinished") + dataset_prompt = restored_dataloader.dataset[ + int(restored_group.prompt_ref.sample_id) + ] + restored_ledger.bind_runtime_prompt("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)): + RolloutRecoveryLedger().load_state_dict(state) # type: ignore[arg-type] diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index c4691a1c378..7935a4cbc37 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -56,6 +56,7 @@ 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 @@ -132,7 +133,7 @@ def must_survive(self) -> set[str]: Handed out, not trained, not deliberately evicted -- and below the cursor, so the dataloader will never produce them again. This is the - bar the feature should eventually meet, not the bar it meets today. + no-data-loss bar enforced by the recovery matrix. """ return { _gid(g.gid) @@ -162,7 +163,7 @@ class Case: Args: scenario: The buffer state at checkpoint time. sampler: Which sampler the run is configured with. - why: For a case that fails today, the line that drops the data. + why: Optional diagnostic context for an expected behavior gap. """ scenario: Scenario @@ -256,10 +257,10 @@ class RoundTrip: ``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 how a future partial-group restore is built. A - group could come back already committed (its missing rollouts regenerated - before the save), or as a reserved slot waiting to be finished -- either - way the run has not lost the prompt, and either way these tests notice. + 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. """ @@ -287,6 +288,25 @@ async def _round_trip( if sampler_a.supports_buffer_checkpoint else None ) + recovery_ledger_a = RolloutRecoveryLedger() + 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( + 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") @@ -305,6 +325,22 @@ async def _round_trip( expected_group_size=ROLLOUTS_PER_GROUP, expected_manifest_digest=sidecar["manifest_digest"], ) + recovery_ledger_b = RolloutRecoveryLedger() + recovery_ledger_b.load_state_dict(recovery_sidecar) + recovery_ledger_b.discard_canonical_groups(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, + ) + recovery_ledger_b.discard_group(group_id) ready = { gid for gid, is_ready in zip(buf_b._group_ids, buf_b.ready_list) if is_ready diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 72aa50f25e1..4c210924138 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -29,7 +29,9 @@ from __future__ import annotations import asyncio +import hashlib import threading +from collections import deque from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -40,6 +42,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( REPLAY_BUFFER_METADATA_FILENAME, + DataPlaneCheckpointBarrier, TQReplayBuffer, ) from nemo_rl.algorithms.async_utils.staleness_sampler import InOrderSampler @@ -47,8 +50,16 @@ from nemo_rl.algorithms.metric_utils import SetupTimingMetrics from nemo_rl.algorithms.single_controller import SingleControllerActor from nemo_rl.data_plane.adapters.noop import NoOpDataPlaneClient +from nemo_rl.data.interfaces import DatumSpec 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, + prompt_payload_sha256, +) from tests.unit.single_controller._checkpoint_scenarios import ( _record, patch_converter, @@ -66,11 +77,16 @@ class _CountingInOrderSampler(InOrderSampler): 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): + self.admission_commits += 1 + return super().commit_admission() + class _BlockingBeforeAdmissionSampler(_CountingInOrderSampler): """Pause after the dataloader advances but before admission mutates state.""" @@ -80,10 +96,10 @@ def __init__(self) -> None: self.admission_entered = asyncio.Event() self.release_admission = asyncio.Event() - async def admit(self, *, trainer_version_fn): + async def wait_until_admissible(self, *, trainer_version_fn): self.admission_entered.set() await self.release_admission.wait() - return await super().admit(trainer_version_fn=trainer_version_fn) + await super().wait_until_admissible(trainer_version_fn=trainer_version_fn) @dataclass(frozen=True) @@ -126,12 +142,21 @@ def assign_target_step(self, group_id: str, target_step: int) -> None: def state_dict(self) -> dict[str, Any]: return { - "schema_version": 1, + "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_payload": group.prompt_payload, + "prompt_ref": { + "sample_id": str(group.prompt_payload.get("idx", "unknown")), + "task_name": group.prompt_payload.get("task_name"), + "payload_sha256": prompt_payload_sha256(group.prompt_payload), + }, + "expected_generations": 2, + "start_weight_version": 7, + "phase": ("reserved" if group.target_step is None else "admitted"), } for group in self._groups ], @@ -142,19 +167,26 @@ def release(self, group_id: str) -> None: class _RecoveryRolloutManager: - def __init__(self, ledger: _PendingLedger) -> None: + def __init__(self, ledger: RolloutRecoveryLedger) -> None: self.recovery_ledger = ledger self.recovered: list[tuple[str, int | None]] = [] - async def recover_group(self, group_id: str) -> bool: - group = next( - group - for group in self.recovery_ledger.groups() - if group.group_id == group_id - ) + async def complete_recovery(self, group_id: str) -> None: + group = self.recovery_ledger.get_group(group_id) self.recovered.append((group.group_id, group.target_step)) - self.recovery_ledger.release(group_id) - return True + self.recovery_ledger.discard_group(group_id) + + def mark_prompt_group_admitted( + self, group_id: str, *, target_step: int | None + ) -> None: + self.recovery_ledger.mark_group_admitted( + group_id, + target_step=target_step, + start_weight_version=7, + ) + + def discard_prompt_group(self, group_id: str) -> None: + self.recovery_ledger.discard_group(group_id) class _BlockingRolloutManager: @@ -170,8 +202,14 @@ def set_weight_version(self, version: int) -> None: self.weight_version = version def reserve_prompt_group( - self, prompt: dict[str, Any], *, target_step: int | None = None + self, + 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(): @@ -184,20 +222,27 @@ def reserve_prompt_group( ) return group_id - def mark_prompt_group_admitted(self, group_id: str, *, target_step: int) -> None: + def mark_prompt_group_admitted( + self, group_id: str, *, target_step: int | None + ) -> None: + if target_step is None: + return self.recovery_ledger.assign_target_step(group_id, target_step) + def discard_prompt_group(self, group_id: str) -> None: + self.recovery_ledger.release(group_id) + async def generate_and_push( self, - prompt: dict[str, Any], + prompt: DatumSpec, *, target_step: int | None = None, inflight_registry: dict[str, Any] | None = None, - recovery_group_id: str | None = None, + lineage_group_id: str | None = None, ) -> RolloutOutcome: del inflight_registry - if recovery_group_id is None: - recovery_group_id = self.reserve_prompt_group( + if lineage_group_id is None: + lineage_group_id = self.reserve_prompt_group( prompt, target_step=target_step, ) @@ -225,17 +270,71 @@ def save_checkpoint( super().save_checkpoint(checkpoint_dir, metadata=metadata) -def _enable_recovery_checkpoint_capture(controller: Any) -> None: - """Install the narrow recovery hooks expected by the future foundation.""" +class _LedgerFacade: + """Minimal RolloutManager ownership surface for reserve-pool cuts.""" + + def __init__(self) -> None: + self.recovery_ledger = RolloutRecoveryLedger() - async def _inventory_is_valid(**_: Any) -> None: - return None + def reserve_prompt_group( + self, + prompt: DatumSpec, + *, + target_step: int | None = None, + admitted: bool = True, + admission_id: str | None = None, + ) -> str: + record = self.recovery_ledger.reserve_group( + 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 - controller._validate_rollout_recovery_inventory = _inventory_is_valid - controller._master_config.__dict__["token_capture"] = SimpleNamespace( - enabled=True, - staging_partition="rollout_staging", + def mark_prompt_group_admitted( + self, group_id: str, *, target_step: int | None + ) -> None: + self.recovery_ledger.mark_group_admitted( + group_id, + target_step=target_step, + start_weight_version=7, + ) + + def discard_prompt_group(self, group_id: str) -> None: + self.recovery_ledger.discard_group(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 _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: @@ -249,13 +348,6 @@ async def exercise() -> int | None: assert asyncio.run(exercise()) == 8 -@pytest.mark.xfail( - strict=True, - reason=( - "The dataloader cursor can advance before unfinished prompt ownership " - "is recorded in the checkpoint bundle." - ), -) def test_checkpoint_after_fetch_before_admit_owns_the_prompt(tmp_path) -> None: """A checkpoint cut inside admit retains the fetched batch for recovery.""" @@ -297,8 +389,6 @@ async def exercise() -> None: sampler = _BlockingBeforeAdmissionSampler() sampler.restore_dispatch_index(6) controller._sampler = sampler - _enable_recovery_checkpoint_capture(controller) - pump = asyncio.create_task(controller._rollout_pump()) await asyncio.wait_for(sampler.admission_entered.wait(), timeout=1.0) assert controller._sampler.dispatch_index == 6 @@ -322,7 +412,8 @@ async def exercise() -> None: ) assert len(recovery_state["groups"]) == 1 assert recovery_state["groups"][0]["target_step"] is None - assert recovery_state["groups"][0]["prompt_payload"]["idx"] == 70 + 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, @@ -331,13 +422,6 @@ async def exercise() -> None: asyncio.run(exercise()) -@pytest.mark.xfail( - strict=True, - reason=( - "An unfinished admitted batch is not yet persisted alongside the native " - "TQ checkpoint." - ), -) 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.""" @@ -379,8 +463,6 @@ async def exercise() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = controller_cls(config, actor_args, SetupTimingMetrics()) - _enable_recovery_checkpoint_capture(controller) - pump = asyncio.create_task(controller._rollout_pump()) await asyncio.wait_for(rollout_manager.started.wait(), timeout=1.0) assert controller._sampler.dispatch_index == 7 @@ -408,13 +490,6 @@ async def exercise() -> None: asyncio.run(exercise()) -@pytest.mark.xfail( - strict=True, - reason=( - "A commit that loses the checkpoint-barrier race is not yet retained in " - "a durable unfinished-group ledger." - ), -) def test_commit_contending_with_checkpoint_has_exactly_one_saved_owner( tmp_path, monkeypatch, @@ -472,8 +547,6 @@ async def exercise() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = controller_cls(config, actor_args, SetupTimingMetrics()) - _enable_recovery_checkpoint_capture(controller) - save_task = asyncio.create_task( controller._save_checkpoint( {"loss": 1.0}, @@ -509,12 +582,8 @@ async def exercise() -> None: 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"] - } + 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 @@ -524,13 +593,6 @@ async def exercise() -> None: asyncio.run(exercise()) -@pytest.mark.xfail( - strict=True, - reason=( - "A canonical replay group is not yet filtered out of the checkpointed " - "unfinished-group ledger." - ), -) def test_canonical_replay_wins_over_stale_ledger_entry( tmp_path, monkeypatch, @@ -591,8 +653,6 @@ async def exercise() -> None: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = controller_cls(config, actor_args, SetupTimingMetrics()) - _enable_recovery_checkpoint_capture(controller) - await buffer.commit( group_id, _record(), @@ -618,12 +678,8 @@ async def exercise() -> None: 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"] - } + 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 @@ -632,48 +688,393 @@ async def exercise() -> None: asyncio.run(exercise()) -@pytest.mark.xfail( - strict=True, - reason=( - "The controller does not yet persist and replay unfinished prompt-group " - "ownership alongside the TQ checkpoint." - ), -) -def test_recovery_replays_step_7_without_readmitting_the_batch() -> None: +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) - ledger = _PendingLedger( - _PendingGroup( - group_id="batch-7-prompt-0", - target_step=7, - prompt_payload={"idx": 70, "message_log": []}, - ) + saved_ledger = RolloutRecoveryLedger() + saved_ledger.reserve_group( + group_id="batch-7-prompt-0", + 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, ) + saved_state = saved_ledger.state_dict() + saved_state["batch_shortfall"] = {6: 1} + saved_state["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_payload_sha256": "checkpoint-cut-digest" + "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) + controller._async_cfg = SimpleNamespace( + max_buffered_rollouts=4, + max_inflight_prompts=2, + ) controller._buffer_capacity = asyncio.Semaphore(4) + controller._trainer_version = 7 + controller._dataloader = SimpleNamespace( + dataset={70: {"idx": 70, "message_log": []}} + ) + controller._buffer = SimpleNamespace( + count_for_target_step=lambda _target_step: 0, + metadata_state_dict=lambda *, saved_capacity: { + "groups": [], + "saved_capacity": saved_capacity, + }, + ) - async def _inventory_is_valid(*, clear_unreferenced: bool) -> None: - assert clear_unreferenced - - controller._validate_rollout_recovery_inventory = _inventory_is_valid + async def _recover( + _prompt: dict[str, Any], + _target_step: int | None, + group_id: str, + ) -> None: + await rollout_manager.complete_recovery(group_id) await controller._maybe_restore_rollout_recovery(restored_replay_groups=0) + await controller._redispatch_restored_rollouts(_recover) - assert ledger.prepare_calls == 1 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_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() + for prompt_idx in (70, 71): + saved_ledger.reserve_group( + 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: + await rollout_manager.complete_recovery(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_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() + for prompt_idx in (70, 71): + saved_ledger.reserve_group( + 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._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 + old_group_id = manager.reserve_prompt_group( + _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(): + replacement = controller._take_replacement(7, 0) + assert replacement is not None + manager.discard_prompt_group(old_group_id) + manager.reserve_prompt_group(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() + for group_id in group_ids: + controller._rollout_manager.mark_prompt_group_admitted( + 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 index fe26792c733..e75c5a11eb8 100644 --- a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -18,10 +18,8 @@ in_order; ready_first is included here because it now advertises the same completed-buffer recovery capability. -Normal test runs keep the unfinished-group rows as strict xfails so #3480 can -land without claiming partial-rollout recovery. During development of the -dispatch ledger, run this file with --runxfail: those rows become the red TDD -contract and must all pass before the xfail marks are removed. +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 @@ -41,36 +39,16 @@ ) ALL_CASES = [ - Case(scenario, sampler) - for sampler in SAMPLERS - for scenario in ALL_SCENARIOS + 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 + Case(scenario, sampler) for sampler in SAMPLERS for scenario in FULLY_GENERATED ] UNFINISHED_CASES = [ - Case( - scenario, - sampler, - "The dataloader has advanced, but metadata_state_dict() omits every " - "ready=False reservation. A dispatch ledger and recovery pump must " - "redispatch the prompt group after restart.", - ) - for sampler in SAMPLERS - for scenario in WITH_IN_FLIGHT + Case(scenario, sampler) for sampler in SAMPLERS for scenario in WITH_IN_FLIGHT ] -def _known_gap(case: Case): - return pytest.param( - case, - id=case.id, - marks=pytest.mark.xfail(strict=True, reason=case.why), - ) - - @pytest.fixture(autouse=True) def _converter(monkeypatch): patch_converter(monkeypatch) @@ -93,9 +71,9 @@ def test_fully_generated_scenarios_have_no_data_loss(case, tmp_path): assert result.recovered == case.scenario.must_survive() -@pytest.mark.parametrize("case", [_known_gap(case) for case in UNFINISHED_CASES]) +@pytest.mark.parametrize("case", UNFINISHED_CASES, ids=lambda case: case.id) def test_unfinished_groups_are_owned_across_restart(case, tmp_path): - """Desired end state: handed-out unfinished groups remain recoverable.""" + """Handed-out unfinished groups remain recoverable.""" assert_no_data_loss(case.scenario, case.sampler, tmp_path) diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index e3b8d98b330..a5cac0e7be2 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -35,6 +35,7 @@ from __future__ import annotations import asyncio +import hashlib import json import os import threading @@ -79,6 +80,11 @@ from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs from nemo_rl.data.utils import load_dataloader_state from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta +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 @@ -351,6 +357,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) @@ -1019,12 +1026,19 @@ def test_saves_authoritative_tq_state_and_metadata_only_replay_index( assert save_call["checkpoint_dir"] == str( tmp_path / "checkpoints" / "tmp_step_1" / "data_plane" ) - assert save_call["metadata"] == _data_plane_checkpoint_metadata( + 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 ( @@ -1032,6 +1046,15 @@ def test_saves_authoritative_tq_state_and_metadata_only_replay_index( == 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( @@ -1186,6 +1209,27 @@ async def _main() -> int: 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() diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 8ce81627b4e..54a65cf3e95 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -91,6 +91,7 @@ def _init_pump_ledgers(ctrl: Any) -> None: ctrl._batch_replacements = {} ctrl._batch_promotions = {} ctrl._replacement_reserve = deque() + ctrl._rollout_recovery_enabled = False class _RecordingBuffer: diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index 7d47453ce5a..fbf60e18763 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -35,6 +35,7 @@ ReadyFirstSampler, ReadyFirstSamplerConfig, SamplerConfig, + TransactionalAdmissionSampler, WeightFifoSampler, WeightFifoSamplerConfig, WindowedSampler, @@ -105,9 +106,19 @@ 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 + assert sampler.commit_admission() == 0 + assert sampler.dispatch_index == 0 + 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. diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 8f84825bfea..451ff9ac735 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -148,6 +148,14 @@ 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) @@ -334,6 +342,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) From 45246ccb8f1c9543af468c37530101461caf9120 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 18:01:46 -0400 Subject: [PATCH 19/32] test(sc): stabilize checkpoint recovery coverage Signed-off-by: Anish Mahishi --- .../test_checkpoint_dispatch_races.py | 36 ++++++++++++++++--- .../single_controller/test_checkpointing.py | 1 + tests/unit/single_controller/test_setup.py | 5 ++- 3 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 4c210924138..33779df8adc 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -70,6 +70,32 @@ _make_actor_args, ) +_ASYNC_TEST_TIMEOUT_S = 10.0 + + +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.""" @@ -390,7 +416,7 @@ async def exercise() -> None: sampler.restore_dispatch_index(6) controller._sampler = sampler pump = asyncio.create_task(controller._rollout_pump()) - await asyncio.wait_for(sampler.admission_entered.wait(), timeout=1.0) + await _wait_for_event_or_pump(sampler.admission_entered, pump) assert controller._sampler.dispatch_index == 6 try: @@ -400,9 +426,9 @@ async def exercise() -> None: ) finally: sampler.release_admission.set() - await asyncio.wait_for(rollout_manager.started.wait(), timeout=1.0) + await _wait_for_event_or_pump(rollout_manager.started, pump) rollout_manager.release.set() - await asyncio.wait_for(pump, timeout=1.0) + await asyncio.wait_for(pump, timeout=_ASYNC_TEST_TIMEOUT_S) controller._checkpointer.shutdown() checkpoint = tmp_path / "checkpoints" / "step_7" @@ -464,7 +490,7 @@ async def exercise() -> None: controller = controller_cls(config, actor_args, SetupTimingMetrics()) pump = asyncio.create_task(controller._rollout_pump()) - await asyncio.wait_for(rollout_manager.started.wait(), timeout=1.0) + await _wait_for_event_or_pump(rollout_manager.started, pump) assert controller._sampler.dispatch_index == 7 try: @@ -474,7 +500,7 @@ async def exercise() -> None: ) finally: rollout_manager.release.set() - await asyncio.wait_for(pump, timeout=1.0) + await asyncio.wait_for(pump, timeout=_ASYNC_TEST_TIMEOUT_S) controller._checkpointer.shutdown() checkpoint = tmp_path / "checkpoints" / "step_7" diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index a5cac0e7be2..19e119afe1e 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -1322,6 +1322,7 @@ 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 diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index a77098f3599..f7d17134b0b 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -392,8 +392,11 @@ def __init__(self, **kwargs): assert teacher_topology is None -def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): +def test_single_controller_mopd_recipe_resolves_to_runtime_contract( + tmp_path, monkeypatch +): """The inherited recipe resolves exactly as the SC entrypoint consumes it.""" + monkeypatch.setenv("HF_HOME", str(tmp_path / "huggingface")) register_omegaconf_resolvers() repo_root = Path(__file__).resolve().parents[3] recipe = repo_root / ( From 7d1cac5f9d90de41bb71cb0c32d471589a1b8850 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 18:05:07 -0400 Subject: [PATCH 20/32] revert: remove MOPD recipe test environment override Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_setup.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unit/single_controller/test_setup.py b/tests/unit/single_controller/test_setup.py index f7d17134b0b..a77098f3599 100644 --- a/tests/unit/single_controller/test_setup.py +++ b/tests/unit/single_controller/test_setup.py @@ -392,11 +392,8 @@ def __init__(self, **kwargs): assert teacher_topology is None -def test_single_controller_mopd_recipe_resolves_to_runtime_contract( - tmp_path, monkeypatch -): +def test_single_controller_mopd_recipe_resolves_to_runtime_contract(): """The inherited recipe resolves exactly as the SC entrypoint consumes it.""" - monkeypatch.setenv("HF_HOME", str(tmp_path / "huggingface")) register_omegaconf_resolvers() repo_root = Path(__file__).resolve().parents[3] recipe = repo_root / ( From fcf1604dc16f601007a21da7579882569eda8607 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 18:21:08 -0400 Subject: [PATCH 21/32] fix(sc): rehydrate unfinished prompts deterministically Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 44 +++++- nemo_rl/experience/rollout_recovery.py | 141 +++++++++++++++++- .../unit/experience/test_rollout_recovery.py | 14 ++ .../test_checkpoint_dispatch_races.py | 20 ++- .../single_controller/test_checkpointing.py | 7 + 5 files changed, 213 insertions(+), 13 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index ce8a19c04e4..0d5b0c80699 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -731,19 +731,55 @@ async def _rehydrate_rollout_recovery_prompts(self) -> None: prompt = resolved_prompts.get(sample_id) if prompt is None: try: - prompt = await asyncio.to_thread(dataset.__getitem__, sample_index) + 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(prompt, dict): + if not isinstance(dataset_prompt, dict): raise TypeError( f"dataset sample_id={sample_id!r} resolved to " - f"{type(prompt).__name__}, expected a DatumSpec dictionary" + f"{type(dataset_prompt).__name__}, expected a DatumSpec " + "dictionary" + ) + + # The ledger fingerprints the prompt after dataloader collation, + # because that is the object actually dispatched to RolloutManager. + # Re-run the same one-row collation here so tensor scalars, optional + # fields, and multimodal wrappers match the original runtime shape. + 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(group.group_id, prompt) + recovery_ledger.bind_runtime_prompt( + group.group_id, + cast(DatumSpec, prompt), + ) async def _admit_reserved_prompt_groups( self, diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 08c3aaafb19..223e1df0b9d 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -18,16 +18,21 @@ import copy import hashlib -import pickle +import os import uuid +from collections.abc import Mapping from dataclasses import dataclass, replace from enum import StrEnum from typing import TYPE_CHECKING, Any, NotRequired, TypedDict +import numpy as np +import torch +from PIL import Image + if TYPE_CHECKING: from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 2 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 3 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" @@ -109,11 +114,135 @@ def _require_int(value: Any, *, field: str, minimum: int) -> int: return value +def _update_fingerprint_bytes( + digest: Any, + kind: bytes, + payload: bytes, +) -> None: + """Add one length-delimited value to a prompt fingerprint.""" + digest.update(len(kind).to_bytes(4, "big")) + digest.update(kind) + digest.update(len(payload).to_bytes(8, "big")) + digest.update(payload) + + +def _update_prompt_fingerprint(digest: Any, value: object, *, path: str) -> None: + """Hash supported prompt values by content rather than object identity.""" + if value is None: + _update_fingerprint_bytes(digest, b"none", b"") + elif isinstance(value, bool): + _update_fingerprint_bytes(digest, b"bool", b"1" if value else b"0") + elif isinstance(value, int): + _update_fingerprint_bytes(digest, b"int", str(value).encode()) + elif isinstance(value, float): + _update_fingerprint_bytes(digest, b"float", value.hex().encode()) + elif isinstance(value, str): + _update_fingerprint_bytes(digest, b"str", value.encode("utf-8")) + elif isinstance(value, (bytes, bytearray, memoryview)): + _update_fingerprint_bytes(digest, b"bytes", bytes(value)) + elif isinstance(value, os.PathLike): + _update_fingerprint_bytes( + digest, + b"path", + os.fsencode(value), + ) + elif isinstance(value, torch.Tensor): + if value.layout is not torch.strided: + raise TypeError( + f"prompt fingerprint at {path} does not support tensor layout " + f"{value.layout}" + ) + tensor = value.detach().cpu().contiguous() + _update_fingerprint_bytes(digest, b"tensor-dtype", str(tensor.dtype).encode()) + _update_fingerprint_bytes( + digest, + b"tensor-shape", + ",".join(str(dimension) for dimension in tensor.shape).encode(), + ) + _update_fingerprint_bytes( + digest, + b"tensor-data", + tensor.reshape(-1).view(torch.uint8).numpy().tobytes(), + ) + elif isinstance(value, np.ndarray): + _update_fingerprint_bytes(digest, b"ndarray-dtype", value.dtype.str.encode()) + _update_fingerprint_bytes( + digest, + b"ndarray-shape", + ",".join(str(dimension) for dimension in value.shape).encode(), + ) + if value.dtype.hasobject: + _update_prompt_fingerprint(digest, value.tolist(), path=f"{path}.tolist") + else: + _update_fingerprint_bytes( + digest, + b"ndarray-data", + np.ascontiguousarray(value).tobytes(), + ) + elif isinstance(value, np.generic): + _update_fingerprint_bytes(digest, b"numpy-dtype", value.dtype.str.encode()) + _update_prompt_fingerprint(digest, value.item(), path=f"{path}.item") + elif isinstance(value, Image.Image): + _update_fingerprint_bytes(digest, b"image-mode", value.mode.encode()) + _update_fingerprint_bytes( + digest, + b"image-size", + f"{value.width},{value.height}".encode(), + ) + _update_fingerprint_bytes(digest, b"image-data", value.tobytes()) + elif isinstance(value, Mapping): + keys = list(value) + if any(not isinstance(key, str) for key in keys): + raise TypeError( + f"prompt fingerprint at {path} requires string mapping keys" + ) + _update_fingerprint_bytes(digest, b"mapping-size", str(len(keys)).encode()) + for key in sorted(keys): + _update_fingerprint_bytes(digest, b"mapping-key", key.encode("utf-8")) + _update_prompt_fingerprint( + digest, + value[key], + path=f"{path}.{key}", + ) + elif isinstance(value, (list, tuple)): + kind = b"list-size" if isinstance(value, list) else b"tuple-size" + _update_fingerprint_bytes(digest, kind, str(len(value)).encode()) + for index, item in enumerate(value): + _update_prompt_fingerprint(digest, item, path=f"{path}[{index}]") + elif ( + type(value).__module__ == "nemo_rl.data.multimodal_utils" + and type(value).__name__ == "PackedTensor" + ): + _update_fingerprint_bytes(digest, b"packed-tensor", b"") + for attribute in ( + "tensors", + "dim_to_pack", + "pad_to_max_shape", + "_row_offsets", + "_segment_indices", + ): + _update_fingerprint_bytes( + digest, + b"packed-tensor-attribute", + attribute.encode(), + ) + _update_prompt_fingerprint( + digest, + getattr(value, attribute), + path=f"{path}.{attribute}", + ) + else: + raise TypeError( + f"prompt fingerprint at {path} does not support " + f"{type(value).__module__}.{type(value).__qualname__}" + ) + + def prompt_payload_sha256(prompt_payload: object) -> str: - """Fingerprint a prompt so dataset rehydration cannot silently change it.""" - return hashlib.sha256( - pickle.dumps(prompt_payload, protocol=pickle.HIGHEST_PROTOCOL) - ).hexdigest() + """Fingerprint prompt content so dataset rehydration detects real changes.""" + digest = hashlib.sha256() + _update_prompt_fingerprint(digest, prompt_payload, path="prompt") + return digest.hexdigest() def _prompt_task_name(prompt_payload: DatumSpec) -> str | None: diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index eef76f3d97d..3a24ff86eb6 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -221,6 +221,20 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: restored.bind_runtime_prompt("g7", _prompt(8)) +def test_prompt_fingerprint_is_stable_for_equivalent_tensor_content() -> None: + first = {"idx": 7, "length": torch.tensor(3), "tokens": torch.tensor([1, 2, 3])} + second = { + "tokens": torch.tensor([1, 2, 3]), + "length": torch.tensor(3), + "idx": 7, + } + + assert prompt_payload_sha256(first) == prompt_payload_sha256(second) + + second["tokens"] = torch.tensor([1, 2, 4]) + assert prompt_payload_sha256(first) != prompt_payload_sha256(second) + + def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: """Shuffle position and prompt identity are independent recovery assets.""" diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 33779df8adc..5030cbb94a8 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -35,7 +35,7 @@ from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace -from typing import Any +from typing import Any, cast import pytest import torch @@ -49,6 +49,7 @@ 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_plane.adapters.noop import NoOpDataPlaneClient from nemo_rl.data.interfaces import DatumSpec from nemo_rl.distributed.batched_data_dict import BatchedDataDict @@ -720,12 +721,24 @@ def test_recovery_replays_step_7_without_readmitting_the_batch(tmp_path) -> None 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() saved_ledger.reserve_group( group_id="batch-7-prompt-0", admission_id="batch-7", prompt_id="70", - prompt_payload={"idx": 70, "message_log": []}, + prompt_payload=dispatched_prompt, expected_generations=2, target_step=7, start_weight_version=7, @@ -758,7 +771,8 @@ async def exercise() -> None: controller._buffer_capacity = asyncio.Semaphore(4) controller._trainer_version = 7 controller._dataloader = SimpleNamespace( - dataset={70: {"idx": 70, "message_log": []}} + dataset={70: dataset_prompt}, + collate_fn=rl_collate_fn, ) controller._buffer = SimpleNamespace( count_for_target_step=lambda _target_step: 0, diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 19e119afe1e..2a887005c17 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -395,6 +395,13 @@ 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, state: dict[str, Any], From cb412ad75e66e3b1cf21378dd8a312bd9c9747d8 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 18:32:59 -0400 Subject: [PATCH 22/32] refactor(sc): simplify prompt recovery fingerprint Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_recovery.py | 137 +----------------- .../unit/experience/test_rollout_recovery.py | 4 +- 2 files changed, 7 insertions(+), 134 deletions(-) diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 223e1df0b9d..d8bff1d38e9 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -18,16 +18,13 @@ import copy import hashlib -import os +import io import uuid -from collections.abc import Mapping from dataclasses import dataclass, replace from enum import StrEnum from typing import TYPE_CHECKING, Any, NotRequired, TypedDict -import numpy as np import torch -from PIL import Image if TYPE_CHECKING: from nemo_rl.data.interfaces import DatumSpec @@ -114,135 +111,11 @@ def _require_int(value: Any, *, field: str, minimum: int) -> int: return value -def _update_fingerprint_bytes( - digest: Any, - kind: bytes, - payload: bytes, -) -> None: - """Add one length-delimited value to a prompt fingerprint.""" - digest.update(len(kind).to_bytes(4, "big")) - digest.update(kind) - digest.update(len(payload).to_bytes(8, "big")) - digest.update(payload) - - -def _update_prompt_fingerprint(digest: Any, value: object, *, path: str) -> None: - """Hash supported prompt values by content rather than object identity.""" - if value is None: - _update_fingerprint_bytes(digest, b"none", b"") - elif isinstance(value, bool): - _update_fingerprint_bytes(digest, b"bool", b"1" if value else b"0") - elif isinstance(value, int): - _update_fingerprint_bytes(digest, b"int", str(value).encode()) - elif isinstance(value, float): - _update_fingerprint_bytes(digest, b"float", value.hex().encode()) - elif isinstance(value, str): - _update_fingerprint_bytes(digest, b"str", value.encode("utf-8")) - elif isinstance(value, (bytes, bytearray, memoryview)): - _update_fingerprint_bytes(digest, b"bytes", bytes(value)) - elif isinstance(value, os.PathLike): - _update_fingerprint_bytes( - digest, - b"path", - os.fsencode(value), - ) - elif isinstance(value, torch.Tensor): - if value.layout is not torch.strided: - raise TypeError( - f"prompt fingerprint at {path} does not support tensor layout " - f"{value.layout}" - ) - tensor = value.detach().cpu().contiguous() - _update_fingerprint_bytes(digest, b"tensor-dtype", str(tensor.dtype).encode()) - _update_fingerprint_bytes( - digest, - b"tensor-shape", - ",".join(str(dimension) for dimension in tensor.shape).encode(), - ) - _update_fingerprint_bytes( - digest, - b"tensor-data", - tensor.reshape(-1).view(torch.uint8).numpy().tobytes(), - ) - elif isinstance(value, np.ndarray): - _update_fingerprint_bytes(digest, b"ndarray-dtype", value.dtype.str.encode()) - _update_fingerprint_bytes( - digest, - b"ndarray-shape", - ",".join(str(dimension) for dimension in value.shape).encode(), - ) - if value.dtype.hasobject: - _update_prompt_fingerprint(digest, value.tolist(), path=f"{path}.tolist") - else: - _update_fingerprint_bytes( - digest, - b"ndarray-data", - np.ascontiguousarray(value).tobytes(), - ) - elif isinstance(value, np.generic): - _update_fingerprint_bytes(digest, b"numpy-dtype", value.dtype.str.encode()) - _update_prompt_fingerprint(digest, value.item(), path=f"{path}.item") - elif isinstance(value, Image.Image): - _update_fingerprint_bytes(digest, b"image-mode", value.mode.encode()) - _update_fingerprint_bytes( - digest, - b"image-size", - f"{value.width},{value.height}".encode(), - ) - _update_fingerprint_bytes(digest, b"image-data", value.tobytes()) - elif isinstance(value, Mapping): - keys = list(value) - if any(not isinstance(key, str) for key in keys): - raise TypeError( - f"prompt fingerprint at {path} requires string mapping keys" - ) - _update_fingerprint_bytes(digest, b"mapping-size", str(len(keys)).encode()) - for key in sorted(keys): - _update_fingerprint_bytes(digest, b"mapping-key", key.encode("utf-8")) - _update_prompt_fingerprint( - digest, - value[key], - path=f"{path}.{key}", - ) - elif isinstance(value, (list, tuple)): - kind = b"list-size" if isinstance(value, list) else b"tuple-size" - _update_fingerprint_bytes(digest, kind, str(len(value)).encode()) - for index, item in enumerate(value): - _update_prompt_fingerprint(digest, item, path=f"{path}[{index}]") - elif ( - type(value).__module__ == "nemo_rl.data.multimodal_utils" - and type(value).__name__ == "PackedTensor" - ): - _update_fingerprint_bytes(digest, b"packed-tensor", b"") - for attribute in ( - "tensors", - "dim_to_pack", - "pad_to_max_shape", - "_row_offsets", - "_segment_indices", - ): - _update_fingerprint_bytes( - digest, - b"packed-tensor-attribute", - attribute.encode(), - ) - _update_prompt_fingerprint( - digest, - getattr(value, attribute), - path=f"{path}.{attribute}", - ) - else: - raise TypeError( - f"prompt fingerprint at {path} does not support " - f"{type(value).__module__}.{type(value).__qualname__}" - ) - - def prompt_payload_sha256(prompt_payload: object) -> str: - """Fingerprint prompt content so dataset rehydration detects real changes.""" - digest = hashlib.sha256() - _update_prompt_fingerprint(digest, prompt_payload, path="prompt") - return digest.hexdigest() + """Fingerprint a prompt for recovery within the same software runtime.""" + payload = io.BytesIO() + torch.save(prompt_payload, payload) + return hashlib.sha256(payload.getbuffer()).hexdigest() def _prompt_task_name(prompt_payload: DatumSpec) -> str | None: diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 3a24ff86eb6..318c38b9fb0 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -224,9 +224,9 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: def test_prompt_fingerprint_is_stable_for_equivalent_tensor_content() -> None: first = {"idx": 7, "length": torch.tensor(3), "tokens": torch.tensor([1, 2, 3])} second = { - "tokens": torch.tensor([1, 2, 3]), - "length": torch.tensor(3), "idx": 7, + "length": torch.tensor(3), + "tokens": torch.tensor([1, 2, 3]), } assert prompt_payload_sha256(first) == prompt_payload_sha256(second) From 3656fdeefc0d4c65b852f9d8b643518953037cad Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 18:51:48 -0400 Subject: [PATCH 23/32] fix(sc): stabilize recovered prompt fingerprints Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 5 ++-- nemo_rl/experience/rollout_recovery.py | 28 ++++++++++++++++-- .../unit/experience/test_rollout_recovery.py | 29 +++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 0d5b0c80699..0404e7637ac 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -748,8 +748,9 @@ async def _rehydrate_rollout_recovery_prompts(self) -> None: # The ledger fingerprints the prompt after dataloader collation, # because that is the object actually dispatched to RolloutManager. - # Re-run the same one-row collation here so tensor scalars, optional - # fields, and multimodal wrappers match the original runtime shape. + # Re-run one-row collation here to reconstruct tensor scalars, + # optional fields, and multimodal wrappers. Batch-derived transport + # fields can legitimately differ and are excluded from the digest. collate_fn = getattr(self._dataloader, "collate_fn", None) if collate_fn is None: prompt = dataset_prompt diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index d8bff1d38e9..26e8c4370b3 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 3 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" @@ -111,10 +111,32 @@ def _require_int(value: Any, *, field: str, minimum: int) -> int: return value +def _clone_tensor_leaves(value: Any) -> Any: + """Detach tensor content from batch-sized backing storage before hashing.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().clone() + if isinstance(value, dict): + return {key: _clone_tensor_leaves(item) for key, item in value.items()} + if isinstance(value, list): + return [_clone_tensor_leaves(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_tensor_leaves(item) for item in value) + return value + + def prompt_payload_sha256(prompt_payload: object) -> str: - """Fingerprint a prompt for recovery within the same software runtime.""" + """Fingerprint stable prompt content within the same software runtime.""" + if not isinstance(prompt_payload, dict): + raise TypeError("prompt payload fingerprint requires a dictionary") + canonical_payload = { + key: _clone_tensor_leaves(value) + for key, value in prompt_payload.items() + # Derived from the other prompts in the original dataloader batch and + # unused after collation; one-row recovery legitimately recomputes it. + if key != "batch_max_length" + } payload = io.BytesIO() - torch.save(prompt_payload, payload) + torch.save(canonical_payload, payload) return hashlib.sha256(payload.getbuffer()).hexdigest() diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 318c38b9fb0..9e6996ec5c1 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -235,6 +235,35 @@ def test_prompt_fingerprint_is_stable_for_equivalent_tensor_content() -> None: assert prompt_payload_sha256(first) != prompt_payload_sha256(second) +def test_prompt_fingerprint_ignores_batch_storage_and_derived_max_length() -> None: + batch_lengths = torch.tensor([3, 9]) + batch_loss_multipliers = torch.tensor([1.0, 1.0]) + dispatched = { + "idx": 7, + "length": batch_lengths[0], + "loss_multiplier": batch_loss_multipliers[0], + "batch_max_length": batch_lengths.max(), + "message_log": [{"role": "user", "token_ids": torch.tensor([1, 2, 3])}], + } + + restored_lengths = torch.tensor([3]) + restored_loss_multipliers = torch.tensor([1.0]) + one_row_recollated = { + "idx": 7, + "length": restored_lengths[0], + "loss_multiplier": restored_loss_multipliers[0], + "batch_max_length": restored_lengths.max(), + "message_log": [{"role": "user", "token_ids": torch.tensor([1, 2, 3])}], + } + + assert batch_lengths[0].untyped_storage().nbytes() != ( + restored_lengths[0].untyped_storage().nbytes() + ) + assert prompt_payload_sha256(dispatched) == prompt_payload_sha256( + one_row_recollated + ) + + def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: """Shuffle position and prompt identity are independent recovery assets.""" From fd1e424e54bf0b89f9224a8f04d9410b40bf397b Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 19:07:27 -0400 Subject: [PATCH 24/32] refactor(sc): remove prompt payload fingerprints Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 7 +- nemo_rl/experience/rollout_recovery.py | 99 ++----------------- .../unit/experience/test_rollout_recovery.py | 54 +--------- .../test_checkpoint_dispatch_races.py | 2 - 4 files changed, 14 insertions(+), 148 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 0404e7637ac..0efde6d50e2 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -746,11 +746,8 @@ async def _rehydrate_rollout_recovery_prompts(self) -> None: "dictionary" ) - # The ledger fingerprints the prompt after dataloader collation, - # because that is the object actually dispatched to RolloutManager. - # Re-run one-row collation here to reconstruct tensor scalars, - # optional fields, and multimodal wrappers. Batch-derived transport - # fields can legitimately differ and are excluded from the digest. + # 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 diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 26e8c4370b3..aa45c895b3e 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -17,19 +17,15 @@ from __future__ import annotations import copy -import hashlib -import io import uuid -from dataclasses import dataclass, replace +from dataclasses import dataclass from enum import StrEnum from typing import TYPE_CHECKING, Any, NotRequired, TypedDict -import torch - if TYPE_CHECKING: from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 4 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 5 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" @@ -45,7 +41,6 @@ class PromptRefState(TypedDict): sample_id: str task_name: str | None - payload_sha256: str class PromptGroupRecoveryState(TypedDict): @@ -72,11 +67,10 @@ class RolloutRecoveryState(TypedDict): @dataclass(frozen=True) class PromptRef: - """Stable dataset identity and integrity check for one prompt.""" + """Stable dataset identity for rebuilding one prompt.""" sample_id: str task_name: str | None - payload_sha256: str | None = None @dataclass(frozen=True) @@ -111,35 +105,6 @@ def _require_int(value: Any, *, field: str, minimum: int) -> int: return value -def _clone_tensor_leaves(value: Any) -> Any: - """Detach tensor content from batch-sized backing storage before hashing.""" - if isinstance(value, torch.Tensor): - return value.detach().cpu().clone() - if isinstance(value, dict): - return {key: _clone_tensor_leaves(item) for key, item in value.items()} - if isinstance(value, list): - return [_clone_tensor_leaves(item) for item in value] - if isinstance(value, tuple): - return tuple(_clone_tensor_leaves(item) for item in value) - return value - - -def prompt_payload_sha256(prompt_payload: object) -> str: - """Fingerprint stable prompt content within the same software runtime.""" - if not isinstance(prompt_payload, dict): - raise TypeError("prompt payload fingerprint requires a dictionary") - canonical_payload = { - key: _clone_tensor_leaves(value) - for key, value in prompt_payload.items() - # Derived from the other prompts in the original dataloader batch and - # unused after collation; one-row recovery legitimately recomputes it. - if key != "batch_max_length" - } - payload = io.BytesIO() - torch.save(canonical_payload, payload) - return hashlib.sha256(payload.getbuffer()).hexdigest() - - 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): @@ -174,25 +139,6 @@ def _validate_prompt_identity( ) -def _validate_prompt_ref( - prompt_ref: PromptRef, - prompt_payload: DatumSpec, - *, - group_id: str, -) -> str: - _validate_prompt_identity(prompt_ref, prompt_payload, group_id=group_id) - payload_sha256 = prompt_payload_sha256(prompt_payload) - if ( - prompt_ref.payload_sha256 is not None - and payload_sha256 != prompt_ref.payload_sha256 - ): - raise ValueError( - f"recovery group {group_id!r} prompt fingerprint mismatch for " - f"sample_id={prompt_ref.sample_id!r}" - ) - return payload_sha256 - - class RolloutRecoveryLedger: """Own prompts after dataloader advance and before canonical TQ commit.""" @@ -216,7 +162,7 @@ def reserve_group( Args: prompt_id: Dataset-level prompt identity used for diagnostics. prompt_payload: Runtime prompt used for whole-group regeneration. Only - its stable dataset reference and fingerprint are checkpointed. + 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. @@ -266,7 +212,7 @@ def reserve_group( # 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 locator and fingerprint. + # dispatch; state_dict() persists only its dataset locator. prompt_ref=PromptRef( sample_id=prompt_id, task_name=_prompt_task_name(prompt_payload), @@ -319,9 +265,9 @@ def bind_runtime_prompt( group_id: str, prompt_payload: DatumSpec, ) -> None: - """Attach and verify a dataset-rehydrated prompt after checkpoint load.""" + """Attach a dataset-rehydrated prompt after identity validation.""" record = self._require_group(group_id) - payload_sha256 = _validate_prompt_ref( + _validate_prompt_identity( record.prompt_ref, prompt_payload, group_id=group_id, @@ -333,7 +279,6 @@ def bind_runtime_prompt( prompt_ref=PromptRef( sample_id=record.prompt_ref.sample_id, task_name=record.prompt_ref.task_name, - payload_sha256=payload_sha256, ), runtime_prompt_payload=prompt_payload, expected_generations=record.expected_generations, @@ -367,7 +312,7 @@ def discard_canonical_groups(self, group_ids: set[str]) -> int: def state_dict(self) -> RolloutRecoveryState: """Return versioned references without serializing full prompt payloads.""" groups: list[PromptGroupRecoveryState] = [] - for group_id, record in list(self._groups.items()): + for record in self._groups.values(): prompt_payload = record.runtime_prompt_payload if prompt_payload is None: raise RuntimeError( @@ -379,20 +324,6 @@ def state_dict(self) -> RolloutRecoveryState: prompt_payload, group_id=record.group_id, ) - payload_sha256 = record.prompt_ref.payload_sha256 - if payload_sha256 is None: - payload_sha256 = prompt_payload_sha256(prompt_payload) - record = replace( - record, - prompt_ref=replace( - record.prompt_ref, - payload_sha256=payload_sha256, - ), - ) - # Prompts are immutable after dataloader processing. Cache the - # first durable fingerprint so repeated checkpoints do not - # serialize the same long prompt merely to hash it again. - self._groups[group_id] = record groups.append( { "group_id": record.group_id, @@ -401,7 +332,6 @@ def state_dict(self) -> RolloutRecoveryState: "prompt_ref": { "sample_id": record.prompt_ref.sample_id, "task_name": record.prompt_ref.task_name, - "payload_sha256": payload_sha256, }, "expected_generations": record.expected_generations, "target_step": record.target_step, @@ -493,7 +423,6 @@ def load_state_dict(self, state: RolloutRecoveryState) -> None: ) sample_id = raw_prompt_ref.get("sample_id") task_name = raw_prompt_ref.get("task_name") - payload_sha256 = raw_prompt_ref.get("payload_sha256") if not isinstance(sample_id, str) or not sample_id: raise ValueError( f"rollout recovery groups[{index}].prompt_ref.sample_id " @@ -509,17 +438,6 @@ def load_state_dict(self, state: RolloutRecoveryState) -> None: f"rollout recovery groups[{index}].prompt_ref.task_name " "must be a string or None" ) - if ( - not isinstance(payload_sha256, str) - or len(payload_sha256) != 64 - or any( - character not in "0123456789abcdef" for character in payload_sha256 - ) - ): - raise ValueError( - f"rollout recovery groups[{index}].prompt_ref.payload_sha256 " - "must be a lowercase SHA-256 digest" - ) restored[group_id] = PromptGroupRecoveryRecord( group_id=group_id, admission_id=admission_id, @@ -527,7 +445,6 @@ def load_state_dict(self, state: RolloutRecoveryState) -> None: prompt_ref=PromptRef( sample_id=sample_id, task_name=task_name, - payload_sha256=payload_sha256, ), runtime_prompt_payload=None, expected_generations=expected_generations, diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 9e6996ec5c1..7d32f158daa 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -9,7 +9,6 @@ ROLLOUT_RECOVERY_SCHEMA_VERSION, PromptGroupPhase, RolloutRecoveryLedger, - prompt_payload_sha256, ) @@ -45,7 +44,6 @@ def _group_state( target_step: int | None = 7, phase: str = "admitted", ) -> dict: - prompt = _prompt(idx) return { "group_id": f"g{idx}", "admission_id": "batch-7", @@ -53,7 +51,6 @@ def _group_state( "prompt_ref": { "sample_id": str(idx), "task_name": None, - "payload_sha256": prompt_payload_sha256(prompt), }, "expected_generations": 2, "target_step": target_step, @@ -173,14 +170,13 @@ def test_state_dict_stores_a_prompt_ref_without_the_full_payload() -> None: assert group_state["prompt_ref"] == { "sample_id": "7", "task_name": None, - "payload_sha256": prompt_payload_sha256(prompt), } group_state["prompt_ref"]["sample_id"] = "100" assert ledger.get_group("g7").prompt_ref.sample_id == "7" -def test_bind_runtime_prompt_rejects_changed_dataset_content() -> None: +def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> None: ledger = RolloutRecoveryLedger() original = _prompt() ledger.reserve_group( @@ -198,8 +194,9 @@ def test_bind_runtime_prompt_rejects_changed_dataset_content() -> None: changed = _prompt() changed["message_log"][0]["content"] = "different prompt" - with pytest.raises(ValueError, match="fingerprint mismatch"): - restored.bind_runtime_prompt("g7", changed) + restored.bind_runtime_prompt("g7", changed) + + assert restored.get_group("g7").prompt_payload == changed def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: @@ -221,49 +218,6 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: restored.bind_runtime_prompt("g7", _prompt(8)) -def test_prompt_fingerprint_is_stable_for_equivalent_tensor_content() -> None: - first = {"idx": 7, "length": torch.tensor(3), "tokens": torch.tensor([1, 2, 3])} - second = { - "idx": 7, - "length": torch.tensor(3), - "tokens": torch.tensor([1, 2, 3]), - } - - assert prompt_payload_sha256(first) == prompt_payload_sha256(second) - - second["tokens"] = torch.tensor([1, 2, 4]) - assert prompt_payload_sha256(first) != prompt_payload_sha256(second) - - -def test_prompt_fingerprint_ignores_batch_storage_and_derived_max_length() -> None: - batch_lengths = torch.tensor([3, 9]) - batch_loss_multipliers = torch.tensor([1.0, 1.0]) - dispatched = { - "idx": 7, - "length": batch_lengths[0], - "loss_multiplier": batch_loss_multipliers[0], - "batch_max_length": batch_lengths.max(), - "message_log": [{"role": "user", "token_ids": torch.tensor([1, 2, 3])}], - } - - restored_lengths = torch.tensor([3]) - restored_loss_multipliers = torch.tensor([1.0]) - one_row_recollated = { - "idx": 7, - "length": restored_lengths[0], - "loss_multiplier": restored_loss_multipliers[0], - "batch_max_length": restored_lengths.max(), - "message_log": [{"role": "user", "token_ids": torch.tensor([1, 2, 3])}], - } - - assert batch_lengths[0].untyped_storage().nbytes() != ( - restored_lengths[0].untyped_storage().nbytes() - ) - assert prompt_payload_sha256(dispatched) == prompt_payload_sha256( - one_row_recollated - ) - - def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: """Shuffle position and prompt identity are independent recovery assets.""" diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 5030cbb94a8..b154c56bed9 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -59,7 +59,6 @@ ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, RolloutRecoveryLedger, - prompt_payload_sha256, ) from tests.unit.single_controller._checkpoint_scenarios import ( _record, @@ -179,7 +178,6 @@ def state_dict(self) -> dict[str, Any]: "prompt_ref": { "sample_id": str(group.prompt_payload.get("idx", "unknown")), "task_name": group.prompt_payload.get("task_name"), - "payload_sha256": prompt_payload_sha256(group.prompt_payload), }, "expected_generations": 2, "start_weight_version": 7, From 6e5aedb315709e20bf5bd55156482f50c52f2725 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 19:49:00 -0400 Subject: [PATCH 25/32] fix(sc): prevent recovery deadlocks and ledger leaks Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 82 +++++++++++++------ .../test_checkpoint_dispatch_races.py | 69 ++++++++++++++++ .../single_controller/test_rollout_pump.py | 29 ++++++- 3 files changed, 155 insertions(+), 25 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 0efde6d50e2..131ad5104ce 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -849,9 +849,23 @@ async def _redispatch_restored_rollouts( if not groups_to_recover: return + # 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; - # ADMITTED groups retain their original target step and cursor position. + # 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: @@ -859,18 +873,20 @@ async def _redispatch_restored_rollouts( group.group_id ) for group_ids in reserved_admissions.values(): - await self._admit_reserved_prompt_groups(group_ids) - - refreshed_groups = recovery_ledger.groups() - for group in refreshed_groups: - await launch( - group.prompt_payload, - group.target_step, - group.group_id, - ) + _, 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 {len(refreshed_groups)} unfinished rollout " + f"📦 Redispatched {redispatched} unfinished rollout " "group(s) before new dataloader work", flush=True, ) @@ -2362,20 +2378,38 @@ 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 + 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(): + # 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(group_id) + for _, task in stale_groups: + task.cancel() + else: + stale_groups = _stale_groups() + for _, task in stale_groups: + task.cancel() - for task in stale_tasks: - 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 diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index b154c56bed9..cfa1ba47d64 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -876,6 +876,75 @@ async def _recover( 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() + ledger.reserve_group( + 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( + 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 = DataPlaneCheckpointBarrier() + 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)) + await rollout_manager.complete_recovery(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: diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 54a65cf3e95..4c52377c8d8 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -25,7 +25,10 @@ import ray import torch -from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer +from nemo_rl.algorithms.async_utils.replay_buffer import ( + DataPlaneCheckpointBarrier, + TQReplayBuffer, +) from nemo_rl.algorithms.async_utils.staleness_sampler import ( InOrderSampler, WeightFifoSampler, @@ -43,6 +46,7 @@ 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 # Reuse fixtures from the experience tests; same shape as test_async_rollout_manager. from tests.unit.experience.test_rollout_manager import ( @@ -772,17 +776,39 @@ async def _main() -> None: stale = asyncio.create_task(asyncio.Event().wait()) await asyncio.sleep(0) + ledger = RolloutRecoveryLedger() + for group_id, prompt_idx, start_weight_version in ( + ("fresh", 50, 5), + ("stale", 10, 1), + ): + ledger.reserve_group( + 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 = DataPlaneCheckpointBarrier() + 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): @@ -807,6 +833,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() From df31decca304aa934319a042fb6110eb89911adf Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 20:25:17 -0400 Subject: [PATCH 26/32] refactor(sc): initialize rollout recovery schema at v1 Signed-off-by: Anish Mahishi --- nemo_rl/experience/rollout_recovery.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index aa45c895b3e..96c5d6828fc 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -25,7 +25,7 @@ if TYPE_CHECKING: from nemo_rl.data.interfaces import DatumSpec -ROLLOUT_RECOVERY_SCHEMA_VERSION = 5 +ROLLOUT_RECOVERY_SCHEMA_VERSION = 1 ROLLOUT_RECOVERY_STATE_FILENAME = "rollout_recovery.pt" From c3b86240b450aca0aa4d62dc25b838516a4eccd5 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 21:30:13 -0400 Subject: [PATCH 27/32] test(sc): cover unfinished recovery edge cases Signed-off-by: Anish Mahishi --- .../L1_Functional_Tests_SingleController.sh | 6 +- ...p_single_controller_unfinished_recovery.sh | 6 + .../test_checkpoint_dispatch_races.py | 83 ++++++++++++ .../single_controller/test_rollout_pump.py | 122 +++++++++++++++++- 4 files changed, 213 insertions(+), 4 deletions(-) diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh index 55f33c16eae..9322f0a6df0 100755 --- a/tests/functional/L1_Functional_Tests_SingleController.sh +++ b/tests/functional/L1_Functional_Tests_SingleController.sh @@ -83,9 +83,9 @@ run_test env VICTIM_STATE=serving uv run --no-sync bash ./tests/functional/ 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 -# Full mode: deterministic process restart with an admitted group held before -# canonical TQ commit, followed by exact-once redispatch at its stable group ID. -run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_unfinished_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/grpo_dp_single_controller_unfinished_recovery.sh b/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh index b52f716f440..d400af8db26 100644 --- a/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh +++ b/tests/functional/grpo_dp_single_controller_unfinished_recovery.sh @@ -27,6 +27,12 @@ COMMON_OVERRIDES=( 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 ===" diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index cfa1ba47d64..75b3eafdb4c 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -32,6 +32,7 @@ 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 @@ -343,6 +344,88 @@ def _reserve_prompt(idx: int) -> DatumSpec: } +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() + saved_ledger.reserve_group( + 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() + restored_ledger.load_state_dict(saved_ledger.state_dict()) + + controller_cls = SingleControllerActor.__ray_metadata__.modified_class + controller = object.__new__(controller_cls) + controller._rollout_manager = SimpleNamespace(recovery_ledger=restored_ledger) + controller._dataloader = SimpleNamespace( + dataset={7: dataset_prompt}, + collate_fn=collate_fn, + ) + return controller, restored_ledger + + +def test_recovery_rehydration_accepts_an_identity_dict_collator() -> None: + controller, ledger = _rehydration_controller(_identity_dict_collator) + + asyncio.run(controller._rehydrate_rollout_recovery_prompts()) + + 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): + asyncio.run(controller._rehydrate_rollout_recovery_prompts()) + + def _reserve_controller() -> Any: controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 4c52377c8d8..1f1d3e44f3c 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 @@ -46,7 +48,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 +from nemo_rl.experience.rollout_recovery import ( + RolloutRecoveryLedger, + RolloutRecoveryState, +) # Reuse fixtures from the experience tests; same shape as test_async_rollout_manager. from tests.unit.experience.test_rollout_manager import ( @@ -98,6 +103,22 @@ def _init_pump_ledgers(ctrl: Any) -> None: 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[None]: + async with super().mutation(): + yield + self.mutation_applied.set() + await self.release_mutation.wait() + + class _RecordingBuffer: """TQReplayBuffer stand-in recording the target_step of each reserve. @@ -817,6 +838,105 @@ 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() + ledger.reserve_group( + 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 = DataPlaneCheckpointBarrier() + ctrl._rollout_manager = SimpleNamespace( + recovery_ledger=ledger, + discard_prompt_group=ledger.discard_group, + ) + + async with ctrl._data_plane_checkpoint_barrier.checkpoint(): + 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("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() + ledger.reserve_group( + 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() -> RolloutRecoveryState: + 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: From bca9a721034845998aa1146c5c46cb919949c063 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Thu, 27 Aug 2026 21:34:45 -0400 Subject: [PATCH 28/32] test(sc): add rollout recovery copyright header Signed-off-by: Anish Mahishi --- tests/unit/experience/test_rollout_recovery.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 7d32f158daa..71676dff9a9 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -1,3 +1,17 @@ +# 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 pytest From 47dc1489b80bb910f8ffbdca9da3b5be899fc34a Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 12:24:30 -0400 Subject: [PATCH 29/32] fix(sc): harden rollout recovery checkpointing Signed-off-by: Anish Mahishi --- docs/guides/single-controller.md | 6 +- .../algorithms/async_utils/replay_buffer.py | 42 ++- .../async_utils/staleness_sampler.py | 23 +- nemo_rl/algorithms/single_controller.py | 149 +++++---- nemo_rl/experience/rollout_manager.py | 28 +- nemo_rl/experience/rollout_recovery.py | 137 +++++++- tests/unit/experience/test_rollout_manager.py | 60 +++- .../unit/experience/test_rollout_recovery.py | 188 +++++++++-- .../_checkpoint_scenarios.py | 132 ++++++-- .../test_checkpoint_dispatch_races.py | 304 ++++++++++++------ .../test_checkpoint_recovery_matrix.py | 47 ++- .../single_controller/test_rollout_pump.py | 89 ++--- .../test_sampler_interface.py | 21 +- .../test_tq_replay_buffer.py | 16 + 14 files changed, 945 insertions(+), 297 deletions(-) diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index a8ff1b7c17b..20078620280 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -76,6 +76,8 @@ With `checkpointing.save_data_plane: true`, each Single-Controller checkpoint co - 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. @@ -84,8 +86,8 @@ On resume, Single-Controller validates the TQ snapshot against the trainer check Replay recovery is supported by all built-in samplers: `in_order`, `weight_fifo`, `ready_first`, and `windowed`. Custom samplers must explicitly declare `supports_buffer_checkpoint = True`. Otherwise, setup emits a warning and completed buffered groups are not restored. -:::{warning} -This checkpointing path recovers completed groups that have been committed to TQ. It does not recover generations that were still in flight at the checkpoint boundary. +:::{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. diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index ad23230a49c..2cf2d09df08 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -54,6 +54,7 @@ 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" @@ -182,6 +183,24 @@ def replay_manifest_digest(groups: list[TQReplayGroupMetadata]) -> str: 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. @@ -197,22 +216,24 @@ def __init__(self) -> None: self._active_mutations = 0 @asynccontextmanager - async def mutation(self) -> AsyncIterator[None]: - """Enter a commit/clear section, waiting only for an active checkpoint.""" + 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 + 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[None]: - """Block new mutations and wait for active ones before snapshotting.""" + 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 @@ -222,9 +243,11 @@ async def checkpoint(self) -> AsyncIterator[None]: self._checkpoint_active = False self._condition.notify_all() raise + cut = DataPlaneMutationCut(self) try: - yield + yield cut finally: + cut._invalidate() async with self._condition: self._checkpoint_active = False self._condition.notify_all() @@ -989,6 +1012,13 @@ def set_data_plane_checkpoint_barrier( 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]], diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py index 90874620a19..ff927ff3493 100644 --- a/nemo_rl/algorithms/async_utils/staleness_sampler.py +++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py @@ -54,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. @@ -141,8 +144,8 @@ async def wait_until_admissible( """Wait until one admission can commit without mutating sampler state.""" ... - def commit_admission(self) -> Optional[int]: - """Advance the admission cursor and return the batch target-step stamp.""" + def commit_admission(self, cut: DataPlaneMutationCut) -> Optional[int]: + """Advance the cursor under a live data-plane mutation cut.""" ... @@ -332,13 +335,14 @@ async def wait_until_admissible( """Return immediately because buffer capacity is this policy's gate.""" del trainer_version_fn - def commit_admission(self) -> Optional[int]: + 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]: await self.wait_until_admissible(trainer_version_fn=trainer_version_fn) - return self.commit_admission() + return None async def select( self, @@ -410,14 +414,19 @@ async def wait_until_admissible( while self._dispatch_index >= trainer_version_fn() + self._gate_window: await asyncio.sleep(_GATE_POLL_SECONDS) - def commit_admission(self) -> Optional[int]: + 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() + return self._commit_admission() def _stamp(self) -> Optional[int]: return None diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 131ad5104ce..86459061f05 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -61,8 +61,10 @@ LEGACY_REPLAY_BUFFER_FILENAME, REPLAY_BUFFER_METADATA_FILENAME, REPLAY_BUFFER_METADATA_SCHEMA_VERSION, + REPLACEMENT_RESERVE_FILENAME, DataPlaneCheckpointBarrier, DataPlaneCheckpointMetadata, + DataPlaneMutationCut, TQReplayMetadataState, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( @@ -105,6 +107,8 @@ ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, RolloutRecoveryState, + build_rollout_recovery_state, + parse_rollout_recovery_state, ) from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration from nemo_rl.models.generation.vllm import VllmGeneration @@ -633,53 +637,28 @@ async def _maybe_restore_rollout_recovery( io.BytesIO(payload), weights_only=True, ) - if not isinstance(state, dict): - raise TypeError( - "rollout recovery sidecar must contain a dictionary, got " - f"{type(state).__name__}" - ) - groups = state.get("groups") - if not isinstance(groups, list) or len(groups) != expected_group_count: + 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 - recovery_ledger.load_state_dict(cast(RolloutRecoveryState, state)) - raw_batch_shortfall = state.get("batch_shortfall", {}) - if not isinstance(raw_batch_shortfall, dict): - raise TypeError("rollout recovery batch_shortfall must be a dictionary") - restored_batch_shortfall: dict[int, int] = {} - for step, count in raw_batch_shortfall.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}" - ) - restored_batch_shortfall[step] = count - 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" + 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 ) - self._batch_shortfall = restored_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(canonical_group_ids) - await self._rehydrate_rollout_recovery_prompts() + 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 = ( - raw_sampler_stamps - if raw_sampler_stamps is not None + 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() ) @@ -698,8 +677,15 @@ async def _maybe_restore_rollout_recovery( flush=True, ) - async def _rehydrate_rollout_recovery_prompts(self) -> None: - """Resolve durable prompt references against the restored dataset.""" + 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: @@ -775,6 +761,7 @@ async def _rehydrate_rollout_recovery_prompts(self) -> None: ) resolved_prompts[sample_id] = cast(DatumSpec, prompt) recovery_ledger.bind_runtime_prompt( + cut, group.group_id, cast(DatumSpec, prompt), ) @@ -793,11 +780,15 @@ async def _admit_reserved_prompt_groups( if not group_ids: raise ValueError("sampler admission requires at least one prompt group") - def _commit(target_step: Optional[int]) -> tuple[Optional[int], list[str], int]: + 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, ) @@ -810,26 +801,28 @@ def _commit(target_step: Optional[int]) -> tuple[Optional[int], list[str], int]: 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(group_id) + 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(): - target_step = self._sampler.commit_admission() - return _commit(target_step) + 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; custom implementations can - # opt into TransactionalAdmissionSampler to avoid delaying checkpoints - # while their gate waits. - async with self._data_plane_checkpoint_barrier.mutation(): + # 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(target_step) + return _commit(cut, target_step) async def _redispatch_restored_rollouts( self, @@ -849,6 +842,21 @@ async def _redispatch_restored_rollouts( 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 @@ -947,7 +955,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 @@ -1143,7 +1151,9 @@ async def _dispatch_one_prompt( if self._rollout_recovery_enabled: assert lineage_group_id is not None - async with self._data_plane_checkpoint_barrier.mutation(): + async with ( + self._data_plane_checkpoint_barrier.mutation() + ) as cut: replacement = self._take_replacement( target_step, replacements ) @@ -1151,13 +1161,16 @@ async def _dispatch_one_prompt( # 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(lineage_group_id) + 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, ) @@ -1296,7 +1309,9 @@ async def _launch( dataloader_iterator = iter(self._dataloader) while True: prompt_dispatches: list[tuple[DatumSpec, str]] = [] - async with self._data_plane_checkpoint_barrier.mutation(): + async with ( + self._data_plane_checkpoint_barrier.mutation() + ) as cut: try: prompt_batch = next(dataloader_iterator) except StopIteration: @@ -1310,6 +1325,7 @@ async def _launch( 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, @@ -1430,7 +1446,7 @@ async def _drain_reserve_into_steps( 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(): + async with self._data_plane_checkpoint_barrier.mutation() as cut: step_prompts = [ self._replacement_reserve.popleft() for _ in range(num_prompts_per_step) @@ -1438,6 +1454,7 @@ async def _drain_reserve_into_steps( 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, @@ -2390,7 +2407,7 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: return stale_groups if self._rollout_recovery_enabled: - async with self._data_plane_checkpoint_barrier.mutation(): + 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() @@ -2398,7 +2415,7 @@ def _stale_groups() -> list[tuple[str, asyncio.Task[None]]]: # 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(group_id) + self._rollout_manager.discard_prompt_group(cut, group_id) for _, task in stale_groups: task.cancel() else: @@ -2512,14 +2529,12 @@ async def _save_checkpoint( await self._validate_replay_inventory(replay_metadata) if self._rollout_recovery_enabled: - rollout_recovery_state = ( - self._rollout_manager.recovery_ledger.state_dict() - ) - rollout_recovery_state["batch_shortfall"] = ( - self._batch_shortfall.copy() - ) - rollout_recovery_state["sampler_stamps_target_steps"] = ( - self._sampler_stamps_target_steps + 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 = { @@ -2602,7 +2617,7 @@ 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( diff --git a/nemo_rl/experience/rollout_manager.py b/nemo_rl/experience/rollout_manager.py index 13c8c9ff324..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, ) @@ -1225,6 +1226,7 @@ def recovery_ledger(self) -> RolloutRecoveryLedger: def reserve_prompt_group( self, + cut: DataPlaneMutationCut, input_sample: DatumSpec, *, target_step: Optional[int], @@ -1239,6 +1241,7 @@ def reserve_prompt_group( 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, @@ -1251,20 +1254,26 @@ def reserve_prompt_group( 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, group_id: str) -> None: + 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(group_id) + self._recovery_ledger.discard_group(cut, group_id) def set_weight_version(self, version: int) -> None: """Set the weight_version used for rollout tags. @@ -1392,8 +1401,14 @@ async def generate_and_push( flush=True, ) if cleanup_failed: - # A retry cannot safely reuse the stable ID while the previous - # slot may still exist. Re-raise the original rollout failure. + # 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), @@ -1461,7 +1476,10 @@ async def generate_and_push( # succeeded on a retry also counts -- the fleet recovered either way. self._consecutive_infra_drops = 0 if lineage_group_id is not None: - self._recovery_ledger.discard_group(lineage_group_id) + 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 diff --git a/nemo_rl/experience/rollout_recovery.py b/nemo_rl/experience/rollout_recovery.py index 96c5d6828fc..98afc450ecb 100644 --- a/nemo_rl/experience/rollout_recovery.py +++ b/nemo_rl/experience/rollout_recovery.py @@ -23,6 +23,7 @@ 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 @@ -56,11 +57,16 @@ class PromptGroupRecoveryState(TypedDict): phase: str -class RolloutRecoveryState(TypedDict): - """Versioned checkpoint sidecar for unfinished prompt groups.""" +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] @@ -98,6 +104,15 @@ def prompt_payload(self) -> DatumSpec: 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: @@ -140,13 +155,18 @@ def _validate_prompt_identity( class RolloutRecoveryLedger: - """Own prompts after dataloader advance and before canonical TQ commit.""" + """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, @@ -160,6 +180,7 @@ def reserve_group( """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. @@ -175,6 +196,7 @@ def reserve_group( 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") @@ -230,12 +252,14 @@ def reserve_group( 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( @@ -262,10 +286,17 @@ def mark_group_admitted( def bind_runtime_prompt( self, + cut: DataPlaneMutationCut, group_id: str, prompt_payload: DatumSpec, ) -> None: - """Attach a dataset-rehydrated prompt after identity validation.""" + """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, @@ -295,13 +326,19 @@ 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, group_id: str) -> None: + 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, group_ids: set[str]) -> int: + 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: @@ -309,7 +346,7 @@ def discard_canonical_groups(self, group_ids: set[str]) -> int: discarded += 1 return discarded - def state_dict(self) -> RolloutRecoveryState: + def state_dict(self) -> RolloutRecoveryLedgerState: """Return versioned references without serializing full prompt payloads.""" groups: list[PromptGroupRecoveryState] = [] for record in self._groups.values(): @@ -324,6 +361,9 @@ def state_dict(self) -> RolloutRecoveryState: 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, @@ -344,8 +384,13 @@ def state_dict(self) -> RolloutRecoveryState: "groups": groups, } - def load_state_dict(self, state: RolloutRecoveryState) -> None: + 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" @@ -472,3 +517,79 @@ def _require_group(self, group_id: str) -> PromptGroupRecoveryRecord: 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/tests/unit/experience/test_rollout_manager.py b/tests/unit/experience/test_rollout_manager.py index 65b2770f30c..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 @@ -75,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] = [] @@ -353,9 +365,13 @@ async def _assert_ledger_owns_inflight_prompt(_sample): _FakeImpl(on_run=_assert_ledger_owns_inflight_prompt), ) prompt = {"idx": 0, "message_log": [], "prompt": "p"} - group_id = mgr.reserve_prompt_group( - prompt, - target_step=None, + group_id = _with_cut( + buf, + lambda cut: mgr.reserve_prompt_group( + cut, + prompt, + target_step=None, + ), ) _run( @@ -373,14 +389,19 @@ 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( - _FakeBuffer(), + buf, _FakeImpl(on_run=_fail_rollout), RolloutRetryPolicy.single_attempt(max_skipped_prompts=1), ) - group_id = mgr.reserve_prompt_group( - {"idx": 7, "message_log": []}, - target_step=7, + group_id = _with_cut( + buf, + lambda cut: mgr.reserve_prompt_group( + cut, + {"idx": 7, "message_log": []}, + target_step=7, + ), ) outcome = _run( @@ -395,15 +416,20 @@ async def _fail_rollout(_sample): assert mgr.recovery_ledger.get_group(group_id).target_step == 7 def test_tracked_dispatch_rejects_changed_generations_per_prompt(self): - mgr = _make_manager(_FakeBuffer(), _FakeImpl()) - mgr.recovery_ledger.reserve_group( - group_id="g0", - prompt_id="0", - prompt_payload={"idx": 0, "message_log": []}, - expected_generations=2, - target_step=0, - start_weight_version=0, - admitted=True, + 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"): diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 71676dff9a9..606c4ac08a1 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -14,17 +14,53 @@ 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 { @@ -75,7 +111,8 @@ def _group_state( def test_ledger_round_trip_preserves_group_ownership() -> None: ledger = RolloutRecoveryLedger() - ledger.reserve_group( + _reserve( + ledger, group_id="g7", admission_id="batch-7", prompt_id="7", @@ -88,19 +125,120 @@ def test_ledger_round_trip_preserves_group_ownership() -> None: state = ledger.state_dict() restored = RolloutRecoveryLedger() - restored.load_state_dict(state) + _load(restored, state) with pytest.raises(RuntimeError, match="has not rehydrated prompt"): _ = restored.get_group("g7").prompt_payload - restored.bind_runtime_prompt("g7", _prompt()) + _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 = ledger.reserve_group( + record = _reserve( + ledger, group_id="windowed", admission_id="batch-windowed", prompt_id="7", @@ -117,7 +255,8 @@ def test_target_step_none_does_not_mean_unadmitted() -> None: def test_reserved_group_can_be_admitted_exactly_once() -> None: ledger = RolloutRecoveryLedger() - ledger.reserve_group( + _reserve( + ledger, group_id="g7", admission_id="batch-7", prompt_id="7", @@ -128,7 +267,8 @@ def test_reserved_group_can_be_admitted_exactly_once() -> None: admitted=False, ) - ledger.mark_group_admitted( + _mark( + ledger, "g7", target_step=7, start_weight_version=7, @@ -139,7 +279,8 @@ def test_reserved_group_can_be_admitted_exactly_once() -> None: assert record.target_step == 7 assert record.start_weight_version == 7 with pytest.raises(ValueError, match="already admitted"): - ledger.mark_group_admitted( + _mark( + ledger, "g7", target_step=8, start_weight_version=8, @@ -149,7 +290,8 @@ def test_reserved_group_can_be_admitted_exactly_once() -> None: def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> None: ledger = RolloutRecoveryLedger() for idx, group_id in enumerate(("canonical", "unfinished"), start=7): - ledger.reserve_group( + _reserve( + ledger, group_id=group_id, admission_id="batch-7", prompt_id=str(idx), @@ -160,14 +302,17 @@ def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> admitted=True, ) - assert ledger.discard_canonical_groups({"canonical"}) == 1 + 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() - ledger.reserve_group( + _reserve( + ledger, group_id="g7", admission_id="batch-7", prompt_id="7", @@ -193,7 +338,8 @@ def test_state_dict_stores_a_prompt_ref_without_the_full_payload() -> None: def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> None: ledger = RolloutRecoveryLedger() original = _prompt() - ledger.reserve_group( + _reserve( + ledger, group_id="g7", admission_id="batch-7", prompt_id="7", @@ -204,18 +350,19 @@ def test_bind_runtime_prompt_accepts_changed_content_with_the_same_identity() -> admitted=True, ) restored = RolloutRecoveryLedger() - restored.load_state_dict(ledger.state_dict()) + _load(restored, ledger.state_dict()) changed = _prompt() changed["message_log"][0]["content"] = "different prompt" - restored.bind_runtime_prompt("g7", changed) + _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() - ledger.reserve_group( + _reserve( + ledger, group_id="g7", admission_id="batch-7", prompt_id="7", @@ -226,10 +373,10 @@ def test_bind_runtime_prompt_rejects_the_wrong_dataset_sample() -> None: admitted=True, ) restored = RolloutRecoveryLedger() - restored.load_state_dict(ledger.state_dict()) + _load(restored, ledger.state_dict()) with pytest.raises(ValueError, match="expected '7'"): - restored.bind_runtime_prompt("g7", _prompt(8)) + _bind(restored, "g7", _prompt(8)) def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: @@ -241,7 +388,8 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: owned_prompt = fetched[-1] ledger = RolloutRecoveryLedger() - ledger.reserve_group( + _reserve( + ledger, group_id="unfinished", admission_id="shuffled-batch", prompt_id=str(owned_prompt["idx"]), @@ -260,12 +408,12 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: assert next(iter(restored_dataloader)) == expected_next_prompt restored_ledger = RolloutRecoveryLedger() - restored_ledger.load_state_dict(ledger_state) + _load(restored_ledger, ledger_state) restored_group = restored_ledger.get_group("unfinished") dataset_prompt = restored_dataloader.dataset[ int(restored_group.prompt_ref.sample_id) ] - restored_ledger.bind_runtime_prompt("unfinished", dataset_prompt) + _bind(restored_ledger, "unfinished", dataset_prompt) assert restored_ledger.get_group("unfinished").prompt_payload == owned_prompt @@ -293,4 +441,4 @@ def test_prompt_ref_rehydrates_through_a_restored_shuffled_dataloader() -> None: ) def test_restore_rejects_incompatible_or_malformed_state(state: dict) -> None: with pytest.raises((TypeError, ValueError)): - RolloutRecoveryLedger().load_state_dict(state) # type: ignore[arg-type] + _load(RolloutRecoveryLedger(), state) # type: ignore[arg-type] diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index 7935a4cbc37..e49713bb66c 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -151,6 +151,14 @@ def committed_outstanding(self) -> set[str]: 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}" @@ -262,7 +270,10 @@ class RoundTrip: 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. + 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] @@ -271,10 +282,17 @@ class RoundTrip: 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 + 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) @@ -289,23 +307,25 @@ async def _round_trip( else None ) recovery_ledger_a = RolloutRecoveryLedger() - 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( - 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, - ) + 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") @@ -326,8 +346,9 @@ async def _round_trip( expected_manifest_digest=sidecar["manifest_digest"], ) recovery_ledger_b = RolloutRecoveryLedger() - recovery_ledger_b.load_state_dict(recovery_sidecar) - recovery_ledger_b.discard_canonical_groups(set(buf_b._group_ids)) + 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, @@ -340,24 +361,64 @@ async def _round_trip( start_weight_version=group.start_weight_version, end_weight_version=group.start_weight_version, ) - recovery_ledger_b.discard_group(group_id) + 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=set(buf_b._group_ids), + recovered=recovered, ready=ready, - pending=set(buf_b._group_ids) - ready, + pending=pending, saved_sidecar=sidecar is not None, rows_before=rows_before, - rows_after=set(dp_b.list_sample_ids(PARTITION)), + rows_after=rows_after, + stamps=stamps, + selected=selected, + selected_count=selected_count, ) -def round_trip(scenario: Scenario, sampler_name: str, tmp_path: Path) -> RoundTrip: +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)) + return asyncio.run( + _round_trip( + scenario, + sampler_name, + tmp_path, + select_current_train_weight=select_current_train_weight, + ) + ) def assert_no_data_loss( @@ -410,6 +471,21 @@ def assert_completed_groups_survive( 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=( @@ -485,7 +561,7 @@ def assert_completed_groups_survive( ) # Everything fully generated -- the case this PR set out to recover. -FULLY_GENERATED = (S_ALL_COMPLETE, S_STALE_ONLY) +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/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 75b3eafdb4c..a9377fe515f 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -36,7 +36,7 @@ from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace -from typing import Any, cast +from typing import Any, TypeVar, cast import pytest import torch @@ -44,6 +44,7 @@ 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 @@ -60,6 +61,7 @@ ROLLOUT_RECOVERY_STATE_FILENAME, PromptGroupPhase, RolloutRecoveryLedger, + build_rollout_recovery_state, ) from tests.unit.single_controller._checkpoint_scenarios import ( _record, @@ -72,6 +74,15 @@ ) _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( @@ -110,9 +121,9 @@ 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): + def commit_admission(self, cut: DataPlaneMutationCut): self.admission_commits += 1 - return super().commit_admission() + return super().commit_admission(cut) class _BlockingBeforeAdmissionSampler(_CountingInOrderSampler): @@ -197,22 +208,35 @@ def __init__(self, ledger: RolloutRecoveryLedger) -> None: self.recovery_ledger = ledger self.recovered: list[tuple[str, int | None]] = [] - async def complete_recovery(self, group_id: str) -> 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(group_id) + self.recovery_ledger.discard_group(cut, group_id) def mark_prompt_group_admitted( - self, group_id: str, *, target_step: int | None + 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, group_id: str) -> None: - self.recovery_ledger.discard_group(group_id) + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + self.recovery_ledger.discard_group(cut, group_id) class _BlockingRolloutManager: @@ -229,6 +253,7 @@ def set_weight_version(self, version: int) -> None: def reserve_prompt_group( self, + cut: DataPlaneMutationCut | None, prompt: DatumSpec, *, target_step: int | None = None, @@ -249,13 +274,23 @@ def reserve_prompt_group( return group_id def mark_prompt_group_admitted( - self, group_id: str, *, target_step: int | None + 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, group_id: str) -> None: + def discard_prompt_group( + self, + cut: DataPlaneMutationCut, + group_id: str, + ) -> None: + del cut self.recovery_ledger.release(group_id) async def generate_and_push( @@ -269,6 +304,7 @@ async def generate_and_push( del inflight_registry if lineage_group_id is None: lineage_group_id = self.reserve_prompt_group( + None, prompt, target_step=target_step, ) @@ -304,6 +340,7 @@ def __init__(self) -> None: def reserve_prompt_group( self, + cut: DataPlaneMutationCut, prompt: DatumSpec, *, target_step: int | None = None, @@ -311,6 +348,7 @@ def reserve_prompt_group( admission_id: str | None = None, ) -> str: record = self.recovery_ledger.reserve_group( + cut, prompt_id=str(prompt["idx"]), prompt_payload=prompt, expected_generations=2, @@ -322,16 +360,25 @@ def reserve_prompt_group( return record.group_id def mark_prompt_group_admitted( - self, group_id: str, *, target_step: int | None + 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, group_id: str) -> None: - self.recovery_ledger.discard_group(group_id) + 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: @@ -368,20 +415,26 @@ def _rehydration_controller( """Build a restored ledger whose prompt must be resolved from the dataset.""" dataset_prompt = _reserve_prompt(7) saved_ledger = RolloutRecoveryLedger() - saved_ledger.reserve_group( - group_id="rehydrate-7", - prompt_id="7", - prompt_payload=dataset_prompt, - expected_generations=2, - target_step=7, - start_weight_version=7, - admitted=True, + _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() - restored_ledger.load_state_dict(saved_ledger.state_dict()) + _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}, @@ -390,10 +443,18 @@ def _rehydration_controller( 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) - asyncio.run(controller._rehydrate_rollout_recovery_prompts()) + _run_rehydration(controller) assert ledger.get_group("rehydrate-7").prompt_payload["length"] == 99 @@ -423,7 +484,7 @@ def test_recovery_rehydration_rejects_invalid_collator_results( controller, _ = _rehydration_controller(collate_fn) with pytest.raises(expected_error, match=match): - asyncio.run(controller._rehydrate_rollout_recovery_prompts()) + _run_rehydration(controller) def _reserve_controller() -> Any: @@ -815,19 +876,23 @@ async def exercise() -> None: {key: value[0] for key, value in prompt_batch.items()}, ) saved_ledger = RolloutRecoveryLedger() - saved_ledger.reserve_group( - 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, + 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, ) - saved_state = saved_ledger.state_dict() - saved_state["batch_shortfall"] = {6: 1} - saved_state["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() @@ -851,6 +916,7 @@ async def exercise() -> None: ) 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, @@ -868,7 +934,8 @@ async def _recover( _target_step: int | None, group_id: str, ) -> None: - await rollout_manager.complete_recovery(group_id) + 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) @@ -883,22 +950,59 @@ async def _recover( 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() - for prompt_idx in (70, 71): - saved_ledger.reserve_group( - 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, - ) + 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) @@ -943,7 +1047,8 @@ async def _recover( _target_step: int | None, group_id: str, ) -> None: - await rollout_manager.complete_recovery(group_id) + 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) @@ -966,26 +1071,30 @@ async def exercise() -> None: sampler = _CountingInOrderSampler() sampler.restore_dispatch_index(7) ledger = RolloutRecoveryLedger() - ledger.reserve_group( - 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( - 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, - ) + 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 @@ -993,7 +1102,7 @@ async def exercise() -> None: controller._sampler = sampler controller._rollout_manager = rollout_manager controller._trainer_version = 6 - controller._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + controller._data_plane_checkpoint_barrier = barrier controller._buffer = SimpleNamespace( count_for_target_step=lambda _target_step: 0, ) @@ -1006,7 +1115,8 @@ async def _recover( group_id: str, ) -> None: launched.append((group_id, target_step)) - await rollout_manager.complete_recovery(group_id) + 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. @@ -1035,17 +1145,19 @@ def test_recovery_load_does_not_require_every_unfinished_group_to_fit_at_once( async def exercise() -> None: saved_ledger = RolloutRecoveryLedger() - for prompt_idx in (70, 71): - saved_ledger.reserve_group( - 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, - ) + 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) @@ -1136,21 +1248,23 @@ def test_checkpoint_waits_for_replacement_pop_and_reownership() -> None: async def exercise() -> None: controller = _reserve_controller() manager = controller._rollout_manager - old_group_id = manager.reserve_prompt_group( - _reserve_prompt(20), - target_step=7, - ) + 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(): + 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(old_group_id) - manager.reserve_prompt_group(replacement, target_step=7) + manager.discard_prompt_group(cut, old_group_id) + manager.reserve_prompt_group(cut, replacement, target_step=7) mutation_applied.set() await release_mutation.wait() @@ -1198,11 +1312,13 @@ async def block_admission( ) -> tuple[int, list[str], int]: admission_started.set() await release_admission.wait() - for group_id in group_ids: - controller._rollout_manager.mark_prompt_group_admitted( - group_id, - target_step=7, - ) + 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( diff --git a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py index e75c5a11eb8..1777da91514 100644 --- a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -15,8 +15,8 @@ """Checkpoint recovery contract across the built-in async samplers. The six scenarios come from #3827. That PR covered windowed, weight_fifo, and -in_order; ready_first is included here because it now advertises the same -completed-buffer recovery capability. +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. @@ -29,7 +29,11 @@ from tests.unit.single_controller._checkpoint_scenarios import ( ALL_SCENARIOS, FULLY_GENERATED, + GROUPS_PER_STEP, SAMPLERS, + S_ALL_COMPLETE, + S_LAG2, + S_ZERO_LAG_ALL_COMPLETE, WITH_IN_FLIGHT, Case, assert_completed_groups_survive, @@ -47,6 +51,11 @@ 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) @@ -77,6 +86,14 @@ def test_unfinished_groups_are_owned_across_restart(case, tmp_path): 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.""" @@ -94,3 +111,29 @@ def test_intentionally_evicted_group_is_not_resurrected(sampler, tmp_path): 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_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py index 1f1d3e44f3c..5c953e85e39 100644 --- a/tests/unit/single_controller/test_rollout_pump.py +++ b/tests/unit/single_controller/test_rollout_pump.py @@ -29,6 +29,7 @@ from nemo_rl.algorithms.async_utils.replay_buffer import ( DataPlaneCheckpointBarrier, + DataPlaneMutationCut, TQReplayBuffer, ) from nemo_rl.algorithms.async_utils.staleness_sampler import ( @@ -50,7 +51,7 @@ from nemo_rl.experience.rollout_manager import RolloutManager, RolloutOutcome from nemo_rl.experience.rollout_recovery import ( RolloutRecoveryLedger, - RolloutRecoveryState, + RolloutRecoveryLedgerState, ) # Reuse fixtures from the experience tests; same shape as test_async_rollout_manager. @@ -112,9 +113,9 @@ def __init__(self) -> None: self.release_mutation = asyncio.Event() @asynccontextmanager - async def mutation(self) -> AsyncIterator[None]: - async with super().mutation(): - yield + async def mutation(self) -> AsyncIterator[DataPlaneMutationCut]: + async with super().mutation() as cut: + yield cut self.mutation_applied.set() await self.release_mutation.wait() @@ -798,19 +799,22 @@ async def _main() -> None: await asyncio.sleep(0) ledger = RolloutRecoveryLedger() - for group_id, prompt_idx, start_weight_version in ( - ("fresh", 50, 5), - ("stale", 10, 1), - ): - ledger.reserve_group( - 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, - ) + 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) @@ -818,7 +822,7 @@ async def _main() -> None: 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 = DataPlaneCheckpointBarrier() + ctrl._data_plane_checkpoint_barrier = barrier ctrl._rollout_manager = SimpleNamespace( recovery_ledger=ledger, discard_prompt_group=ledger.discard_group, @@ -845,15 +849,18 @@ async def _main() -> None: completed = asyncio.create_task(asyncio.Event().wait()) await asyncio.sleep(0) ledger = RolloutRecoveryLedger() - ledger.reserve_group( - group_id="completed", - prompt_id="10", - prompt_payload={"idx": 10, "message_log": []}, - expected_generations=2, - target_step=None, - start_weight_version=1, - admitted=True, - ) + 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) @@ -861,18 +868,18 @@ async def _main() -> None: ctrl._trainer_version = 5 ctrl._inflight_by_group_id = {"completed": (completed, 1)} ctrl._rollout_recovery_enabled = True - ctrl._data_plane_checkpoint_barrier = DataPlaneCheckpointBarrier() + 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(): + 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("completed") + ledger.discard_group(cut, "completed") assert await asyncio.wait_for(abort_task, timeout=1.0) == 0 assert not completed.cancelled() @@ -891,15 +898,17 @@ async def _main() -> None: stale = asyncio.create_task(asyncio.Event().wait()) await asyncio.sleep(0) ledger = RolloutRecoveryLedger() - ledger.reserve_group( - group_id="stale", - prompt_id="10", - prompt_payload={"idx": 10, "message_log": []}, - expected_generations=2, - target_step=None, - start_weight_version=1, - admitted=True, - ) + 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 @@ -919,7 +928,7 @@ async def _main() -> None: checkpoint_entered = asyncio.Event() - async def checkpoint_snapshot() -> RolloutRecoveryState: + async def checkpoint_snapshot() -> RolloutRecoveryLedgerState: async with barrier.checkpoint(): checkpoint_entered.set() return ledger.state_dict() diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py index fbf60e18763..a62c567d9ad 100644 --- a/tests/unit/single_controller/test_sampler_interface.py +++ b/tests/unit/single_controller/test_sampler_interface.py @@ -27,6 +27,7 @@ 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, @@ -116,9 +117,27 @@ def test_wait_does_not_advance_gated_dispatch_cursor(self): _run(sampler.wait_until_admissible(trainer_version_fn=lambda: 0)) assert sampler.dispatch_index == -1 - assert sampler.commit_admission() == 0 + + 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. diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 3806434cb67..edfa15a0b48 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -211,6 +211,22 @@ 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() From 2dd821232e2ade73dcc787d3b2a567a77279aebb Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 13:13:03 -0400 Subject: [PATCH 30/32] test(sc): initialize recovery fixture barrier Signed-off-by: Anish Mahishi --- tests/unit/single_controller/test_checkpoint_dispatch_races.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index a9377fe515f..927e7651152 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -1164,6 +1164,7 @@ async def exercise() -> None: 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 = { From 54d9e82fcec40d01a42b81565b8aee7627d33de9 Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 11:09:03 -0700 Subject: [PATCH 31/32] fix: lint issues Signed-off-by: Anish Mahishi --- nemo_rl/algorithms/single_controller.py | 15 +++++---------- tests/unit/experience/test_rollout_recovery.py | 4 +--- .../single_controller/_checkpoint_scenarios.py | 3 +-- .../test_checkpoint_dispatch_races.py | 8 +++----- .../test_checkpoint_recovery_matrix.py | 2 +- 5 files changed, 11 insertions(+), 21 deletions(-) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 86459061f05..80c3e1f4609 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -59,9 +59,9 @@ 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, - REPLACEMENT_RESERVE_FILENAME, DataPlaneCheckpointBarrier, DataPlaneCheckpointMetadata, DataPlaneMutationCut, @@ -847,9 +847,7 @@ async def _redispatch_restored_rollouts( PromptGroupPhase.RESERVED, ) unhandled_groups = [ - group - for group in groups_to_recover - if group.phase not in recognized_phases + group for group in groups_to_recover if group.phase not in recognized_phases ] if unhandled_groups: details = ", ".join( @@ -1309,9 +1307,7 @@ async def _launch( dataloader_iterator = iter(self._dataloader) while True: prompt_dispatches: list[tuple[DatumSpec, str]] = [] - async with ( - self._data_plane_checkpoint_barrier.mutation() - ) as cut: + async with self._data_plane_checkpoint_barrier.mutation() as cut: try: prompt_batch = next(dataloader_iterator) except StopIteration: @@ -2395,6 +2391,7 @@ 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.""" + 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(): @@ -2532,9 +2529,7 @@ async def _save_checkpoint( 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 - ), + sampler_stamps_target_steps=(self._sampler_stamps_target_steps), ) if replay_metadata is not None: canonical_group_ids = { diff --git a/tests/unit/experience/test_rollout_recovery.py b/tests/unit/experience/test_rollout_recovery.py index 606c4ac08a1..ef96bbf0b93 100644 --- a/tests/unit/experience/test_rollout_recovery.py +++ b/tests/unit/experience/test_rollout_recovery.py @@ -302,9 +302,7 @@ def test_canonical_groups_are_discarded_without_touching_unfinished_groups() -> admitted=True, ) - assert _mutate( - lambda cut: ledger.discard_canonical_groups(cut, {"canonical"}) - ) == 1 + assert _mutate(lambda cut: ledger.discard_canonical_groups(cut, {"canonical"})) == 1 assert [group.group_id for group in ledger.groups()] == ["unfinished"] diff --git a/tests/unit/single_controller/_checkpoint_scenarios.py b/tests/unit/single_controller/_checkpoint_scenarios.py index e49713bb66c..5e657615746 100644 --- a/tests/unit/single_controller/_checkpoint_scenarios.py +++ b/tests/unit/single_controller/_checkpoint_scenarios.py @@ -387,8 +387,7 @@ async def _round_trip( ) if selected_meta is not None: selected = { - sample_id.rpartition("_g")[0] - for sample_id in selected_meta.sample_ids + sample_id.rpartition("_g")[0] for sample_id in selected_meta.sample_ids } return RoundTrip( recovered=recovered, diff --git a/tests/unit/single_controller/test_checkpoint_dispatch_races.py b/tests/unit/single_controller/test_checkpoint_dispatch_races.py index 927e7651152..062b2595d62 100644 --- a/tests/unit/single_controller/test_checkpoint_dispatch_races.py +++ b/tests/unit/single_controller/test_checkpoint_dispatch_races.py @@ -52,8 +52,8 @@ 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_plane.adapters.noop import NoOpDataPlaneClient 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 ( @@ -68,8 +68,8 @@ patch_converter, ) from tests.unit.single_controller.test_checkpointing import ( - _FakeDataloader, _actor_master_config, + _FakeDataloader, _make_actor_args, ) @@ -961,9 +961,7 @@ async def exercise() -> None: ) controller_cls = SingleControllerActor.__ray_metadata__.modified_class controller = object.__new__(controller_cls) - controller._rollout_manager = SimpleNamespace( - recovery_ledger=recovery_ledger - ) + controller._rollout_manager = SimpleNamespace(recovery_ledger=recovery_ledger) launched = False async def _recover( diff --git a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py index 1777da91514..c99e9cb15a4 100644 --- a/tests/unit/single_controller/test_checkpoint_recovery_matrix.py +++ b/tests/unit/single_controller/test_checkpoint_recovery_matrix.py @@ -30,10 +30,10 @@ ALL_SCENARIOS, FULLY_GENERATED, GROUPS_PER_STEP, - SAMPLERS, S_ALL_COMPLETE, S_LAG2, S_ZERO_LAG_ALL_COMPLETE, + SAMPLERS, WITH_IN_FLIGHT, Case, assert_completed_groups_survive, From 2017b6006550bc6ca91fc36151c5a9b3a3aab63d Mon Sep 17 00:00:00 2001 From: Anish Mahishi Date: Fri, 28 Aug 2026 17:48:43 -0400 Subject: [PATCH 32/32] fix(sc): repair checkpointing CI fixtures Signed-off-by: Anish Mahishi --- tests/functional/ppo_async_single_controller.sh | 2 +- tests/unit/single_controller/test_single_controller_actor.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/functional/ppo_async_single_controller.sh b/tests/functional/ppo_async_single_controller.sh index b3df65ab72c..2861a038f60 100755 --- a/tests/functional/ppo_async_single_controller.sh +++ b/tests/functional/ppo_async_single_controller.sh @@ -65,7 +65,7 @@ TRAIN_CMD=( checkpointing.checkpoint_dir="${CKPT_DIR}" checkpointing.metric_name=null checkpointing.save_period=1 - checkpointing.save_data_plane=true + +checkpointing.save_data_plane=true ) cd "${PROJECT_ROOT}" diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index f8ce41a7157..d9c7d359b21 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -455,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={}) @@ -484,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={})