@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0335316a9b

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 11:15 UTC

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

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>

Aaron Hao committed Apr 29, 2026 at 14:51 UTC 0335316a9ba245e5e82a20ef1b53ba3da108afd5

3 files changed +232 -20

tests/v1/distributed/test_async_llm_dp.py

+122 -7

@@ -365,10 +365,13 @@ async def test_dp_pause_keep_race_staggered_engines():

365 async def staggered_pause_keep(method: str, *args) -> Any:

366 if method != "pause_scheduler" or not args or args[0] != "keep":

367 return await original_call_utility(method, *args)

368 - # Send pause(keep) to engine 0 first

369 - await client._call_utility_async(

370 - method, *args, engine=client.core_engines[0]

368 + # Fire pause(keep) to engine 0 (don't await — with DP

369 + # two-phase pause, consensus requires all ranks).

370 + pause_0 = asyncio.create_task(

371 + client._call_utility_async(method, *args, engine=client.core_engines[0])

372 )

373 + # Let the event loop send the message to engine 0.

374 + await asyncio.sleep(0.5)

375 # In the middle: send two requests (race window)

376 sp = SamplingParams(max_tokens=5, ignore_eos=True)

377

@@ -384,11 +387,13 @@ async def test_dp_pause_keep_race_staggered_engines():

387 t2 = asyncio.create_task(consume_gen("race-2"))

388 mid_pause_tasks.extend([t1, t2])

389 await asyncio.sleep(3)

387 - # Then send pause(keep) to engine 1

388 - result = await client._call_utility_async(

389 - method, *args, engine=client.core_engines[1]

390 + # Fire pause(keep) to engine 1, then await both so

391 + # consensus can be reached.

392 + pause_1 = asyncio.create_task(

393 + client._call_utility_async(method, *args, engine=client.core_engines[1])

394 )

391 - return result

395 + results = await asyncio.gather(pause_0, pause_1)

396 + return results[0]

397

398 client.call_utility_async = staggered_pause_keep

399

@@ -398,3 +403,113 @@ async def test_dp_pause_keep_race_staggered_engines():

403 assert not await engine.is_paused()

404 # Let the two requests we sent mid-pause complete

405 await asyncio.gather(*mid_pause_tasks)

406 +

407 +

408 +@pytest.mark.asyncio

409 +async def test_dp_pause_barrier_request_deadlock():

410 + """

411 + Test that start_dp_wave is ignored while paused.

412 +

413 + Sequence:

414 + 1. Pause all engines (PAUSED_ALL).

415 + 2. Send barrier to engine 0 only — blocks in dist.barrier(dp_group).

416 + 3. Send a request routed to engine 1.

417 + 4. Wait for any (buggy) START_DP_WAVE propagation.

418 + 5. Send barrier to engine 1 — completes in fixed code, deadlocks

419 + in buggy code because engine 1 is stuck in EP all-to-all.

420 + """

421 + if DP_SIZE != 2:

422 + pytest.skip("requires DP_SIZE=2")

423 +

424 + with ExitStack() as after:

425 + engine_args = _get_dp_pause_engine_args(expert_parallel=True)

426 + engine = AsyncLLM.from_engine_args(engine_args)

427 + after.callback(engine.shutdown)

428 +

429 + client = engine.engine_core

430 +

431 + # Cache get_supported_tasks so that generate() won't need to

432 + # send a utility call to all engines (which would hang once

433 + # engine 0 is blocked in the barrier).

434 + await engine.get_supported_tasks()

435 +

436 + # Pause all engines normally — no staggering.

437 + await engine.pause_generation(mode="keep")

438 + assert await engine.is_paused()

439 +

440 + original_call_utility = client.call_utility_async

441 + mid_barrier_tasks: list[asyncio.Task] = []

442 +

443 + async def staggered_barrier(method: str, *args) -> Any:

444 + if method != "barrier":

445 + return await original_call_utility(method, *args)

446 +

447 + # Send barrier to engine 0 only — it blocks in

448 + # dist.barrier(dp_group) waiting for engine 1.

449 + barrier_0 = asyncio.create_task(

450 + client._call_utility_async(method, *args, engine=client.core_engines[0])

451 + )

452 + await asyncio.sleep(1)

453 +

454 + # While engine 0 is blocked, send a request routed

455 + # specifically to engine 1.

456 + sp = SamplingParams(max_tokens=5, ignore_eos=True)

457 +

458 + engine_1 = client.core_engines[1]

459 + original_get_engine = client.get_core_engine_for_request

460 +

461 + def route_to_engine_1(req):

462 + client.reqs_in_flight[req.request_id] = engine_1

463 + return engine_1

464 +

465 + client.get_core_engine_for_request = route_to_engine_1

466 +

467 + async def consume_gen(req_id: str) -> None:

468 + async for _ in engine.generate(

469 + request_id=req_id,

470 + prompt=DP_PAUSE_PROMPT,

471 + sampling_params=sp,

472 + ):

473 + pass

474 +

475 + t1 = asyncio.create_task(consume_gen("race-1"))

476 + mid_barrier_tasks.append(t1)

477 +

478 + # Yield so generate() preprocessing completes and

479 + # add_request_async is called (which, in buggy code,

480 + # would send FIRST_REQ and wake engine 1).

481 + for _ in range(200):

482 + await asyncio.sleep(0)

483 +

484 + client.get_core_engine_for_request = original_get_engine

485 +

486 + # Wait for any START_DP_WAVE to propagate and for

487 + # engine 1 to potentially enter execute_dummy_batch.

488 + await asyncio.sleep(5)

489 +

490 + # Now send barrier to engine 1. In buggy code engine 1

491 + # is stuck in execute_dummy_batch (EP all-to-all) while

492 + # engine 0 is stuck in dist.barrier(dp_group) — deadlock.

493 + result = await client._call_utility_async(

494 + method, *args, engine=client.core_engines[1]

495 + )

496 + await barrier_0

497 + return result

498 +

499 + client.call_utility_async = staggered_barrier

500 +

501 + # Drive the staggered barrier. Old code deadlocks here.

502 + try:

503 + await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30)

504 + except asyncio.TimeoutError:

505 + for t in mid_barrier_tasks:

506 + t.cancel()

507 + pytest.fail(

508 + "Staggered barrier deadlocked — FIRST_REQ sent while "

509 + "paused caused collective-op mismatch between engines"

510 + )

511 +

512 + await engine.resume_generation()

513 + assert not await engine.is_paused()

514 + # Let the two requests we sent mid-barrier complete.

515 + await asyncio.gather(*mid_barrier_tasks)

vllm/config/parallel.py

+27

@@ -663,6 +663,33 @@ class ParallelConfig:

663 aggregated_has_unfinished = bool(tensor.item())

664 return aggregated_has_unfinished

665

666 + @staticmethod

667 + def sync_dp_state(

668 + dp_group: ProcessGroup, has_unfinished: bool, pending_pause: bool

669 + ) -> tuple[bool, bool]:

670 + """Combined all-reduce for DP state synchronization.

671 +

672 + Uses a single SUM all-reduce on a 2-element tensor:

673 + [0] = 1 if this rank has unfinished work, else 0.

674 + SUM > 0 ≡ logical OR across ranks → any rank has work.

675 + [1] = 1 if this rank has a pending pause request, else 0.

676 + SUM == dp_size ≡ all ranks reached pause consensus.

677 +

678 + has_unfinished_global is true if any rank has unfinished work,

679 + or if some ranks are waiting for a pause consensus.

680 +

681 + Returns:

682 + (has_unfinished_global, pause_consensus)

683 + """

684 + tensor = torch.tensor(

685 + [int(has_unfinished), int(pending_pause)], dtype=torch.int32, device="cpu"

686 + )

687 + torch.distributed.all_reduce(tensor, op=ReduceOp.SUM, group=dp_group)

688 + dp_size = dp_group.size()

689 + pause_count = tensor[1].item()

690 + has_unfinished_global = tensor[0].item() > 0 or pause_count % dp_size != 0

691 + return has_unfinished_global, pause_count == dp_size

692 +

693 @staticmethod

694 def sync_kv_cache_memory_size(dp_group: ProcessGroup, kv_cache_memory: int) -> int:

695 if kv_cache_memory == -1:

vllm/v1/engine/core.py

+83 -13

@@ -1571,7 +1571,8 @@ class EngineCoreProc(EngineCore):

1571

1572 pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW

1573 self.scheduler.set_pause_state(pause_state)

1574 - if not self.has_work():

1574 +

1575 + if self._pause_complete():

1576 if clear_cache:

1577 self._reset_caches()

1578 return None

@@ -1580,6 +1581,13 @@ class EngineCoreProc(EngineCore):

1581 self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))

