feat: pause SingleController refits and enable KV recomputation by default

docs/guides/single-controller.md

@@ -345,7 +345,7 @@ The SC path splits the async-GRPO loop across a rollout pump and a train pump th
 2. **Actor startup**: `SingleControllerActor` launches `_rollout_pump` and `_train_pump` concurrently as asyncio tasks; both share the same `TQReplayBuffer` and `StalenessSampler`.
 3. **Rollout pump loop**: `sampler.admit` gates dispatch against the current trainer version (returning a `target_step` for `in_order`); the pump then reserves a buffer slot, drives `RolloutManager.generate_and_push`, and commits with the observed `start_weight` / `end_weight`.
 4. **Train pump loop**: `sampler.evict` drops out-of-window groups and `sampler.select` picks the next batch. On PPO, `_value_stage` runs the critic forward, `_advantage_stage` computes advantages, and `_value_train_epochs` runs `ppo.critic_ppo_epochs` critic updates. The TQPolicy split API then runs one optimizer step per RL step on GRPO, or `ppo.ppo_epochs` policy updates on PPO.
-5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, calls the weight synchronizer, and re-opens the rollout pump for the next version.
+5. **Weight sync**: after each optimizer step the pump bumps the trainer version, clears rollout permission, asks the generation backend to pause in-flight work, calls the weight synchronizer, resumes generation, and re-opens the rollout pump for the next version. vLLM's native pause preserves active request state; backends without pause support warn and retain their existing refit behavior.
 
 ## Relation to Legacy Async GRPO
 
