vLLM/tests/v1/distributed/test_async_llm_dp.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import asyncio
import os
import time
from contextlib import ExitStack
from dataclasses import dataclass
from typing import Any

import pytest

from vllm import SamplingParams
from vllm.config import VllmConfig
from vllm.config.parallel import DataParallelBackend
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.inputs import PromptType
from vllm.outputs import RequestOutput
from vllm.platforms import current_platform
from vllm.sampling_params import RequestOutputKind
from vllm.v1.engine import EngineCoreRequestType
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.engine.core_client import DPLBAsyncMPClient
from vllm.v1.metrics.loggers import StatLoggerBase
from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats

DP_SIZE = int(os.getenv("DP_SIZE", 2))


async def generate(
    engine: AsyncLLM,
    request_id: str,
    prompt: PromptType,
    output_kind: RequestOutputKind,
    max_tokens: int,
    prompt_logprobs: int | None = None,
    data_parallel_rank: int | None = None,
) -> tuple[int, str]:
    # Ensure generate doesn't complete too fast for cancellation test.
    await asyncio.sleep(0.2)

    count = 0
    sampling_params = SamplingParams(
        max_tokens=max_tokens,
        ignore_eos=True,
        output_kind=output_kind,
        temperature=0,
        prompt_logprobs=prompt_logprobs,
    )
    async for out in engine.generate(
        request_id=request_id,
        prompt=prompt,
        sampling_params=sampling_params,
        data_parallel_rank=data_parallel_rank,
    ):
        num_tokens = len(out.outputs[0].token_ids)
        if output_kind == RequestOutputKind.DELTA:
            count += num_tokens
        else:
            count = num_tokens

        await asyncio.sleep(0.0)

    return count, request_id


@pytest.mark.parametrize(
    "model",
    [
        "ibm-research/PowerMoE-3b",
        "hmellor/tiny-random-LlamaForCausalLM",
    ],
)
@pytest.mark.parametrize(
    "output_kind",
    [
        RequestOutputKind.DELTA,
        RequestOutputKind.FINAL_ONLY,
    ],
)
@pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"])
@pytest.mark.parametrize("async_scheduling", [True, False])
@pytest.mark.asyncio
async def test_load(
    model: str,
    output_kind: RequestOutputKind,
    data_parallel_backend: DataParallelBackend,
    async_scheduling: bool,
):
    if async_scheduling and data_parallel_backend == "ray":
        # TODO(NickLucche) Re-enable when async scheduling is supported
        pytest.skip("Async scheduling is not supported with ray")
    elif data_parallel_backend == "ray" and current_platform.is_rocm():
        pytest.skip(
            "Ray as the distributed executor backend is not supported with ROCm."
        )
    stats_loggers = {}

    @dataclass
    class SimpleStatsLogger(StatLoggerBase):
        init_count: int = 0
        finished_req_count: int = 0

        def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
            stats_loggers[engine_index] = self

        def record(
            self,
            scheduler_stats: SchedulerStats | None,
            iteration_stats: IterationStats | None,
            mm_cache_stats: MultiModalCacheStats | None = None,
            engine_idx: int = 0,
        ):
            if iteration_stats:
                self.finished_req_count += len(iteration_stats.finished_requests)

        def log_engine_initialized(self):
            self.init_count += 1

    with ExitStack() as after:
        prompt = "This is a test of data parallel"

        engine_args = AsyncEngineArgs(
            model=model,
            enforce_eager=True,
            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
            data_parallel_size=DP_SIZE,
            data_parallel_backend=data_parallel_backend,
            async_scheduling=async_scheduling,
        )
        engine = AsyncLLM.from_engine_args(
            engine_args, stat_loggers=[SimpleStatsLogger]
        )
        after.callback(engine.shutdown)

        # The first request an engine serves pays a few hundred milliseconds of
        # one-off setup, and while it does the engine holds the requests already
        # routed to it while reporting itself idle, which is enough to skew the
        # split asserted on below. Pay it on every engine first, then measure.
        await asyncio.gather(
            *(
                generate(
                    engine,
                    f"warmup-{rank}",
                    prompt,
                    output_kind,
                    max_tokens=1,
                    data_parallel_rank=rank,
                )
                for rank in range(DP_SIZE)
            )
        )
        for sl in stats_loggers.values():
            sl.finished_req_count = 0

        NUM_REQUESTS = 100
        NUM_EXPECTED_TOKENS = 10

        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

        # Create concurrent requests.
        tasks = []
        for request_id in request_ids:
            tasks.append(
                asyncio.create_task(
                    generate(
                        engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS
                    )
                )
            )
            # Short sleep to ensure that requests are distributed.
            await asyncio.sleep(0.01)
        # Confirm that we got all the EXPECTED tokens from the requests.
        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
        for task in pending:
            task.cancel()
        for task in done:
            num_generated_tokens, request_id = await task
            assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
                f"{request_id} generated {num_generated_tokens} but "
                f"expected {NUM_EXPECTED_TOKENS}"
            )

        assert not engine.output_processor.has_unfinished_requests()

        # testing internals here which may break
        core_client = engine.engine_core
        assert isinstance(core_client, DPLBAsyncMPClient)
        # the engines only synchronize stopping every N steps so
        # allow a small amount of time here.
        for _ in range(10):
            if not core_client.engines_running:
                break
            await asyncio.sleep(0.5)

        assert not core_client.engines_running
        assert not core_client.reqs_in_flight

        # Check that requests were distributed between the engines
        print(f"Stats loggers after test: {stats_loggers}")
        assert len(stats_loggers) == DP_SIZE
        assert stats_loggers[0].init_count == 1

        for sl in stats_loggers.values():
            slogger: SimpleStatsLogger = sl

            assert slogger.finished_req_count > NUM_REQUESTS // (DP_SIZE + 1), (
                f"requests are imbalanced: {stats_loggers}"
            )