1582 return future

1583

1584 + def _pause_complete(self) -> bool:

1585 + """Returns True if the pause has fully completed and the caller can

1586 + return ``None`` synchronously; False if the pause is still pending

1587 + and the caller should register an idle-state callback to finish it.

1588 + """

1589 + return not self.has_work()

1590 +

1591 def _send_finish_outputs_to_client(

1592 self, req_ids: list[str], client_index: int, finish_reason: FinishReason

1593 ) -> None:

@@ -1635,6 +1643,14 @@ class DPEngineCoreProc(EngineCoreProc):

1643 self.current_wave = 0

1644 self.last_counts = (0, 0)

1645

1646 + # Two-phase pause protocol state. When pending_pause is True, the

1647 + # engine keeps stepping (dummy batches) while waiting for all DP

1648 + # ranks to also set pending_pause. Once all ranks agree via

1649 + # all-reduce, ignore_start_dp_wave is set so that stale

1650 + # START_DP_WAVE messages cannot re-wake the engines.

1651 + self.pending_pause = False

1652 + self.ignore_start_dp_wave = False

1653 +

1654 from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState

1655

1656 self.eep_scaling_state: ElasticEPScalingState | None = None

@@ -1664,6 +1680,7 @@ class DPEngineCoreProc(EngineCoreProc):

