# 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.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.async_llm import AsyncLLM from vllm.v1.engine.core_client import DPAsyncMPClient 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: str, 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) 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: DPAsyncMPClient = engine.engine_core # 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 @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 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)