@pytest.mark.parametrize("prefill_schedule_interval", [1, 4])
@pytest.mark.asyncio
async def test_dp_prefill_schedule_interval(prefill_schedule_interval: int):
    """Throttling new prefills to every Nth step (DP balancing) must not break
    generation: a stream of staggered requests should still all complete with
    the expected number of tokens.

    The throttle only engages in the DP MoE/EP engine-core path
    (`DPEngineCoreProc`), so this uses an MoE model with expert parallel.
    """
    with ExitStack() as after:
        prompt = "This is a test of data parallel"

        engine_args = AsyncEngineArgs(
            model="ibm-research/PowerMoE-3b",
            enforce_eager=True,
            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
            data_parallel_size=DP_SIZE,
            data_parallel_backend="mp",
            enable_expert_parallel=True,
            prefill_schedule_interval=prefill_schedule_interval,
        )
        engine = AsyncLLM.from_engine_args(engine_args)
        after.callback(engine.shutdown)

        NUM_REQUESTS = 50
        NUM_EXPECTED_TOKENS = 10

        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]

        # Create requests with a small stagger so they arrive across many
        # steps and (with interval > 1) accumulate in the waiting queue
        # before being admitted together on cadence-aligned steps.
        tasks = []
        for request_id in request_ids:
            tasks.append(
                asyncio.create_task(
                    generate(
                        engine,
                        request_id,
                        prompt,
                        RequestOutputKind.DELTA,
                        NUM_EXPECTED_TOKENS,
                    )
                )
            )
            await asyncio.sleep(0.01)

        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
        for task in pending:
            task.cancel()
        for task in done:
            num_generated_tokens, request_id = await task
            assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
                f"{request_id} generated {num_generated_tokens} but "
                f"expected {NUM_EXPECTED_TOKENS}"
            )

        assert not engine.output_processor.has_unfinished_requests()


# =============================================================================
# DP Pause/Resume Tests
# =============================================================================
# When expert_parallel=False: uses non-MoE model (DP replicas as separate engines).
# When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path).

DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b"
DP_PAUSE_PROMPT = "This is a test of data parallel pause"