1680 assert 0 <= local_dp_rank <= dp_rank < dp_size

1681

1682 self.dp_rank = dp_rank

1683 + self.dp_size = dp_size

1684 dp_group, dp_store = parallel_config.stateless_init_dp_group(return_store=True)

1685 self.dp_group, self.dp_store = dp_group, dp_store

1686

@@ -1672,6 +1689,24 @@ class DPEngineCoreProc(EngineCoreProc):

1689 if dp_group := getattr(self, "dp_group", None):

1690 stateless_destroy_torch_distributed_process_group(dp_group)

1691

1692 + def _pause_complete(self) -> bool:

1693 + """Two-phase DP-aware pause.

1694 +

1695 + Phase 1: Set local pause state and ``pending_pause`` flag. If the

1696 + engines are idle, kick-start them by setting ``engines_running`` to

1697 + True so ranks enter the stepping loop and reach the all-reduce

1698 + consensus checkpoint in ``_has_global_unfinished_reqs``.

1699 +

1700 + Phase 2 (in ``_has_global_unfinished_reqs``): Once the all-reduce

1701 + confirms that **all** ranks have ``pending_pause`` set, collectively

1702 + stop stepping and set ``ignore_start_dp_wave`` so that stale

1703 + ``START_DP_WAVE`` messages cannot re-wake any engine.

1704 + """

1705 + self.pending_pause = True

1706 + self.engines_running = True

1707 +

1708 + return False

1709 +

1710 def add_request(self, request: Request, request_wave: int = 0):

