From 17ba4f69ef4ba6695e18cbc4fbab4efbed2aaa5c Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Wed, 29 Jul 2026 16:14:05 -0700
Subject: [PATCH 01/15] feat(sc): add TQReplayBuffer state_dict/load_state_dict
 for checkpointing

Port the replay-buffer checkpoint state capture from PR #3138 onto the
split single-controller path. state_dict snapshots ready slots on the
event loop, then fetches each group's DataPlane rows; unready
reservations (in-flight rollouts) are dropped. load_state_dict validates
the envelope (partition, group size, sample_id uniqueness) before any
DataPlane write, truncates to the current capacity keeping the freshest
groups, and re-puts rows while rebuilding the parallel slot lists.
Staleness filtering is intentionally left to the sampler's first evict.

Covered by 9 new unit tests (round-trip, preflight rejection,
capacity truncation); 20/20 pass in tests/unit/single_controller/
test_tq_replay_buffer.py.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 .../algorithms/async_utils/replay_buffer.py   | 188 +++++++++++++++
 .../test_tq_replay_buffer.py                  | 216 ++++++++++++++++++
 2 files changed, 404 insertions(+)

diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py
index 1bf9c5ae93..b4d1b39c2d 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -823,6 +823,194 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int:
 
         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"}, ...]}``.
+        """
+        snapshot: list[tuple[KVBatchMeta, int, int, Optional[int], str]] = []
+        for i, ready in enumerate(self.ready_list):
+            if not ready:
+                continue
+            meta = self.meta_list[i]
+            assert meta is not None  # commit sets meta before ready=True
+            snapshot.append(
+                (
+                    meta,
+                    self.start_weight_list[i],
+                    self.end_weight_list[i],
+                    self.target_step_list[i],
+                    self._group_ids[i],
+                )
+            )
+
+        groups: list[dict[str, Any]] = []
+        for meta, start_weight, end_weight, target_step, group_id in snapshot:
+            fields_data = await self._call_dp(
+                "get_samples",
+                sample_ids=meta.sample_ids,
+                partition_id=self._partition_id,
+                select_fields=meta.fields,
+            )
+            groups.append(
+                {
+                    "meta": meta,
+                    "start_weight": start_weight,
+                    "end_weight": end_weight,
+                    "target_step": target_step,
+                    "group_id": group_id,
+                    "fields_data": fields_data,
+                }
+            )
+        return {
+            "partition_id": self._partition_id,
+            "saved_capacity": saved_capacity,
+            "groups": groups,
+        }
+
+    async def load_state_dict(
+        self,
+        state: dict[str, Any],
+        *,
+        max_groups: int,
+        expected_partition_id: str,
+        expected_group_size: int,
+    ) -> int:
+        """Validate and re-put checkpointed groups into the buffer.
+
+        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.
+
+        Staleness is intentionally NOT handled here — load only loads. The
+        train pump's first ``sampler.evict`` drops any restored group that is
+        outside the staleness window and releases its capacity permit, keeping
+        eviction in one place.
+
+        Args:
+            state: Envelope produced by ``state_dict``.
+            max_groups: Current max_buffered_rollouts; the restored count
+                never exceeds it.
+            expected_partition_id: Partition this buffer writes to; must
+                match the envelope.
+            expected_group_size: num_generations_per_prompt; every group must
+                hold exactly this many rows (a changed group size silently
+                breaks the group-relative baseline).
+
+        Returns:
+            Number of groups restored into the buffer.
+
+        Raises:
+            ValueError: If the envelope is malformed (missing keys, partition
+                mismatch, misaligned or wrongly sized groups, duplicate
+                sample_ids).
+        """
+        required_keys = {"partition_id", "saved_capacity", "groups"}
+        missing_keys = required_keys - set(state)
+        if missing_keys:
+            raise ValueError(
+                f"Replay buffer checkpoint missing required keys: {missing_keys}"
+            )
+        if state["partition_id"] != expected_partition_id:
+            raise ValueError(
+                "Replay buffer checkpoint partition_id mismatch: "
+                f"checkpoint={state['partition_id']!r}, "
+                f"expected={expected_partition_id!r}"
+            )
+
+        groups = list(state["groups"])
+        group_keys = {
+            "meta",
+            "start_weight",
+            "end_weight",
+            "target_step",
+            "group_id",
+            "fields_data",
+        }
+        seen_sample_ids: set[str] = set()
+        for group in groups:
+            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"]
+            num_tags = len(meta.tags) if meta.tags is not None else -1
+            num_lengths = (
+                len(meta.sequence_lengths) if meta.sequence_lengths is not None else -1
+            )
+            if not (
+                len(meta.sample_ids) == num_tags == num_lengths == expected_group_size
+            ):
+                raise ValueError(
+                    "Replay buffer checkpoint group misaligned: "
+                    f"sample_ids={len(meta.sample_ids)}, tags={num_tags}, "
+                    f"sequence_lengths={num_lengths}, "
+                    f"expected_group_size={expected_group_size}"
+                )
+            for sid in meta.sample_ids:
+                if sid in seen_sample_ids:
+                    raise ValueError(
+                        f"Replay buffer checkpoint has duplicate sample_id: {sid!r}"
+                    )
+                seen_sample_ids.add(sid)
+
+        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:
+            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),
+            )
+            indices_to_keep = sorted(prioritized[num_truncated:])
+            groups = [groups[i] for i in indices_to_keep]
+
+        for group in groups:
+            meta = group["meta"]
+            await self._call_dp(
+                "put_samples",
+                sample_ids=list(meta.sample_ids),
+                partition_id=self._partition_id,
+                fields=group["fields_data"],
+                tags=[dict(t) for t in meta.tags],
+            )
+            self.meta_list.append(meta)
+            self.start_weight_list.append(group["start_weight"])
+            self.end_weight_list.append(group["end_weight"])
+            self.target_step_list.append(group["target_step"])
+            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)
+        return len(groups)
+
     def size(self) -> int:
         """Return the number of prompt-group entries currently held."""
         return len(self.meta_list)
diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py
index 76cb919849..f8ba5b663a 100644
--- a/tests/unit/single_controller/test_tq_replay_buffer.py
+++ b/tests/unit/single_controller/test_tq_replay_buffer.py
@@ -63,6 +63,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.get_calls: list[dict[str, Any]] = []
 
     def put_samples(
         self,
@@ -98,6 +99,24 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None
         for sid in ids:
             self._rows.pop(sid, None)
 
+    def get_samples(
+        self,
+        sample_ids: list[str],
+        partition_id: str,
+        select_fields: list[str] | None = None,
+    ) -> dict[str, Any]:
+        assert partition_id == self._partition_id
+        self.get_calls.append(
+            {
+                "sample_ids": list(sample_ids),
+                "select_fields": (
+                    list(select_fields) if select_fields is not None else None
+                ),
+            }
+        )
+        # Opaque per-group payload; load_state_dict must re-put it verbatim.
+        return {"payload_for": list(sample_ids)}
+
     def depth(self) -> int:
         return len(self._rows)
 
@@ -363,3 +382,200 @@ def test_size_and_len(self):
         _run(buf.remove([0], remove_in_dp=True))
         assert buf.size() == 1
         assert len(buf) == 1
+
+
+# ── state_dict / load_state_dict (checkpointing) ─────────────────────────────
+
+
+def _group_id_of(meta: KVBatchMeta) -> str:
+    head, _, _ = meta.sample_ids[0].rpartition("_g")
+    return head
+
+
+def _make_group_entry(
+    group_id: str,
+    weight: int,
+    *,
+    n: int = _N_GENS,
+    target_step: int | None = None,
+    sample_ids: list[str] | None = None,
+    sequence_lengths: list[int] | None = None,
+    partition_id: str = "rollout_data",
+) -> dict[str, Any]:
+    """Hand-built envelope group (bypasses commit) for preflight tests."""
+    sids = (
+        list(sample_ids)
+        if sample_ids is not None
+        else [f"{group_id}_g{i}" for i in range(n)]
+    )
+    meta = KVBatchMeta(
+        partition_id=partition_id,
+        task_name="train",
+        sample_ids=sids,
+        fields=["input_ids", "input_lengths", "total_reward"],
+        sequence_lengths=(
+            sequence_lengths if sequence_lengths is not None else [3] * len(sids)
+        ),
+        tags=[{"weight_version": weight}] * len(sids),
+    )
+    return {
+        "meta": meta,
+        "start_weight": weight,
+        "end_weight": weight,
+        "target_step": target_step,
+        "group_id": group_id,
+        "fields_data": {"payload_for": sids},
+    }
+
+
+def _make_envelope(
+    groups: list[dict[str, Any]],
+    *,
+    partition_id: str = "rollout_data",
+    saved_capacity: int = 8,
+) -> dict[str, Any]:
+    return {
+        "partition_id": partition_id,
+        "saved_capacity": saved_capacity,
+        "groups": list(groups),
+    }
+
+
+def _load(
+    buf: TQReplayBuffer,
+    state: dict[str, Any],
+    *,
+    max_groups: int = 8,
+    expected_partition_id: str = "rollout_data",
+    expected_group_size: int = _N_GENS,
+) -> int:
+    return _run(
+        buf.load_state_dict(
+            state,
+            max_groups=max_groups,
+            expected_partition_id=expected_partition_id,
+            expected_group_size=expected_group_size,
+        )
+    )
+
+
+class TestTQReplayBufferStateDict:
+    def test_state_dict_serializes_ready_and_skips_unready(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
+
+        state = _run(buf.state_dict(saved_capacity=8))
+
+        assert state["partition_id"] == "rollout_data"
+        assert state["saved_capacity"] == 8
+        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 dp.get_calls[0]["select_fields"] == list(metas[0].fields)
+        assert state["groups"][0]["fields_data"] == {
+            "payload_for": list(metas[0].sample_ids)
+        }
+
+    def test_round_trip_restores_lists_and_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))
+
+        dp2 = FakeDataPlaneClient()
+        buf2 = _make_buffer(dp2)
+        restored = _load(buf2, state)
+
+        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),
+        ]
+        # 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]
+
+
+class TestTQReplayBufferLoadPreflight:
+    """Malformed envelopes raise ValueError before any DataPlane write."""
+
+    def _assert_rejected(self, state: dict[str, Any], match: str, **load_kwargs):
+        dp = FakeDataPlaneClient()
+        buf = _make_buffer(dp)
+        with pytest.raises(ValueError, match=match):
+            _load(buf, state, **load_kwargs)
+        assert dp.put_calls == []
+        assert buf.size() == 0
+
+    def test_missing_envelope_keys(self):
+        self._assert_rejected({"groups": []}, match="missing required keys")
+
+    def test_partition_id_mismatch(self):
+        state = _make_envelope([], partition_id="other_partition")
+        self._assert_rejected(state, match="partition_id mismatch")
+
+    def test_group_missing_keys(self):
+        group = _make_group_entry("g0", weight=1)
+        del group["fields_data"]
+        self._assert_rejected(_make_envelope([group]), match="group missing keys")
+
+    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")
+
+    def test_group_size_mismatch(self):
+        state = _make_envelope([_make_group_entry("g0", weight=1, n=2)])
+        self._assert_rejected(state, match="misaligned", expected_group_size=3)
+
+    def test_duplicate_sample_ids_across_groups(self):
+        g0 = _make_group_entry("g0", weight=1)
+        g1 = _make_group_entry(
+            "g1", weight=2, sample_ids=["g0_g0", "g1_g1"]
+        )  # g0_g0 collides
+        self._assert_rejected(_make_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)),
+        )
+
+        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)

From 19d6797daf213d830d0f7b226cf31133c22a9f2b Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Wed, 29 Jul 2026 16:27:09 -0700
Subject: [PATCH 02/15] feat(sc): support checkpoint resume in prompt-group
 samplers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two additions the SC checkpointing path needs from the sampler layer:

- resume_from_step: BaseSampler (and the three built-in policies +
  create_sampler) now accept the trainer step the run starts from — 0
  for a fresh run, the restored current_step on resume. It seeds the
  dispatch cursor to preserve the fresh-start invariant
  _dispatch_index == trainer_version - 1. Without it a restored
  InOrderSampler stamps target_steps from 0 and every dispatched batch
  is instantly evicted (target < trainer_version), livelocking the
  train pump. create_sampler forwards the kwarg to custom samplers
  only on resume, so fresh starts don't constrain their constructors
  and an unsupported class fails loudly instead of silently running
  with an unseeded cursor.

- supports_buffer_checkpoint: new PromptGroupSampler property gating
  replay-buffer save/restore. Only the ungated WindowedSampler returns
  True — gated policies dispatch a fixed quota per trainer step, so
  restored groups could never complete an already-consumed window.

Covered by 8 new unit tests in test_sampler_interface.py (cursor
seeding, gate behavior after resume, factory forwarding, custom
fail-loud, checkpoint-support matrix).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 .../async_utils/staleness_sampler.py          | 118 +++++++++++++++---
 .../test_sampler_interface.py                 |  88 +++++++++++++
 2 files changed, 190 insertions(+), 16 deletions(-)

diff --git a/nemo_rl/algorithms/async_utils/staleness_sampler.py b/nemo_rl/algorithms/async_utils/staleness_sampler.py
index 69959de6e8..6810f2531e 100644
--- a/nemo_rl/algorithms/async_utils/staleness_sampler.py
+++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py
@@ -100,6 +100,16 @@ def is_on_policy(self) -> bool:
         """True when the policy admits zero staleness (sync mode)."""
         ...
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        """True when buffered groups may be checkpointed and restored.
+
+        Gated policies dispatch a fixed quota per trainer step, so groups
+        restored from a checkpoint can never complete an already-consumed
+        quota window — only ungated (over-sampled) policies can consume them.
+        """
+        ...
+
     def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
         """Buffer-capacity the policy needs, or ``None`` if unconstrained."""
         ...
@@ -113,11 +123,26 @@ class BaseSampler(abc.ABC):
     select-finalize / weight-window-evict helpers.
     """
 
-    def __init__(self, buffer: TQReplayBuffer) -> None:
+    def __init__(self, buffer: TQReplayBuffer, *, resume_from_step: int = 0) -> None:
+        """Initialize shared sampler state.
+
+        Args:
+            buffer: Shared TQReplayBuffer holding the candidate slots.
+            resume_from_step: Trainer step this run starts from — 0 for a
+                fresh run, the restored ``current_step`` when resuming from a
+                checkpoint. Seeds the dispatch cursor so gated ``admit`` and
+                ``InOrderSampler``'s target_step stamps line up with the
+                restored trainer version exactly as they would at step 0 of
+                a fresh run.
+        """
+        if resume_from_step < 0:
+            raise ValueError(
+                f"resume_from_step must be non-negative, got {resume_from_step}"
+            )
         self._buffer = buffer
-        # Pre-incremented before each admitted batch, so -1 lets the first
-        # batch through a zero-staleness gate.
-        self._dispatch_index: int = -1
+        # Pre-incremented before each admitted batch, so the cursor trails
+        # the run's starting step by one.
+        self._dispatch_index: int = resume_from_step - 1
 
     # ── rollout-pump side ────────────────────────────────────────────────
     @abc.abstractmethod
@@ -158,6 +183,10 @@ async def evict(self, *, current_train_weight: int) -> int:
     def is_on_policy(self) -> bool:
         return self._eviction_window() == 0
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        return False
+
     def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
         return None
 
@@ -215,8 +244,9 @@ def __init__(
         *,
         max_staleness_versions: int,
         sample_freshest_first: bool = False,
+        resume_from_step: int = 0,
     ) -> None:
