[Bugfix][DP] Synchronize the device on pause completion (#52914)

Signed-off-by: aoshen02 <aoshen@inferact.ai>

tests/v1/distributed/test_async_llm_dp.py

@@ -713,3 +713,113 @@ async def test_dp_pause_barrier_request_deadlock():
         assert not await engine.is_paused()
         # Let the two requests we sent mid-barrier complete.
         await asyncio.gather(*mid_barrier_tasks)
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_wait_mode_drains_in_flight():
+    """mode="wait" through the DP consensus path: the pause lets the
+    in-flight request run to completion while the consensus is pending,
+    resolves only once drained, and the engine generates again on resume."""
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))
+        after.callback(engine.shutdown)
+
+        collector = await engine.add_request(
+            request_id="inflight",
+            prompt=DP_PAUSE_PROMPT,
+            params=SamplingParams(max_tokens=64),
+        )
+        await engine.pause_generation(mode="wait")
+        assert await engine.is_paused()
+        while True:
+            out = await asyncio.wait_for(collector.get(), timeout=30)
+            if out.finished:
+                break
+
+        await engine.resume_generation()
+        async for out in engine.generate(
+            request_id="after-wait",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=SamplingParams(max_tokens=5),
+        ):
+            pass
+        assert out.finished
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_while_asleep():
+    """Pausing a sleeping DP engine: sleeping ranks skip dummy batches yet
+    must still reach the pause consensus, and the completion barrier must
+    be harmless on workers whose memory is unmapped."""
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))
+        after.callback(engine.shutdown)
+
+        async for _ in engine.generate(
+            request_id="warmup",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=SamplingParams(max_tokens=5),
+        ):
+            pass
+        await engine.sleep(level=1)
+        assert await engine.is_sleeping()
+
+        await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=30)
+        assert await engine.is_paused()
+
+        # A full wake makes the memory resident and resumes the scheduler.
+        await engine.wake_up()
+        assert not await engine.is_sleeping()
+        async for out in engine.generate(
+            request_id="after-wake",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=SamplingParams(max_tokens=5),
+        ):
+            pass
+        assert out.finished
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_completion_implies_device_idle(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """A resolved pause future promises an idle device. Inflate the dummy
+    batches the pause consensus manufactures so any work the pause fails to
+    wait on stays visible, then require a quiet worker stream the moment
+    pause_generation() returns."""
+    with ExitStack() as after:
+        # The probes below ship functions through collective_rpc.
+        monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+        engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True))
+        after.callback(engine.shutdown)
+
+        async for _ in engine.generate(
+            request_id="warmup",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=SamplingParams(max_tokens=5),
+        ):
+            pass
+
+        def inflate_dummy_batches(worker: Any) -> bool:
+            import torch
+
+            inner = worker.execute_dummy_batch
+
+            def slow_dummy_batch() -> None:
+                inner()
+                # Keep the stream busy after launch, like a large-MoE forward.
+                torch.cuda._sleep(200_000_000)
+
+            worker.execute_dummy_batch = slow_dummy_batch
+            return True
+
+        assert all(await engine.engine_core.collective_rpc_async(inflate_dummy_batches))
+
+        await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=60)
+
+        def device_idle(worker: Any) -> bool:
+            import torch
+
+            return bool(torch.cuda.current_stream().query())
+
+        assert all(await engine.engine_core.collective_rpc_async(device_idle))

tests/v1/engine/test_engine_core.py

@@ -5,7 +5,7 @@ import copy
 import time
 import uuid
 from concurrent.futures import Future, ThreadPoolExecutor
-from unittest.mock import PropertyMock, patch
+from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 from transformers import AutoTokenizer
@@ -23,7 +23,7 @@ from vllm.engine.arg_utils import EngineArgs
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_default_torch_num_threads
 from vllm.v1.engine import EngineCoreRequest