@@ -372,7 +372,7 @@ Do not carry `max_num_epochs: -1` across either. [ppo.md](./ppo.md#asynchronous-
 | `enabled: true` | Implicit — SC is always async; use `sampler.max_lookahead_versions: 0` for sync semantics, `>= 1` for async |
 | `max_trajectory_age_steps: N` | `sampler.name: in_order` with `sampler.max_lookahead_versions: N` |
 | `warmup_generation_lead_steps` (PPO) | `sampler.warmup_lookahead_versions` — the lookahead to use while critic warmup is in progress |
-| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same) |
+| `recompute_kv_cache_after_weight_updates` | `recompute_kv_cache_after_weight_updates` (same); when vLLM pauses natively, the preserved requests recompute their KV cache after refit |
 | `in_flight_weight_updates` | Always effectively true; `false`-equivalent behavior is not yet supported (drain-gate tracked in [issue #2625](https://github.com/NVIDIA-NeMo/RL/issues/2625)) |
 | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run |
 | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` |

nemo_rl/algorithms/single_controller.py

@@ -4635,17 +4635,22 @@ async def _sync_weights(
         *,
         calibration_data: Optional[BatchedDataDict[Any]] = None,
     ) -> int:
-        """Pause new rollout dispatches, synchronize weights, resume.
+        """Pause rollout dispatch and generation, synchronize weights, resume.
 
-        SC owns the pause gate. vLLM serves through the refit; it supports live weight updates.
-        A colocated engine is instead already stood down: the synchronizer's sync is the wake,
-        and step 4 resumes dispatch.
+        SC owns the rollout-dispatch gate and asks every generation backend to
+        pause in-flight work through the common refit lifecycle hooks. Backends
+        without native pause support warn and retain their existing in-flight
+        update behavior. Continuous-serving vLLM preserves request state while
+        paused. A colocated engine is already stood down before training; the
+        synchronizer's sync remains its weight-carrying wake.
 
         Flow:
           1. _rollout_permitted.clear()  — no new dispatches
-          2. Optionally calibrate FP8 KV-cache scales.
-          3. weight_synchronizer.sync_weights(kv_scales=...)
-          4. _rollout_permitted.set()   — resume
+          2. Reconcile refit membership and optionally calibrate FP8 KV-cache scales.
+          3. generation.pause_generation_for_refit(clear_cache=...)
+          4. weight_synchronizer.sync_weights(kv_scales=...)
+          5. generation.resume_generation_after_refit()
+          6. _rollout_permitted.set()   — resume dispatch
 
         Args:
             calibration_data: Optional data used to calibrate FP8 KV-cache
@@ -4704,7 +4709,21 @@ async def _sync_weights(
         # shard the communicator was deliberately built without.
         participants = self._refit_participants()
 
+        clear_cache = self._async_cfg.recompute_kv_cache_after_weight_updates
+        timeout_s = self._refit_await_budget_s()
+        generation_paused_for_refit = False
+        print("⏸️ Requesting generation pause before refit", flush=True)
         try:
+            generation_paused_for_refit = await asyncio.to_thread(
+                self._gen.pause_generation_for_refit,
+                clear_cache=clear_cache,
+                timeout_s=timeout_s,
+            )
+            if generation_paused_for_refit:
+                print(
+                    f"   {len(self._inflight_by_group_id)} in-flight rollout group(s) paused",
+                    flush=True,
+                )
             await self._sync_weights_within(kv_scales, "first")
         except (RefitAborted, RayActorError) as failure:
             # DETECT AND FAIL FAST, because this one cannot be recovered from.
@@ -4732,11 +4751,14 @@ async def _sync_weights(
                 raise
             with self._recovery_window():
                 await self._recover_from_failed_refit(failure)
-                # Re-read: the recovery condemns the silent participant and rebuilds over
-                # the survivors, so the retry's membership is not the first attempt's. This
-                # is the likelier of the two windows -- the shard is restarting precisely
-                # because this refit just failed.
                 participants = self._refit_participants()
+                # A restarted engine may have joined the rebuilt membership without
+                # receiving the initial pause, even if it reuses the same shard index.
+                generation_paused_for_refit = await asyncio.to_thread(
+                    self._gen.pause_generation_for_refit,
+                    clear_cache=clear_cache,
+                    timeout_s=timeout_s,
+                )
                 # Once only: a second failure is a real fault, not a membership problem,
                 # and retrying forever would recreate the wedge this exists to remove.
                 await self._sync_weights_within(kv_scales, "retry")
@@ -4750,12 +4772,28 @@ async def _sync_weights(
             # promoted inside its window, and everything below this must still run on
             # both paths.
             self._record_refit_landed(participants)
-        if self._async_cfg.recompute_kv_cache_after_weight_updates:
-            # to_thread, like every other call into the workers here. Run directly on
-            # the loop this is a blocking Ray call, and a wedged generation worker would
-            # freeze the event loop itself -- taking the watchdog, which is an asyncio
-            # task on that same loop, down with it.
-            await asyncio.to_thread(self._gen.invalidate_kv_cache)
+
+        print("▶️ Requesting generation resume after refit", flush=True)
+        generation_resumed_after_refit = await asyncio.to_thread(
+            self._gen.resume_generation_after_refit, timeout_s=timeout_s
+        )
+        if generation_paused_for_refit and not generation_resumed_after_refit:
+            raise RuntimeError(
+                "Failed to resume generation after a successful refit pause"
+            )
+
+        if clear_cache and not generation_paused_for_refit:
+            # A backend that paused natively already received clear_cache=True at pause
+            # time. Other backends must explicitly confirm equivalent engine-side cache
+            # handling; merely clearing reusable prefix entries does not invalidate KV
+            # blocks still owned by running requests.
+            invalidated = await asyncio.to_thread(self._gen.invalidate_kv_cache)
+            if not invalidated:
+                raise NotImplementedError(
+                    "recompute_kv_cache_after_weight_updates=True requires the "
+                    "generation backend to clear active-request KV state during a "
+                    "native refit pause or confirm equivalent engine-side handling"
+                )
         elapsed = time.monotonic() - t0
 
         print(f"  _sync_weights: sync done in {elapsed:.3f}s", flush=True)

nemo_rl/models/generation/interfaces.py

@@ -611,12 +611,22 @@ def attach_fleet_health(self, monitor: Any, selector: Any) -> None:
             f"{type(self).__name__} generation backend"
         )
 
-    # Optional hook; backends may override to invalidate any reusable caches
-    # (e.g., vLLM prefix/KV caches) after weight updates.
     def invalidate_kv_cache(self) -> bool:
+        """Confirm stale KV state cannot be reused after a weight update.
+
+        Backends may implement this by invalidating active-request KV state directly or
+        by performing equivalent engine-side handling atomically with refit. Clearing
+        only unreferenced prefix-cache entries is insufficient while requests are active.
+
+        Returns:
+            True only when the backend guarantees the recompute contract; False when it
+            is unsupported or could not be completed.
+        """
         return False
 
-    def pause_generation_for_refit(self, *, clear_cache: bool) -> bool:
+    def pause_generation_for_refit(
+        self, *, clear_cache: bool, timeout_s: Optional[float] = None
+    ) -> bool:
         """Pause in-flight generation while preserving request state.
 
         Backends with native in-flight refit support override this hook. The default
@@ -628,6 +638,7 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool:
         Args:
             clear_cache: Also clear the engine's reusable caches at pause time so
                 preserved requests recompute their KV after the weight update.
+            timeout_s: Maximum time to wait for engines to pause; None waits indefinitely.
 
         Returns:
             True if every engine paused; False when the backend has no native pause
@@ -636,13 +647,18 @@ def pause_generation_for_refit(self, *, clear_cache: bool) -> bool:
         _warn_unsupported_in_flight_refit_pause_once(type(self).__name__)
         return False
 
-    def resume_generation_after_refit(self) -> bool:
+    def resume_generation_after_refit(
+        self, *, timeout_s: Optional[float] = None
+    ) -> bool:
         """Resume generation paused by :meth:`pause_generation_for_refit`.
 
         The default implementation shares the once-per-backend warning emitted by
         :meth:`pause_generation_for_refit` and lets the refit continue for backends
         without native pause/resume support.
 
+        Args:
+            timeout_s: Maximum time to wait for engines to resume; None waits indefinitely.
+
         Returns:
             True if every engine resumed; False when the backend has no native resume
             support. Backends with native support raise when resuming fails.

nemo_rl/models/generation/vllm/vllm_generation.py

@@ -33,6 +33,7 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict, SlicedDataDict
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env
+from nemo_rl.distributed.refit_watchdog import RefitAborted
 from nemo_rl.distributed.virtual_cluster import NVLINK_DOMAIN_UNKNOWN, RayVirtualCluster
 from nemo_rl.distributed.worker_groups import RayWorkerBuilder, RayWorkerGroup
 from nemo_rl.models.generation.fleet_health import (
@@ -1623,23 +1624,35 @@ def invalidate_kv_cache(self) -> bool:
             print(f"Error invalidating vLLM caches: {e}")
             return False
 
-    def pause_generation_for_refit(self, *, clear_cache: bool) -> bool:
+    def pause_generation_for_refit(
+        self, *, clear_cache: bool, timeout_s: Optional[float] = None
+    ) -> bool:
         """Pause every async vLLM engine while preserving in-flight requests."""
         if not self.cfg["vllm_cfg"]["async_engine"]:
             raise RuntimeError("pause_generation_for_refit requires async_engine=True")
         if not self.worker_group or not self.worker_group.workers:
             raise RuntimeError("Worker group is not initialized")
 
-        futures = self.worker_group.run_all_workers_single_data(
-            "pause_generation_async",
-            clear_cache=clear_cache,
-            run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
-        )
-        if not all(ray.get(futures)):
+        # Address surviving DP leaders directly. SingleController may have rebuilt the
+        # refit communicator after losing a shard; walking the whole worker group here
+        # would call the dead Ray actor and undo that recovery before the refit starts.
+        futures = [
+            worker.pause_generation_async.remote(clear_cache=clear_cache)
+            for worker in self._refit_leader_workers()
+        ]
+        try:
+            results = ray.get(futures, timeout=timeout_s)
+        except ray.exceptions.GetTimeoutError as exc:
+            raise RefitAborted(
+                f"vLLM generation pause did not return within {timeout_s}s"
+            ) from exc
+        if not all(results):
             raise RuntimeError("Failed to pause every async vLLM engine")
         return True
 
-    def resume_generation_after_refit(self) -> bool:
+    def resume_generation_after_refit(
+        self, *, timeout_s: Optional[float] = None
+    ) -> bool:
         """Resume every async vLLM engine paused for refit."""
         if not self.cfg["vllm_cfg"]["async_engine"]:
             raise RuntimeError(
@@ -1648,11 +1661,17 @@ def resume_generation_after_refit(self) -> bool:
         if not self.worker_group or not self.worker_group.workers:
             raise RuntimeError("Worker group is not initialized")
 
-        futures = self.worker_group.run_all_workers_single_data(
-            "resume_generation_async",
-            run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
-        )
-        if not all(ray.get(futures)):
+        futures = [
+            worker.resume_generation_async.remote()
+            for worker in self._refit_leader_workers()
+        ]
+        try:
+            results = ray.get(futures, timeout=timeout_s)
+        except ray.exceptions.GetTimeoutError as exc:
+            raise RefitAborted(
+                f"vLLM generation resume did not return within {timeout_s}s"
+            ) from exc
+        if not all(results):
             raise RuntimeError("Failed to resume every async vLLM engine")
         return True
 

tests/unit/models/generation/test_vllm_generation.py

@@ -21,7 +21,7 @@
 import types
 from copy import deepcopy
 from pathlib import Path
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, call, patch
 
 import pytest
 import ray
@@ -32,6 +32,7 @@
 from nemo_rl.algorithms.loss import NLLLossFn
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.distributed.refit_watchdog import RefitAborted
 from nemo_rl.distributed.virtual_cluster import RayVirtualCluster
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.models.generation.interfaces import (
@@ -491,43 +492,69 @@ def shutdown(self) -> None:
     assert worker.llm is None
 
 
+@pytest.mark.parametrize("timeout_s", [None, 0.25])
 def test_vllm_generation_broadcasts_native_refit_pause_and_resume(
     monkeypatch: pytest.MonkeyPatch,
+    timeout_s: float | None,
 ) -> None:
     generation = VllmGeneration.__new__(VllmGeneration)
     generation.cfg = {"vllm_cfg": {"async_engine": True}}
-    generation.worker_group = MagicMock()
-    generation.worker_group.workers = [object(), object()]
-    futures = [object(), object()]
-    generation.worker_group.run_all_workers_single_data.return_value = futures
+    generation.worker_group = MagicMock(workers=[object(), object()])
+    leaders = [MagicMock(), MagicMock()]
+    pause_futures = [object(), object()]
+    resume_futures = [object(), object()]
+    for worker, pause_future, resume_future in zip(
+        leaders, pause_futures, resume_futures
+    ):
+        worker.pause_generation_async.remote.return_value = pause_future
+        worker.resume_generation_async.remote.return_value = resume_future
+    generation._refit_leader_workers = MagicMock(return_value=leaders)
     ray_get = MagicMock(side_effect=[[True, True], [True, True]])
     monkeypatch.setattr(ray, "get", ray_get)
 
-    assert generation.pause_generation_for_refit(clear_cache=True)
-    assert generation.resume_generation_after_refit()
+    assert generation.pause_generation_for_refit(clear_cache=True, timeout_s=timeout_s)
+    assert generation.resume_generation_after_refit(timeout_s=timeout_s)
 
-    generation.worker_group.run_all_workers_single_data.assert_any_call(
-        "pause_generation_async",
-        clear_cache=True,
-        run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
-    )
-    generation.worker_group.run_all_workers_single_data.assert_any_call(
-        "resume_generation_async",
-        run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
-    )
+    for worker in leaders:
+        worker.pause_generation_async.remote.assert_called_once_with(clear_cache=True)
+        worker.resume_generation_async.remote.assert_called_once_with()
+    assert ray_get.call_args_list == [
+        call(pause_futures, timeout=timeout_s),
+        call(resume_futures, timeout=timeout_s),
+    ]
+
+
+@pytest.mark.parametrize("operation", ["pause", "resume"])
+def test_vllm_generation_refit_pause_and_resume_raise_on_timeout(
+    monkeypatch: pytest.MonkeyPatch,
+    operation: str,
+) -> None:
+    generation = VllmGeneration.__new__(VllmGeneration)
+    generation.cfg = {"vllm_cfg": {"async_engine": True}}
+    generation.worker_group = MagicMock(workers=[object()])
+    worker = MagicMock()
+    generation._refit_leader_workers = MagicMock(return_value=[worker])
+    ray_get = MagicMock(side_effect=ray.exceptions.GetTimeoutError("timed out"))
+    monkeypatch.setattr(ray, "get", ray_get)
+
+    with pytest.raises(RefitAborted, match=rf"{operation} did not return within 0.25s"):
+        if operation == "pause":
+            generation.pause_generation_for_refit(clear_cache=True, timeout_s=0.25)
+        else:
+            generation.resume_generation_after_refit(timeout_s=0.25)
+
+    future = getattr(worker, f"{operation}_generation_async").remote.return_value
+    ray_get.assert_called_once_with([future], timeout=0.25)
 
 
 def test_vllm_generation_rejects_partial_refit_pause_and_resume(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     generation = VllmGeneration.__new__(VllmGeneration)
     generation.cfg = {"vllm_cfg": {"async_engine": True}}
-    generation.worker_group = MagicMock()
-    generation.worker_group.workers = [object(), object()]
-    generation.worker_group.run_all_workers_single_data.return_value = [
-        object(),
-        object(),
-    ]
+    generation.worker_group = MagicMock(workers=[object(), object()])
+    leaders = [MagicMock(), MagicMock()]
+    generation._refit_leader_workers = MagicMock(return_value=leaders)
     monkeypatch.setattr(
         ray,
         "get",

tests/unit/single_controller/test_checkpointing.py

@@ -437,6 +437,18 @@ def snapshot_step_metrics(self) -> None:
     def get_step_metrics(self) -> dict[str, float]:
         return {}
 
+    def pause_generation_for_refit(
+        self, *, clear_cache: bool, timeout_s: float | None = None
+    ) -> bool:
+        del clear_cache, timeout_s
+        return False
+
+    def resume_generation_after_refit(self, *, timeout_s: float | None = None) -> bool:
+        return False
+
+    def invalidate_kv_cache(self) -> bool:
+        return False
+
     def wake_carries_weight_updates(self) -> bool:
         return True