iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[BUG] Two phase pause to prevent deadlock (#39366)

Browse Source 
Signed-off-by: ahao-anyscale <ahao@anyscale.com>
Signed-off-by: Aaron Hao <ahao@anyscale.com>
Co-authored-by: Junjie Zhang <junj.jay.zhang@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Aaron Hao
2026-04-29 17:51:03 -04:00
committed by GitHub
co-authored by Junjie Zhang Nick Hill
parent 944e138bcf
commit 0335316a9b
3 changed files with 232 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/distributed/test_async_llm_dp.py
+122 -7
View File
@@ -365,10 +365,13 @@ async def test_dp_pause_keep_race_staggered_engines():
async def staggered_pause_keep(method: str, *args) -> Any:
if method != "pause_scheduler" or not args or args[0] != "keep":
return await original_call_utility(method, *args)
# Send pause(keep) to engine 0 first
await client._call_utility_async(
method, *args, engine=client.core_engines[0]
# Fire pause(keep) to engine 0 (don't await — with DP
# two-phase pause, consensus requires all ranks).
pause_0 = asyncio.create_task(
client._call_utility_async(method, *args, engine=client.core_engines[0])
)
# Let the event loop send the message to engine 0.
await asyncio.sleep(0.5)
# In the middle: send two requests (race window)
sp = SamplingParams(max_tokens=5, ignore_eos=True)
@@ -384,11 +387,13 @@ async def test_dp_pause_keep_race_staggered_engines():
t2 = asyncio.create_task(consume_gen("race-2"))
mid_pause_tasks.extend([t1, t2])
await asyncio.sleep(3)
# Then send pause(keep) to engine 1
result = await client._call_utility_async(
method, *args, engine=client.core_engines[1]
# Fire pause(keep) to engine 1, then await both so
# consensus can be reached.
pause_1 = asyncio.create_task(
client._call_utility_async(method, *args, engine=client.core_engines[1])
)
return result
results = await asyncio.gather(pause_0, pause_1)
return results[0]
client.call_utility_async = staggered_pause_keep
@@ -398,3 +403,113 @@ async def test_dp_pause_keep_race_staggered_engines():
assert not await engine.is_paused()
# Let the two requests we sent mid-pause complete
await asyncio.gather(*mid_pause_tasks)
@pytest.mark.asyncio
async def test_dp_pause_barrier_request_deadlock():
"""
Test that start_dp_wave is ignored while paused.
Sequence:
1. Pause all engines (PAUSED_ALL).
2. Send barrier to engine 0 only — blocks in dist.barrier(dp_group).
3. Send a request routed to engine 1.
4. Wait for any (buggy) START_DP_WAVE propagation.
5. Send barrier to engine 1 — completes in fixed code, deadlocks
in buggy code because engine 1 is stuck in EP all-to-all.
"""
if DP_SIZE != 2:
pytest.skip("requires DP_SIZE=2")
with ExitStack() as after:
engine_args = _get_dp_pause_engine_args(expert_parallel=True)
engine = AsyncLLM.from_engine_args(engine_args)
after.callback(engine.shutdown)
client = engine.engine_core
# Cache get_supported_tasks so that generate() won't need to
# send a utility call to all engines (which would hang once
# engine 0 is blocked in the barrier).
await engine.get_supported_tasks()
# Pause all engines normally — no staggering.
await engine.pause_generation(mode="keep")
assert await engine.is_paused()
original_call_utility = client.call_utility_async
mid_barrier_tasks: list[asyncio.Task] = []
async def staggered_barrier(method: str, *args) -> Any:
if method != "barrier":
return await original_call_utility(method, *args)
# Send barrier to engine 0 only — it blocks in
# dist.barrier(dp_group) waiting for engine 1.
barrier_0 = asyncio.create_task(
client._call_utility_async(method, *args, engine=client.core_engines[0])
)
await asyncio.sleep(1)
# While engine 0 is blocked, send a request routed
# specifically to engine 1.
sp = SamplingParams(max_tokens=5, ignore_eos=True)
engine_1 = client.core_engines[1]
original_get_engine = client.get_core_engine_for_request
def route_to_engine_1(req):
client.reqs_in_flight[req.request_id] = engine_1
return engine_1
client.get_core_engine_for_request = route_to_engine_1
async def consume_gen(req_id: str) -> None:
async for _ in engine.generate(
request_id=req_id,
prompt=DP_PAUSE_PROMPT,
sampling_params=sp,
):
pass
t1 = asyncio.create_task(consume_gen("race-1"))
mid_barrier_tasks.append(t1)
# Yield so generate() preprocessing completes and
# add_request_async is called (which, in buggy code,
# would send FIRST_REQ and wake engine 1).
for _ in range(200):
await asyncio.sleep(0)
client.get_core_engine_for_request = original_get_engine
# Wait for any START_DP_WAVE to propagate and for
# engine 1 to potentially enter execute_dummy_batch.
await asyncio.sleep(5)
# Now send barrier to engine 1. In buggy code engine 1
# is stuck in execute_dummy_batch (EP all-to-all) while
# engine 0 is stuck in dist.barrier(dp_group) — deadlock.
result = await client._call_utility_async(
method, *args, engine=client.core_engines[1]
)
await barrier_0
return result
client.call_utility_async = staggered_barrier
# Drive the staggered barrier. Old code deadlocks here.
try:
await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30)
except asyncio.TimeoutError:
for t in mid_barrier_tasks:
t.cancel()
pytest.fail(
"Staggered barrier deadlocked — FIRST_REQ sent while "
"paused caused collective-op mismatch between engines"
)
await engine.resume_generation()
assert not await engine.is_paused()
# Let the two requests we sent mid-barrier complete.
await asyncio.gather(*mid_barrier_tasks)
vllm/config/parallel.py
+27
View File
@@ -663,6 +663,33 @@ class ParallelConfig:
aggregated_has_unfinished = bool(tensor.item())
return aggregated_has_unfinished
@staticmethod
def sync_dp_state(
dp_group: ProcessGroup, has_unfinished: bool, pending_pause: bool
) -> tuple[bool, bool]:
"""Combined all-reduce for DP state synchronization.
Uses a single SUM all-reduce on a 2-element tensor:
[0] = 1 if this rank has unfinished work, else 0.
SUM > 0 ≡ logical OR across ranks → any rank has work.
[1] = 1 if this rank has a pending pause request, else 0.
SUM == dp_size ≡ all ranks reached pause consensus.
has_unfinished_global is true if any rank has unfinished work,
or if some ranks are waiting for a pause consensus.
Returns:
(has_unfinished_global, pause_consensus)
"""
tensor = torch.tensor(
[int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu"
)
torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=dp_group)
dp_size = dp_group.size()
pause_count = tensor[1].item()
has_unfinished_global = tensor[0].item() > 0 or pause_count % dp_size != 0
return has_unfinished_global, pause_count == dp_size
@staticmethod
def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> int:
if kv_cache_memory == -1:
vllm/v1/engine/core.py
+83 -13
View File
@@ -1571,7 +1571,8 @@ class EngineCoreProc(EngineCore):
pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
self.scheduler.set_pause_state(pause_state)
if not self.has_work():
if self._pause_complete():
if clear_cache:
self._reset_caches()
return None
@@ -1580,6 +1581,13 @@ class EngineCoreProc(EngineCore):
self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
return future
def _pause_complete(self) -> bool:
"""Returns True if the pause has fully completed and the caller can
return ``None`` synchronously; False if the pause is still pending
and the caller should register an idle-state callback to finish it.
"""
return not self.has_work()
def _send_finish_outputs_to_client(
self, req_ids: list[str], client_index: int, finish_reason: FinishReason
) -> None:
@@ -1635,6 +1643,14 @@ class DPEngineCoreProc(EngineCoreProc):
self.current_wave = 0
self.last_counts = (0, 0)
# Two-phase pause protocol state. When pending_pause is True, the
# engine keeps stepping (dummy batches) while waiting for all DP
# ranks to also set pending_pause. Once all ranks agree via
# all-reduce, ignore_start_dp_wave is set so that stale
# START_DP_WAVE messages cannot re-wake the engines.
self.pending_pause = False
self.ignore_start_dp_wave = False
from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState
self.eep_scaling_state: ElasticEPScalingState | None = None
@@ -1664,6 +1680,7 @@ class DPEngineCoreProc(EngineCoreProc):
assert 0 <= local_dp_rank <= dp_rank < dp_size
self.dp_rank = dp_rank
self.dp_size = dp_size
dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True)
self.dp_group, self.dp_store = dp_group, dp_store
@@ -1672,6 +1689,24 @@ class DPEngineCoreProc(EngineCoreProc):
if dp_group := getattr(self, "dp_group", None):
stateless_destroy_torch_distributed_process_group(dp_group)
def _pause_complete(self) -> bool:
"""Two-phase DP-aware pause.
Phase 1: Set local pause state and ``pending_pause`` flag. If the
engines are idle, kick-start them by setting ``engines_running`` to
True so ranks enter the stepping loop and reach the all-reduce
consensus checkpoint in ``_has_global_unfinished_reqs``.
Phase 2 (in ``_has_global_unfinished_reqs``): Once the all-reduce
confirms that **all** ranks have ``pending_pause`` set, collectively
stop stepping and set ``ignore_start_dp_wave`` so that stale
``START_DP_WAVE`` messages cannot re-wake any engine.
"""
self.pending_pause = True
self.engines_running = True
return False
def add_request(self, request: Request, request_wave: int = 0):
super().add_request(request, request_wave)
if self.has_coordinator and request_wave != self.current_wave:
@@ -1681,36 +1716,60 @@ class DPEngineCoreProc(EngineCoreProc):
not self.engines_running
and self.scheduler.pause_state == PauseState.UNPAUSED
):
self.engines_running = True
# Request received for an already-completed wave, notify
# front-end that we need to start the next one.
self.engines_running = True
self.output_queue.put_nowait(
(-1, EngineCoreOutputs(start_wave=self.current_wave))
)
def resume_scheduler(self):
super().resume_scheduler()
if (
self.has_coordinator
and not self.engines_running
and self.scheduler.has_unfinished_requests()
):
# Wake up other DP engines.
self.output_queue.put_nowait(
(-1, EngineCoreOutputs(start_wave=self.current_wave))
if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave):
raise RuntimeError(
"resume_scheduler called while pause is still in "
"flight. Wait for the pause future to resolve before "
"resuming."
)
if self.engines_running:
logger.debug("Resume called while engines are not paused, ignoring.")
return
super().resume_scheduler()
self.ignore_start_dp_wave = False
# Barrier: wait for all DP ranks to have resumed (and cleared
# ignore_start_dp_wave) before any rank starts stepping. Uses
# the existing all-reduce which is safe because engines are
# stopped.
has_global_unfinished = ParallelConfig.has_unfinished_dp(
self.dp_group, self.scheduler.has_unfinished_requests()
)
if has_global_unfinished:
self.engines_running = True
def barrier(self):
"""Blocking barrier on the DP process group (test-only utility)."""
import torch.distributed as dist
dist.barrier(group=self.dp_group)
def _handle_client_request(
self, request_type: EngineCoreRequestType, request: Any
) -> None:
if request_type == EngineCoreRequestType.START_DP_WAVE:
if self.ignore_start_dp_wave:
return
new_wave, exclude_eng_index = request
if exclude_eng_index != self.engine_index and (
new_wave >= self.current_wave
):
self.current_wave = new_wave
if not self.engines_running:
logger.debug("EngineCore starting idle loop for wave %d.", new_wave)
logger.debug(
"EngineCore starting idle loop for wave %d.",
new_wave,
)
self.engines_running = True
else:
super()._handle_client_request(request_type, request)
@@ -1790,7 +1849,18 @@ class DPEngineCoreProc(EngineCoreProc):
if self.step_counter % 32 != 0:
return True
return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)
has_unfinished, pause_consensus = ParallelConfig.sync_dp_state(
self.dp_group,
has_unfinished=local_unfinished,
pending_pause=self.pending_pause,
)
if pause_consensus:
self.ignore_start_dp_wave = True
self.pending_pause = False
logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.")
return has_unfinished
def reinitialize_distributed(
self, reconfig_request: ReconfigureDistributedRequest