-        super().__init__(buffer)
+        super().__init__(buffer, resume_from_step=resume_from_step)
         if max_staleness_versions < 0:
             raise ValueError(
                 f"max_staleness_versions must be non-negative, got "
@@ -228,6 +258,12 @@ def __init__(
     def _eviction_window(self) -> int:
         return self.max_staleness_versions
 
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        # Ungated: restored groups are ordinary in-window candidates, so the
+        # buffer can round-trip through a checkpoint.
+        return True
+
     async def admit(self, *, trainer_version_fn: Callable[[], int]) -> Optional[int]:
         # Over-sampled: dispatch is bounded by buffer capacity, not by version.
         return None
@@ -275,8 +311,14 @@ class _GatedSampler(BaseSampler):
     (``gate_window`` versions of lookahead).
     """
 
-    def __init__(self, buffer: TQReplayBuffer, *, gate_window: int) -> None:
-        super().__init__(buffer)
+    def __init__(
+        self,
+        buffer: TQReplayBuffer,
+        *,
+        gate_window: int,
+        resume_from_step: int = 0,
+    ) -> None:
+        super().__init__(buffer, resume_from_step=resume_from_step)
         if gate_window < 0:
             raise ValueError(f"gate_window must be non-negative, got {gate_window}")
         self._gate_window = gate_window
@@ -308,8 +350,18 @@ class WeightFifoSampler(_GatedSampler):
     that weight's batch to fill. Evict uses the weight window (default).
     """
 
-    def __init__(self, buffer: TQReplayBuffer, *, max_staleness_versions: int) -> None:
-        super().__init__(buffer, gate_window=max_staleness_versions)
+    def __init__(
+        self,
+        buffer: TQReplayBuffer,
+        *,
+        max_staleness_versions: int,
+        resume_from_step: int = 0,
+    ) -> None:
+        super().__init__(
+            buffer,
+            gate_window=max_staleness_versions,
+            resume_from_step=resume_from_step,
+        )
         self.max_staleness_versions = max_staleness_versions
 
     async def select(
@@ -349,8 +401,18 @@ class InOrderSampler(_GatedSampler):
     upcoming is never dropped early, and evict/select can't disagree.
     """
 
-    def __init__(self, buffer: TQReplayBuffer, *, max_lookahead_versions: int) -> None:
-        super().__init__(buffer, gate_window=max_lookahead_versions)
+    def __init__(
+        self,
+        buffer: TQReplayBuffer,
+        *,
+        max_lookahead_versions: int,
+        resume_from_step: int = 0,
+    ) -> None:
+        super().__init__(
+            buffer,
+            gate_window=max_lookahead_versions,
+            resume_from_step=resume_from_step,
+        )
         self.max_lookahead_versions = max_lookahead_versions
 
     def _stamp(self) -> Optional[int]:
@@ -454,20 +516,40 @@ def required_buffer_capacity_for_config(
 def create_sampler(
     buffer: TQReplayBuffer,
     cfg: SamplerConfig,
+    *,
+    resume_from_step: int = 0,
 ) -> PromptGroupSampler:
-    """Build a sampler from its config (or import one by FQN)."""
+    """Build a sampler from its config (or import one by FQN).
+
+    Args:
+        buffer: Shared TQReplayBuffer holding the candidate slots.
+        cfg: Discriminated sampler config selecting the policy.
+        resume_from_step: Trainer step this run starts from; 0 for a fresh
+            run, the restored ``current_step`` on checkpoint resume (see
+            ``BaseSampler.__init__``). Custom samplers receive it only on
+            resume, so fresh starts don't constrain their constructor
+            signature — a custom class that doesn't accept the kwarg fails
+            loudly the first time a run actually resumes with it.
+    """
     if isinstance(cfg, WindowedSamplerConfig):
         return WindowedSampler(
             buffer,
             max_staleness_versions=cfg.max_staleness_versions,
             sample_freshest_first=cfg.sample_freshest_first,
+            resume_from_step=resume_from_step,
         )
     if isinstance(cfg, WeightFifoSamplerConfig):
         return WeightFifoSampler(
-            buffer, max_staleness_versions=cfg.max_staleness_versions
+            buffer,
+            max_staleness_versions=cfg.max_staleness_versions,
+            resume_from_step=resume_from_step,
         )
     if isinstance(cfg, InOrderSamplerConfig):
-        return InOrderSampler(buffer, max_lookahead_versions=cfg.max_lookahead_versions)
+        return InOrderSampler(
+            buffer,
+            max_lookahead_versions=cfg.max_lookahead_versions,
+            resume_from_step=resume_from_step,
+        )
     if isinstance(cfg, CustomSamplerConfig):
         module_name, sep, class_name = cfg.target.partition(":")
         if not sep:
@@ -475,11 +557,15 @@ def create_sampler(
                 f"custom sampler target must be 'module:ClassName', got {cfg.target!r}"
             )
         sampler_cls = getattr(importlib.import_module(module_name), class_name)
-        sampler = sampler_cls(buffer, **(cfg.model_extra or {}))
+        custom_kwargs = dict(cfg.model_extra or {})
+        if resume_from_step != 0:
+            custom_kwargs["resume_from_step"] = resume_from_step
+        sampler = sampler_cls(buffer, **custom_kwargs)
         if not isinstance(sampler, PromptGroupSampler):
             raise TypeError(
                 f"{cfg.target} does not implement the PromptGroupSampler "
-                f"interface (needs admit/select/evict)"
+                f"interface (needs admit/select/evict, is_on_policy, "
+                f"supports_buffer_checkpoint, required_buffer_capacity)"
             )
         return sampler
     raise ValueError(f"unknown sampler config {type(cfg).__name__}")
diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py
index 8e0ef3e20c..c99be7fb8b 100644
--- a/tests/unit/single_controller/test_sampler_interface.py
+++ b/tests/unit/single_controller/test_sampler_interface.py
@@ -304,5 +304,93 @@ def test_windowed_evict_skips_unready_stale(self):
         assert _run(s.evict(current_train_weight=5)) == 0
 
 
+class TestSupportsBufferCheckpoint:
+    """Only ungated (over-sampled) policies may checkpoint buffered groups."""
+
+    def test_windowed_supports_buffer_checkpoint(self):
+        assert WindowedSampler(
+            FakeBuffer(), max_staleness_versions=1
+        ).supports_buffer_checkpoint
+
+    def test_gated_samplers_do_not(self):
+        assert not WeightFifoSampler(
+            FakeBuffer(), max_staleness_versions=1
+        ).supports_buffer_checkpoint
+        assert not InOrderSampler(
+            FakeBuffer(), max_lookahead_versions=1
+        ).supports_buffer_checkpoint
+
+
+class TestDispatchCursorRestore:
+    """Checkpoint resume passes resume_from_step=current_step, restoring the
+    fresh-start invariant _dispatch_index == trainer_version - 1. Without it,
+    a restored InOrderSampler would stamp target_steps starting at 0 and every
+    dispatched batch would be instantly evicted (target < trainer_version)."""
+
+    def test_resumed_in_order_stamps_from_trainer_version(self):
+        s = InOrderSampler(FakeBuffer(), max_lookahead_versions=1, resume_from_step=7)
+        assert _run(s.admit(trainer_version_fn=lambda: 7)) == 7
+        assert _run(s.admit(trainer_version_fn=lambda: 8)) == 8
+
+    def test_resumed_gate_admits_window_then_blocks(self):
+        s = WeightFifoSampler(
+            FakeBuffer(), max_staleness_versions=0, resume_from_step=7
+        )
+        # Resumed at step 7, window 0: one batch admitted, then the gate
+        # closes exactly as it would on a fresh run at step 0.
+        assert _run(s.admit(trainer_version_fn=lambda: 7)) is None
+        with pytest.raises(asyncio.TimeoutError):
+            _run(asyncio.wait_for(s.admit(trainer_version_fn=lambda: 7), timeout=0.05))
+
+    def test_negative_resume_step_rejected(self):
+        with pytest.raises(ValueError, match="resume_from_step"):
+            WindowedSampler(FakeBuffer(), max_staleness_versions=1, resume_from_step=-1)
+
+
+class TestFactorySeeding:
+    def test_builtin_receives_resume_step(self):
+        s = create_sampler(
+            FakeBuffer(),
+            InOrderSamplerConfig(max_lookahead_versions=1),
+            resume_from_step=7,
+        )
+        assert _run(s.admit(trainer_version_fn=lambda: 7)) == 7
+
+    def test_custom_receives_resume_step_on_restore(self):
+        from nemo_rl.algorithms.async_utils.staleness_sampler import (
+            CustomSamplerConfig,
+        )
+
+        s = create_sampler(
+            FakeBuffer(),
+            CustomSamplerConfig(
+                target=f"{__name__}:EchoSampler", max_lookahead_versions=1
+            ),
+            resume_from_step=6,
+        )
+        assert _run(s.admit(trainer_version_fn=lambda: 6)) == 6
+
+    def test_custom_without_restore_support_fails_only_on_restore(self):
+        from nemo_rl.algorithms.async_utils.staleness_sampler import (
+            CustomSamplerConfig,
+        )
+
+        cfg = CustomSamplerConfig(
+            target=f"{__name__}:NoRestoreSampler", max_lookahead_versions=1
+        )
+        # Fresh start: the kwarg is withheld, existing custom samplers keep working.
+        assert isinstance(create_sampler(FakeBuffer(), cfg), NoRestoreSampler)
+        # Resume: loud TypeError instead of a silently unseeded cursor.
+        with pytest.raises(TypeError):
+            create_sampler(FakeBuffer(), cfg, resume_from_step=5)
+
+
 class EchoSampler(InOrderSampler):
     """Stand-in for a user-defined sampler loaded by FQN."""
+
+
+class NoRestoreSampler(InOrderSampler):
+    """User sampler whose constructor predates resume_from_step."""
+
+    def __init__(self, buffer, *, max_lookahead_versions: int) -> None:
+        super().__init__(buffer, max_lookahead_versions=max_lookahead_versions)

From d9eac70738e3625c9ce9a37e129d9c04e5daf922 Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Wed, 29 Jul 2026 17:06:31 -0700
Subject: [PATCH 03/15] feat(sc): wire checkpoint resume through
 setup_single_controller
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Replace the checkpointing NotImplementedError guard with the actual
driver-side resume wiring, following the grpo.py setup pattern:

- Build a CheckpointManager unconditionally and resolve the latest
  checkpoint: load_training_info() populates save_state (default
  GRPOSaveState when starting fresh) and get_resume_paths() yields the
  weights/optimizer paths.
- _build_trainer takes kw-only weights_path/optimizer_path and forwards
  them to TQPolicy (previously hardcoded to None) on both the colocated
  and non-colocated build paths.
- Restore the dataloader position from train_dataloader.pt when present
  (load_dataloader_state, with its dataset-swap guard); warn and start
  fresh otherwise. Runs before _clamp_max_num_steps as before.
- Forward checkpointing.pretrained_checkpoint into the policy config.
- SingleControllerActorArgs carries two new fields, save_state and
  last_checkpoint_path, for the actor-side restore (next step).

Saving itself is not wired yet — that lands in the SingleControllerActor
train pump next. Existing tests updated for the new surface: the setup
tests' hand-built checkpointing block now carries the keys
CheckpointManager indexes, and the pump tests pass the two new
ActorArgs fields.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 .../single_controller_utils/setup.py          | 71 ++++++++++++++++---
 .../single_controller/test_rollout_pump.py    |  3 +
 .../test_single_controller_setup.py           | 13 +++-
 .../unit/single_controller/test_train_pump.py |  3 +
 4 files changed, 78 insertions(+), 12 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index 6ad8d107a3..579ede5d87 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -21,8 +21,10 @@
 
 from __future__ import annotations
 
+import os
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Optional, cast
 
 from torchdata.stateful_dataloader import StatefulDataLoader
@@ -32,7 +34,9 @@
 from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer
 from nemo_rl.algorithms.grpo import MasterConfig as GrpoMasterConfig
 from nemo_rl.algorithms.grpo import (
+    GRPOSaveState,
     _create_advantage_estimator,
+    _default_grpo_save_state,
     _should_use_nemo_gym,
 )
 from nemo_rl.algorithms.loss import ClippedPGLossFn
@@ -43,7 +47,7 @@
 )
 from nemo_rl.algorithms.utils import set_seed
 from nemo_rl.data.collate_fn import rl_collate_fn
-from nemo_rl.data.utils import setup_response_data
+from nemo_rl.data.utils import load_dataloader_state, setup_response_data
 from nemo_rl.data_plane import DataPlaneClient, build_data_plane_client
 from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
 from nemo_rl.environments.interfaces import EnvironmentInterface
@@ -61,6 +65,7 @@
     router_replay_enabled,
 )
 from nemo_rl.models.policy.tq_policy import TQPolicy
+from nemo_rl.utils.checkpoint import CheckpointManager
 from nemo_rl.weight_sync import WeightSynchronizer, create_weight_synchronizer
 
 
@@ -85,6 +90,8 @@ class SingleControllerActorArgs:
     rollout_manager: RolloutManager
     tq_buffer: TQReplayBuffer
     partition_id: str
+    save_state: GRPOSaveState
+    last_checkpoint_path: Optional[str]
 
 
 def _build_clusters(
@@ -197,6 +204,9 @@ def _build_trainer(
     master_config: MasterConfig,
     tokenizer,
     processor,
+    *,
+    weights_path: Optional[Path],
+    optimizer_path: Optional[Path],
 ):
     """Build the TQ-mediated trainer (driver-side TQPolicy).
 
@@ -212,8 +222,8 @@ def _build_trainer(
         config=master_config.policy,
         tokenizer=tokenizer,
         processor=processor,
-        weights_path=None,
-        optimizer_path=None,
+        weights_path=weights_path,
+        optimizer_path=optimizer_path,
         init_optimizer=True,
         init_reference_model=init_reference_model,
         dp_cfg=master_config.data_plane,
@@ -296,12 +306,6 @@ def setup_single_controller(
             "SingleController doesn't support validation now, will support "
             "later. Set grpo.val_period=0, val_at_start=false, val_at_end=false."
         )
-    if master_config.checkpointing["enabled"]:
-        raise NotImplementedError(
-            "SingleController doesn't support checkpointing now, will support "
-            "later. Set checkpointing.enabled=false."
-        )
-
     if dp_config is None or not dp_config.get("enabled", False):
         raise ValueError(
             "single_controller_utils.setup requires "
@@ -319,8 +323,24 @@ def setup_single_controller(
             "data.use_multiple_dataloader=True yet."
         )
 
+    checkpointing_pretrained = master_config.checkpointing.get("pretrained_checkpoint")
+    if checkpointing_pretrained is not None:
+        policy_config["pretrained_checkpoint"] = checkpointing_pretrained
+
     set_seed(grpo_config["seed"])
 
+    # ==========================
+    # Checkpointing
+    # ==========================
+    checkpointer = CheckpointManager(master_config.checkpointing)
+    last_checkpoint_path = checkpointer.get_latest_checkpoint_path()
+    save_state = cast(
+        Optional[GRPOSaveState], checkpointer.load_training_info(last_checkpoint_path)
+    )
+    if save_state is None:
+        save_state = _default_grpo_save_state()
+    weights_path, optimizer_path = checkpointer.get_resume_paths(last_checkpoint_path)
+
     # ==========================
     # Setup Dataset & Environments
     # ==========================
@@ -352,6 +372,20 @@ def setup_single_controller(
         drop_last=True,
         num_workers=data_config["num_workers"],
     )
+    if last_checkpoint_path is not None:
+        dataloader_state_path = os.path.join(
+            last_checkpoint_path, "train_dataloader.pt"
+        )
+        if os.path.exists(dataloader_state_path):
+            print(
+                f"📦 Restoring dataloader state from checkpoint: {dataloader_state_path}"
+            )
+            load_dataloader_state(dataloader, last_checkpoint_path, data_config)
+        else:
+            print(
+                f"⚠️ No dataloader state found at {dataloader_state_path}. "
+                "Starting with a fresh dataloader position."
+            )
 
     _clamp_max_num_steps(master_config, dataloader)
     _maybe_inject_megatron_train_iters(master_config)
@@ -365,7 +399,14 @@ def setup_single_controller(
         # Colocated: vLLM prefers a clean GPU at load time, so generation
         # comes up before the policy.
         generation = _build_generation(inference_cluster, master_config)
-        policy = _build_trainer(train_cluster, master_config, tokenizer, processor)
+        policy = _build_trainer(
+            train_cluster,
+            master_config,
+            tokenizer,
+            processor,
+            weights_path=weights_path,
+            optimizer_path=optimizer_path,
+        )
     else:
         # Non-colocated: generation + policy run on disjoint GPUs, so
         # bring them up in parallel.
@@ -374,7 +415,13 @@ def setup_single_controller(
                 _build_generation, inference_cluster, master_config
             )
             policy_future = executor.submit(
-                _build_trainer, train_cluster, master_config, tokenizer, processor
+                _build_trainer,
+                train_cluster,
+                master_config,
+                tokenizer,
+                processor,
+                weights_path=weights_path,
+                optimizer_path=optimizer_path,
             )
             generation = gen_future.result()
             policy = policy_future.result()
@@ -458,4 +505,6 @@ def setup_single_controller(
         rollout_manager=rollout_manager,
         tq_buffer=tq_buffer,
         partition_id=partition_id,
+        save_state=save_state,
+        last_checkpoint_path=last_checkpoint_path,
     )
diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py
index cfeda7b1f5..08a56ae426 100644
--- a/tests/unit/single_controller/test_rollout_pump.py
+++ b/tests/unit/single_controller/test_rollout_pump.py
@@ -31,6 +31,7 @@
     WindowedSampler,
     WindowedSamplerConfig,
 )
+from nemo_rl.algorithms.grpo import _default_grpo_save_state
 from nemo_rl.algorithms.single_controller import SingleControllerActor
 from nemo_rl.algorithms.single_controller_utils.config import (
     AsyncRLConfig,
@@ -343,6 +344,8 @@ def test_rollout_pump_writes_expected_tq_data(
         rollout_manager=rollout_manager,
         tq_buffer=tq_buffer,
         partition_id=_PARTITION_ID,
+        save_state=_default_grpo_save_state(),
+        last_checkpoint_path=None,
     )
     ctrl = SingleControllerActor.remote(
         master_config=master_config, actor_args=actor_args
diff --git a/tests/unit/single_controller/test_single_controller_setup.py b/tests/unit/single_controller/test_single_controller_setup.py
index 3a71a23b6a..5b5bac11b2 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -77,7 +77,18 @@ def _make_master_config(
                 "colocated": {"enabled": colocated, "resources": {}},
             },
         },
-        checkpointing={"enabled": False},
+        # Full block: setup builds a CheckpointManager unconditionally (resume
+        # lookup), which indexes these keys directly. Nothing is written while
+        # enabled=False and the dir doesn't exist.
+        checkpointing={
+            "enabled": False,
+            "checkpoint_dir": "results/_sc_setup_test_ckpt",
+            "metric_name": None,
+            "higher_is_better": False,
+            "keep_top_k": None,
+            "save_period": 10,
+            "save_optimizer": False,
+        },
         loss_fn=ClippedPGLossConfig(),
         env=env if env is not None else {},
         async_rl=AsyncRLConfig(
diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py
index fe2b59ebbd..9c8ec46b39 100644
--- a/tests/unit/single_controller/test_train_pump.py
+++ b/tests/unit/single_controller/test_train_pump.py
@@ -28,6 +28,7 @@
 
 from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer
 from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig
+from nemo_rl.algorithms.grpo import _default_grpo_save_state
 from nemo_rl.algorithms.single_controller import SingleControllerActor
 from nemo_rl.algorithms.single_controller_utils.config import (
     AsyncRLConfig,
@@ -336,6 +337,8 @@ def test_train_pump_drives_mcore_training_step(
             rollout_manager=rollout_manager,
             tq_buffer=tq_buffer,
             partition_id=_PARTITION_ID,
+            save_state=_default_grpo_save_state(),
+            last_checkpoint_path=None,
         )
         ctrl = _RecordingSingleControllerActor.remote(
             metric_log_handle=log,

From 9148ef4931fe687d14a5cacbeebf482f37e2b460 Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Thu, 30 Jul 2026 10:00:28 -0700
Subject: [PATCH 04/15] feat(sc): save and restore checkpoints in
 SingleControllerActor
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wire the actor side of SC checkpointing (the setup/resume half landed in
the previous commit), with Megatron async_save supported end to end:

- Restore: __init__ rebuilds counters (train_steps/trainer_version/
  current_epoch/consumed_samples/total_valid_tokens) from the save_state
  loaded by setup, seeds the sampler with resume_from_step, and run()
  reloads replay-buffer groups (ungated samplers only, one capacity
  permit per restored group) before the pumps start.
- Save: after each weight sync, _save_checkpoint mirrors
  async_grpo_train's block — finalize_pending flushes the previous
  background finalization, save_checkpoint returns after D2H staging
  under async_save, aux state (training info, dataloader position,
  replay buffer when the sampler supports it) is written synchronously,
  then begin_finalization defers the tmp->step rename until the async
  weight writes complete. run() flushes the last checkpoint via
  checkpointer.shutdown() on every exit path.
- TimeoutChecker drives checkpoint_must_save_by: a timeout save also
  stops training early, matching the legacy loops.
- latest_checkpoint_status.json is refreshed after each save for
  external watchdogs (reuses grpo's _write_latest_checkpoint_status).

The pump tests' hand-built configs gain the checkpointing block the
actor now reads (enabled=false keeps them write-free).

Validated end to end on GB200 (Qwen3-0.6B, megatron async_save=true):
4-step run saves step_2/step_4 with no tmp_step_* leftovers; a
checkpoint_must_save_by run stops early with a complete checkpoint; the
resume run restores dataloader + 4 replay groups and continues from
step 2 to step 4.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 nemo_rl/algorithms/single_controller.py       | 203 +++++++++++++++++-
 .../single_controller/test_rollout_pump.py    |  12 ++
 .../unit/single_controller/test_train_pump.py |  12 ++
 3 files changed, 221 insertions(+), 6 deletions(-)

diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index 9effa5168c..899683b9a2 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -35,7 +35,9 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import time
+import warnings
 from functools import partial
 from typing import Any, Optional, Union
 
@@ -43,6 +45,7 @@
 import torch
 
 from nemo_rl.algorithms.async_utils.staleness_sampler import create_sampler
+from nemo_rl.algorithms.grpo import GRPOSaveState, _write_latest_checkpoint_status
 from nemo_rl.algorithms.single_controller_utils.config import (
     AdvantageConfig,
     MasterConfig,
@@ -64,8 +67,9 @@
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
 from nemo_rl.models.policy.tq_policy import TQPolicy
+from nemo_rl.utils.checkpoint import CheckpointManager
 from nemo_rl.utils.logger import Logger
-from nemo_rl.utils.timer import Timer
+from nemo_rl.utils.timer import TimeoutChecker, Timer
 
 Generation = Union[VllmGeneration, SGLangGeneration]
 
@@ -126,6 +130,26 @@ def __init__(
         self._logger = Logger(master_config.logger)  # type: ignore
         self._timer = Timer()
 
+        # Also built here, not on the driver: TimeoutChecker must capture
+        # wall-clock start times inside the actor, not at driver setup time.
+        # actor_args only carries the driver-side restore products
+        # (save_state, last_checkpoint_path).
+        self._checkpointer = CheckpointManager(master_config.checkpointing)
+        self._timeout = TimeoutChecker(
+            timeout=master_config.checkpointing["checkpoint_must_save_by"],
+            fit_last_save_time=True,
+        )
+        self._timeout.start_iterations()
+
+        # Loaded (or default) GRPOSaveState; keys SC does not own
+        # (val_reward, ...) pass through to saved checkpoints untouched.
+        self._save_state: GRPOSaveState = actor_args.save_state
+        self._last_checkpoint_path: Optional[str] = actor_args.last_checkpoint_path
+        self._consumed_samples: int = actor_args.save_state["consumed_samples"]
+        self._total_valid_tokens: int = actor_args.save_state.get(
+            "total_valid_tokens", 0
+        )  # Default to 0 for backward compatibility with older checkpoints
+
         # Pin clusters so RayVirtualCluster.__del__ doesn't remove the PGs.
         self._train_cluster = actor_args.train_cluster
         self._inference_cluster = actor_args.inference_cluster
@@ -134,6 +158,7 @@ def __init__(
         self._sampler = create_sampler(
             self._buffer,
             self._async_cfg.sampler,
+            resume_from_step=actor_args.save_state["current_step"],
         )
         required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step)
         validate_sampler_buffer_capacity(
@@ -165,9 +190,9 @@ def __init__(
             self._async_cfg.max_buffered_rollouts
         )
 
-        self._trainer_version: int = 0
-        self._train_steps: int = 0
-        self._current_epoch: int = 0
+        self._trainer_version: int = actor_args.save_state["current_step"]
+        self._train_steps: int = actor_args.save_state["current_step"]
+        self._current_epoch: int = actor_args.save_state["current_epoch"]
         self._step_log_dict: dict[str, list] = {
             "rewards": [],
             "masked_advantages": [],
@@ -190,6 +215,8 @@ async def run(self) -> dict[str, Any]:
         # Synchronize weights before starting the pumps
         await self._sync_weights()
 
+        await self._maybe_restore_replay_buffer()
+
         # Start the rollout and train pumps
         rollout_task = asyncio.create_task(self._rollout_pump())
         train_task = asyncio.create_task(self._train_pump())
@@ -206,6 +233,8 @@ async def run(self) -> dict[str, Any]:
             rollout_task.cancel()
             train_task.cancel()
             await asyncio.gather(rollout_task, train_task, return_exceptions=True)
+            # Flush the last checkpoint's background finalization before exit.
+            await asyncio.to_thread(self._checkpointer.shutdown)
             self._logger.finish()
 
         return {
@@ -226,6 +255,44 @@ 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.
+
+        No-op unless the sampler supports_buffer_checkpoint (ungated only).
+        """
+        if not (
+            self._sampler.supports_buffer_checkpoint
+            and self._last_checkpoint_path is not None
+        ):
+            return
+        buffer_path = os.path.join(self._last_checkpoint_path, "replay_buffer.pt")
+        if not os.path.exists(buffer_path):
+            print(
+                f"⚠️ No replay buffer checkpoint found at {buffer_path}. "
+                "Starting with an empty replay buffer.",
+                flush=True,
+            )
+            return
+        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.
+        buffer_state = await asyncio.to_thread(
+            torch.load, buffer_path, weights_only=False
+        )
+        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"
+            ],
+        )
+        # Each buffered group holds one _buffer_capacity permit; the load
+        # truncation guarantees restored <= capacity, so this never blocks.
+        assert restored <= self._async_cfg.max_buffered_rollouts
+        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
@@ -524,6 +591,30 @@ async def _train_pump(self) -> None:
                     )
                     await self._sync_weights(calibration_data=calibration_data)
 
+                # Checkpointing (mirrors async_grpo_train's save block).
+                self._consumed_samples += grpo_cfg["num_prompts_per_step"]
+                self._total_valid_tokens += step_metrics.get(
+                    "global_valid_toks", 0
+                )
+                self._timeout.mark_iteration()
+
+                is_last_step = self._train_steps >= grpo_cfg["max_num_steps"]
+                # _train_steps was already incremented above, so it equals
+                # the legacy loop's 1-indexed `step + 1`.
+                should_save_by_step = (
+                    is_last_step
+                    or self._train_steps
+                    % self._master_config.checkpointing["save_period"]
+                    == 0
+                )
+                should_save_by_timeout = self._timeout.check_save()
+
+                if self._master_config.checkpointing["enabled"] and (
+                    should_save_by_step or should_save_by_timeout
+                ):
+                    with self._timer.time("checkpointing"):
+                        await self._save_checkpoint(step_metrics)
+
             timing_metrics: dict[str, float] = self._timer.get_timing_metrics(
                 reduction_op="sum"
             )  # type: ignore
@@ -549,8 +640,6 @@ async def _train_pump(self) -> None:
                 percent = (v / total_time * 100) if total_time > 0 else 0.0
                 print(f"  • {k}: {v:.2f}s ({percent:.1f}%)")
 
-            # TODO: checkpointing (save_period/top-k metric_name,
-            #   policy.save_checkpoint, dataloader state, TQReplayBuffer state).
             # TODO: per-step train_data jsonl dump, vllm metrics logger,
             #   histogram log, rollout_metrics, seq_logprob_error_metrics,
             #   pretty-print "Training Results" block, print_performance_metrics.
@@ -573,6 +662,108 @@ async def _train_pump(self) -> None:
                 flush=True,
             )
 
+            if should_save_by_timeout:
+                print("Timeout has been reached, stopping training early", flush=True)
+                break
+
+    async def _save_checkpoint(
+        self,
+        step_metrics: dict[str, Any],
+        val_metrics: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """Write a full checkpoint for the just-finished train step.
+
+        Everything except the (possibly async) policy weight write must be
+        on disk before begin_finalization; rollouts keep running throughout.
+        val_metrics stays None until SC grows a validation loop.
+        """
+        save_state = self._save_state
+        save_state["current_step"] = self._train_steps
+        save_state["total_steps"] = self._train_steps
+        save_state["current_epoch"] = self._current_epoch
+        save_state["consumed_samples"] = self._consumed_samples
+        save_state["total_valid_tokens"] = self._total_valid_tokens
+        # Snapshot before any await so it can't interleave with
+        # _rollout_pump iterating this same dataloader.
+        dataloader_state = self._dataloader.state_dict()
+        if val_metrics is not None:
+            save_state["val_reward"] = val_metrics["accuracy"]
+        elif "val_reward" in save_state:
+            del save_state["val_reward"]
+
+        full_metric_name = self._master_config.checkpointing["metric_name"]
+        if full_metric_name is not None:
+            assert full_metric_name.startswith(
+                "train:"
+            ) or full_metric_name.startswith("val:"), (
+                f"metric_name={full_metric_name} must start with 'val:' or 'train:',\n"
+                f'followed by the corresponding name in the "val" or "train" metrics dictionary.'
+                f"  If you are using an old config, please updated checkpointing.metric_name to the new format, "
+                f" e.g. 'val_reward --> 'val:accuracy'"
+            )
+            prefix, metric_name = full_metric_name.split(":", 1)
+            metrics_source = step_metrics if prefix == "train" else val_metrics
+            if not metrics_source:
+                warnings.warn(
+                    f"You asked to save checkpoints based on {metric_name} but no {prefix} metrics were collected. "
+                    "This checkpoint will not be saved as top-k.",
+                    stacklevel=2,
+                )
+                if full_metric_name in save_state:
+                    del save_state[full_metric_name]
+            elif metric_name not in metrics_source:
+                raise ValueError(
+                    f"Metric {metric_name} not found in {prefix} metrics"
+                )
+            else:
+                save_state[full_metric_name] = metrics_source[metric_name]
+
+        # Flush the previous checkpoint's background finalization first;
+        # re-raises a failure from it.
+        await asyncio.to_thread(self._checkpointer.finalize_pending)
+
+        print(f"Saving checkpoint for step {self._train_steps}...")
+        checkpoint_path = await asyncio.to_thread(
+            self._checkpointer.init_tmp_checkpoint,
+            self._train_steps,
+            save_state,
+            self._master_config,
+        )
+        # With async_save this returns after D2H staging; disk writes finish
+        # in the background.
+        await asyncio.to_thread(
+            self._trainer.save_checkpoint,
+            weights_path=os.path.join(checkpoint_path, "policy", "weights"),
+            optimizer_path=os.path.join(checkpoint_path, "policy", "optimizer")
+            if self._checkpointer.save_optimizer
+            else None,
+            tokenizer_path=os.path.join(checkpoint_path, "policy", "tokenizer"),
+            checkpointing_cfg=self._master_config.checkpointing,
+        )
+        await asyncio.to_thread(
+            torch.save,
+            dataloader_state,
+            os.path.join(checkpoint_path, "train_dataloader.pt"),
+        )
+        if self._sampler.supports_buffer_checkpoint:
+            buffer_state = await self._buffer.state_dict(
+                saved_capacity=self._async_cfg.max_buffered_rollouts
+            )
+            await asyncio.to_thread(
+                torch.save,
+                buffer_state,
+                os.path.join(checkpoint_path, "replay_buffer.pt"),
+            )
+        # Rename happens in the background once the async weight writes
+        # finish; flushed at the next save or on exit.
+        self._checkpointer.begin_finalization(
+            checkpoint_path,
+            wait_fn=self._trainer.finalize_async_save,
+        )
+        _write_latest_checkpoint_status(
+            self._checkpointer, last_checkpoint_step=self._train_steps
+        )
+
     async def _sync_weights(
         self,
         *,
diff --git a/tests/unit/single_controller/test_rollout_pump.py b/tests/unit/single_controller/test_rollout_pump.py
index 08a56ae426..ad337c75fb 100644
--- a/tests/unit/single_controller/test_rollout_pump.py
+++ b/tests/unit/single_controller/test_rollout_pump.py
@@ -310,6 +310,18 @@ def test_rollout_pump_writes_expected_tq_data(
             "mlflow_enabled": False,
             "monitor_gpus": False,
         },
+        # Actor __init__ builds a CheckpointManager + TimeoutChecker from
+        # this block; enabled=False keeps the run write-free.
+        checkpointing={
+            "enabled": False,
+            "checkpoint_dir": str(tmp_path / "checkpoints"),
+            "metric_name": None,
+            "higher_is_better": False,
+            "keep_top_k": None,
+            "save_period": 10_000,
+            "save_optimizer": False,
+            "checkpoint_must_save_by": None,
+        },
     )
     # Wrap each value in a single-element list so size==1 and v[0] returns the original field.
     batched_sample = BatchedDataDict({k: [v] for k, v in input_sample.items()})
diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py
index 9c8ec46b39..2d0998fd85 100644
--- a/tests/unit/single_controller/test_train_pump.py
+++ b/tests/unit/single_controller/test_train_pump.py
@@ -321,6 +321,18 @@ def test_train_pump_drives_mcore_training_step(
                 "mlflow_enabled": False,
                 "monitor_gpus": False,
             },
+            # Actor __init__ builds a CheckpointManager + TimeoutChecker from
+            # this block; enabled=False keeps the run write-free.
+            checkpointing={
+                "enabled": False,
+                "checkpoint_dir": str(tmp_path / "checkpoints"),
+                "metric_name": None,
+                "higher_is_better": False,
+                "keep_top_k": None,
+                "save_period": 10_000,
+                "save_optimizer": False,
+                "checkpoint_must_save_by": None,
+            },
         )
 
         actor_args = SingleControllerActorArgs(

From a96e5f84d0e12fa8811f0ca571ce9f2c90539645 Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Thu, 30 Jul 2026 10:21:29 -0700
Subject: [PATCH 05/15] test(sc): add unit tests for SC checkpointing
 save/restore

Port the checkpointing test suite from PR #3138 onto the split SC
architecture (actor_args, the PromptGroupSampler protocol, the split
trainer step API) and extend it for the async-save path:

- counter/sampler-cursor restore, save triggers (period boundary, last
  step, checkpoint_must_save_by timeout, disabled, save_optimizer),
  metric_name handling, dataloader state round-trip with the
  dataset-swap guard, and setup resume wiring (get_resume_paths
  forwarded to the trainer factory, training_info.json loaded).
- replay-buffer persistence is asserted against
  sampler.supports_buffer_checkpoint (windowed saves/restores with one
  capacity permit per group; gated samplers skip both sides).
- new async-save coverage: the tmp->step rename stays deferred until
  finalize_async_save completes and is flushed by shutdown; a failed
  background finalization re-raises at the next save; _save_checkpoint
  records val_metrics into val_reward and a val:* metric_name.

32 tests, in-process actor with fakes (ray.cluster_resources patched);
108 passed together with the existing single_controller suite.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 .../test_sc_checkpointing.py                  | 1083 +++++++++++++++++
 1 file changed, 1083 insertions(+)
 create mode 100644 tests/unit/single_controller/test_sc_checkpointing.py

diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py
new file mode 100644
index 0000000000..c7e3b96ece
--- /dev/null
+++ b/tests/unit/single_controller/test_sc_checkpointing.py
@@ -0,0 +1,1083 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for SC checkpointing.
+
+Covers:
+  - counter restore from save_state (train_steps / trainer_version / sampler
+    dispatch cursor, current_epoch);
+  - save trigger + write path through _train_pump with fakes (period
+    boundary, last step, timeout, disabled);
+  - async-save finalization (rename deferred until finalize_async_save,
+    background failure re-raised at the next save);
+  - metric_name behavior (val:* warn-and-save, train:* value recorded);
+  - dataloader state: train_dataloader.pt written at save, position
+    round-trip through a real StatefulDataLoader, dataset-swap guard,
+    setup restore wiring + missing-file fresh-position fallback;
+  - replay buffer persistence gated on sampler.supports_buffer_checkpoint;
+  - setup_single_controller resume-path wiring (get_resume_paths forwarded
+    to the trainer factory, save_state loaded from training_info.json).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import threading
+from pathlib import Path
+from typing import Any, Optional
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import yaml
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from nemo_rl.algorithms.async_utils.staleness_sampler import (
+    InOrderSamplerConfig,
+    WindowedSamplerConfig,
+)
+from nemo_rl.algorithms.grpo import _default_grpo_save_state
+from nemo_rl.algorithms.loss import ClippedPGLossConfig
+from nemo_rl.algorithms.single_controller import SingleControllerActor
+from nemo_rl.algorithms.single_controller_utils import (
+    AsyncRLConfig,
+    MasterConfig,
+    setup_single_controller,
+)
+from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs
+from nemo_rl.data.utils import load_dataloader_state
+from nemo_rl.data_plane import KVBatchMeta
+from nemo_rl.utils.checkpoint import CheckpointManager
+
+# Reuse the factory patches from the setup tests (same cross-module fixture
+# import pattern as test_rollout_pump.py).
+from tests.unit.single_controller.test_single_controller_setup import (
+    patched_factories,  # noqa: F401
+)
+
+# Instantiate the underlying class in-process (same pattern as
+# tests/unit/algorithms/test_async_utils.py for AsyncTrajectoryCollector).
+_ACTOR_CLS = SingleControllerActor.__ray_metadata__.modified_class
+
+_PARTITION_ID = "rollout_data"
+
+
+# ── fakes ────────────────────────────────────────────────────────────────────
+
+
+class _FakeTrainer:
+    """TQPolicy stand-in: train methods are no-ops, save_checkpoint records calls."""
+
+    def __init__(self, step_metrics: Optional[dict[str, Any]] = None) -> None:
+        self._step_metrics = dict(step_metrics or {})
+        self.save_calls: list[dict[str, Any]] = []
+        self.finalize_calls: int = 0
+
+    def prepare_for_lp_inference(self) -> None:
+        pass
+
+    def get_logprobs_from_meta(self, meta: KVBatchMeta) -> None:
+        pass
+
+    def get_reference_policy_logprobs_from_meta(self, meta: KVBatchMeta) -> None:
+        pass
+
+    def prepare_for_training(self) -> None:
+        pass
+
+    def begin_train_step(self, loss_fn: Any) -> None:
+        pass
+
+    def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None:
+        pass
+
+    def finish_train_step(self) -> dict[str, Any]:
+        return dict(self._step_metrics)
+
+    def save_checkpoint(
+        self,
+        *,
+        weights_path: str,
+        optimizer_path: Optional[str],
+        tokenizer_path: str,
+        checkpointing_cfg: dict[str, Any],
+    ) -> None:
+        self.save_calls.append(
+            {
+                "weights_path": weights_path,
+                "optimizer_path": optimizer_path,
+                "tokenizer_path": tokenizer_path,
+                "checkpointing_cfg": checkpointing_cfg,
+            }
+        )
+        # Mimic the real Policy: materialize the checkpoint subdirs.
+        os.makedirs(weights_path, exist_ok=True)
+        if optimizer_path is not None:
+            os.makedirs(optimizer_path, exist_ok=True)
+        os.makedirs(tokenizer_path, exist_ok=True)
+
+    def finalize_async_save(self) -> None:
+        self.finalize_calls += 1
+
+
+class _GatedFinalizeTrainer(_FakeTrainer):
+    """Async-save stand-in: finalize_async_save blocks until released."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.release = threading.Event()
+
+    def finalize_async_save(self) -> None:
+        assert self.release.wait(timeout=30.0), "test never released the writer"
+        super().finalize_async_save()
+
+
+class _FailingFinalizeTrainer(_FakeTrainer):
+    def finalize_async_save(self) -> None:
+        raise RuntimeError("injected async-writer failure")
+
+
+class _FakeSampler:
+    """PromptGroupSampler stand-in: always returns a full, fresh batch."""
+
+    def __init__(self, supports_buffer_checkpoint: bool = True) -> None:
+        self._supports_buffer_checkpoint = supports_buffer_checkpoint
+        self._step = 0
+
+    async def admit(self, *, trainer_version_fn) -> Optional[int]:
+        return None
+
+    async def evict(self, *, current_train_weight: int) -> int:
+        return 0
+
+    async def select(
+        self,
+        *,
+        current_train_weight: int,
+        min_prompt_groups: int,
+        max_prompt_groups: int,
+    ) -> tuple[KVBatchMeta, int]:
+        n = max_prompt_groups
+        sample_ids = [f"s{self._step}-{i}" for i in range(n)]
+        self._step += 1
+        meta = KVBatchMeta(
+            partition_id=_PARTITION_ID,
+            task_name=None,
+            sample_ids=sample_ids,
+            sequence_lengths=[16] * n,
+            tags=[{"weight_version": current_train_weight}] * n,
+        )
+        return meta, n
+
+    @property
+    def is_on_policy(self) -> bool:
+        return False
+
+    @property
+    def supports_buffer_checkpoint(self) -> bool:
+        return self._supports_buffer_checkpoint
+
+    def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
+        return None
+
+
+class _FakeDPClient:
+    def __init__(self) -> None:
+        self.clear_calls: list[tuple[list[str], str]] = []
+
+    def clear_samples(self, sample_ids: list[str], partition_id: str) -> None:
+        self.clear_calls.append((list(sample_ids), partition_id))
+
+
+class _FakeWeightSynchronizer:
+    def __init__(self) -> None:
+        self.sync_count = 0
+
+    def sync_weights(self, *, kv_scales: Any = None) -> None:
+        self.sync_count += 1
+
+
+class _FakeRolloutManager:
+    def __init__(self) -> None:
+        self.weight_versions: list[int] = []
+        self._tq_buffer = None
+
+    def set_weight_version(self, version: int) -> None:
+        self.weight_versions.append(version)
+
+
+class _FakeTQBuffer:
+    """TQReplayBuffer stand-in for the SC save/restore integration tests."""
+
+    def __init__(
+        self,
+        state: Optional[dict[str, Any]] = None,
+        load_return: int = 0,
+    ) -> None:
+        self._state = state if state is not None else {"fake_buffer_envelope": 1}
+        self.load_return = load_return
+        self.state_dict_calls: list[int] = []
+        self.load_calls: list[dict[str, Any]] = []
+
+    async def state_dict(self, *, saved_capacity: int) -> dict[str, Any]:
+        self.state_dict_calls.append(saved_capacity)
+        return dict(self._state)
+
+    async def load_state_dict(
+        self,
+        state: dict[str, Any],
+        *,
+        max_groups: int,
+        expected_partition_id: str,
+        expected_group_size: int,
+    ) -> int:
+        self.load_calls.append(
+            {
+                "state": state,
+                "max_groups": max_groups,
+                "expected_partition_id": expected_partition_id,
+                "expected_group_size": expected_group_size,
+            }
+        )
+        return self.load_return
+
+
+# Default position sentinel the fake dataloader reports via state_dict().
+_SENTINEL_DL_STATE = {"fake_position": 42}
+
+
+class _FakeDataloader(list):
+    """List-backed dataloader with the StatefulDataLoader state_dict surface.
+
+    The save block snapshots ``self._dataloader.state_dict()``; return a
+    sentinel dict so tests can assert the exact object written to
+    train_dataloader.pt.
+    """
+
+    def __init__(self, batches: Any = (), state: Optional[dict[str, Any]] = None):
+        super().__init__(batches)
+        self._state = dict(state) if state is not None else dict(_SENTINEL_DL_STATE)
+
+    def state_dict(self) -> dict[str, Any]:
+        return dict(self._state)
+
+
+# ── builders ─────────────────────────────────────────────────────────────────
+
+
+def _actor_master_config(
+    tmp_path: Path,
+    *,
+    max_num_steps: int = 4,
+    save_period: int = 2,
+    enabled: bool = True,
+    metric_name: Optional[str] = None,
+    save_optimizer: bool = True,
+    checkpoint_must_save_by: Optional[str] = None,
+    num_prompts_per_step: int = 2,
+    max_num_epochs: int = 1,
+    buffer_checkpoint: bool = True,
+) -> MasterConfig:
+    """MasterConfig for in-process SingleControllerActor tests.
+
+    All fields are populated (init_tmp_checkpoint dumps the whole config to
+    config.yaml); values satisfy validate_single_controller_config.
+    buffer_checkpoint selects the sampler: windowed supports replay-buffer
+    checkpointing, the gated in_order sampler does not.
+    """
+    sampler_cfg = (
+        WindowedSamplerConfig(max_staleness_versions=1)
+        if buffer_checkpoint
+        else InOrderSamplerConfig(max_lookahead_versions=1)
+    )
+    return MasterConfig.model_construct(
+        policy={
+            # One optimizer.step per RL step: prompts * generations == gbs.
+            "train_global_batch_size": num_prompts_per_step * 2,
+        },
+        loss_fn=ClippedPGLossConfig(),
+        env={},
+        data={"shuffle": False, "num_workers": 0},
+        grpo={
+            "max_num_steps": max_num_steps,
+            "max_num_epochs": max_num_epochs,
+            "num_prompts_per_step": num_prompts_per_step,
+            "num_generations_per_prompt": 2,
+            "seed": 42,
+        },
+        logger={
+            "log_dir": str(tmp_path / "logs"),
+            "wandb_enabled": False,
+            "swanlab_enabled": False,
+            "tensorboard_enabled": False,
+            "mlflow_enabled": False,
+            "monitor_gpus": False,
+        },
+        cluster={"num_nodes": 1, "gpus_per_node": 1},
+        checkpointing={
+            "enabled": enabled,
+            "checkpoint_dir": str(tmp_path / "checkpoints"),
+            "metric_name": metric_name,
+            "higher_is_better": True,
+            "keep_top_k": None,
+            "save_period": save_period,
+            "save_optimizer": save_optimizer,
+            "checkpoint_must_save_by": checkpoint_must_save_by,
+        },
+        data_plane={"enabled": True, "impl": "transfer_queue"},
+        async_rl=AsyncRLConfig(
+            sampler=sampler_cfg,
+            min_groups_for_streaming_train=1,
+            max_inflight_prompts=4,
+            max_buffered_rollouts=4,
+        ),
+    )
+
+
+def _make_actor_args(
+    *,
+    trainer: Optional[_FakeTrainer] = None,
+    save_state: Optional[dict[str, Any]] = None,
+    dataloader: Optional[_FakeDataloader] = None,
+    tq_buffer: Optional[_FakeTQBuffer] = None,
+    last_checkpoint_path: Optional[str] = None,
+) -> SingleControllerActorArgs:
+    return SingleControllerActorArgs(
+        gen_handle=object(),
+        trainer_handle=trainer if trainer is not None else _FakeTrainer(),
+        env_handles={},
+        train_cluster=None,  # type: ignore[arg-type]
+        inference_cluster=None,  # type: ignore[arg-type]
+        dp_client=_FakeDPClient(),
+        dataloader=dataloader if dataloader is not None else _FakeDataloader(),
+        weight_synchronizer=_FakeWeightSynchronizer(),  # type: ignore[arg-type]
+        advantage_estimator=None,
+        loss_fn=object(),  # type: ignore[arg-type]
+        rollout_manager=_FakeRolloutManager(),  # type: ignore[arg-type]
+        tq_buffer=tq_buffer if tq_buffer is not None else _FakeTQBuffer(),  # type: ignore[arg-type]
+        partition_id=_PARTITION_ID,
+        save_state=(
+            save_state if save_state is not None else _default_grpo_save_state()
+        ),
+        last_checkpoint_path=last_checkpoint_path,
+    )
+
+
+def _run_train_pump(
+    mc: MasterConfig,
+    actor_args: SingleControllerActorArgs,
+    *,
+    flush: bool = True,
+):
+    """Construct the actor in-process and drive _train_pump to completion.
+
+    flush=True joins the (possibly async) checkpoint finalization afterwards,
+    like run()'s exit path does, so step_N dirs are visible to assertions.
+    """
+
+    async def _main():
+        actor = _ACTOR_CLS(mc, actor_args)
+        actor._sampler = _FakeSampler(
+            supports_buffer_checkpoint=(mc.async_rl.sampler.name == "windowed")
+        )
+        # In-process runs have no Ray runtime; the pump only reads the GPU
+        # count for a throughput metric.
+        with patch("ray.cluster_resources", return_value={"GPU": 0}):
+            await actor._train_pump()
+        if flush:
+            actor._checkpointer.shutdown()
+        return actor
+
+    return asyncio.run(_main())
+
+
+def _run_actor_run(mc: MasterConfig, actor_args: SingleControllerActorArgs):
+    """Construct the actor in-process and drive run() to completion.
+
+    max_num_steps=0 makes _train_pump exit immediately, so run() executes
+    only the restore block + pump startup/teardown. The wait_for bounds the
+    would-be-deadlock cases (an over-capacity permit acquisition would hang
+    run() forever).
+    """
+
+    async def _main():
+        actor = _ACTOR_CLS(mc, actor_args)
+        result = await asyncio.wait_for(actor.run(), timeout=60.0)
+        return actor, result
+
+    return asyncio.run(_main())
+
+
+def _step_dir_names(ckpt_dir: Path) -> set[str]:
+    if not ckpt_dir.exists():
+        return set()
+    return {p.name for p in ckpt_dir.iterdir() if p.name != "latest_checkpoint_status.json"}
+
+
+def _training_info(ckpt_dir: Path, step: int) -> dict[str, Any]:
+    with open(ckpt_dir / f"step_{step}" / "training_info.json") as f:
+        return json.load(f)
+
+
+# ── counter restore ──────────────────────────────────────────────────────────
+
+
+class TestCounterRestore:
+    def test_restore_from_step_n(self, tmp_path):
+        save_state = _default_grpo_save_state()
+        save_state["current_step"] = 7
+        save_state["current_epoch"] = 2
+        save_state["consumed_samples"] = 42
+        save_state["total_valid_tokens"] = 1234
+
+        actor = _ACTOR_CLS(
+            _actor_master_config(tmp_path), _make_actor_args(save_state=save_state)
+        )
+
+        assert actor._train_steps == 7
+        assert actor._trainer_version == 7
+        # The sampler dispatch cursor is seeded to preserve the fresh-start
+        # invariant _dispatch_index == trainer_version - 1.
+        assert actor._sampler._dispatch_index == 6
+        assert actor._consumed_samples == 42
+        assert actor._current_epoch == 2
+        assert actor._total_valid_tokens == 1234
+
+    def test_fresh_start_defaults(self, tmp_path):
+        actor = _ACTOR_CLS(_actor_master_config(tmp_path), _make_actor_args())
+
+        assert actor._train_steps == 0
+        assert actor._trainer_version == 0
+        assert actor._sampler._dispatch_index == -1
+        assert actor._consumed_samples == 0
+        assert actor._current_epoch == 0
+        assert actor._total_valid_tokens == 0
+
+    def test_old_checkpoint_without_total_valid_tokens(self, tmp_path):
+        # Older checkpoints may predate the total_valid_tokens key.
+        save_state = {
+            "consumed_samples": 10,
+            "current_step": 5,
+            "current_epoch": 0,
+            "total_steps": 5,
+        }
+
+        actor = _ACTOR_CLS(
+            _actor_master_config(tmp_path), _make_actor_args(save_state=save_state)
+        )
+
+        assert actor._train_steps == 5
+        assert actor._sampler._dispatch_index == 4
+        assert actor._total_valid_tokens == 0
+
+
+# ── save trigger + write path ────────────────────────────────────────────────
+
+
+class TestSaveTrigger:
+    def test_saves_on_period_boundary_and_last_step(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=4, save_period=2)
+        trainer = _FakeTrainer()
+
+        actor = _run_train_pump(mc, _make_actor_args(trainer=trainer))
+
+        assert actor._train_steps == 4
+        ckpt_dir = tmp_path / "checkpoints"
+        # Finalized exactly at steps 2 and 4; no tmp_step_* leftovers.
+        assert _step_dir_names(ckpt_dir) == {"step_2", "step_4"}
+
+        info_2 = _training_info(ckpt_dir, 2)
+        assert info_2["current_step"] == 2
+        assert info_2["total_steps"] == 2
+        assert info_2["consumed_samples"] == 4  # 2 prompts/step * 2 steps
+        # No validation ran, so the default val_reward is dropped.
+        assert "val_reward" not in info_2
+
+        info_4 = _training_info(ckpt_dir, 4)
+        assert info_4["current_step"] == 4
+        assert info_4["consumed_samples"] == 8
+
+        # config.yaml is dumped next to training_info.json.
+        assert (ckpt_dir / "step_2" / "config.yaml").exists()
+
+        # save_checkpoint was called into the tmp dir with all three paths.
+        assert len(trainer.save_calls) == 2
+        first = trainer.save_calls[0]
+        assert first["weights_path"] == str(
+            ckpt_dir / "tmp_step_2" / "policy" / "weights"
+        )
+        assert first["optimizer_path"] == str(
+            ckpt_dir / "tmp_step_2" / "policy" / "optimizer"
+        )
+        assert first["tokenizer_path"] == str(
+            ckpt_dir / "tmp_step_2" / "policy" / "tokenizer"
+        )
+        assert first["checkpointing_cfg"] is mc.checkpointing
+        assert trainer.save_calls[1]["weights_path"] == str(
+            ckpt_dir / "tmp_step_4" / "policy" / "weights"
+        )
+
+        # The async writers were waited on before each rename.
+        assert trainer.finalize_calls == 2
+
+        # The tmp dirs were finalized: policy/* survive under step_*.
+        assert (ckpt_dir / "step_2" / "policy" / "weights").is_dir()
+        assert (ckpt_dir / "step_2" / "policy" / "optimizer").is_dir()
+        assert (ckpt_dir / "step_4" / "policy" / "tokenizer").is_dir()
+
+    def test_last_step_saves_off_period_boundary(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=3, save_period=2)
+
+        _run_train_pump(mc, _make_actor_args())
+
+        # step 2 (boundary) + step 3 (last step), no step_1.
+        assert _step_dir_names(tmp_path / "checkpoints") == {"step_2", "step_3"}
+
+    def test_save_optimizer_false_gates_optimizer_path(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, save_optimizer=False
+        )
+        trainer = _FakeTrainer()
+
+        _run_train_pump(mc, _make_actor_args(trainer=trainer))
+
+        assert len(trainer.save_calls) == 1
+        assert trainer.save_calls[0]["optimizer_path"] is None
+        ckpt_dir = tmp_path / "checkpoints"
+        assert (ckpt_dir / "step_2" / "policy" / "weights").is_dir()
+        assert not (ckpt_dir / "step_2" / "policy" / "optimizer").exists()
+
+    def test_no_save_when_disabled(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=1, enabled=False)
+        trainer = _FakeTrainer()
+
+        actor = _run_train_pump(mc, _make_actor_args(trainer=trainer))
+
+        assert actor._train_steps == 2
+        assert trainer.save_calls == []
+        assert _step_dir_names(tmp_path / "checkpoints") == set()
+
+    def test_timeout_saves_and_stops_training_early(self, tmp_path):
+        # 0-second budget: the first check_save() fires; the pump must save
+        # at step 1 (off the period boundary) and break out of the loop.
+        mc = _actor_master_config(
+            tmp_path,
+            max_num_steps=4,
+            save_period=100,
+            checkpoint_must_save_by="00:00:00:00",
+        )
+        trainer = _FakeTrainer()
+
+        actor = _run_train_pump(mc, _make_actor_args(trainer=trainer))
+
+        assert actor._train_steps == 1
+        assert len(trainer.save_calls) == 1
+        assert _step_dir_names(tmp_path / "checkpoints") == {"step_1"}
+
+
+# ── async-save finalization ──────────────────────────────────────────────────
+
+
+class TestAsyncSaveFinalization:
+    def test_rename_deferred_until_async_writes_finish(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2)
+        trainer = _GatedFinalizeTrainer()
+
+        actor = _run_train_pump(mc, _make_actor_args(trainer=trainer), flush=False)
+
+        # The async writer hasn't finished: the checkpoint must still be a
+        # tmp dir, invisible to resume lookups.
+        ckpt_dir = tmp_path / "checkpoints"
+        assert (ckpt_dir / "tmp_step_2").is_dir()
+        assert not (ckpt_dir / "step_2").exists()
+
+        trainer.release.set()
+        actor._checkpointer.shutdown()
+
+        assert (ckpt_dir / "step_2").is_dir()
+        assert not (ckpt_dir / "tmp_step_2").exists()
+        assert trainer.finalize_calls == 1
+
+    def test_failed_background_finalization_raises_at_next_save(self, tmp_path):
+        # The step-2 checkpoint's background finalization fails; the failure
+        # must surface at the step-4 save's finalize_pending, not vanish.
+        mc = _actor_master_config(tmp_path, max_num_steps=4, save_period=2)
+        trainer = _FailingFinalizeTrainer()
+
+        with pytest.raises(RuntimeError, match="finalization failed"):
+            _run_train_pump(mc, _make_actor_args(trainer=trainer), flush=False)
+
+    def test_save_checkpoint_records_val_metrics(self, tmp_path):
+        # Direct _save_checkpoint call: the val_metrics parameter feeds both
+        # val_reward and a val:* metric_name (wired by the future validation
+        # loop; the pump passes None today).
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, metric_name="val:accuracy"
+        )
+
+        async def _main():
+            actor = _ACTOR_CLS(mc, _make_actor_args())
+            actor._train_steps = 2
+            await actor._save_checkpoint({"loss": 1.0}, val_metrics={"accuracy": 0.75})
+            actor._checkpointer.shutdown()
+            return actor
+
+        actor = asyncio.run(_main())
+
+        info = _training_info(tmp_path / "checkpoints", 2)
+        assert info["val_reward"] == 0.75
+        assert info["val:accuracy"] == 0.75
+        assert actor._save_state["val_reward"] == 0.75
+
+
+# ── metric_name behavior ─────────────────────────────────────────────────────
+
+
+class TestMetricName:
+    def test_val_metric_without_validation_warns_and_still_saves(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, metric_name="val:accuracy"
+        )
+
+        with pytest.warns(UserWarning, match="no val metrics were collected"):
+            _run_train_pump(mc, _make_actor_args())
+
+        ckpt_dir = tmp_path / "checkpoints"
+        assert _step_dir_names(ckpt_dir) == {"step_2"}
+        assert "val:accuracy" not in _training_info(ckpt_dir, 2)
+
+    def test_train_metric_lands_in_training_info(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, metric_name="train:loss"
+        )
+        trainer = _FakeTrainer(step_metrics={"loss": 0.5})
+
+        _run_train_pump(mc, _make_actor_args(trainer=trainer))
+
+        info = _training_info(tmp_path / "checkpoints", 2)
+        assert info["train:loss"] == 0.5
+
+    def test_train_metric_missing_key_raises(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, metric_name="train:not_a_metric"
+        )
+
+        with pytest.raises(ValueError, match="not found in train metrics"):
+            _run_train_pump(mc, _make_actor_args())
+
+    def test_metric_name_requires_train_or_val_prefix(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, metric_name="reward"
+        )
+
+        with pytest.raises(AssertionError, match="must start with"):
+            _run_train_pump(mc, _make_actor_args())
+
+
+# ── setup resume-path wiring ─────────────────────────────────────────────────
+
+
+_STEP_3_SAVE_STATE = {
+    "consumed_samples": 24,
+    "current_step": 3,
+    "current_epoch": 1,
+    "total_steps": 3,
+    "total_valid_tokens": 999,
+}
+
+
+def _write_checkpoint(
+    ckpt_dir: Path,
+    step: int,
+    save_state: dict[str, Any],
+    *,
+    with_optimizer: bool = True,
+    dataloader_state: Optional[dict[str, Any]] = None,
+    config: Optional[dict[str, Any]] = None,
+) -> Path:
+    step_dir = ckpt_dir / f"step_{step}"
+    (step_dir / "policy" / "weights").mkdir(parents=True)
+    if with_optimizer:
+        (step_dir / "policy" / "optimizer").mkdir(parents=True)
+    with open(step_dir / "training_info.json", "w") as f:
+        json.dump(save_state, f)
+    if dataloader_state is not None:
+        torch.save(dataloader_state, step_dir / "train_dataloader.pt")
+    if config is not None:
+        with open(step_dir / "config.yaml", "w") as f:
+            yaml.safe_dump(config, f)
+    return step_dir
+
+
+def _setup_master_config(checkpoint_dir: str) -> MasterConfig:
+    """Partially-populated MasterConfig for setup_single_controller tests.
+
+    Same shape as test_single_controller_setup._make_master_config, plus the
+    checkpointing block setup now reads.
+    """
+    return MasterConfig.model_construct(
+        data_plane={"enabled": True, "impl": "transfer_queue"},
+        data={
+            "use_multiple_dataloader": False,
+            "shuffle": False,
+            "num_workers": 0,
+            "train": [{"env_name": "math"}],
+        },
+        grpo={
+            "max_num_steps": 100,
+            "max_num_epochs": 1,
+            "num_prompts_per_step": 4,
+            "num_generations_per_prompt": 2,
+            "max_rollout_turns": 1,
+            "seed": 42,
+            "val_period": 0,
+            "val_at_start": False,
+            "val_at_end": False,
+        },
+        policy={
+            "train_global_batch_size": 8,
+            "max_total_sequence_length": 32,
+            "tokenizer": {"use_fastokens": False},
+            "megatron_cfg": {"enabled": False},
+            "generation": {
+                "backend": "vllm",
+                "colocated": {"enabled": True, "resources": {}},
+            },
+        },
+        loss_fn=ClippedPGLossConfig(),
+        env={},
+        async_rl=AsyncRLConfig(
+            min_groups_for_streaming_train=4,
+            max_buffered_rollouts=8,
+        ),
+        checkpointing={
+            "enabled": True,
+            "checkpoint_dir": checkpoint_dir,
+            "metric_name": None,
+            "higher_is_better": True,
+            "keep_top_k": None,
+            "save_period": 2,
+            "save_optimizer": True,
+            "checkpoint_must_save_by": None,
+        },
+    )
+
+
+class TestGetResumePaths:
+    def test_resume_paths_from_fixture_layout(self, tmp_path):
+        step_dir = _write_checkpoint(tmp_path, 3, _STEP_3_SAVE_STATE)
+
+        weights_path, optimizer_path = CheckpointManager.get_resume_paths(str(step_dir))
+
+        assert weights_path == step_dir / "policy" / "weights"
+        assert optimizer_path == step_dir / "policy" / "optimizer"
+
+    def test_resume_paths_without_optimizer_state(self, tmp_path):
+        step_dir = _write_checkpoint(
+            tmp_path, 3, _STEP_3_SAVE_STATE, with_optimizer=False
+        )
+
+        with pytest.warns(UserWarning, match="Optimizer state not found"):
+            weights_path, optimizer_path = CheckpointManager.get_resume_paths(
+                str(step_dir)
+            )
+
+        assert weights_path == step_dir / "policy" / "weights"
+        assert optimizer_path is None
+
+    def test_no_checkpoint_gives_none(self):
+        assert CheckpointManager.get_resume_paths(None) == (None, None)
+
+
+class TestSetupResumeWiring:
+    def test_setup_forwards_latest_resume_paths(
+        self,
+        patched_factories,  # noqa: F811
+        tmp_path,
+    ):
+        ckpt_dir = tmp_path / "ckpts"
+        _write_checkpoint(ckpt_dir, 1, {**_STEP_3_SAVE_STATE, "current_step": 1})
+        step_3 = _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE)
+        mc = _setup_master_config(str(ckpt_dir))
+
+        actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        # Latest checkpoint (step_3) wins; its paths reach the trainer factory.
+        trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs
+        assert trainer_kwargs["weights_path"] == step_3 / "policy" / "weights"
+        assert trainer_kwargs["optimizer_path"] == step_3 / "policy" / "optimizer"
+        # training_info.json is loaded into the actor args for the actor.
+        assert actor_args.save_state == _STEP_3_SAVE_STATE
+        assert actor_args.last_checkpoint_path == str(step_3)
+
+    def test_setup_fresh_start_passes_none_paths(
+        self,
+        patched_factories,  # noqa: F811
+        tmp_path,
+    ):
+        ckpt_dir = tmp_path / "empty_ckpts"
+        ckpt_dir.mkdir()
+        mc = _setup_master_config(str(ckpt_dir))
+
+        actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs
+        assert trainer_kwargs["weights_path"] is None
+        assert trainer_kwargs["optimizer_path"] is None
+        assert actor_args.save_state == _default_grpo_save_state()
+        assert actor_args.last_checkpoint_path is None
+
+    def test_setup_forwards_pretrained_checkpoint(
+        self,
+        patched_factories,  # noqa: F811
+        tmp_path,
+    ):
+        mc = _setup_master_config(str(tmp_path / "ckpts"))
+        pretrained = {"path": "/some/ckpt", "format": "megatron_bridge"}
+        mc.checkpointing["pretrained_checkpoint"] = pretrained
+
+        setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        assert mc.policy["pretrained_checkpoint"] == pretrained
+
+
+def _make_int_dataloader() -> StatefulDataLoader:
+    """8 ints, batch_size=2 → batches [0,1], [2,3], [4,5], [6,7]."""
+    return StatefulDataLoader(
+        list(range(8)),
+        batch_size=2,
+        shuffle=False,
+        drop_last=True,
+        num_workers=0,
+    )
+
+
+class TestDataloaderState:
+    def test_save_writes_dataloader_state(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2)
+        save_state = _default_grpo_save_state()
+        save_state["current_epoch"] = 3
+        dataloader = _FakeDataloader(state={"fake_position": 7})
+
+        _run_train_pump(
+            mc, _make_actor_args(save_state=save_state, dataloader=dataloader)
+        )
+
+        ckpt_dir = tmp_path / "checkpoints"
+        dl_state_path = ckpt_dir / "step_2" / "train_dataloader.pt"
+        assert dl_state_path.exists()
+        # The snapshot taken at save time round-trips through torch.save.
+        assert torch.load(dl_state_path) == {"fake_position": 7}
+        # current_epoch flows save_state → actor → training_info.json.
+        assert _training_info(ckpt_dir, 2)["current_epoch"] == 3
+
+    def test_stateful_dataloader_position_roundtrip(self, tmp_path):
+        data_config = {"train": [{"dataset_name": "math_train"}]}
+        dataloader = _make_int_dataloader()
+        it = iter(dataloader)
+        assert [next(it).tolist() for _ in range(2)] == [[0, 1], [2, 3]]
+
+        step_dir = _write_checkpoint(
+            tmp_path,
+            5,
+            _default_grpo_save_state(),
+            dataloader_state=dataloader.state_dict(),
+            config={"data": {"train": [{"dataset_name": "math_train"}]}},
+        )
+
+        restored = _make_int_dataloader()
+        load_dataloader_state(restored, str(step_dir), data_config)
+
+        # Resumes at batch k+1, not from the top.
+        assert next(iter(restored)).tolist() == [4, 5]
+
+    def test_dataset_swap_skips_restore(self, tmp_path):
+        dataloader = _make_int_dataloader()
+        it = iter(dataloader)
+        assert [next(it).tolist() for _ in range(2)] == [[0, 1], [2, 3]]
+
+        step_dir = _write_checkpoint(
+            tmp_path,
+            5,
+            _default_grpo_save_state(),
+            dataloader_state=dataloader.state_dict(),
+            config={"data": {"train": [{"dataset_name": "old_dataset"}]}},
+        )
+
+        restored = _make_int_dataloader()
+        load_dataloader_state(
+            restored, str(step_dir), {"train": [{"dataset_name": "new_dataset"}]}
+        )
+
+        # Restore skipped on dataset swap: the new dataset starts from index 0.
+        assert next(iter(restored)).tolist() == [0, 1]
+
+    def test_setup_restores_dataloader_state(
+        self,
+        patched_factories,  # noqa: F811
+        tmp_path,
+    ):
+        ckpt_dir = tmp_path / "ckpts"
+        sentinel = {"fake_position": 123}
+        _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE, dataloader_state=sentinel)
+        mc = _setup_master_config(str(ckpt_dir))
+
+        setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        fake_dataloader = patched_factories["dataloader"]
+        fake_dataloader.load_state_dict.assert_called_once()
+        assert fake_dataloader.load_state_dict.call_args.args[0] == sentinel
+
+    def test_setup_missing_dataloader_state_starts_fresh(
+        self,
+        patched_factories,  # noqa: F811
+        tmp_path,
+        monkeypatch,
+    ):
+        # Checkpoint with training_info.json + policy/ but no
+        # train_dataloader.pt. Setup must warn and keep the fresh position
+        # while still forwarding the weight resume paths.
+        ckpt_dir = tmp_path / "ckpts"
+        step_3 = _write_checkpoint(ckpt_dir, 3, _STEP_3_SAVE_STATE)
+        mc = _setup_master_config(str(ckpt_dir))
+        printed: list[str] = []
+        # Repo addopts run pytest with -s, so capsys sees nothing; record
+        # print calls instead.
+        monkeypatch.setattr(
+            "builtins.print",
+            lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)),
+        )
+
+        setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        patched_factories["dataloader"].load_state_dict.assert_not_called()
+        assert any("No dataloader state found" in line for line in printed)
+        trainer_kwargs = patched_factories["_build_trainer"].call_args.kwargs
+        assert trainer_kwargs["weights_path"] == step_3 / "policy" / "weights"
+
+
+# ── replay buffer persistence ────────────────────────────────────────────────
+
+
+class TestReplayBufferPersistence:
+    def test_save_writes_replay_buffer_when_sampler_supports_it(self, tmp_path):
+        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2)
+        envelope = {"groups": [], "sentinel": "abc"}
+        buffer = _FakeTQBuffer(state=envelope)
+
+        _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.
+        assert buffer.state_dict_calls == [4]
+
+    def test_no_replay_buffer_with_gated_sampler(self, tmp_path):
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=2, buffer_checkpoint=False
+        )
+        buffer = _FakeTQBuffer()
+
+        _run_train_pump(mc, _make_actor_args(tq_buffer=buffer))
+
+        ckpt_dir = tmp_path / "checkpoints"
+        assert (ckpt_dir / "step_2" / "training_info.json").exists()
+        assert not (ckpt_dir / "step_2" / "replay_buffer.pt").exists()
+        assert buffer.state_dict_calls == []
+
+    def test_run_restores_replay_buffer_and_permits(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)
+
+        actor, result = _run_actor_run(
+            mc,
+            _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)),
+        )
+
+        assert buffer.load_calls == [
+            {
+                "state": envelope,
+                "max_groups": 4,
+                "expected_partition_id": _PARTITION_ID,
+                "expected_group_size": 2,
+            }
+        ]
+        # Each restored group holds one _buffer_capacity permit.
+        assert actor._buffer_capacity._value == 4 - 3
+        assert result["train_steps"] == 0
+
+    def test_run_restore_at_full_capacity_does_not_hang(self, tmp_path):
+        # K == max_buffered_rollouts: the acquisitions must all complete
+        # without waiting (no pump is running yet to release permits).
+        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)
+        buffer = _FakeTQBuffer(load_return=4)
+
+        actor, result = _run_actor_run(
+            mc,
+            _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)),
+        )
+
+        assert len(buffer.load_calls) == 1
+        assert actor._buffer_capacity._value == 0
+        assert result["train_steps"] == 0
+
+    def test_run_missing_replay_buffer_file_starts_empty(self, tmp_path, monkeypatch):
+        # Resuming from a checkpoint written by a gated-sampler run: no
+        # replay_buffer.pt.
+        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)),
+        )
+
+        actor, _ = _run_actor_run(
+            mc,
+            _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)),
+        )
+
+        assert buffer.load_calls == []
+        assert actor._buffer_capacity._value == 4  # zero permits consumed
+        assert any("No replay buffer checkpoint found" in line for line in printed)
+
+    def test_run_no_restore_with_gated_sampler(self, tmp_path):
+        # File present but the sampler doesn't support buffer checkpointing:
+        # nothing is restored.
+        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, buffer_checkpoint=False)
+        buffer = _FakeTQBuffer(load_return=2)
+
+        actor, _ = _run_actor_run(
+            mc,
+            _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)),
+        )
+
+        assert buffer.load_calls == []
+        assert actor._buffer_capacity._value == 4

