diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 1b7739d2f0..70a5136a57 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -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) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 4f903eeefa..6ba392802e 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -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: diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 36864ba738..11c5ee19a6 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -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