1711 super().add_request(request, request_wave)

1712 if self.has_coordinator and request_wave != self.current_wave:

@@ -1681,36 +1716,60 @@ class DPEngineCoreProc(EngineCoreProc):

1716 not self.engines_running

1717 and self.scheduler.pause_state == PauseState.UNPAUSED

1718 ):

1684 - self.engines_running = True

1719 # Request received for an already-completed wave, notify

1720 # front-end that we need to start the next one.

1721 + self.engines_running = True

1722 self.output_queue.put_nowait(

1723 (-1, EngineCoreOutputs(start_wave=self.current_wave))

1724 )

1725

1726 def resume_scheduler(self):

1692 - super().resume_scheduler()

1693 - if (

1694 - self.has_coordinator

1695 - and not self.engines_running

1696 - and self.scheduler.has_unfinished_requests()

1697 - ):

1698 - # Wake up other DP engines.

1699 - self.output_queue.put_nowait(

1700 - (-1, EngineCoreOutputs(start_wave=self.current_wave))

1727 + if self.pending_pause or (self.engines_running and self.ignore_start_dp_wave):

1728 + raise RuntimeError(

1729 + "resume_scheduler called while pause is still in "

1730 + "flight. Wait for the pause future to resolve before "

1731 + "resuming."

1732 )

1733 + if self.engines_running:

1734 + logger.debug("Resume called while engines are not paused, ignoring.")

1735 + return

1736 +

1737 + super().resume_scheduler()

1738 + self.ignore_start_dp_wave = False

1739 +

1740 + # Barrier: wait for all DP ranks to have resumed (and cleared

1741 + # ignore_start_dp_wave) before any rank starts stepping. Uses

1742 + # the existing all-reduce which is safe because engines are

1743 + # stopped.

1744 + has_global_unfinished = ParallelConfig.has_unfinished_dp(

1745 + self.dp_group, self.scheduler.has_unfinished_requests()

1746 + )

1747 +

1748 + if has_global_unfinished:

1749 + self.engines_running = True

1750 +

1751 + def barrier(self):

1752 + """Blocking barrier on the DP process group (test-only utility)."""

1753 + import torch.distributed as dist

1754 +

1755 + dist.barrier(group=self.dp_group)

1756

1757 def _handle_client_request(

1758 self, request_type: EngineCoreRequestType, request: Any

1759 ) -> None:

1760 if request_type == EngineCoreRequestType.START_DP_WAVE:

1761 + if self.ignore_start_dp_wave:

1762 + return

1763 new_wave, exclude_eng_index = request

1764 if exclude_eng_index != self.engine_index and (

1765 new_wave >= self.current_wave

1766 ):

1767 self.current_wave = new_wave

1768 if not self.engines_running:

1713 - logger.debug("EngineCore starting idle loop for wave %d.", new_wave)

1769 + logger.debug(

1770 + "EngineCore starting idle loop for wave %d.",

1771 + new_wave,

1772 + )

1773 self.engines_running = True

1774 else:

1775 super()._handle_client_request(request_type, request)

@@ -1790,7 +1849,18 @@ class DPEngineCoreProc(EngineCoreProc):

1849 if self.step_counter % 32 != 0:

1850 return True

1851

1793 - return ParallelConfig.has_unfinished_dp(self.dp_group, local_unfinished)

1852 + has_unfinished, pause_consensus = ParallelConfig.sync_dp_state(

1853 + self.dp_group,

1854 + has_unfinished=local_unfinished,

1855 + pending_pause=self.pending_pause,

1856 + )

1857 +

1858 + if pause_consensus:

1859 + self.ignore_start_dp_wave = True

1860 + self.pending_pause = False

1861 + logger.debug("DP pause consensus reached, ignoring START_DP_WAVE.")

1862 +

1863 + return has_unfinished

1864

1865 def reinitialize_distributed(

1866 self, reconfig_request: ReconfigureDistributedRequest