def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs:
    """Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama."""
    model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL
    return AsyncEngineArgs(
        model=model,
        enforce_eager=True,
        tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
        data_parallel_size=DP_SIZE,
        data_parallel_backend="mp",
        enable_expert_parallel=expert_parallel,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("expert_parallel", [False, True])
async def test_dp_pause_resume_basic(expert_parallel: bool):
    """Pausing from the client (one call) pauses all DP ranks; resume clears it."""
    with ExitStack() as after:
        engine_args = _get_dp_pause_engine_args(expert_parallel)
        engine = AsyncLLM.from_engine_args(engine_args)
        after.callback(engine.shutdown)

        assert not await engine.is_paused()
        await engine.pause_generation(mode="abort")
        assert await engine.is_paused()
        await engine.resume_generation()
        assert not await engine.is_paused()

        # Engine still works after resume
        sampling_params = SamplingParams(max_tokens=5)
        async for out in engine.generate(
            request_id="after-resume",
            prompt=DP_PAUSE_PROMPT,
            sampling_params=sampling_params,
        ):
            pass
        assert out.finished


async def _consume(generator) -> None:
    async for _ in generator:
        pass


async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool:
    """Wait for the front-end's view of the DP engines to reach ``want``."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if engine.engine_core.dp_engines_running() == want:
            return True
        await asyncio.sleep(0.05)
    return False


@pytest.mark.asyncio
async def test_dp_pause_late_request_does_not_block_drain():
    """A request arriving after pause must not make the coordinator believe
    the engines are running.

    Paused engines start no wave for it, so nothing could report wave
    completion afterwards; a running flag set on its behalf would leave drain
    no way back but a timeout.

    MoE only: wave coordination is enabled iff the model is MoE, so a dense
    model never reaches the coordinator state this exercises.
    """
    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)

        # Run a wave first, so the engines are quiesced by the pause rather
        # than by never having started, and drain has something to observe.
        long_request = asyncio.create_task(
            _consume(
                engine.generate(
                    request_id="warmup",
                    prompt=DP_PAUSE_PROMPT,
                    sampling_params=SamplingParams(max_tokens=400, ignore_eos=True),
                )
            )
        )

        # A design that never reports the engines as running would satisfy
        # every drain assertion below while destroying the signal, so pin it
        # down first. Rank 0 reports the wave only after its first step, which
        # can outlast the 2s on a cold device, so wait for the report rather
        # than sample it.
        await asyncio.sleep(2)
        assert not long_request.done(), "the warmup request was too short to sample"
        assert await _poll_flag(engine, True, timeout=30), (
            "the coordinator does not report the engines as running while they are"
        )

        await long_request
        assert await _poll_flag(engine, False, timeout=30)

        await engine.pause_generation(mode="abort")
        await engine.wait_for_requests_to_drain(drain_timeout=30)

        collector = await engine.add_request(
            request_id="late",
            prompt=DP_PAUSE_PROMPT,
            params=SamplingParams(max_tokens=5),
        )
        assert not await _poll_flag(engine, True, timeout=5)

        # The late request was held rather than dropped: it completes on resume.
        await engine.resume_generation()
        while True:
            out = await asyncio.wait_for(collector.get(), timeout=60)
            if out.finished:
                break


@pytest.mark.asyncio
async def test_dp_request_reaching_rank_after_resume_wakes_peers():
    """A request can reach its rank only after resume (e.g. from another API
    server), with no other work anywhere. That rank must wake its peers itself,
    or it steps alone. Modelled here by holding the request's ADD on the client
    side across pause and resume."""
    with ExitStack() as after:
        engine = AsyncLLM.from_engine_args(
            _get_dp_pause_engine_args(expert_parallel=True)
        )
        after.callback(engine.shutdown)
        client = engine.engine_core
        params = SamplingParams(max_tokens=5)

        await _consume(engine.generate(DP_PAUSE_PROMPT, params, "warmup"))
        await engine.pause_generation(mode="abort")
        assert await _poll_flag(engine, False, timeout=30)

        held: list[tuple] = []
        send_input = client._send_input

        def hold_add(request_type, request, engine_id=None):
            if request_type == EngineCoreRequestType.ADD:
                held.append((request_type, request, engine_id))
                return asyncio.sleep(0)
            return send_input(request_type, request, engine_id)

        client._send_input = hold_add
        task = asyncio.create_task(
            _consume(engine.generate(DP_PAUSE_PROMPT, params, "late"))
        )
        await asyncio.sleep(2)
        assert held

        await engine.resume_generation()
        client._send_input = send_input
        await send_input(*held[0])
        await asyncio.wait_for(task, timeout=60)


@pytest.mark.asyncio
async def test_dp_sleep_late_request_does_not_block_drain():
    """The same latch, reached through sleep rather than pause.

    Sleep stops the engines stepping just as pause does, so a request arriving
    while they are asleep can leave the coordinator believing they are running
    with nothing able to say otherwise. This is worth pinning separately from
    the pause case because it is the shape that reaches
    `_drain_requests_for_elastic_ep`, which decides from the same signal
    whether it is safe to scale.
    """
    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)

        async for _ in engine.generate(
            request_id="warmup",
            prompt=DP_PAUSE_PROMPT,
            sampling_params=SamplingParams(max_tokens=5),
        ):
            pass
        assert await _poll_flag(engine, False, timeout=30)

        await engine.sleep(level=1)
        assert await engine.is_sleeping()

        collector = await engine.add_request(
            request_id="while-asleep",
            prompt=DP_PAUSE_PROMPT,
            params=SamplingParams(max_tokens=5),
        )
        assert not await _poll_flag(engine, True, timeout=5)

        await engine.wake_up()
        assert not await engine.is_sleeping()
        while True:
            out = await asyncio.wait_for(collector.get(), timeout=60)
            if out.finished:
                break


@pytest.mark.asyncio
@pytest.mark.parametrize("expert_parallel", [False, True])
async def test_dp_pause_abort(expert_parallel: bool):
    """Pause with abort from one client aborts in-flight requests on all DP ranks."""
    with ExitStack() as after:
        engine_args = _get_dp_pause_engine_args(expert_parallel)
        engine = AsyncLLM.from_engine_args(engine_args)
        after.callback(engine.shutdown)

        # Start several requests so they are distributed across ranks
        sampling_params = SamplingParams(max_tokens=500, ignore_eos=True)
        num_requests = 4
        outputs_by_id: dict[str, list[RequestOutput]] = {}

        async def gen(rid: str):
            out_list: list[RequestOutput] = []
            outputs_by_id[rid] = out_list
            async for out in engine.generate(
                request_id=rid,
                prompt=DP_PAUSE_PROMPT,
                sampling_params=sampling_params,
            ):
                out_list.append(out)
            return out_list[-1] if out_list else None

        tasks = [asyncio.create_task(gen(f"req-{i}")) for i in range(num_requests)]
        # Wait for some tokens on at least one request
        while not any(len(o) >= 2 for o in outputs_by_id.values()):
            await asyncio.sleep(0.02)

        await engine.pause_generation(mode="abort")

        finals = await asyncio.gather(*tasks)
        for i, final in enumerate(finals):
            assert final is not None, f"req-{i} had no output"
            assert final.finished
            assert final.outputs[0].finish_reason == "abort"

        assert await engine.is_paused()
        await engine.resume_generation()
        assert not await engine.is_paused()

        # New request completes after resume
        async for out in engine.generate(
            request_id="after-abort",
            prompt=DP_PAUSE_PROMPT,
            sampling_params=SamplingParams(max_tokens=5),
        ):
            pass
        assert out.finished
        assert not engine.output_processor.has_unfinished_requests()


@pytest.mark.asyncio
@pytest.mark.parametrize("expert_parallel", [False, True])
async def test_dp_pause_keep_then_resume(expert_parallel: bool):
    """Start generation, pause after a few tokens (keep mode), resume; verify gap."""
    pause_duration = 2.0
    min_tokens_before_pause = 3

    with ExitStack() as after:
        engine_args = _get_dp_pause_engine_args(expert_parallel)
        engine = AsyncLLM.from_engine_args(engine_args)
        after.callback(engine.shutdown)

        sampling_params = SamplingParams(max_tokens=15, ignore_eos=True)
        token_times: list[tuple[int, float]] = []
        pause_token_idx = 0

        async def generator_task():
            nonlocal pause_token_idx
            out = None
            async for output in engine.generate(
                request_id="keep-resume-req",
                prompt=DP_PAUSE_PROMPT,
                sampling_params=sampling_params,
            ):
                token_count = len(output.outputs[0].token_ids)
                token_times.append((token_count, time.monotonic()))
                out = output
            return out

        async def controller_task():
            nonlocal pause_token_idx
            while len(token_times) < min_tokens_before_pause:
                await asyncio.sleep(0.01)
            await engine.pause_generation(mode="keep")
            await asyncio.sleep(pause_duration)
            pause_token_idx = len(token_times)
            await engine.resume_generation()

        gen_task = asyncio.create_task(generator_task())
        ctrl_task = asyncio.create_task(controller_task())
        final_output, _ = await asyncio.gather(gen_task, ctrl_task)

        assert final_output is not None and final_output.finished
        assert await engine.is_paused() is False
        assert pause_token_idx >= min_tokens_before_pause
        if pause_token_idx > 0 and pause_token_idx < len(token_times):
            pause_gap = (
                token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
            )
            assert pause_gap >= pause_duration * 0.8, (
                f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s"
            )


@pytest.mark.asyncio
async def test_dp_pause_keep_race_staggered_engines():
    """Race: send pause(keep) to engine 0, then add two requests,
    then pause(keep) to engine 1. Ensures no deadlock when pause
    requests are staggered and requests arrive in between."""
    if DP_SIZE != 2:
        pytest.skip("test_dp_pause_keep_race_staggered_engines 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

        original_call_utility = client.call_utility_async
        mid_pause_tasks: list[asyncio.Task] = []

        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)
            # 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)

            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"))
            t2 = asyncio.create_task(consume_gen("race-2"))
            mid_pause_tasks.extend([t1, t2])
            await asyncio.sleep(3)
            # 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])
            )
            results = await asyncio.gather(pause_0, pause_1)
            return results[0]

        client.call_utility_async = staggered_pause_keep

        await engine.pause_generation(mode="keep")
        assert await engine.is_paused()
        await engine.resume_generation()
        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 a request sent while paused does not wake any engine.

    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 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 TimeoutError:
            for t in mid_barrier_tasks:
                t.cancel()
            pytest.fail(
                "Staggered barrier deadlocked — a request 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)


@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),
        )
        # Wait for the first output, so drain has something to observe: a
        # request still queued holds no blocks and is not waited on.
        out = await asyncio.wait_for(collector.get(), timeout=30)
        await engine.pause_generation(mode="wait")
        assert await engine.is_paused()
        while not out.finished:
            out = await asyncio.wait_for(collector.get(), timeout=30)

        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))