From 17da48596c98946d3e3e6896e2ebd341e809f3bd Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 27 Aug 2026 00:50:33 +0800 Subject: [PATCH] [Bugfix][DP] Synchronize the device on pause completion (#52914) Signed-off-by: aoshen02 --- tests/v1/distributed/test_async_llm_dp.py | 110 ++++++++++++++++++++++ tests/v1/engine/test_engine_core.py | 38 +++++++- vllm/v1/engine/core.py | 16 ++-- vllm/v1/worker/gpu_worker.py | 2 + vllm/v1/worker/worker_base.py | 6 ++ 5 files changed, 164 insertions(+), 8 deletions(-) diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index f51fde9053..e1adc948a8 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/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)) diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index aa2a70559d..d0663d0293 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/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"] diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 0ee3435b3d..c734e67229 100644 --- a/vllm/v1/engine/core.py +++ b/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]() diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 7ac9a2c6c9..990fae6db2 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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() diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 7a17e10586..cf5b0e7ce3 100644 --- a/vllm/v1/worker/worker_base.py +++ b/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.