-from vllm.v1.engine.core import EngineCore
+from vllm.v1.engine.core import EngineCore, EngineCoreProc
 from vllm.v1.executor.abstract import Executor
 from vllm.v1.executor.uniproc_executor import UniProcExecutor
 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -608,3 +608,37 @@ def test_encoder_instance_zero_kv_cache(
         assert not engine_core.scheduler.ec_connector.is_producer, (
             "Consumer instance EC connector should be consumer"
         )
+
+
+def _pausable_engine_core_proc() -> EngineCoreProc:
+    """A bare EngineCoreProc holding just the state pause_scheduler touches."""
+    core = object.__new__(EngineCoreProc)
+    core.model_executor = MagicMock()
+    core.scheduler = MagicMock()
+    core.scheduler.has_requests.return_value = False
+    core.batch_queue = None
+    core.engines_running = False
+    core._idle_state_callbacks = []
+    return core
+
+
+@pytest.mark.parametrize("deferred", [False, True])
+def test_pause_synchronizes_device_before_cache_reset(deferred: bool):
+    """A resolved pause promises an idle device: the barrier must run before
+    caches are cleared and before the caller is unblocked."""
+    core = _pausable_engine_core_proc()
+    core.engines_running = deferred
+    order: list[str] = []
+    core.model_executor.collective_rpc.side_effect = lambda method: order.append(method)
+    core._reset_caches = lambda: order.append("reset_caches")
+
+    result = EngineCoreProc.pause_scheduler(core, mode="keep", clear_cache=True)
+    if deferred:
+        assert isinstance(result, Future)
+        assert not result.done() and order == []
+        core.engines_running = False
+        core._notify_idle_state_callbacks()
+        assert result.result(timeout=0) is None
+    else:
+        assert result is None
+    assert order == ["synchronize_device", "reset_caches"]

vllm/v1/engine/core.py

@@ -839,6 +839,13 @@ class EngineCore:
         self.reset_mm_cache()
         self.reset_encoder_cache()
 
+    def _finish_pause(self, clear_cache: bool) -> None:
+        # A completed pause promises an idle device: nothing else waits on
+        # the last dummy batch an idle DP rank launches.
+        self.model_executor.collective_rpc("synchronize_device")
+        if clear_cache:
+            self._reset_caches()
+
     def pause_scheduler(
         self, mode: PauseMode = "abort", clear_cache: bool = True
     ) -> Future | None:
@@ -865,8 +872,7 @@ class EngineCore:
 
         pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
         self.scheduler.set_pause_state(pause_state)
-        if clear_cache:
-            self._reset_caches()
+        self._finish_pause(clear_cache)
 
         return None
 
@@ -1937,8 +1943,7 @@ class EngineCoreProc(EngineCore):
             raise ValueError(f"Invalid pause mode: {mode}")
 
         def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
-            if clear_cache:
-                engine._reset_caches()
+            engine._finish_pause(clear_cache)
             future.set_result(None)
 
         if mode == "abort":
@@ -1951,8 +1956,7 @@ class EngineCoreProc(EngineCore):
         self.scheduler.set_pause_state(pause_state)
 
         if self._pause_complete():
-            if clear_cache:
-                self._reset_caches()
+            self._finish_pause(clear_cache)
             return None
 
         future = Future[Any]()

vllm/v1/worker/gpu_worker.py

@@ -259,6 +259,8 @@ class Worker(WorkerBase):
         if tags is None or "kv_cache" in tags:
             self.model_runner.post_kv_cache_wake_up()
 
+        self.synchronize_device()
+
     def checkpoint_prepare(self) -> None:
         checkpoint_prepare_distributed_state()
 

vllm/v1/worker/worker_base.py

@@ -125,6 +125,12 @@ class WorkerBase:
         """Basic health check (override for device-specific checks)."""
         return
 
+    def synchronize_device(self) -> None:
+        """Block until in-flight device work completes; backends outside
+        ``torch.accelerator`` must override with their own wait."""
+        if torch.accelerator.is_available():
+            torch.accelerator.synchronize()
+
     def init_device(self) -> None:
         """Initialize device state, such as loading the model or other on-device
         memory allocations.