From 1b7cb0a5e5d565c5b4a26aef1dbca803d8a13c59 Mon Sep 17 00:00:00 2001
From: haitian-nvidia <haitianj@nvidia.com>
Date: Thu, 30 Jul 2026 10:45:07 -0700
Subject: [PATCH 06/15] fix(sc): keep training failures visible through
 checkpoint-flush errors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

On exit, run() now propagates a failed checkpoint finalization only on
the clean path; when an exception is already propagating the flush is
best-effort (warning), so the original training failure stays the
raised exception — matching async_grpo_train's guarded cleanup
shutdown. logger.finish() moves into its own finally so it runs either
way.

Also drop the stale "SC does not support checkpointing yet." comment
from the SC exemplar config.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: haitian-nvidia <haitianj@nvidia.com>
---
 ...po_math_1B_megatron_single_controller.yaml |  1 -
 nemo_rl/algorithms/single_controller.py       | 19 ++++++++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml
index fe4a6b3f31..14be5604e1 100644
--- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml
+++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml
@@ -21,7 +21,6 @@ async_rl:
   # Enable per-rollout diagnostic prints (prompt content / completion previews).
   diagnostics: false
 
-# SC does not support checkpointing yet.
 checkpointing:
   enabled: false
   checkpoint_dir: results/grpo-single-controller
diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index 899683b9a2..7938b9586b 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -36,6 +36,7 @@
 
 import asyncio
 import os
+import sys
 import time
 import warnings
 from functools import partial
@@ -233,9 +234,21 @@ async def run(self) -> dict[str, Any]:
             rollout_task.cancel()
             train_task.cancel()
             await asyncio.gather(rollout_task, train_task, return_exceptions=True)
-            # Flush the last checkpoint's background finalization before exit.
-            await asyncio.to_thread(self._checkpointer.shutdown)
-            self._logger.finish()
+            # Flush the last checkpoint's background finalization; a failure
+            # raises on a clean exit but never masks a propagating exception.
+            propagating = sys.exc_info()[0] is not None
+            try:
+                await asyncio.to_thread(self._checkpointer.shutdown)
+            except Exception:
+                if not propagating:
+                    raise
+                warnings.warn(
+                    "Checkpoint finalization failed while handling an "
+                    "exception; the original exception will be re-raised.",
+                    stacklevel=2,
+                )
+            finally:
+                self._logger.finish()
 
         return {
             "train_steps": self._train_steps,

From 0495de9030d3113aad5731f5a6e761e7ddf96ca3 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Wed, 29 Jul 2026 18:40:23 -0400
Subject: [PATCH 07/15] build(data-plane): bump TransferQueue to v0.1.9

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 nemo_rl/data_plane/adapters/transfer_queue.py |  7 ++++++-
 pyproject.toml                                | 15 ++++++---------
 uv.lock                                       |  7 ++++---
 3 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py
index ce293b4e26..6bb89e6c58 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -173,7 +173,12 @@ def patched(*args, **kwargs):
 
     patched_any = False
     try:
-        from transfer_queue.storage.simple_backend import SimpleStorageUnit
+        try:
+            # TQ v0.1.9+
+            from transfer_queue.storage.simple_storage import SimpleStorageUnit
+        except ImportError:
+            # Compatibility with older TQ revisions.
+            from transfer_queue.storage.simple_backend import SimpleStorageUnit
 
         patched_any |= _install(SimpleStorageUnit)
     except ImportError:
diff --git a/pyproject.toml b/pyproject.toml
index 5fd74de631..1136e927a8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,11 +84,10 @@ dependencies = [
   # automatically include them. Removes the need for a `[data-plane]`
   # extra and the corresponding plumbing in the per-worker venv builder.
   "tensordict",
-  # Pinned to b266d39 (post-0.1.6, pre-0.1.7) for PR #77's MooncakeStore
-  # refactor: `clear` switched from unanchored `remove_by_regex` to
-  # exact-key `batch_remove`, which fixes a collateral-key-deletion bug
-  # that breaks DAPO + mooncake_cpu. Bump to the 0.1.7 tag when released.
-  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@b266d39",
+  # TransferQueue v0.1.9 adds full-system checkpoint save/load APIs and
+  # retains the exact-key MooncakeStore clear behavior required by DAPO.
+  # Pin the immutable release commit rather than the mutable version tag.
+  "TransferQueue @ git+https://github.com/Ascend/TransferQueue.git@c51614308b68c8d7a87c9b3ef62d59e14c69bde2",
   # Backs data_plane.backend="mooncake_cpu". Default backend is "simple"
   # (in-process), but the mooncake_cpu path needs the `mooncake_master`
   # binary that ships in this wheel at <site-packages>/mooncake/. Bundled
@@ -379,10 +378,8 @@ override-dependencies = [
   "pytest>=9.0.3",
   "langchain>=0.3.28",      # Address CVE-2025-65106
   "langchain-core>=0.3.80", # Address CVE-2025-65106
-  # TransferQueue (data-plane extra) pins numpy<2.0.0; megatron-core needs
-  # numpy>=2.1.0 via onnx → ml-dtypes. Override globally so the data-plane
-  # extra composes with mcore/automodel without version-mirroring TQ's
-  # requirements.txt. Forward-compatible across TQ minor bumps.
+  # Keep the NumPy floor compatible with megatron-core's onnx → ml-dtypes
+  # dependency while composing the data-plane stack with mcore/automodel.
   "numpy>=2.1.0",
   # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame).
   # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it
diff --git a/uv.lock b/uv.lock
index 9654360842..146f7104e6 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4421,7 +4421,7 @@ requires-dist = [
     { name = "torchdata" },
     { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.26.0", index = "https://download.pytorch.org/whl/cu130" },
     { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.26.0", index = "https://pypi.org/simple" },
-    { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39" },
+    { name = "transferqueue", git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2" },
     { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" },
     { name = "transformers", specifier = ">=5.5.0,<5.9.0" },
     { name = "transformers", marker = "extra == 'automodel'", specifier = ">=5.5.0,<5.6.0" },
@@ -7874,14 +7874,15 @@ wheels = [
 
 [[package]]
 name = "transferqueue"
-version = "0.1.7.dev0"
-source = { git = "https://github.com/Ascend/TransferQueue.git?rev=b266d39#b266d39a15aae114730de36cf8317b6285436f7f" }
+version = "0.1.9"
+source = { git = "https://github.com/Ascend/TransferQueue.git?rev=c51614308b68c8d7a87c9b3ef62d59e14c69bde2#c51614308b68c8d7a87c9b3ef62d59e14c69bde2" }
 dependencies = [
     { name = "hydra-core", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-trtllm' and extra == 'extra-7-nemo-rl-vllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-automodel' and extra != 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "hydra-core", version = "1.3.4", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-7-nemo-rl-automodel' or extra != 'extra-7-nemo-rl-mcore' or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-trtllm') or (extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm')" },
     { name = "msgspec" },
     { name = "numpy" },
     { name = "omegaconf" },
+    { name = "prometheus-client" },
     { name = "psutil" },
     { name = "pyzmq" },
     { name = "ray", extra = ["default"] },

From 46219100944c8fd576111f4c9dad9986edb59c04 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Wed, 29 Jul 2026 19:10:54 -0400
Subject: [PATCH 08/15] fix(data-plane): normalize uniform TQ nested reads

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 nemo_rl/data_plane/adapters/transfer_queue.py | 15 +++++++-
 tests/unit/data_plane/test_codec_mooncake.py  | 38 +++++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py
index 6bb89e6c58..9d1af6c363 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -361,13 +361,26 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict:
 
 
 def _from_wire(td: TensorDict) -> TensorDict:
-    """Inverse of `_promote_1d_leaves`: squeeze trailing 1 back to (N,)."""
+    """Normalize Mooncake reads and invert :func:`_promote_1d_leaves`.
+
+    TQ v0.1.9 reconstructs every non-scalar field as a nested tensor before
+    attempting a dense representation, including fields whose rows all have
+    the same shape. Densify those uniform nested tensors first so regular
+    batched inputs retain their dense representation. Truly ragged fields
+    remain nested. Finally, squeeze the singleton dimension introduced by
+    :func:`_promote_1d_leaves`.
+    """
     # Same top-level iteration as `_promote_1d_leaves`: NonTensorData /
     # NonTensorStack leaves are only visible via td.keys(), not leaves_only.
     new_dict: dict[str, Any] = {}
     changed = False
     for k in td.keys():
         v = td.get(k)
+        if isinstance(v, torch.Tensor) and v.is_nested:
+            rows = list(v.unbind())
+            if rows and all(row.shape == rows[0].shape for row in rows[1:]):
+                v = torch.stack(rows)
+                changed = True
         if (
             isinstance(v, torch.Tensor)
             and not v.is_nested
diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py
index c9b71820d7..584a51ac81 100644
--- a/tests/unit/data_plane/test_codec_mooncake.py
+++ b/tests/unit/data_plane/test_codec_mooncake.py
@@ -73,6 +73,44 @@ def test_promote_1d_roundtrip_via_from_wire() -> None:
     assert torch.equal(back["reward"], original)
 
 
+def test_from_wire_densifies_uniform_nested_rows() -> None:
+    """TQ v0.1.9's uniform nested reads are restored to dense tensors."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import _from_wire
+
+    rows = [torch.tensor([i], dtype=torch.float32) for i in range(4)]
+    wire = TensorDict(
+        {"reward": torch.nested.as_nested_tensor(rows, layout=torch.jagged)},
+        batch_size=[len(rows)],
+    )
+
+    back = _from_wire(wire)
+
+    assert not back["reward"].is_nested
+    assert back["reward"].shape == (len(rows),)
+    assert torch.equal(back["reward"], torch.arange(len(rows), dtype=torch.float32))
+
+
+def test_from_wire_preserves_ragged_nested_rows() -> None:
+    """Variable-length rollout fields must remain nested."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import _from_wire
+
+    rows = [torch.arange(i + 1) for i in range(3)]
+    nested = torch.nested.as_nested_tensor(rows, layout=torch.jagged)
+    wire = TensorDict({"token_ids": nested}, batch_size=[len(rows)])
+
+    back = _from_wire(wire)
+
+    assert back["token_ids"].is_nested
+    assert all(
+        torch.equal(actual, expected)
+        for actual, expected in zip(back["token_ids"].unbind(), rows, strict=True)
+    )
+
+
 # ── P2: pack_per_token_field — tolerates SP padding ──────────────────────────
 
 

From 53c95b992f8a5d5ad5865d0eff3d818f009f1871 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Fri, 31 Jul 2026 12:15:50 -0400
Subject: [PATCH 09/15] fix(data-plane): preserve TQ wire shapes

Track Mooncake 1D promotion in durable TQ metadata so genuine singleton token columns retain their rank. Remove the obsolete pre-0.1.9 storage import fallback and add focused regression coverage.

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 nemo_rl/data_plane/adapters/transfer_queue.py | 136 +++++++++++++++---
 tests/unit/data_plane/test_codec_mooncake.py  | 120 ++++++++++++++++
 2 files changed, 234 insertions(+), 22 deletions(-)

diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py
index 9d1af6c363..5abc1b1a49 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -173,12 +173,7 @@ def patched(*args, **kwargs):
 
     patched_any = False
     try:
-        try:
-            # TQ v0.1.9+
-            from transfer_queue.storage.simple_storage import SimpleStorageUnit
-        except ImportError:
-            # Compatibility with older TQ revisions.
-            from transfer_queue.storage.simple_backend import SimpleStorageUnit
+        from transfer_queue.storage.simple_storage import SimpleStorageUnit
 
         patched_any |= _install(SimpleStorageUnit)
     except ImportError:
@@ -311,6 +306,77 @@ def _init_tq(cfg: DataPlaneConfig) -> None:
 # ──────────────────────────────────────────────────────────────────────────
 
 
+_PROMOTED_1D_TAG_PREFIX = "__nemo_rl_promoted_1d__:"
+_WIRE_SHAPE_VERSION_TAG = "__nemo_rl_wire_shape_version__"
+_WIRE_SHAPE_VERSION = 1
+
+
+def _add_wire_shape_tags(
+    tags: list[dict[str, Any]], td: TensorDict
+) -> list[dict[str, Any]]:
+    """Record which tensor fields the Mooncake workaround will promote.
+
+    TQ metadata is durable and visible to readers in other processes, unlike
+    client-local state. The version marker distinguishes new unpromoted fields
+    from legacy rows; promoted fields get an additional boolean marker.
+    """
+    field_markers: dict[str, int | bool] = {
+        _WIRE_SHAPE_VERSION_TAG: _WIRE_SHAPE_VERSION
+    }
+    for k in td.keys():
+        value = td.get(k)
+        if isinstance(value, torch.Tensor) and not value.is_nested and value.dim() == 1:
+            field_markers[f"{_PROMOTED_1D_TAG_PREFIX}{k}"] = True
+    return [{**tag, **field_markers} for tag in tags]
+
+
+def _promoted_1d_fields_from_tags(
+    tags: list[dict[str, Any]], field_names: list[str]
+) -> set[str] | None:
+    """Recover promoted fields, or ``None`` for legacy unmarked rows.
+
+    Every selected tensor field must have the same marker on every row. Mixed
+    metadata means rows were written by incompatible codecs; fail instead of
+    silently returning tensors with the wrong rank.
+    """
+    versions = [tag.get(_WIRE_SHAPE_VERSION_TAG) for tag in tags]
+    if not any(version is not None for version in versions):
+        return None
+    if not versions or any(version != _WIRE_SHAPE_VERSION for version in versions):
+        raise RuntimeError(
+            "Inconsistent Mooncake shape metadata: expected wire-shape "
+            f"version {_WIRE_SHAPE_VERSION} on every sample."
+        )
+
+    promoted: set[str] = set()
+    for field in field_names:
+        marker = f"{_PROMOTED_1D_TAG_PREFIX}{field}"
+        values = [tag.get(marker) for tag in tags]
+        if any(value is not None for value in values) and any(
+            value is not True for value in values
+        ):
+            raise RuntimeError(
+                "Inconsistent Mooncake shape metadata for "
+                f"field {field!r}: expected the promotion marker on every sample."
+            )
+        if values and values[0] is True:
+            promoted.add(field)
+    return promoted
+
+
+def _strip_wire_shape_tags(tags: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Remove adapter-private shape metadata from user-visible tags."""
+    return [
+        {
+            key: value
+            for key, value in tag.items()
+            if key != _WIRE_SHAPE_VERSION_TAG
+            and not key.startswith(_PROMOTED_1D_TAG_PREFIX)
+        }
+        for tag in tags
+    ]
+
+
 def _assert_no_key_loss(src_dict: dict, new_td: TensorDict, fn: str) -> None:
     """Guard against silent leaf drops through TensorDict constructor rebuild.
 
@@ -360,15 +426,18 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict:
     return new_td
 
 
-def _from_wire(td: TensorDict) -> TensorDict:
+def _from_wire(
+    td: TensorDict, promoted_1d_fields: set[str] | None = None
+) -> TensorDict:
     """Normalize Mooncake reads and invert :func:`_promote_1d_leaves`.
 
     TQ v0.1.9 reconstructs every non-scalar field as a nested tensor before
     attempting a dense representation, including fields whose rows all have
     the same shape. Densify those uniform nested tensors first so regular
     batched inputs retain their dense representation. Truly ragged fields
-    remain nested. Finally, squeeze the singleton dimension introduced by
-    :func:`_promote_1d_leaves`.
+    remain nested. Finally, squeeze only singleton dimensions known to have
+    been introduced by :func:`_promote_1d_leaves`. ``None`` retains the old
+    shape-only behavior for checkpoints written before shape tags existed.
     """
     # Same top-level iteration as `_promote_1d_leaves`: NonTensorData /
     # NonTensorStack leaves are only visible via td.keys(), not leaves_only.
@@ -386,6 +455,7 @@ def _from_wire(td: TensorDict) -> TensorDict:
             and not v.is_nested
             and v.dim() >= 2
             and v.shape[-1] == 1
+            and (promoted_1d_fields is None or str(k) in promoted_1d_fields)
         ):
             new_dict[str(k)] = v.squeeze(-1).contiguous()
             changed = True
@@ -550,7 +620,10 @@ def claim_meta(
         # because shard_meta_for_dp reads it directly), but the rest
         # of the tag dict travels through unchanged so consumers can
         # filter on it without fetching data.
-        tags = list(tq_meta.custom_meta) if tq_meta.custom_meta else [{} for _ in keys]
+        wire_tags = (
+            list(tq_meta.custom_meta) if tq_meta.custom_meta else [{} for _ in keys]
+        )
+        tags = _strip_wire_shape_tags(wire_tags)
         seqlens: list[int] | None = None
         if tags and any("input_lengths" in t for t in tags):
             seqlens = [int(t.get("input_lengths", 0)) for t in tags]
@@ -601,8 +674,10 @@ def put_samples(
             return KVBatchMeta(
                 partition_id=partition_id, task_name=None, sample_ids=[], fields=None
             )
-        if tags is None:
-            tags = [{} for _ in sample_ids]
+        user_tags = (
+            [{} for _ in sample_ids] if tags is None else [dict(tag) for tag in tags]
+        )
+        wire_tags = user_tags
 
         wire_fields: TensorDict | None = None
         field_names: list[str] | None = None
@@ -614,6 +689,7 @@ def put_samples(
             # destructive for non-tensors.
             wire_fields = fields.detach()  # type: ignore[bad-assignment,missing-argument]
             if self._promote_1d:
+                wire_tags = _add_wire_shape_tags(wire_tags, wire_fields)
                 wire_fields = _promote_1d_leaves(wire_fields)  # type: ignore[bad-argument-type]
             field_names = list(wire_fields.keys())
 
@@ -622,7 +698,7 @@ def put_samples(
             keys=list(sample_ids),
             partition_id=partition_id,
             fields=wire_fields,
-            tags=tags,
+            tags=wire_tags,
         )
 
         return KVBatchMeta(
@@ -630,7 +706,7 @@ def put_samples(
             task_name=None,
             sample_ids=list(sample_ids),
             fields=field_names,
-            tags=[dict(t) for t in tags] if tags else None,
+            tags=user_tags if user_tags else None,
         )
 
     def get_samples(
@@ -641,15 +717,31 @@ def get_samples(
     ) -> TensorDict:
         if not sample_ids:
             return TensorDict({}, batch_size=(0,))
-        # TQ's wire vocabulary is `keys=` — translation point.
-        td = tq.kv_batch_get(
-            keys=list(sample_ids),
-            partition_id=partition_id,
-            select_fields=select_fields,
+        if not self._promote_1d:
+            return tq.kv_batch_get(
+                keys=list(sample_ids),
+                partition_id=partition_id,
+                select_fields=select_fields,
+            )
+
+        # Inline TQ's public kv_batch_get flow so the adapter can also inspect
+        # the durable per-row shape markers carried by BatchMeta.custom_meta.
+        client = tq.get_client()
+        tq_meta = client.kv_retrieve_meta(
+            keys=list(sample_ids), partition_id=partition_id, create=False
         )
-        if self._promote_1d:
-            td = _from_wire(td)
-        return td
+        if tq_meta.size == 0:
+            raise ValueError("keys or partition were not found!")
+        tq_meta = tq_meta.select_fields(list(select_fields))
+        if not tq_meta.is_ready:
+            raise ValueError("Some fields are not ready in all the requested keys!")
+        td = client.get_data(tq_meta)
+        tensor_fields = [
+            str(k) for k in td.keys() if isinstance(td.get(k), torch.Tensor)
+        ]
+        wire_tags = list(tq_meta.custom_meta) if tq_meta.custom_meta else []
+        promoted_1d_fields = _promoted_1d_fields_from_tags(wire_tags, tensor_fields)
+        return _from_wire(td, promoted_1d_fields)
 
     def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None:
         cleared_via_none = sample_ids is None
diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py
index 584a51ac81..b0a8868377 100644
--- a/tests/unit/data_plane/test_codec_mooncake.py
+++ b/tests/unit/data_plane/test_codec_mooncake.py
@@ -92,6 +92,126 @@ def test_from_wire_densifies_uniform_nested_rows() -> None:
     assert torch.equal(back["reward"], torch.arange(len(rows), dtype=torch.float32))
 
 
+def test_from_wire_preserves_genuine_length_one_token_column() -> None:
+    """Only fields promoted from ``(N,)`` are squeezed after a TQ read."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import _from_wire
+
+    n = 4
+    wire = TensorDict(
+        {
+            "reward": torch.nested.as_nested_tensor(
+                [torch.tensor([float(i)]) for i in range(n)], layout=torch.jagged
+            ),
+            "input_ids": torch.nested.as_nested_tensor(
+                [torch.tensor([i]) for i in range(n)], layout=torch.jagged
+            ),
+        },
+        batch_size=[n],
+    )
+
+    back = _from_wire(wire, promoted_1d_fields={"reward"})
+
+    assert back["reward"].shape == (n,)
+    assert back["input_ids"].shape == (n, 1)
+    assert torch.equal(back["input_ids"], torch.arange(n).unsqueeze(-1))
+
+
+def test_wire_shape_tags_roundtrip_promoted_field_names() -> None:
+    """Durable tags distinguish scalar and length-one tensor columns."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import (
+        _add_wire_shape_tags,
+        _promoted_1d_fields_from_tags,
+        _strip_wire_shape_tags,
+    )
+
+    n = 3
+    fields = TensorDict(
+        {
+            "reward": torch.arange(n, dtype=torch.float32),
+            "input_ids": torch.arange(n).unsqueeze(-1),
+        },
+        batch_size=[n],
+    )
+    user_tags = [{"weight_version": 7} for _ in range(n)]
+
+    wire_tags = _add_wire_shape_tags(user_tags, fields)
+
+    assert _promoted_1d_fields_from_tags(wire_tags, ["reward", "input_ids"]) == {
+        "reward"
+    }
+    assert _strip_wire_shape_tags(wire_tags) == user_tags
+    assert all(
+        not any(key.startswith("__nemo_rl_promoted_1d__:") for key in tag)
+        for tag in user_tags
+    )
+
+
+def test_get_samples_uses_persisted_shape_tags(monkeypatch) -> None:
+    """The Mooncake adapter restores ranks using TQ's durable row metadata."""
+    from tensordict import TensorDict
+
+    import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter
+
+    n = 3
+    original = TensorDict(
+        {
+            "reward": torch.arange(n, dtype=torch.float32),
+            "input_ids": torch.arange(n).unsqueeze(-1),
+        },
+        batch_size=[n],
+    )
+    wire_tags = tq_adapter._add_wire_shape_tags([{} for _ in range(n)], original)
+    wire_data = TensorDict(
+        {
+            "reward": torch.nested.as_nested_tensor(
+                [row for row in original["reward"].unsqueeze(-1)],
+                layout=torch.jagged,
+            ),
+            "input_ids": torch.nested.as_nested_tensor(
+                [row for row in original["input_ids"]], layout=torch.jagged
+            ),
+        },
+        batch_size=[n],
+    )
+
+    class FakeMeta:
+        size = n
+        is_ready = True
+        custom_meta = wire_tags
+
+        def select_fields(self, fields):
+            assert fields == ["reward", "input_ids"]
+            return self
+
+    class FakeClient:
+        def kv_retrieve_meta(self, *, keys, partition_id, create):
+            assert keys == ["a", "b", "c"]
+            assert partition_id == "train"
+            assert create is False
+            return FakeMeta()
+
+        def get_data(self, meta):
+            assert isinstance(meta, FakeMeta)
+            return wire_data
+
+    monkeypatch.setattr(
+        tq_adapter.tq, "get_client", lambda: FakeClient(), raising=False
+    )
+    client = object.__new__(tq_adapter.TQDataPlaneClient)
+    client._promote_1d = True
+
+    restored = client.get_samples(["a", "b", "c"], "train", ["reward", "input_ids"])
+
+    assert restored["reward"].shape == (n,)
+    assert restored["input_ids"].shape == (n, 1)
+    assert torch.equal(restored["reward"], original["reward"])
+    assert torch.equal(restored["input_ids"], original["input_ids"])
+
+
 def test_from_wire_preserves_ragged_nested_rows() -> None:
     """Variable-length rollout fields must remain nested."""
     from tensordict import TensorDict

From 2f2261934a90637909603f95ae612fc33782ad31 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Fri, 31 Jul 2026 13:40:11 -0400
Subject: [PATCH 10/15] fix(data-plane): normalize TQ reads across backends

Densify uniform nested reads for the simple backend without applying Mooncake's singleton squeeze. Make wire-field typing explicit, test invalid shape provenance, and warn when either TQ actor runtime patch is unavailable.

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 nemo_rl/data_plane/adapters/transfer_queue.py | 55 +++++++------
 tests/unit/data_plane/test_codec_mooncake.py  | 77 +++++++++++++++++++
 tests/unit/data_plane/test_tq_lifecycle.py    |  7 +-
 3 files changed, 113 insertions(+), 26 deletions(-)

diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py
index 5abc1b1a49..e9dbe25d1f 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -27,8 +27,9 @@
 import socket
 import subprocess
 import time
+import warnings
 from importlib import resources
-from typing import Any
+from typing import Any, cast
 
 import torch
 import transfer_queue as tq
@@ -171,27 +172,28 @@ def patched(*args, **kwargs):
         cls.options = patched  # type: ignore[method-assign]
         return True
 
-    patched_any = False
+    unpatched_classes: list[str] = []
     try:
         from transfer_queue.storage.simple_storage import SimpleStorageUnit
 
-        patched_any |= _install(SimpleStorageUnit)
+        if not _install(SimpleStorageUnit):
+            unpatched_classes.append("SimpleStorageUnit")
     except ImportError:
-        pass
+        unpatched_classes.append("SimpleStorageUnit")
     try:
         from transfer_queue.controller import TransferQueueController
 
-        patched_any |= _install(TransferQueueController)
+        if not _install(TransferQueueController):
+            unpatched_classes.append("TransferQueueController")
     except ImportError:
-        pass
+        unpatched_classes.append("TransferQueueController")
 
-    if not patched_any:
+    if unpatched_classes:
         # Soft-fail: TQ may have moved its actor classes. The driver will
         # still work; multi-node TQ may need the per-node `uv sync` workaround.
-        import warnings
-
         warnings.warn(
-            "Could not patch TQ actor classes for runtime_env injection. "
+            "Could not patch every TQ actor class for runtime_env injection: "
+            f"unpatched={unpatched_classes}. "
             "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' "
             "on worker nodes. Workaround: run `uv sync` inside each node's "
             "container before the driver runs.",
@@ -429,15 +431,15 @@ def _promote_1d_leaves(td: TensorDict) -> TensorDict:
 def _from_wire(
     td: TensorDict, promoted_1d_fields: set[str] | None = None
 ) -> TensorDict:
-    """Normalize Mooncake reads and invert :func:`_promote_1d_leaves`.
-
-    TQ v0.1.9 reconstructs every non-scalar field as a nested tensor before
-    attempting a dense representation, including fields whose rows all have
-    the same shape. Densify those uniform nested tensors first so regular
-    batched inputs retain their dense representation. Truly ragged fields
-    remain nested. Finally, squeeze only singleton dimensions known to have
-    been introduced by :func:`_promote_1d_leaves`. ``None`` retains the old
-    shape-only behavior for checkpoints written before shape tags existed.
+    """Normalize TQ reads and invert :func:`_promote_1d_leaves` when needed.
+
+    Both TQ v0.1.9 storage managers reconstruct every non-scalar field as a
+    nested tensor, including fields whose rows all have the same shape.
+    Densify those uniform nested tensors first so regular batched inputs retain
+    their dense representation. Truly ragged fields remain nested. Finally,
+    squeeze only singleton dimensions known to have been introduced by
+    :func:`_promote_1d_leaves`. ``None`` retains the old shape-only behavior
+    for checkpoints written before shape tags existed.
     """
     # Same top-level iteration as `_promote_1d_leaves`: NonTensorData /
     # NonTensorStack leaves are only visible via td.keys(), not leaves_only.
@@ -687,11 +689,15 @@ def put_samples(
             # TDs. TQ's encoder forces ``.contiguous()`` per tensor leaf
             # itself, so the call here was redundant for tensors and
             # destructive for non-tensors.
-            wire_fields = fields.detach()  # type: ignore[bad-assignment,missing-argument]
+            detached_fields = cast(
+                TensorDict,
+                fields.detach(),  # type: ignore[missing-argument]
+            )
             if self._promote_1d:
-                wire_tags = _add_wire_shape_tags(wire_tags, wire_fields)
-                wire_fields = _promote_1d_leaves(wire_fields)  # type: ignore[bad-argument-type]
-            field_names = list(wire_fields.keys())
+                wire_tags = _add_wire_shape_tags(wire_tags, detached_fields)
+                detached_fields = _promote_1d_leaves(detached_fields)
+            wire_fields = detached_fields
+            field_names = [str(key) for key in detached_fields.keys()]
 
         # TQ's wire vocabulary is `keys=` — translation point.
         tq.kv_batch_put(
@@ -718,11 +724,12 @@ def get_samples(
         if not sample_ids:
             return TensorDict({}, batch_size=(0,))
         if not self._promote_1d:
-            return tq.kv_batch_get(
+            td = tq.kv_batch_get(
                 keys=list(sample_ids),
                 partition_id=partition_id,
                 select_fields=select_fields,
             )
+            return _from_wire(td, promoted_1d_fields=set())
 
         # Inline TQ's public kv_batch_get flow so the adapter can also inspect
         # the durable per-row shape markers carried by BatchMeta.custom_meta.
diff --git a/tests/unit/data_plane/test_codec_mooncake.py b/tests/unit/data_plane/test_codec_mooncake.py
index b0a8868377..27ffecb82b 100644
--- a/tests/unit/data_plane/test_codec_mooncake.py
+++ b/tests/unit/data_plane/test_codec_mooncake.py
@@ -22,6 +22,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 
 from nemo_rl.data_plane.codec import pack_per_token_field, to_nested_by_length
@@ -150,6 +151,51 @@ def test_wire_shape_tags_roundtrip_promoted_field_names() -> None:
     )
 
 
+def test_wire_shape_tags_legacy_rows_have_no_promotion_provenance() -> None:
+    """Rows written before wire-shape tags use the legacy decode path."""
+    from nemo_rl.data_plane.adapters.transfer_queue import (
+        _promoted_1d_fields_from_tags,
+    )
+
+    assert _promoted_1d_fields_from_tags([], ["reward"]) is None
+
+
+def test_wire_shape_tags_reject_mixed_codec_versions() -> None:
+    """A batch cannot safely mix tagged and legacy rows."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import (
+        _WIRE_SHAPE_VERSION_TAG,
+        _add_wire_shape_tags,
+        _promoted_1d_fields_from_tags,
+    )
+
+    fields = TensorDict({"input_ids": torch.arange(4).reshape(2, 2)}, batch_size=[2])
+    wire_tags = _add_wire_shape_tags([{}, {}], fields)
+    wire_tags[1].pop(_WIRE_SHAPE_VERSION_TAG)
+
+    with pytest.raises(RuntimeError, match="wire-shape version"):
+        _promoted_1d_fields_from_tags(wire_tags, ["input_ids"])
+
+
+def test_wire_shape_tags_reject_partial_promotion_marker() -> None:
+    """All rows must agree that a field was promoted from one dimension."""
+    from tensordict import TensorDict
+
+    from nemo_rl.data_plane.adapters.transfer_queue import (
+        _PROMOTED_1D_TAG_PREFIX,
+        _add_wire_shape_tags,
+        _promoted_1d_fields_from_tags,
+    )
+
+    fields = TensorDict({"reward": torch.arange(2)}, batch_size=[2])
+    wire_tags = _add_wire_shape_tags([{}, {}], fields)
+    wire_tags[1].pop(f"{_PROMOTED_1D_TAG_PREFIX}reward")
+
+    with pytest.raises(RuntimeError, match="promotion marker"):
+        _promoted_1d_fields_from_tags(wire_tags, ["reward"])
+
+
 def test_get_samples_uses_persisted_shape_tags(monkeypatch) -> None:
     """The Mooncake adapter restores ranks using TQ's durable row metadata."""
     from tensordict import TensorDict
@@ -212,6 +258,37 @@ def get_data(self, meta):
     assert torch.equal(restored["input_ids"], original["input_ids"])
 
 
+def test_get_samples_densifies_uniform_rows_without_1d_promotion(monkeypatch) -> None:
+    """The simple backend normalizes uniform nested rows without squeezing."""
+    from tensordict import TensorDict
+
+    import nemo_rl.data_plane.adapters.transfer_queue as tq_adapter
+
+    rows = [torch.tensor([1, 2]), torch.tensor([3, 4])]
+    wire_data = TensorDict(
+        {"input_ids": torch.nested.as_nested_tensor(rows, layout=torch.jagged)},
+        batch_size=[len(rows)],
+    )
+
+    def fake_kv_batch_get(
+        *, keys: list[str], partition_id: str, select_fields: list[str]
+    ) -> TensorDict:
+        assert keys == ["a", "b"]
+        assert partition_id == "train"
+        assert select_fields == ["input_ids"]
+        return wire_data
+
+    monkeypatch.setattr(tq_adapter.tq, "kv_batch_get", fake_kv_batch_get, raising=False)
+    client = object.__new__(tq_adapter.TQDataPlaneClient)
+    client._promote_1d = False
+
+    restored = client.get_samples(["a", "b"], "train", ["input_ids"])
+
+    assert not restored["input_ids"].is_nested
+    assert restored["input_ids"].shape == (2, 2)
+    assert torch.equal(restored["input_ids"], torch.stack(rows))
+
+
 def test_from_wire_preserves_ragged_nested_rows() -> None:
     """Variable-length rollout fields must remain nested."""
     from tensordict import TensorDict
diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py
index a349c27c2e..51d9f8d1a1 100644
--- a/tests/unit/data_plane/test_tq_lifecycle.py
+++ b/tests/unit/data_plane/test_tq_lifecycle.py
@@ -134,10 +134,11 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None:
         consumer_tasks=["read"],
     )
     keys = ["a", "b", "c", "d"]
+    values = torch.arange(12).reshape(4, 3)
     client.put_samples(
         sample_ids=keys,
         partition_id="smoke-backend",
-        fields=TensorDict({"x": torch.arange(4)}, batch_size=[4]),
+        fields=TensorDict({"x": values}, batch_size=[4]),
     )
 
     meta = client.claim_meta(
@@ -150,7 +151,9 @@ def test_smoke_round_trip_backends(tq_client_backends) -> None:
     assert meta.size == 4
 
     data = client.get_data(meta)
-    expected = torch.tensor([keys.index(k) for k in meta.sample_ids])
+    expected = torch.stack([values[keys.index(k)] for k in meta.sample_ids])
+    assert not data["x"].is_nested
+    assert data["x"].shape == expected.shape
     assert torch.equal(data["x"], expected)
 
     client.clear_samples(sample_ids=None, partition_id="smoke-backend")

From 08f592eafc8fcd46539e6520418ea56916f51721 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Fri, 31 Jul 2026 13:18:26 -0700
Subject: [PATCH 11/15] docs: correct the numpy override rationale in pyproject

TransferQueue v0.1.9 dropped its numpy<2.0.0 pin, so the override's
original justification no longer applies. The constraint it actually
bypasses now is tensorrt-llm's numpy>=2.0.0,<2.4.

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 pyproject.toml | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 1136e927a8..3b2da988c2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -378,8 +378,9 @@ override-dependencies = [
   "pytest>=9.0.3",
   "langchain>=0.3.28",      # Address CVE-2025-65106
   "langchain-core>=0.3.80", # Address CVE-2025-65106
-  # Keep the NumPy floor compatible with megatron-core's onnx → ml-dtypes
-  # dependency while composing the data-plane stack with mcore/automodel.
+  # Forces numpy past tensorrt-llm's `numpy>=2.0.0,<2.4` cap (resolves to
+  # 2.5.x). The original driver — TransferQueue pinning `numpy<2.0.0` — is
+  # gone as of TQ v0.1.9; drop this override once the trtllm cap lifts.
   "numpy>=2.1.0",
   # av (PyAV) carries CVE-bundled codec libs (libx264, libx265, libopenh264, libmp3lame).
   # It is only needed by megatron-bridge's optional WAN diffusion path, which installs it

From b95e993ade895b67296b0c1c30d190991a98340b Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Sun, 2 Aug 2026 14:00:17 -0400
Subject: [PATCH 12/15] feat(sc): save native TQ state in checkpoints

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 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       | 107 +++++++--
 .../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                  | 144 +++++++++++-
 .../test_single_controller_setup.py           |  11 +
 .../test_tq_replay_buffer.py                  |  47 +++-
 tools/verify_tq_data_plane_checkpoint.py      | 210 ++++++++++++++++++
 19 files changed, 860 insertions(+), 32 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 8c45f089b9..f7fe7a2b4b 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -496,6 +496,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 14be5604e1..3020b4f664 100644
--- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml
+++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml
@@ -54,6 +54,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 b4d1b39c2d..95fad17634 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -659,6 +659,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,
@@ -749,10 +756,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):
@@ -815,10 +820,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)
@@ -1028,3 +1031,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 7938b9586b..c25b1c7bd4 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -74,6 +74,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()
@@ -296,9 +309,7 @@ async def _maybe_restore_replay_buffer(self) -> None:
             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_group_size=self._master_config.grpo["num_generations_per_prompt"],
         )
         # Each buffered group holds one _buffer_capacity permit; the load
         # truncation guarantees restored <= capacity, so this never blocks.
@@ -321,6 +332,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:
@@ -577,11 +643,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
 
@@ -606,9 +668,7 @@ async def _train_pump(self) -> None:
 
                 # Checkpointing (mirrors async_grpo_train's save block).
                 self._consumed_samples += grpo_cfg["num_prompts_per_step"]
-                self._total_valid_tokens += step_metrics.get(
-                    "global_valid_toks", 0
-                )
+                self._total_valid_tokens += step_metrics.get("global_valid_toks", 0)
                 self._timeout.mark_iteration()
 
                 is_last_step = self._train_steps >= grpo_cfg["max_num_steps"]
@@ -706,9 +766,9 @@ async def _save_checkpoint(
 
         full_metric_name = self._master_config.checkpointing["metric_name"]
         if full_metric_name is not None:
-            assert full_metric_name.startswith(
-                "train:"
-            ) or full_metric_name.startswith("val:"), (
+            assert full_metric_name.startswith("train:") or full_metric_name.startswith(
+                "val:"
+            ), (
                 f"metric_name={full_metric_name} must start with 'val:' or 'train:',\n"
                 f'followed by the corresponding name in the "val" or "train" metrics dictionary.'
                 f"  If you are using an old config, please updated checkpointing.metric_name to the new format, "
@@ -725,9 +785,7 @@ async def _save_checkpoint(
                 if full_metric_name in save_state:
                     del save_state[full_metric_name]
             elif metric_name not in metrics_source:
-                raise ValueError(
-                    f"Metric {metric_name} not found in {prefix} metrics"
-                )
+                raise ValueError(f"Metric {metric_name} not found in {prefix} metrics")
             else:
                 save_state[full_metric_name] = metrics_source[metric_name]
 
@@ -758,10 +816,23 @@ async def _save_checkpoint(
             dataloader_state,
             os.path.join(checkpoint_path, "train_dataloader.pt"),
         )
-        if self._sampler.supports_buffer_checkpoint:
+        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,
diff --git a/nemo_rl/algorithms/single_controller_utils/setup.py b/nemo_rl/algorithms/single_controller_utils/setup.py
index 579ede5d87..0450c21933 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -312,6 +312,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 1c5b00a5e4..f01b40c098 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 e9dbe25d1f..070f388c09 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
@@ -516,6 +518,7 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None:
         # is unaffected). Writer unsqueezes 1D → (N, 1) on put; reader
         # squeezes the trailing 1 back on get. Drop when upstream TQ
         # unifies the schema/data shapes for 1D fields.
+        self._backend = cfg["backend"]
         self._promote_1d = cfg["backend"] == "mooncake_cpu"
 
         if bootstrap:
@@ -779,6 +782,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 41a98f0c0e..6f57bcade6 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 63e551dc20..d7bcd88fca 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 65929bb208..49fc4f2ef5 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -244,6 +244,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 a105892bc0..7ec4ac0984 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 3426c3b506..3a4009b164 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 0d471bc266..2cdbdb1a62 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 51d9f8d1a1..d08629e7af 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 2124588dac..8fe46c942d 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -485,6 +485,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 c7e3b96ece..5caacf9fd1 100644
--- a/tests/unit/single_controller/test_sc_checkpointing.py
+++ b/tests/unit/single_controller/test_sc_checkpointing.py
@@ -195,12 +195,49 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[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:
@@ -231,6 +268,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)
@@ -290,6 +331,7 @@ def _actor_master_config(
     num_prompts_per_step: int = 2,
     max_num_epochs: int = 1,
     buffer_checkpoint: bool = True,
+    data_plane_checkpoint: bool = False,
 ) -> MasterConfig:
     """MasterConfig for in-process SingleControllerActor tests.
 
@@ -337,7 +379,12 @@ def _actor_master_config(
             "save_optimizer": save_optimizer,
             "checkpoint_must_save_by": checkpoint_must_save_by,
         },
-        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,
@@ -353,6 +400,7 @@ def _make_actor_args(
     save_state: Optional[dict[str, Any]] = None,
     dataloader: Optional[_FakeDataloader] = None,
     tq_buffer: Optional[_FakeTQBuffer] = None,
+    dp_client: Optional[_FakeDPClient] = None,
     last_checkpoint_path: Optional[str] = None,
 ) -> SingleControllerActorArgs:
     return SingleControllerActorArgs(
@@ -361,7 +409,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,
@@ -424,7 +472,9 @@ async def _main():
 def _step_dir_names(ckpt_dir: Path) -> set[str]:
     if not ckpt_dir.exists():
         return set()
-    return {p.name for p in ckpt_dir.iterdir() if p.name != "latest_checkpoint_status.json"}
+    return {
+        p.name for p in ckpt_dir.iterdir() if p.name != "latest_checkpoint_status.json"
+    }
 
 
 def _training_info(ckpt_dir: Path, step: int) -> dict[str, Any]:
@@ -561,7 +611,9 @@ def test_save_optimizer_false_gates_optimizer_path(self, tmp_path):
         assert not (ckpt_dir / "step_2" / "policy" / "optimizer").exists()
 
     def test_no_save_when_disabled(self, tmp_path):
-        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=1, enabled=False)
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=2, save_period=1, enabled=False
+        )
         trainer = _FakeTrainer()
 
         actor = _run_train_pump(mc, _make_actor_args(trainer=trainer))
@@ -588,6 +640,88 @@ def test_timeout_saves_and_stops_training_early(self, tmp_path):
         assert _step_dir_names(tmp_path / "checkpoints") == {"step_1"}
 
 
+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 5b5bac11b2..51cc4d94c3 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -218,6 +218,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 f8ba5b663a..0e2f83a3a9 100644
--- a/tests/unit/single_controller/test_tq_replay_buffer.py
+++ b/tests/unit/single_controller/test_tq_replay_buffer.py
@@ -151,10 +151,16 @@ def _make_record() -> PromptGroupRecord:
     )
 
 
-def _make_buffer(dp: FakeDataPlaneClient) -> TQReplayBuffer:
-    return TQReplayBuffer(
+def _make_buffer(
+    dp: FakeDataPlaneClient,
+    *,
+    checkpoint_lock: asyncio.Lock | None = None,
+) -> TQReplayBuffer:
+    buffer = TQReplayBuffer(
         dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0}
     )
+    buffer.set_data_plane_checkpoint_lock(checkpoint_lock or asyncio.Lock())
+    return buffer
 
 
 def _add_group(
@@ -304,6 +310,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 0000000000..0f58c6d5d7
--- /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 912085d6dbb4b91f10f6c3d52f8885050e9c824c Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Sun, 2 Aug 2026 14:55:31 -0400
Subject: [PATCH 13/15] fix(sc): harden TQ checkpoint lifecycle

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
---
 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 f7fe7a2b4b..697e10eba1 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -496,7 +496,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 3020b4f664..ba841363c7 100644
--- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml
+++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml
@@ -54,6 +54,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 95fad17634..4884115f00 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -662,7 +662,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 c25b1c7bd4..54440c8b08 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(
@@ -321,13 +324,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
@@ -342,7 +365,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,
@@ -360,20 +389,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 070f388c09..0d9043b0af 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -527,6 +527,22 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None:
             _connect_existing()
         self._poll_interval_s = cfg["claim_meta_poll_interval_s"]
         self._closed = False
+        # TQ restore is non-transactional and requires a globally clean system.
+        # This process-local guard catches incorrect ordering through this
+        # adapter; setup must still ensure no other client has touched TQ.
+        self._data_operations_started = False
+
+    def _mark_data_operation_started(self) -> None:
+        """Make a later checkpoint load fail instead of mixing TQ states."""
+        self._data_operations_started = True
+
+    def _require_clean_for_load(self) -> None:
+        """Reject restore after this client has performed a data operation."""
+        if self._data_operations_started:
+            raise RuntimeError(
+                "load_checkpoint requires a clean TQ client before any "
+                "register, claim, get, put, clear, or consumption operation"
+            )
 
     # ── (A) task-mediated ───────────────────────────────────────────────
 
@@ -559,6 +575,7 @@ def register_partition(
         # stale metadata from a previous registration.
         if not fields:
             return
+        self._mark_data_operation_started()
         schema_key = (
             f"__schema__:{partition_id}:{os.getpid()}:{id(self)}:{time.time_ns()}"
         )
@@ -584,6 +601,7 @@ def claim_meta(
         blocking: bool = True,
         timeout_s: float = 60.0,
     ) -> KVBatchMeta:
+        self._mark_data_operation_started()
         client = tq.get_client()
         deadline = time.time() + max(0.0, timeout_s)
         sampling_config: dict[str, Any] = {}
@@ -658,6 +676,7 @@ def get_data(
     def check_consumption_status(
         self, partition_id: str, task_names: list[str]
     ) -> bool:
+        self._mark_data_operation_started()
         client = tq.get_client()
         for t in task_names:
             if not client.check_consumption_status(
@@ -702,6 +721,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),
@@ -726,6 +746,7 @@ def get_samples(
     ) -> TensorDict:
         if not sample_ids:
             return TensorDict({}, batch_size=(0,))
+        self._mark_data_operation_started()
         if not self._promote_1d:
             td = tq.kv_batch_get(
                 keys=list(sample_ids),
@@ -756,6 +777,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
@@ -777,6 +799,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)
 
@@ -799,20 +822,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 6f57bcade6..6bb23d1cbd 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 27ffecb82b..19907c79a5 100644
--- a/tests/unit/data_plane/test_codec_mooncake.py
+++ b/tests/unit/data_plane/test_codec_mooncake.py
@@ -249,6 +249,7 @@ def get_data(self, meta):
     )
     client = object.__new__(tq_adapter.TQDataPlaneClient)
     client._promote_1d = True
+    client._data_operations_started = False
 
     restored = client.get_samples(["a", "b", "c"], "train", ["reward", "input_ids"])
 
@@ -281,6 +282,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 d08629e7af..1c88a46e5d 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 8fe46c942d..4813178e39 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -485,7 +485,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 0000000000..25c2455fe7
--- /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 0f58c6d5d7..3525424fff 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_<bundle>/data_plane``, then renames the parent bundle to its
+final path just like ``CheckpointManager``. The load phase starts a fresh TQ
+instance, restores from ``<bundle>/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 a9f52b73435a0948f88f44ef27ae1d2a3ab70a0f Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Sun, 2 Aug 2026 21:01:48 -0400
Subject: [PATCH 14/15] fix(sc): harden TQ checkpoint I/O

---
 .../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 4884115f00..f7338c3070 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -24,6 +24,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.experience.interfaces import PromptGroupRecord
 from nemo_rl.experience.payload import pack_payload, record_to_train_batch
 
@@ -733,7 +734,8 @@ async def commit(
             train_batch, weight_version=start_weight_version, group_id=group_id
         )
         try:
-            await self._call_dp(
+            await call_data_plane(
+                self._dp_client,
                 "put_samples",
                 sample_ids=sample_ids,
                 partition_id=self._partition_id,
@@ -869,7 +871,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,
@@ -999,7 +1002,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,
@@ -1026,17 +1030,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:
@@ -1045,8 +1038,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 54440c8b08..b8be90d53e 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -63,6 +63,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
@@ -320,46 +321,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,
             )
@@ -389,7 +357,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,
@@ -929,7 +898,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,
@@ -979,7 +949,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 0000000000..aee3ef9b53
--- /dev/null
+++ b/nemo_rl/data_plane/async_utils.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Async dispatch helpers for local and Ray data-plane clients."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+
+async def call_data_plane(
+    client: Any,
+    method_name: str,
+    *,
+    offload_sync: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """Call a local data-plane client or a Ray actor exposing its methods.
+
+    Synchronous offloading is opt-in because it allows the actor event loop to
+    issue other calls while this one is running. Callers should enable it only
+    when that concurrency is supported or externally serialized.
+
+    Args:
+        client: Local ``DataPlaneClient`` or Ray actor handle.
+        method_name: Data-plane method to invoke.
+        offload_sync: Run a synchronous local implementation in a worker
+            thread. Ray methods are already asynchronous and ignore this flag.
+        **kwargs: Keyword arguments forwarded to the data-plane method.
+
+    Returns:
+        The method result after awaiting Ray or coroutine results.
+    """
+    method = getattr(client, method_name)
+    remote = getattr(method, "remote", None)
+    if remote is not None:
+        return await remote(**kwargs)
+    if offload_sync:
+        result = await asyncio.to_thread(method, **kwargs)
+    else:
+        result = method(**kwargs)
+    if asyncio.iscoroutine(result):
+        return await result
+    return result
diff --git a/tests/unit/data_plane/test_async_utils.py b/tests/unit/data_plane/test_async_utils.py
new file mode 100644
index 0000000000..25d29b5d71
--- /dev/null
+++ b/tests/unit/data_plane/test_async_utils.py
@@ -0,0 +1,77 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for local and Ray-style async data-plane dispatch."""
+
+import asyncio
+import threading
+
+from nemo_rl.data_plane.async_utils import call_data_plane
+
+
+class _LocalClient:
+    def thread_id(self) -> int:
+        return threading.get_ident()
+
+    async def async_value(self, *, value: int) -> int:
+        return value
+
+
+class _RemoteMethod:
+    def __init__(self) -> None:
+        self.calls: list[int] = []
+
+    async def remote(self, *, value: int) -> int:
+        self.calls.append(value)
+        return value
+
+
+class _RemoteClient:
+    def __init__(self) -> None:
+        self.value = _RemoteMethod()
+
+
+def test_sync_call_stays_inline_by_default() -> None:
+    caller_thread_id = threading.get_ident()
+
+    result = asyncio.run(call_data_plane(_LocalClient(), "thread_id"))
+
+    assert result == caller_thread_id
+
+
+def test_sync_call_can_be_offloaded() -> None:
+    caller_thread_id = threading.get_ident()
+
+    result = asyncio.run(
+        call_data_plane(_LocalClient(), "thread_id", offload_sync=True)
+    )
+
+    assert result != caller_thread_id
+
+
+def test_local_coroutine_result_is_awaited() -> None:
+    result = asyncio.run(
+        call_data_plane(_LocalClient(), "async_value", value=7)
+    )
+
+    assert result == 7
+
+
+def test_ray_style_remote_result_is_awaited() -> None:
+    client = _RemoteClient()
+
+    result = asyncio.run(call_data_plane(client, "value", value=11))
+
+    assert result == 11
+    assert client.value.calls == [11]
diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py
index 5caacf9fd1..0e2f09f5a8 100644
--- a/tests/unit/single_controller/test_sc_checkpointing.py
+++ b/tests/unit/single_controller/test_sc_checkpointing.py
@@ -197,10 +197,12 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[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(
@@ -721,6 +723,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 0e2f83a3a9..cbad17862a 100644
--- a/tests/unit/single_controller/test_tq_replay_buffer.py
+++ b/tests/unit/single_controller/test_tq_replay_buffer.py
@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import asyncio
+import threading
 from typing import Any
 
 import pytest
@@ -63,6 +64,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(
@@ -94,6 +96,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:
@@ -347,6 +350,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 25c2455fe7..6359fc672f 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 3525424fff..33c9f50ada 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 3af936cd1c972ba86386625aa36656e6b8122fc5 Mon Sep 17 00:00:00 2001
From: Anish Mahishi <amahishi@nvidia.com>
Date: Mon, 3 Aug 2026 15:28:53 -0400
Subject: [PATCH 15/15] feat(sc): recover replay buffer from native TQ
 checkpoints

---
 examples/configs/grpo_math_1B.yaml            |   5 +-
 .../algorithms/async_utils/replay_buffer.py   | 434 ++++++++++++------
 .../async_utils/staleness_sampler.py          |  35 +-
 nemo_rl/algorithms/grpo.py                    |   1 +
 nemo_rl/algorithms/single_controller.py       | 219 +++++++--
 .../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            |   7 +
 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 +
 .../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                 |  39 +-
 .../test_sc_checkpointing.py                  | 420 ++++++++++++++---
 .../test_single_controller_setup.py           | 197 ++++++++
 .../test_tq_replay_buffer.py                  | 254 ++++++----
 26 files changed, 1514 insertions(+), 380 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 697e10eba1..ee38971cfb 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -496,7 +496,10 @@ data_plane:
   storage_capacity: 1000000          # max samples retained per partition
   num_storage_units: 2               # storage shards
   claim_meta_poll_interval_s: 0.5    # blocking-claim poll cadence
-  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 f7338c3070..61da9684d9 100644
--- a/nemo_rl/algorithms/async_utils/replay_buffer.py
+++ b/nemo_rl/algorithms/async_utils/replay_buffer.py
@@ -13,12 +13,15 @@
 # limitations under the License.
 
 import asyncio
+import hashlib
+import json
 import statistics
 import threading as _threading
 import uuid
 from collections import Counter
-from collections.abc import Mapping
-from typing import Any, Iterable, Optional
+from collections.abc import AsyncIterator, Mapping
+from contextlib import asynccontextmanager
+from typing import Any, Iterable, Literal, Optional, TypedDict
 
 import ray
 
@@ -29,6 +32,111 @@
 from nemo_rl.experience.payload import pack_payload, record_to_train_batch
 
 
+DATA_PLANE_CHECKPOINT_DIR = "data_plane"
+DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2
+REPLAY_BUFFER_METADATA_FILENAME = "replay_buffer_metadata.pt"
+LEGACY_REPLAY_BUFFER_FILENAME = "replay_buffer.pt"
+REPLAY_BUFFER_METADATA_SCHEMA_VERSION = 1
+REPLAY_BUFFER_METADATA_STORAGE: Literal["tq_checkpoint"] = "tq_checkpoint"
+
+
+class TQReplayGroupMetadata(TypedDict):
+    """Controller-local index for one training-ready group stored in TQ."""
+
+    meta: KVBatchMeta
+    start_weight: int
+    end_weight: int
+    target_step: Optional[int]
+    group_id: str
+
+
+class TQReplayMetadataState(TypedDict):
+    """Versioned metadata-only replay sidecar paired with a TQ snapshot."""
+
+    schema_version: int
+    storage: Literal["tq_checkpoint"]
+    partition_id: str
+    saved_capacity: int
+    manifest_digest: str
+    groups: list[TQReplayGroupMetadata]
+
+
+def replay_manifest_digest(groups: list[TQReplayGroupMetadata]) -> str:
+    """Return a stable digest binding replay metadata to a TQ checkpoint."""
+    digest_input = [
+        {
+            "group_id": group["group_id"],
+            "start_weight": group["start_weight"],
+            "end_weight": group["end_weight"],
+            "target_step": group["target_step"],
+            "meta": {
+                "partition_id": group["meta"].partition_id,
+                "task_name": group["meta"].task_name,
+                "sample_ids": list(group["meta"].sample_ids),
+                "fields": (
+                    list(group["meta"].fields)
+                    if group["meta"].fields is not None
+                    else None
+                ),
+                "sequence_lengths": (
+                    list(group["meta"].sequence_lengths)
+                    if group["meta"].sequence_lengths is not None
+                    else None
+                ),
+                "tags": group["meta"].tags,
+                "extra_info": group["meta"].extra_info,
+            },
+        }
+        for group in groups
+    ]
+    encoded = json.dumps(
+        digest_input, sort_keys=True, separators=(",", ":")
+    ).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()
+
+
+class DataPlaneCheckpointBarrier:
+    """Allow concurrent mutations while giving checkpoints exclusive access."""
+
+    def __init__(self) -> None:
+        self._condition = asyncio.Condition()
+        self._checkpoint_active = False
+        self._active_mutations = 0
+
+    @asynccontextmanager
+    async def mutation(self) -> AsyncIterator[None]:
+        """Enter a commit/clear section, waiting only for an active checkpoint."""
+        async with self._condition:
+            await self._condition.wait_for(lambda: not self._checkpoint_active)
+            self._active_mutations += 1
+        try:
+            yield
+        finally:
+            async with self._condition:
+                self._active_mutations -= 1
+                if self._active_mutations == 0:
+                    self._condition.notify_all()
+
+    @asynccontextmanager
+    async def checkpoint(self) -> AsyncIterator[None]:
+        """Block new mutations and wait for active ones before snapshotting."""
+        async with self._condition:
+            await self._condition.wait_for(lambda: not self._checkpoint_active)
+            self._checkpoint_active = True
+            try:
+                await self._condition.wait_for(lambda: self._active_mutations == 0)
+            except BaseException:
+                self._checkpoint_active = False
+                self._condition.notify_all()
+                raise
+        try:
+            yield
+        finally:
+            async with self._condition:
+                self._checkpoint_active = False
+                self._condition.notify_all()
+
+
 # Classes with @ray.remote can't be inherited from, so we split the implementation out.
 class ReplayBufferImpl(ReplayBufferProtocol):
     """Replay buffer storing per-prompt groups.
@@ -660,18 +768,22 @@ def __init__(
         self.target_step_list: list[Optional[int]] = []
         self.ready_list: list[bool] = []
         self._group_ids: list[str] = []
-        self._data_plane_checkpoint_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,
@@ -729,51 +841,57 @@ async def commit(
                 f"commit called with unknown group_id={group_id!r}; "
                 f"reserve() must precede commit() (or the slot was already removed)"
             )
+        if self._data_plane_checkpoint_barrier is None:
+            raise RuntimeError(
+                "TQReplayBuffer must be bound to the controller data-plane "
+                "checkpoint barrier before committing samples"
+            )
         train_batch = record_to_train_batch(record, pad_value_dict=self._pad_value_dict)
         sample_ids, fields, tags = pack_payload(
             train_batch, weight_version=start_weight_version, group_id=group_id
         )
-        try:
-            await 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``.
@@ -788,11 +906,19 @@ async def remove_group(self, group_id: str, *, remove_in_dp: bool = False) -> in
         Raises:
             ValueError: ``group_id`` has no live slot.
         """
-        try:
-            idx = self._group_ids.index(group_id)
-        except ValueError as error:
-            raise ValueError(f"unknown group_id={group_id!r}") from error
-        return await self.remove([idx], remove_in_dp=remove_in_dp)
+        if self._data_plane_checkpoint_barrier is None:
+            raise RuntimeError(
+                "TQReplayBuffer must be bound to the controller data-plane "
+                "checkpoint barrier before removing a group"
+            )
+        async with self._data_plane_checkpoint_barrier.mutation():
+            try:
+                idx = self._group_ids.index(group_id)
+            except ValueError as error:
+                raise ValueError(f"unknown group_id={group_id!r}") from error
+            return await self._remove_unlocked(
+                [idx], clear_data_plane=remove_in_dp
+            )
 
     async def remove(self, idxs: list[int], remove_in_dp: bool) -> int:
         """Drop entries at the given indices and optionally clear them from DataPlane.
@@ -806,14 +932,26 @@ async def remove(self, idxs: list[int], remove_in_dp: bool) -> int:
         """
         if len(idxs) == 0:
             return 0
-
-        drop_idxs = sorted(idxs, reverse=True)
-        if drop_idxs[0] >= len(self.meta_list):
-            raise IndexError(
-                f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; "
-                f"size={len(self.meta_list)}"
+        if self._data_plane_checkpoint_barrier is None:
+            raise RuntimeError(
+                "TQReplayBuffer must be bound to the controller data-plane "
+                "checkpoint barrier before removing groups"
+            )
+        async with self._data_plane_checkpoint_barrier.mutation():
+            drop_idxs = sorted(idxs, reverse=True)
+            if drop_idxs[0] >= len(self.meta_list):
+                raise IndexError(
+                    f"TQReplayBuffer.remove: indices out of range: {drop_idxs[0]}; "
+                    f"size={len(self.meta_list)}"
+                )
+            return await self._remove_unlocked(
+                drop_idxs, clear_data_plane=remove_in_dp
             )
 
+    async def _remove_unlocked(
+        self, drop_idxs: list[int], *, clear_data_plane: bool
+    ) -> int:
+        """Remove validated indices while the caller owns any required lock."""
         dropped_sample_ids: list[str] = []
         for i in drop_idxs:
             meta = self.meta_list[i]
@@ -826,71 +964,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,
         }
 
@@ -901,14 +1016,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.
+        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
@@ -916,7 +1031,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
@@ -933,12 +1048,33 @@ async def load_state_dict(
                 mismatch, misaligned or wrongly sized groups, duplicate
                 sample_ids).
         """
-        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(
+                "Replay-buffer metadata has unsupported storage: "
+                f"{state['storage']!r}"
+            )
         if state["partition_id"] != expected_partition_id:
             raise ValueError(
                 "Replay buffer checkpoint partition_id mismatch: "
@@ -953,16 +1089,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
@@ -983,33 +1128,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:
-            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"])
@@ -1017,10 +1159,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 size(self) -> int:
@@ -1032,16 +1174,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 6810f2531e..2896192f22 100644
--- a/nemo_rl/algorithms/async_utils/staleness_sampler.py
+++ b/nemo_rl/algorithms/async_utils/staleness_sampler.py
@@ -44,6 +44,7 @@
 from typing import (
     Annotated,
     Callable,
+    ClassVar,
     Literal,
     Optional,
     Protocol,
@@ -455,6 +456,7 @@ async def evict(self, *, current_train_weight: int) -> int:
 
 
 class WindowedSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = True
     name: Literal["windowed"] = "windowed"
     # Max weight-version gap a selected group may have from the trainer.
     max_staleness_versions: NonNegativeInt = 1
@@ -463,18 +465,22 @@ class WindowedSamplerConfig(BaseModel, extra="allow"):
 
 
 class WeightFifoSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = False
     name: Literal["weight_fifo"] = "weight_fifo"
     # Lookahead + selectable weight window, in trainer versions.
     max_staleness_versions: NonNegativeInt = 1
 
 
 class InOrderSamplerConfig(BaseModel, extra="allow"):
+    supports_buffer_checkpoint: ClassVar[bool] = False
     name: Literal["in_order"] = "in_order"
     # How far generation may run ahead of the trainer, in dispatch batches.
     max_lookahead_versions: NonNegativeInt = 1
 
 
 class CustomSamplerConfig(BaseModel, extra="allow"):
+    # A custom implementation's capability is known only after construction.
+    supports_buffer_checkpoint: ClassVar[Optional[bool]] = None
     name: Literal["custom"] = "custom"
     # "module:ClassName" of a PromptGroupSampler defined outside this repo.
     # Extra keys are forwarded to the constructor (after ``buffer``).
@@ -531,26 +537,27 @@ def create_sampler(
             signature — a custom class that doesn't accept the kwarg fails
             loudly the first time a run actually resumes with it.
     """
+    sampler: PromptGroupSampler
     if isinstance(cfg, WindowedSamplerConfig):
-        return WindowedSampler(
+        sampler = WindowedSampler(
             buffer,
             max_staleness_versions=cfg.max_staleness_versions,
             sample_freshest_first=cfg.sample_freshest_first,
             resume_from_step=resume_from_step,
         )
-    if isinstance(cfg, WeightFifoSamplerConfig):
-        return WeightFifoSampler(
+    elif isinstance(cfg, WeightFifoSamplerConfig):
+        sampler = WeightFifoSampler(
             buffer,
             max_staleness_versions=cfg.max_staleness_versions,
             resume_from_step=resume_from_step,
         )
-    if isinstance(cfg, InOrderSamplerConfig):
-        return InOrderSampler(
+    elif isinstance(cfg, InOrderSamplerConfig):
+        sampler = InOrderSampler(
             buffer,
             max_lookahead_versions=cfg.max_lookahead_versions,
             resume_from_step=resume_from_step,
         )
-    if isinstance(cfg, CustomSamplerConfig):
+    elif isinstance(cfg, CustomSamplerConfig):
         module_name, sep, class_name = cfg.target.partition(":")
         if not sep:
             raise ValueError(
@@ -567,5 +574,17 @@ def create_sampler(
                 f"interface (needs admit/select/evict, is_on_policy, "
                 f"supports_buffer_checkpoint, required_buffer_capacity)"
             )
-        return sampler
-    raise ValueError(f"unknown sampler config {type(cfg).__name__}")
+    else:
+        raise ValueError(f"unknown sampler config {type(cfg).__name__}")
+
+    expected_capability = cfg.supports_buffer_checkpoint
+    if (
+        expected_capability is not None
+        and sampler.supports_buffer_checkpoint != expected_capability
+    ):
+        raise RuntimeError(
+            f"{type(cfg).__name__}.supports_buffer_checkpoint="
+            f"{expected_capability} disagrees with {type(sampler).__name__}."
+            f"supports_buffer_checkpoint={sampler.supports_buffer_checkpoint}"
+        )
+    return sampler
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index 85bac9836c..0b3228f61c 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -295,6 +295,7 @@ class GRPOSaveState(TypedDict):
     current_epoch: int
     total_steps: int
     total_valid_tokens: int  # Track total number of non-padding tokens during training
+    trainer_version: NotRequired[int]
     val_reward: NotRequired[
         float
     ]  # Optional field - may not be present during training
diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py
index b8be90d53e..6350739296 100644
--- a/nemo_rl/algorithms/single_controller.py
+++ b/nemo_rl/algorithms/single_controller.py
@@ -45,6 +45,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.single_controller_utils.config import (
@@ -75,10 +84,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:
     """CPU-only Ray actor that orchestrates the RL training loop.
@@ -150,6 +155,9 @@ def __init__(
         # (val_reward, ...) pass through to saved checkpoints untouched.
         self._save_state: GRPOSaveState = actor_args.save_state
         self._last_checkpoint_path: Optional[str] = actor_args.last_checkpoint_path
+        self._data_plane_checkpoint_metadata: Optional[dict[str, Any]] = (
+            actor_args.data_plane_checkpoint_metadata
+        )
         self._consumed_samples: int = actor_args.save_state["consumed_samples"]
         self._total_valid_tokens: int = actor_args.save_state.get(
             "total_valid_tokens", 0
@@ -163,8 +171,20 @@ def __init__(
         self._sampler = create_sampler(
             self._buffer,
             self._async_cfg.sampler,
-            resume_from_step=actor_args.save_state["current_step"],
+            resume_from_step=actor_args.save_state.get(
+                "trainer_version", actor_args.save_state["current_step"]
+            ),
         )
+        if (
+            self._master_config.checkpointing["enabled"]
+            and self._sampler.supports_buffer_checkpoint
+            and not self._master_config.data_plane.get("checkpointing_enabled")
+        ):
+            raise ValueError(
+                "SingleController checkpointing with a replay-checkpoint-capable "
+                "sampler requires data_plane.checkpointing_enabled=true so "
+                "completed, unconsumed rollouts are recoverable."
+            )
         required_capacity = self._sampler.required_buffer_capacity(num_prompts_per_step)
         validate_sampler_buffer_capacity(
             self._async_cfg,
@@ -173,17 +193,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
@@ -208,7 +228,9 @@ def __init__(
             self._async_cfg.max_buffered_rollouts
         )
 
-        self._trainer_version: int = actor_args.save_state["current_step"]
+        self._trainer_version: int = actor_args.save_state.get(
+            "trainer_version", actor_args.save_state["current_step"]
+        )
         self._train_steps: int = actor_args.save_state["current_step"]
         self._current_epoch: int = actor_args.save_state["current_epoch"]
         self._step_log_dict: dict[str, list] = {
@@ -290,40 +312,122 @@ async def _maybe_restore_replay_buffer(self) -> None:
 
         No-op unless the sampler supports_buffer_checkpoint (ungated only).
         """
-        if not (
-            self._sampler.supports_buffer_checkpoint
-            and self._last_checkpoint_path is not None
+        if self._last_checkpoint_path is None:
+            return
+        metadata_path = os.path.join(
+            self._last_checkpoint_path, REPLAY_BUFFER_METADATA_FILENAME
+        )
+        if (
+            os.path.exists(metadata_path)
+            and not self._sampler.supports_buffer_checkpoint
         ):
+            raise RuntimeError(
+                "The checkpoint contains native replay state, but the configured "
+                f"sampler {self._async_cfg.sampler.name!r} does not support "
+                "replay-buffer recovery"
+            )
+        if not self._sampler.supports_buffer_checkpoint:
             return
-        buffer_path = os.path.join(self._last_checkpoint_path, "replay_buffer.pt")
-        if not os.path.exists(buffer_path):
+        if not os.path.exists(metadata_path):
+            legacy_path = os.path.join(
+                self._last_checkpoint_path, LEGACY_REPLAY_BUFFER_FILENAME
+            )
+            if os.path.exists(legacy_path):
+                raise RuntimeError(
+                    "Checkpoint contains legacy replay_buffer.pt state, which "
+                    "predates authoritative native TQ replay recovery. Resume it "
+                    "with the older implementation or explicitly start without "
+                    "restoring buffered rollouts."
+                )
             print(
-                f"⚠️ No replay buffer checkpoint found at {buffer_path}. "
+                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",
@@ -332,13 +436,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,
@@ -352,8 +461,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:
@@ -743,6 +863,7 @@ async def _save_checkpoint(
         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
@@ -806,27 +927,27 @@ async def _save_checkpoint(
             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 0450c21933..e9de1ba7d4 100644
--- a/nemo_rl/algorithms/single_controller_utils/setup.py
+++ b/nemo_rl/algorithms/single_controller_utils/setup.py
@@ -31,7 +31,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 MasterConfig as GrpoMasterConfig
 from nemo_rl.algorithms.grpo import (
     GRPOSaveState,
@@ -92,6 +99,95 @@ 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.get(
+            "trainer_version", save_state["current_step"]
+        ),
+        "single_controller_epoch": save_state["current_epoch"],
+        "partition_id": partition_id,
+        "sampler_name": sampler_name,
+        "mode": "authoritative",
+        "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
+    }
+    mismatches = {
+        key: {"checkpoint": metadata.get(key), "expected": expected}
+        for key, expected in expected_values.items()
+        if metadata.get(key) != expected
+    }
+    if mismatches:
+        raise ValueError(
+            "Native TQ checkpoint metadata does not match the trainer "
+            f"checkpoint: {mismatches}"
+        )
+    manifest_digest = metadata.get("replay_manifest_digest")
+    if not isinstance(manifest_digest, str) or not manifest_digest:
+        raise ValueError(
+            "Native TQ checkpoint metadata is missing replay_manifest_digest"
+        )
+    group_count = metadata.get("replay_group_count")
+    if not isinstance(group_count, int) or group_count < 0:
+        raise ValueError(
+            "Native TQ checkpoint metadata has invalid replay_group_count: "
+            f"{group_count!r}"
+        )
+    print(
+        "📦 Native TQ checkpoint restored and validated: "
+        f"groups={group_count}",
+        flush=True,
+    )
+    return metadata
 
 
 def _build_clusters(
@@ -318,6 +414,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"
@@ -432,6 +538,16 @@ def setup_single_controller(
             generation = gen_future.result()
             policy = policy_future.result()
 
+    # Native TQ restore must run through the bootstrap client before the SC
+    # client is created and before any rollout/train data-plane operation.
+    data_plane_checkpoint_metadata = _maybe_restore_native_data_plane_checkpoint(
+        policy,
+        last_checkpoint_path=last_checkpoint_path,
+        save_state=save_state,
+        partition_id=partition_id,
+        sampler_name=master_config.async_rl.sampler.name,
+    )
+
     # ==========================
     # NeMo-Gym actor (after generation is up so OpenAI URLs are available)
     # ==========================
@@ -513,4 +629,5 @@ def setup_single_controller(
         partition_id=partition_id,
         save_state=save_state,
         last_checkpoint_path=last_checkpoint_path,
+        data_plane_checkpoint_metadata=data_plane_checkpoint_metadata,
     )
diff --git a/nemo_rl/data_plane/adapters/noop.py b/nemo_rl/data_plane/adapters/noop.py
index f01b40c098..34d2f76c7d 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 0d9043b0af..d2d8038919 100644
--- a/nemo_rl/data_plane/adapters/transfer_queue.py
+++ b/nemo_rl/data_plane/adapters/transfer_queue.py
@@ -541,7 +541,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 ───────────────────────────────────────────────
@@ -774,6 +774,12 @@ def get_samples(
         promoted_1d_fields = _promoted_1d_fields_from_tags(wire_tags, tensor_fields)
         return _from_wire(td, promoted_1d_fields)
 
+    def list_sample_ids(self, partition_id: str) -> list[str]:
+        """List TQ keys in ``partition_id`` without fetching tensor payloads."""
+        self._mark_data_operation_started()
+        listing = tq.kv_list(partition_id=partition_id)
+        return sorted(listing.get(partition_id, {}).keys())
+
     def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None:
         cleared_via_none = sample_ids is None
         if sample_ids is None:
diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py
index 6bb23d1cbd..465f42b706 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 d7bcd88fca..0e66d4c1d9 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 61da43511c..43ca603cdb 100644
--- a/nemo_rl/models/policy/tq_policy.py
+++ b/nemo_rl/models/policy/tq_policy.py
@@ -33,6 +33,7 @@
 from collections import defaultdict
 from contextlib import nullcontext
 from dataclasses import replace
+from pathlib import Path
 from typing import Any, Optional
 
 import ray
@@ -134,6 +135,12 @@ def __init__(
 
     # ── lifecycle ──────────────────────────────────────────────────────
 
+    def load_data_plane_checkpoint(
+        self, checkpoint_dir: str | Path
+    ) -> dict[str, Any]:
+        """Restore TQ through the clean bootstrap client during SC setup."""
+        return self.dp_client.load_checkpoint(checkpoint_dir)
+
     def shutdown(self) -> bool:  # type: ignore[override]
         """Close the TQ client before shutting down the worker group."""
         try:
diff --git a/pyrefly.toml b/pyrefly.toml
index 49fc4f2ef5..1b7fc28002 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -114,6 +114,7 @@ project-includes = [
   "nemo_rl/data_plane/adapters/__init__.py",
   "nemo_rl/data_plane/adapters/noop.py",
   "nemo_rl/data_plane/adapters/transfer_queue.py",
+  "nemo_rl/data_plane/async_utils.py",
   "nemo_rl/data_plane/codec.py",
   "nemo_rl/data_plane/column_io.py",
   "nemo_rl/data_plane/factory.py",
diff --git a/tests/functional/L1_Functional_Tests_SingleController.sh b/tests/functional/L1_Functional_Tests_SingleController.sh
index 2a6fad63ee..ea53d7fe6b 100755
--- a/tests/functional/L1_Functional_Tests_SingleController.sh
+++ b/tests/functional/L1_Functional_Tests_SingleController.sh
@@ -36,6 +36,7 @@ run_test() {
 
 run_test fast uv run --no-sync bash ./tests/functional/grpo_dp_single_controller.sh
 run_test fast uv run --no-sync bash ./tests/functional/grpo_async_gym_single_controller.sh
+run_test uv run --no-sync bash ./tests/functional/grpo_dp_single_controller_tq_recovery.sh
 
 cd ${PROJECT_ROOT}/tests
 if compgen -G ".coverage*" > /dev/null; then
diff --git a/tests/functional/grpo_dp_single_controller.sh b/tests/functional/grpo_dp_single_controller.sh
index 855760445d..f62ff6aac0 100755
--- a/tests/functional/grpo_dp_single_controller.sh
+++ b/tests/functional/grpo_dp_single_controller.sh
@@ -22,7 +22,7 @@ rm -rf $EXP_DIR $LOG_DIR
 mkdir -p $EXP_DIR $LOG_DIR
 
 cd $PROJECT_ROOT
-uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+uv run --group test coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
     $PROJECT_ROOT/examples/run_grpo_single_controller.py \
     policy.model_name=Qwen/Qwen3-0.6B \
     grpo.num_prompts_per_step=2 \
@@ -47,11 +47,13 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE
     $@ \
     2>&1 | tee $RUN_LOG
 
-uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+if [[ "${RUN_CONVERGENCE_CHECKS:-1}" == "1" ]]; then
+    uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
 
-uv run tests/check_metrics.py $JSON_METRICS \
-    'max(data["train/gen_kl_error"]) < 0.002' \
-    'min(data["train/probs_ratio_clamped_min"]) > 0.79' \
-    'max(data["train/probs_ratio_clamped_min"]) < 1.21' \
-    'min(data["train/probs_ratio_clamped_max"]) > 0.79' \
-    'max(data["train/probs_ratio_clamped_max"]) < 1.21'
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'max(data["train/gen_kl_error"]) < 0.002' \
+        'min(data["train/probs_ratio_clamped_min"]) > 0.79' \
+        'max(data["train/probs_ratio_clamped_min"]) < 1.21' \
+        'min(data["train/probs_ratio_clamped_max"]) > 0.79' \
+        'max(data["train/probs_ratio_clamped_max"]) < 1.21'
+fi
diff --git a/tests/functional/grpo_dp_single_controller_tq_recovery.sh b/tests/functional/grpo_dp_single_controller_tq_recovery.sh
new file mode 100755
index 0000000000..671b67dd4e
--- /dev/null
+++ b/tests/functional/grpo_dp_single_controller_tq_recovery.sh
@@ -0,0 +1,53 @@
+#!/bin/bash
+# Two-process functional test for native TQ + metadata-only replay recovery.
+
+set -eou pipefail
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
+BASE_TEST=$SCRIPT_DIR/grpo_dp_single_controller.sh
+TEST_DIR=$SCRIPT_DIR/grpo_dp_single_controller_tq_recovery
+CHECKPOINT_DIR=$TEST_DIR/checkpoints
+BASE_RUN_LOG=$SCRIPT_DIR/grpo_dp_single_controller/run.log
+
+rm -rf "$TEST_DIR"
+mkdir -p "$TEST_DIR"
+
+COMMON_OVERRIDES=(
+    checkpointing.enabled=true
+    checkpointing.checkpoint_dir="$CHECKPOINT_DIR"
+    checkpointing.save_period=1
+    data_plane.checkpointing_enabled=true
+    async_rl.sampler.name=windowed
+    '~async_rl.sampler.max_lookahead_versions'
+    '+async_rl.sampler.max_staleness_versions=1'
+    async_rl.max_inflight_prompts=8
+    async_rl.max_buffered_rollouts=8
+)
+
+echo "=== Phase 1: save an authoritative native TQ checkpoint ==="
+# Keep the two-step training horizon identical across both processes so the
+# Megatron optimizer scheduler can be restored. The timeout makes phase 1 save
+# after its first completed step and exit early, simulating an interrupted job.
+RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \
+    "${COMMON_OVERRIDES[@]}" \
+    grpo.max_num_steps=2 \
+    checkpointing.checkpoint_must_save_by=0:0:0:1
+
+test -d "$CHECKPOINT_DIR/step_1/data_plane"
+test -f "$CHECKPOINT_DIR/step_1/replay_buffer_metadata.pt"
+test ! -f "$CHECKPOINT_DIR/step_1/replay_buffer.pt"
+uv run --no-sync python -c \
+    'import json, sys; metadata = json.load(open(sys.argv[1]))["user_metadata"]; assert metadata["mode"] == "authoritative"; assert metadata["replay_group_count"] > 0, metadata' \
+    "$CHECKPOINT_DIR/step_1/data_plane/metadata.json"
+
+echo "=== Phase 2: start a fresh process, restore TQ, and train one more step ==="
+RUN_CONVERGENCE_CHECKS=0 bash "$BASE_TEST" \
+    "${COMMON_OVERRIDES[@]}" grpo.max_num_steps=2
+
+grep -q "Native TQ checkpoint restored and validated" "$BASE_RUN_LOG"
+grep -q "Native TQ replay inventory validated" "$BASE_RUN_LOG"
+grep -Eq "Restored [1-9][0-9]* replay group" "$BASE_RUN_LOG"
+test -d "$CHECKPOINT_DIR/step_2/data_plane"
+test -f "$CHECKPOINT_DIR/step_2/replay_buffer_metadata.pt"
+
+echo "Native TQ recovery functional test passed."
diff --git a/tests/unit/data_plane/test_architecture_invariants.py b/tests/unit/data_plane/test_architecture_invariants.py
index 7ec4ac0984..a5167f053f 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_interface_contract.py b/tests/unit/data_plane/test_interface_contract.py
index 3a4009b164..7c9c900259 100644
--- a/tests/unit/data_plane/test_interface_contract.py
+++ b/tests/unit/data_plane/test_interface_contract.py
@@ -65,8 +65,10 @@ def test_register_put_get_clear(client: DataPlaneClient):
 
     out = client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"])
     assert torch.equal(out["x"], torch.arange(4))
+    assert client.list_sample_ids("p") == keys
 
     client.clear_samples(sample_ids=None, partition_id="p")
+    assert client.list_sample_ids("p") == []
     with pytest.raises(KeyError):
         client.get_samples(sample_ids=keys, partition_id="p", select_fields=["x"])
 
diff --git a/tests/unit/data_plane/test_observability.py b/tests/unit/data_plane/test_observability.py
index 2cdbdb1a62..13a5f59a38 100644
--- a/tests/unit/data_plane/test_observability.py
+++ b/tests/unit/data_plane/test_observability.py
@@ -89,6 +89,22 @@ def test_register_and_clear_recorded(wrapped_client):
     assert ops.count("clear") == 1
 
 
+def test_list_sample_ids_is_forwarded_and_recorded(wrapped_client):
+    client, events = wrapped_client
+    client.register_partition(
+        partition_id="p", fields=["x"], num_samples=2, consumer_tasks=["r"]
+    )
+    client.put_samples(
+        sample_ids=["b", "a"],
+        partition_id="p",
+        fields=TensorDict({"x": torch.ones(2)}, batch_size=[2]),
+    )
+
+    assert client.list_sample_ids("p") == ["a", "b"]
+    assert events[-1]["op"] == "list_sample_ids"
+    assert events[-1]["status"] == "ok"
+
+
 def test_error_status_recorded_and_reraised(wrapped_client):
     """Decorator does NOT swallow errors — re-raise after recording."""
     client, events = wrapped_client
diff --git a/tests/unit/data_plane/test_smoke.py b/tests/unit/data_plane/test_smoke.py
index 25505d1053..e373adf076 100644
--- a/tests/unit/data_plane/test_smoke.py
+++ b/tests/unit/data_plane/test_smoke.py
@@ -89,6 +89,7 @@ def test_dataplane_client_abc_surface() -> None:
         # direct-by-key
         "put_samples",
         "get_samples",
+        "list_sample_ids",
         "clear_samples",
         # lifecycle
         "close",
diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py
index 1c88a46e5d..bc0fa2dd22 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 4813178e39..05ee6fee7d 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -485,7 +485,9 @@ data_plane:
   storage_capacity: 1000000          # max samples retained per partition
   num_storage_units: 2               # storage shards
   claim_meta_poll_interval_s: 0.5    # blocking-claim poll cadence
-  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 1c5d9a8205..3d63f5915e 100644
--- a/tests/unit/single_controller/_dp_fakes.py
+++ b/tests/unit/single_controller/_dp_fakes.py
@@ -147,6 +147,9 @@ def clear_samples(self, sample_ids: list[str], partition_id: str) -> Any:
             )
         )
 
+    def list_sample_ids(self, partition_id: str) -> list[str]:
+        return ray.get(self._handle.list_sample_ids.remote(partition_id))
+
     @staticmethod
     def _padded(td: TensorDict) -> TensorDict:
         out: dict[str, torch.Tensor] = {}
diff --git a/tests/unit/single_controller/test_sampler_interface.py b/tests/unit/single_controller/test_sampler_interface.py
index c99be7fb8b..35f1ca0c10 100644
--- a/tests/unit/single_controller/test_sampler_interface.py
+++ b/tests/unit/single_controller/test_sampler_interface.py
@@ -27,10 +27,12 @@
 import pytest
 
 from nemo_rl.algorithms.async_utils.staleness_sampler import (
+    CustomSamplerConfig,
     InOrderSampler,
     InOrderSamplerConfig,
     PromptGroupSampler,
     WeightFifoSampler,
+    WeightFifoSamplerConfig,
     WindowedSampler,
     WindowedSamplerConfig,
     create_sampler,
@@ -151,6 +153,24 @@ def test_unready_slot_is_never_evicted(self):
 
 
 class TestFactory:
+    @pytest.mark.parametrize(
+        ("config", "expected"),
+        [
+            (WindowedSamplerConfig(), True),
+            (WeightFifoSamplerConfig(), False),
+            (InOrderSamplerConfig(), False),
+            (
+                CustomSamplerConfig(target=f"{__name__}:EchoSampler"),
+                None,
+            ),
+        ],
+    )
+    def test_config_declares_checkpoint_capability_without_serializing_it(
+        self, config, expected
+    ):
+        assert config.supports_buffer_checkpoint is expected
+        assert "supports_buffer_checkpoint" not in config.model_dump()
+
     def test_windowed_config_builds_windowed(self):
         s = create_sampler(
             FakeBuffer(), WindowedSamplerConfig(max_staleness_versions=3)
@@ -164,24 +184,27 @@ def test_in_order_config_builds_in_order(self):
         assert s.max_lookahead_versions == 2
 
     def test_weight_fifo_config_builds_weight_fifo(self):
-        from nemo_rl.algorithms.async_utils.staleness_sampler import (
-            WeightFifoSamplerConfig,
-        )
-
         s = create_sampler(
             FakeBuffer(), WeightFifoSamplerConfig(max_staleness_versions=4)
         )
         assert isinstance(s, WeightFifoSampler)
         assert s.max_staleness_versions == 4
 
+    def test_factory_rejects_config_implementation_capability_drift(
+        self, monkeypatch
+    ):
+        monkeypatch.setattr(
+            WindowedSampler,
+            "supports_buffer_checkpoint",
+            property(lambda _self: False),
+        )
+        with pytest.raises(RuntimeError, match="disagrees with"):
+            create_sampler(FakeBuffer(), WindowedSamplerConfig())
+
 
 class TestCustomFqnSampler:
     def test_custom_target_loads_out_of_repo_sampler(self):
         # A user sampler defined anywhere importable; here, this test module.
-        from nemo_rl.algorithms.async_utils.staleness_sampler import (
-            CustomSamplerConfig,
-        )
-
         s = create_sampler(
             FakeBuffer(),
             CustomSamplerConfig(
diff --git a/tests/unit/single_controller/test_sc_checkpointing.py b/tests/unit/single_controller/test_sc_checkpointing.py
index 0e2f09f5a8..6073f3a132 100644
--- a/tests/unit/single_controller/test_sc_checkpointing.py
+++ b/tests/unit/single_controller/test_sc_checkpointing.py
@@ -25,7 +25,7 @@
   - dataloader state: train_dataloader.pt written at save, position
     round-trip through a real StatefulDataLoader, dataset-swap guard,
     setup restore wiring + missing-file fresh-position fallback;
-  - replay buffer persistence gated on sampler.supports_buffer_checkpoint;
+  - 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).
 """
@@ -45,6 +45,14 @@
 import yaml
 from torchdata.stateful_dataloader import StatefulDataLoader
 
+from nemo_rl.algorithms.async_utils.replay_buffer import (
+    DATA_PLANE_CHECKPOINT_SCHEMA_VERSION,
+    DataPlaneCheckpointBarrier,
+    LEGACY_REPLAY_BUFFER_FILENAME,
+    REPLAY_BUFFER_METADATA_FILENAME,
+    REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
+    REPLAY_BUFFER_METADATA_STORAGE,
+)
 from nemo_rl.algorithms.async_utils.staleness_sampler import (
     InOrderSamplerConfig,
     WindowedSamplerConfig,
@@ -195,11 +203,21 @@ def required_buffer_capacity(self, groups_per_step: int) -> Optional[int]:
 
 
 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())
@@ -263,21 +281,30 @@ class _FakeTQBuffer:
 
     def __init__(
         self,
-        state: Optional[dict[str, Any]] = None,
+        metadata_state: Optional[dict[str, Any]] = None,
         load_return: int = 0,
     ) -> None:
-        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,
@@ -286,6 +313,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(
             {
@@ -293,6 +321,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
@@ -332,7 +361,7 @@ def _actor_master_config(
     checkpoint_must_save_by: Optional[str] = None,
     num_prompts_per_step: int = 2,
     max_num_epochs: int = 1,
-    buffer_checkpoint: bool = True,
+    buffer_checkpoint: bool = False,
     data_plane_checkpoint: bool = False,
 ) -> MasterConfig:
     """MasterConfig for in-process SingleControllerActor tests.
@@ -404,6 +433,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(),
@@ -423,6 +453,7 @@ def _make_actor_args(
             save_state if save_state is not None else _default_grpo_save_state()
         ),
         last_checkpoint_path=last_checkpoint_path,
+        data_plane_checkpoint_metadata=data_plane_checkpoint_metadata,
     )
 
 
@@ -508,6 +539,19 @@ 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 = _default_grpo_save_state()
+        save_state["current_step"] = 7
+        save_state["trainer_version"] = 11
+
+        actor = _ACTOR_CLS(
+            _actor_master_config(tmp_path), _make_actor_args(save_state=save_state)
+        )
+
+        assert actor._train_steps == 7
+        assert actor._trainer_version == 11
+        assert actor._sampler._dispatch_index == 10
+
     def test_fresh_start_defaults(self, tmp_path):
         actor = _ACTOR_CLS(_actor_master_config(tmp_path), _make_actor_args())
 
@@ -553,6 +597,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.
@@ -642,16 +687,46 @@ def test_timeout_saves_and_stops_training_early(self, tmp_path):
         assert _step_dir_names(tmp_path / "checkpoints") == {"step_1"}
 
 
-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,
@@ -664,19 +739,107 @@ 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(
@@ -1124,19 +1287,19 @@ def test_setup_missing_dataloader_state_starts_fresh(
 
 
 class TestReplayBufferPersistence:
-    def test_save_writes_replay_buffer_when_sampler_supports_it(self, tmp_path):
-        mc = _actor_master_config(tmp_path, max_num_steps=2, save_period=2)
-        envelope = {"groups": [], "sentinel": "abc"}
-        buffer = _FakeTQBuffer(state=envelope)
-
-        _run_train_pump(mc, _make_actor_args(tq_buffer=buffer))
+    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,
+        )
 
-        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.
-        assert buffer.state_dict_calls == [4]
+        with pytest.raises(
+            ValueError,
+            match="replay-checkpoint-capable sampler requires",
+        ):
+            _ACTOR_CLS(mc, _make_actor_args())
 
     def test_no_replay_buffer_with_gated_sampler(self, tmp_path):
         mc = _actor_master_config(
@@ -1149,19 +1312,78 @@ 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 buffer.state_dict_calls == []
+        assert not (
+            ckpt_dir / "step_2" / REPLAY_BUFFER_METADATA_FILENAME
+        ).exists()
 
-    def test_run_restores_replay_buffer_and_permits(self, tmp_path):
+    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, last_checkpoint_path=str(ckpt_dir)),
+            _make_actor_args(
+                tq_buffer=buffer,
+                dp_client=_FakeDPClient(sample_ids=sample_ids),
+                last_checkpoint_path=str(ckpt_dir),
+                data_plane_checkpoint_metadata=tq_metadata,
+            ),
         )
 
         assert buffer.load_calls == [
@@ -1170,36 +1392,95 @@ 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_run_restore_at_full_capacity_does_not_hang(self, tmp_path):
-        # K == max_buffered_rollouts: the acquisitions must all complete
-        # without waiting (no pump is running yet to release permits).
+    @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()
-        torch.save({"groups": []}, ckpt_dir / "replay_buffer.pt")
-        mc = _actor_master_config(tmp_path, max_num_steps=0)
-        buffer = _FakeTQBuffer(load_return=4)
+        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, result = _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 len(buffer.load_calls) == 1
-        assert actor._buffer_capacity._value == 0
-        assert result["train_steps"] == 0
+        with pytest.raises(RuntimeError, match="native TQ checkpoint was not restored"):
+            _run_actor_run(
+                mc,
+                _make_actor_args(last_checkpoint_path=str(ckpt_dir)),
+            )
 
-    def test_run_missing_replay_buffer_file_starts_empty(self, tmp_path, monkeypatch):
-        # Resuming from a checkpoint written by a gated-sampler run: no
-        # replay_buffer.pt.
+    def test_run_missing_native_replay_metadata_starts_empty(
+        self, tmp_path, monkeypatch
+    ):
         ckpt_dir = tmp_path / "resume_ckpt"
         ckpt_dir.mkdir()
-        mc = _actor_master_config(tmp_path, max_num_steps=0)
+        mc = _actor_master_config(
+            tmp_path,
+            max_num_steps=0,
+            buffer_checkpoint=True,
+            data_plane_checkpoint=True,
+        )
         buffer = _FakeTQBuffer()
         printed: list[str] = []
         monkeypatch.setattr(
@@ -1214,21 +1495,18 @@ def test_run_missing_replay_buffer_file_starts_empty(self, tmp_path, monkeypatch
 
         assert buffer.load_calls == []
         assert actor._buffer_capacity._value == 4  # zero permits consumed
-        assert any("No replay buffer checkpoint found" in line for line in printed)
+        assert any("No native replay metadata found" in line for line in printed)
 
-    def test_run_no_restore_with_gated_sampler(self, tmp_path):
-        # File present but the sampler doesn't support buffer checkpointing:
-        # nothing is restored.
+    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.pt")
-        mc = _actor_master_config(tmp_path, max_num_steps=0, buffer_checkpoint=False)
-        buffer = _FakeTQBuffer(load_return=2)
-
-        actor, _ = _run_actor_run(
-            mc,
-            _make_actor_args(tq_buffer=buffer, last_checkpoint_path=str(ckpt_dir)),
+        torch.save({"groups": []}, ckpt_dir / REPLAY_BUFFER_METADATA_FILENAME)
+        mc = _actor_master_config(
+            tmp_path, max_num_steps=0, buffer_checkpoint=False
         )
 
-        assert buffer.load_calls == []
-        assert actor._buffer_capacity._value == 4
+        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 51cc4d94c3..e36f724913 100644
--- a/tests/unit/single_controller/test_single_controller_setup.py
+++ b/tests/unit/single_controller/test_single_controller_setup.py
@@ -16,11 +16,21 @@
 
 from __future__ import annotations
 
+from typing import Any, Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 import nemo_rl.algorithms.single_controller_utils.setup as sc_setup_mod
+from nemo_rl.algorithms.async_utils.replay_buffer import (
+    DATA_PLANE_CHECKPOINT_DIR,
+    DATA_PLANE_CHECKPOINT_SCHEMA_VERSION,
+    LEGACY_REPLAY_BUFFER_FILENAME,
+    REPLAY_BUFFER_METADATA_FILENAME,
+    REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
+)
+from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig
+from nemo_rl.algorithms.grpo import _default_grpo_save_state
 from nemo_rl.algorithms.loss import ClippedPGLossConfig
 from nemo_rl.algorithms.single_controller_utils import (
     AsyncRLConfig,
@@ -98,6 +108,27 @@ def _make_master_config(
     )
 
 
+def _native_tq_metadata(
+    *, step: int = 3, trainer_version: Optional[int] = None, epoch: int = 1
+) -> dict[str, Any]:
+    return {
+        "data_plane_checkpoint_schema_version": (
+            DATA_PLANE_CHECKPOINT_SCHEMA_VERSION
+        ),
+        "single_controller_train_steps": step,
+        "single_controller_trainer_version": (
+            step if trainer_version is None else trainer_version
+        ),
+        "single_controller_epoch": epoch,
+        "partition_id": "rollout_data",
+        "sampler_name": "in_order",
+        "mode": "authoritative",
+        "replay_metadata_schema_version": REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
+        "replay_manifest_digest": "digest-1",
+        "replay_group_count": 2,
+    }
+
+
 @pytest.fixture
 def patched_factories():
     """Patch every external factory setup calls.
@@ -229,6 +260,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"):
@@ -465,3 +516,149 @@ def test_nemo_gym_rejects_non_vllm_backend(self, patched_factories, backend):
         ):
             setup_single_controller(mc, MagicMock(pad_token_id=0))
         mock_spinup.assert_not_called()
+
+
+class TestNativeTQRecoverySetup:
+    def test_setup_loads_tq_before_creating_single_controller_client(
+        self, tmp_path, patched_factories
+    ):
+        checkpoint_path = tmp_path / "step_3"
+        (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True)
+        (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch()
+        save_state = _default_grpo_save_state()
+        save_state["current_step"] = 3
+        save_state["current_epoch"] = 1
+        policy = patched_factories["_build_trainer"].return_value
+        events: list[str] = []
+        policy.load_data_plane_checkpoint.side_effect = (
+            lambda checkpoint_dir: events.append("load") or _native_tq_metadata()
+        )
+        patched_factories["build_data_plane_client"].side_effect = (
+            lambda *args, **kwargs: events.append("build")
+            or MagicMock(name="dp_client")
+        )
+        checkpointer = MagicMock()
+        checkpointer.get_latest_checkpoint_path.return_value = str(checkpoint_path)
+        checkpointer.load_training_info.return_value = save_state
+        checkpointer.get_resume_paths.return_value = (None, None)
+        mc = _make_master_config()
+
+        with patch.object(
+            sc_setup_mod, "CheckpointManager", return_value=checkpointer
+        ):
+            actor_args = setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        assert events == ["load", "build"]
+        assert actor_args.data_plane_checkpoint_metadata == _native_tq_metadata()
+
+    def test_loads_authoritative_tq_checkpoint_when_metadata_sidecar_exists(
+        self, tmp_path
+    ):
+        checkpoint_path = tmp_path / "step_3"
+        (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True)
+        (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch()
+        policy = MagicMock()
+        metadata = _native_tq_metadata()
+        policy.load_data_plane_checkpoint.return_value = metadata
+        save_state = {
+            "current_step": 3,
+            "current_epoch": 1,
+        }
+
+        restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+            policy,
+            last_checkpoint_path=str(checkpoint_path),
+            save_state=save_state,
+            partition_id="rollout_data",
+            sampler_name="in_order",
+        )
+
+        assert restored == metadata
+        policy.load_data_plane_checkpoint.assert_called_once_with(
+            checkpoint_path / DATA_PLANE_CHECKPOINT_DIR
+        )
+
+    def test_validates_trainer_version_independently_from_train_step(self, tmp_path):
+        checkpoint_path = tmp_path / "step_3"
+        (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True)
+        (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch()
+        policy = MagicMock()
+        metadata = _native_tq_metadata(step=3, trainer_version=7)
+        policy.load_data_plane_checkpoint.return_value = metadata
+
+        restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+            policy,
+            last_checkpoint_path=str(checkpoint_path),
+            save_state={
+                "current_step": 3,
+                "trainer_version": 7,
+                "current_epoch": 1,
+            },
+            partition_id="rollout_data",
+            sampler_name="in_order",
+        )
+
+        assert restored == metadata
+
+    def test_legacy_replay_checkpoint_is_rejected(self, tmp_path):
+        checkpoint_path = tmp_path / "step_3"
+        checkpoint_path.mkdir()
+        (checkpoint_path / LEGACY_REPLAY_BUFFER_FILENAME).touch()
+        policy = MagicMock()
+
+        with pytest.raises(RuntimeError, match="legacy replay_buffer.pt"):
+            sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+                policy,
+                last_checkpoint_path=str(checkpoint_path),
+                save_state={"current_step": 3, "current_epoch": 1},
+                partition_id="rollout_data",
+                sampler_name="in_order",
+            )
+
+        policy.load_data_plane_checkpoint.assert_not_called()
+
+    def test_checkpoint_without_replay_artifacts_does_not_load_tq(self, tmp_path):
+        checkpoint_path = tmp_path / "step_3"
+        checkpoint_path.mkdir()
+        policy = MagicMock()
+
+        restored = sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+            policy,
+            last_checkpoint_path=str(checkpoint_path),
+            save_state={"current_step": 3, "current_epoch": 1},
+            partition_id="rollout_data",
+            sampler_name="in_order",
+        )
+
+        assert restored is None
+        policy.load_data_plane_checkpoint.assert_not_called()
+
+    def test_metadata_sidecar_requires_matching_tq_directory(self, tmp_path):
+        checkpoint_path = tmp_path / "step_3"
+        checkpoint_path.mkdir()
+        (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch()
+
+        with pytest.raises(FileNotFoundError, match="matching native TQ checkpoint"):
+            sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+                MagicMock(),
+                last_checkpoint_path=str(checkpoint_path),
+                save_state={"current_step": 3, "current_epoch": 1},
+                partition_id="rollout_data",
+                sampler_name="in_order",
+            )
+
+    def test_rejects_tq_checkpoint_from_different_training_step(self, tmp_path):
+        checkpoint_path = tmp_path / "step_3"
+        (checkpoint_path / DATA_PLANE_CHECKPOINT_DIR).mkdir(parents=True)
+        (checkpoint_path / REPLAY_BUFFER_METADATA_FILENAME).touch()
+        policy = MagicMock()
+        policy.load_data_plane_checkpoint.return_value = _native_tq_metadata(step=2)
+
+        with pytest.raises(ValueError, match="does not match the trainer checkpoint"):
+            sc_setup_mod._maybe_restore_native_data_plane_checkpoint(
+                policy,
+                last_checkpoint_path=str(checkpoint_path),
+                save_state={"current_step": 3, "current_epoch": 1},
+                partition_id="rollout_data",
+                sampler_name="in_order",
+            )
diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py
index cbad17862a..1e1336b10f 100644
--- a/tests/unit/single_controller/test_tq_replay_buffer.py
+++ b/tests/unit/single_controller/test_tq_replay_buffer.py
@@ -24,7 +24,13 @@
 import torch
 
 import nemo_rl.algorithms.async_utils.replay_buffer as _replay_buffer_module
-from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer
+from nemo_rl.algorithms.async_utils.replay_buffer import (
+    REPLAY_BUFFER_METADATA_SCHEMA_VERSION,
+    REPLAY_BUFFER_METADATA_STORAGE,
+    DataPlaneCheckpointBarrier,
+    TQReplayBuffer,
+    replay_manifest_digest,
+)
 from nemo_rl.data_plane import KVBatchMeta
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.experience.interfaces import PromptGroupRecord
@@ -102,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],
@@ -117,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:
@@ -157,12 +167,14 @@ def _make_record() -> PromptGroupRecord:
 def _make_buffer(
     dp: FakeDataPlaneClient,
     *,
-    checkpoint_lock: asyncio.Lock | None = None,
+    checkpoint_barrier: DataPlaneCheckpointBarrier | None = None,
 ) -> TQReplayBuffer:
     buffer = TQReplayBuffer(
         dp, partition_id="rollout_data", pad_value_dict={"token_ids": 0}
     )
-    buffer.set_data_plane_checkpoint_lock(checkpoint_lock or asyncio.Lock())
+    buffer.set_data_plane_checkpoint_barrier(
+        checkpoint_barrier or DataPlaneCheckpointBarrier()
+    )
     return buffer
 
 
@@ -182,7 +194,87 @@ def _add_group(
     )
 
 
+class TestDataPlaneCheckpointBarrier:
+    def test_mutations_run_concurrently_without_checkpoint(self):
+        async def exercise() -> None:
+            barrier = DataPlaneCheckpointBarrier()
+            both_entered = asyncio.Event()
+            release = asyncio.Event()
+            active = 0
+
+            async def mutate() -> None:
+                nonlocal active
+                async with barrier.mutation():
+                    active += 1
+                    if active == 2:
+                        both_entered.set()
+                    await release.wait()
+                    active -= 1
+
+            tasks = [asyncio.create_task(mutate()) for _ in range(2)]
+            await asyncio.wait_for(both_entered.wait(), timeout=5.0)
+            assert active == 2
+            release.set()
+            await asyncio.gather(*tasks)
+
+        asyncio.run(exercise())
+
+    def test_checkpoint_waits_for_active_mutation(self):
+        async def exercise() -> None:
+            barrier = DataPlaneCheckpointBarrier()
+            mutation_entered = asyncio.Event()
+            release_mutation = asyncio.Event()
+            checkpoint_entered = asyncio.Event()
+
+            async def mutate() -> None:
+                async with barrier.mutation():
+                    mutation_entered.set()
+                    await release_mutation.wait()
+
+            async def checkpoint() -> None:
+                async with barrier.checkpoint():
+                    checkpoint_entered.set()
+
+            mutation_task = asyncio.create_task(mutate())
+            await mutation_entered.wait()
+            checkpoint_task = asyncio.create_task(checkpoint())
+            await asyncio.sleep(0)
+            assert not checkpoint_entered.is_set()
+
+            release_mutation.set()
+            await asyncio.gather(mutation_task, checkpoint_task)
+            assert checkpoint_entered.is_set()
+
+        asyncio.run(exercise())
+
+
 class TestTQReplayBufferReserveCommit:
+    def test_commit_waits_for_active_checkpoint(self):
+        async def exercise() -> None:
+            dp = FakeDataPlaneClient()
+            checkpoint_barrier = DataPlaneCheckpointBarrier()
+            buf = _make_buffer(dp, checkpoint_barrier=checkpoint_barrier)
+            group_id = buf.reserve(weight_version=3)
+
+            async with checkpoint_barrier.checkpoint():
+                commit_task = asyncio.create_task(
+                    buf.commit(
+                        group_id,
+                        _make_record(),
+                        start_weight_version=3,
+                        end_weight_version=3,
+                    )
+                )
+                await asyncio.sleep(0)
+                assert dp.put_calls == []
+                assert buf.ready_list == [False]
+
+            await commit_task
+            assert len(dp.put_calls) == 1
+            assert buf.ready_list == [True]
+
+        asyncio.run(exercise())
+
     def test_commit_clears_rows_when_put_raises_after_writing(self):
         dp = FailAfterPutDataPlaneClient()
         buf = _make_buffer(dp)
@@ -313,7 +405,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,
@@ -326,11 +418,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,
@@ -339,12 +431,13 @@ 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"]]
 
@@ -482,20 +575,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,
     }
 
 
@@ -506,74 +602,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 == []
 
 
 class TestTQReplayBufferLoadPreflight:
@@ -591,20 +674,28 @@ 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):
@@ -612,28 +703,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")
+        self._assert_rejected(
+            _make_metadata_envelope([g0, g1]), match="duplicate sample_id"
+        )
 
+    def test_metadata_only_restore_rejects_tq_digest_mismatch(self):
+        state = _make_metadata_envelope([_make_group_entry("g0", weight=1)])
+        self._assert_rejected(
+            state,
+            match="does not match the loaded TQ checkpoint",
+            expected_manifest_digest="wrong-digest",
+        )
 
-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,
+    def test_metadata_only_restore_rejects_capacity_truncation(self):
+        state = _make_metadata_envelope(
+            [_make_group_entry(f"g{w}", weight=w) for w in (1, 2, 3)]
         )
-        dp = FakeDataPlaneClient()
-        buf = _make_buffer(dp)
-        printed: list[str] = []
-        monkeypatch.setattr(
-            "builtins.print",
-            lambda *args, **kwargs: printed.append(" ".join(str(a) for a in args)),
+        self._assert_rejected(
+            state,
+            match="more replay groups than the current buffer capacity",
+            max_groups=2,
+            expected_manifest_digest=state["manifest_digest"],
         )
-
-        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)