| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | import os |
| 6 | import time |
| 7 | from contextlib import ExitStack |
| 8 | from dataclasses import dataclass |
| 9 | from typing import Any |
| 10 | |
| 11 | import pytest |
| 12 | |
| 13 | from vllm import SamplingParams |
| 14 | from vllm.config import VllmConfig |
| 15 | from vllm.config.parallel import DataParallelBackend |
| 16 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 17 | from vllm.inputs import PromptType |
| 18 | from vllm.outputs import RequestOutput |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.sampling_params import RequestOutputKind |
| 21 | from vllm.v1.engine.async_llm import AsyncLLM |
| 22 | from vllm.v1.engine.core_client import DPLBAsyncMPClient |
| 23 | from vllm.v1.metrics.loggers import StatLoggerBase |
| 24 | from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats |
| 25 | |
| 26 | DP_SIZE = int(os.getenv("DP_SIZE", 2)) |
| 27 | |
| 28 | |
| 29 | async def generate( |
| 30 | engine: AsyncLLM, |
| 31 | request_id: str, |
| 32 | prompt: PromptType, |
| 33 | output_kind: RequestOutputKind, |
| 34 | max_tokens: int, |
| 35 | prompt_logprobs: int | None = None, |
| 36 | data_parallel_rank: int | None = None, |
| 37 | ) -> tuple[int, str]: |
| 38 | # Ensure generate doesn't complete too fast for cancellation test. |
| 39 | await asyncio.sleep(0.2) |
| 40 | |
| 41 | count = 0 |
| 42 | sampling_params = SamplingParams( |
| 43 | max_tokens=max_tokens, |
| 44 | ignore_eos=True, |
| 45 | output_kind=output_kind, |
| 46 | temperature=0, |
| 47 | prompt_logprobs=prompt_logprobs, |
| 48 | ) |
| 49 | async for out in engine.generate( |
| 50 | request_id=request_id, |
| 51 | prompt=prompt, |
| 52 | sampling_params=sampling_params, |
| 53 | data_parallel_rank=data_parallel_rank, |
| 54 | ): |
| 55 | num_tokens = len(out.outputs[0].token_ids) |
| 56 | if output_kind == RequestOutputKind.DELTA: |
| 57 | count += num_tokens |
| 58 | else: |
| 59 | count = num_tokens |
| 60 | |
| 61 | await asyncio.sleep(0.0) |
| 62 | |
| 63 | return count, request_id |
| 64 | |
| 65 | |
| 66 | @pytest.mark.parametrize( |
| 67 | "model", |
| 68 | [ |
| 69 | "ibm-research/PowerMoE-3b", |
| 70 | "hmellor/tiny-random-LlamaForCausalLM", |
| 71 | ], |
| 72 | ) |
| 73 | @pytest.mark.parametrize( |
| 74 | "output_kind", |
| 75 | [ |
| 76 | RequestOutputKind.DELTA, |
| 77 | RequestOutputKind.FINAL_ONLY, |
| 78 | ], |
| 79 | ) |
| 80 | @pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"]) |
| 81 | @pytest.mark.parametrize("async_scheduling", [True, False]) |
| 82 | @pytest.mark.asyncio |
| 83 | async def test_load( |
| 84 | model: str, |
| 85 | output_kind: RequestOutputKind, |
| 86 | data_parallel_backend: DataParallelBackend, |
| 87 | async_scheduling: bool, |
| 88 | ): |
| 89 | if async_scheduling and data_parallel_backend == "ray": |
| 90 | # TODO(NickLucche) Re-enable when async scheduling is supported |
| 91 | pytest.skip("Async scheduling is not supported with ray") |
| 92 | elif data_parallel_backend == "ray" and current_platform.is_rocm(): |
| 93 | pytest.skip( |
| 94 | "Ray as the distributed executor backend is not supported with ROCm." |
| 95 | ) |
| 96 | stats_loggers = {} |
| 97 | |
| 98 | @dataclass |
| 99 | class SimpleStatsLogger(StatLoggerBase): |
| 100 | init_count: int = 0 |
| 101 | finished_req_count: int = 0 |
| 102 | |
| 103 | def __init__(self, vllm_config: VllmConfig, engine_index: int = 0): |
| 104 | stats_loggers[engine_index] = self |
| 105 | |
| 106 | def record( |
| 107 | self, |
| 108 | scheduler_stats: SchedulerStats | None, |
| 109 | iteration_stats: IterationStats | None, |
| 110 | mm_cache_stats: MultiModalCacheStats | None = None, |
| 111 | engine_idx: int = 0, |
| 112 | ): |
| 113 | if iteration_stats: |
| 114 | self.finished_req_count += len(iteration_stats.finished_requests) |
| 115 | |
| 116 | def log_engine_initialized(self): |
| 117 | self.init_count += 1 |
| 118 | |
| 119 | with ExitStack() as after: |
| 120 | prompt = "This is a test of data parallel" |
| 121 | |
| 122 | engine_args = AsyncEngineArgs( |
| 123 | model=model, |
| 124 | enforce_eager=True, |
| 125 | tensor_parallel_size=int(os.getenv("TP_SIZE", 1)), |
| 126 | data_parallel_size=DP_SIZE, |
| 127 | data_parallel_backend=data_parallel_backend, |
| 128 | async_scheduling=async_scheduling, |
| 129 | ) |
| 130 | engine = AsyncLLM.from_engine_args( |
| 131 | engine_args, stat_loggers=[SimpleStatsLogger] |
| 132 | ) |
| 133 | after.callback(engine.shutdown) |
| 134 | |
| 135 | NUM_REQUESTS = 100 |
| 136 | NUM_EXPECTED_TOKENS = 10 |
| 137 | |
| 138 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 139 | |
| 140 | # Create concurrent requests. |
| 141 | tasks = [] |
| 142 | for request_id in request_ids: |
| 143 | tasks.append( |
| 144 | asyncio.create_task( |
| 145 | generate( |
| 146 | engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS |
| 147 | ) |
| 148 | ) |
| 149 | ) |
| 150 | # Short sleep to ensure that requests are distributed. |
| 151 | await asyncio.sleep(0.01) |
| 152 | # Confirm that we got all the EXPECTED tokens from the requests. |
| 153 | done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) |
| 154 | for task in pending: |
| 155 | task.cancel() |
| 156 | for task in done: |
| 157 | num_generated_tokens, request_id = await task |
| 158 | assert num_generated_tokens == NUM_EXPECTED_TOKENS, ( |
| 159 | f"{request_id} generated {num_generated_tokens} but " |
| 160 | f"expected {NUM_EXPECTED_TOKENS}" |
| 161 | ) |
| 162 | |
| 163 | assert not engine.output_processor.has_unfinished_requests() |
| 164 | |
| 165 | # testing internals here which may break |
| 166 | core_client = engine.engine_core |
| 167 | assert isinstance(core_client, DPLBAsyncMPClient) |
| 168 | # the engines only synchronize stopping every N steps so |
| 169 | # allow a small amount of time here. |
| 170 | for _ in range(10): |
| 171 | if not core_client.engines_running: |
| 172 | break |
| 173 | await asyncio.sleep(0.5) |
| 174 | |
| 175 | assert not core_client.engines_running |
| 176 | assert not core_client.reqs_in_flight |
| 177 | |
| 178 | # Check that requests were distributed between the engines |
| 179 | print(f"Stats loggers after test: {stats_loggers}") |
| 180 | assert len(stats_loggers) == DP_SIZE |
| 181 | assert stats_loggers[0].init_count == 1 |
| 182 | |
| 183 | for sl in stats_loggers.values(): |
| 184 | slogger: SimpleStatsLogger = sl |
| 185 | |
| 186 | assert slogger.finished_req_count > NUM_REQUESTS // (DP_SIZE + 1), ( |
| 187 | f"requests are imbalanced: {stats_loggers}" |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | @pytest.mark.parametrize("prefill_schedule_interval", [1, 4]) |
| 192 | @pytest.mark.asyncio |
| 193 | async def test_dp_prefill_schedule_interval(prefill_schedule_interval: int): |
| 194 | """Throttling new prefills to every Nth step (DP balancing) must not break |
| 195 | generation: a stream of staggered requests should still all complete with |
| 196 | the expected number of tokens. |
| 197 | |
| 198 | The throttle only engages in the DP MoE/EP engine-core path |
| 199 | (`DPEngineCoreProc`), so this uses an MoE model with expert parallel. |
| 200 | """ |
| 201 | with ExitStack() as after: |
| 202 | prompt = "This is a test of data parallel" |
| 203 | |
| 204 | engine_args = AsyncEngineArgs( |
| 205 | model="ibm-research/PowerMoE-3b", |
| 206 | enforce_eager=True, |
| 207 | tensor_parallel_size=int(os.getenv("TP_SIZE", 1)), |
| 208 | data_parallel_size=DP_SIZE, |
| 209 | data_parallel_backend="mp", |
| 210 | enable_expert_parallel=True, |
| 211 | prefill_schedule_interval=prefill_schedule_interval, |
| 212 | ) |
| 213 | engine = AsyncLLM.from_engine_args(engine_args) |
| 214 | after.callback(engine.shutdown) |
| 215 | |
| 216 | NUM_REQUESTS = 50 |
| 217 | NUM_EXPECTED_TOKENS = 10 |
| 218 | |
| 219 | request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] |
| 220 | |
| 221 | # Create requests with a small stagger so they arrive across many |
| 222 | # steps and (with interval > 1) accumulate in the waiting queue |
| 223 | # before being admitted together on cadence-aligned steps. |
| 224 | tasks = [] |
| 225 | for request_id in request_ids: |
| 226 | tasks.append( |
| 227 | asyncio.create_task( |
| 228 | generate( |
| 229 | engine, |
| 230 | request_id, |
| 231 | prompt, |
| 232 | RequestOutputKind.DELTA, |
| 233 | NUM_EXPECTED_TOKENS, |
| 234 | ) |
| 235 | ) |
| 236 | ) |
| 237 | await asyncio.sleep(0.01) |
| 238 | |
| 239 | done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) |
| 240 | for task in pending: |
| 241 | task.cancel() |
| 242 | for task in done: |
| 243 | num_generated_tokens, request_id = await task |
| 244 | assert num_generated_tokens == NUM_EXPECTED_TOKENS, ( |
| 245 | f"{request_id} generated {num_generated_tokens} but " |
| 246 | f"expected {NUM_EXPECTED_TOKENS}" |
| 247 | ) |
| 248 | |
| 249 | assert not engine.output_processor.has_unfinished_requests() |
| 250 | |
| 251 | |
| 252 | # ============================================================================= |
| 253 | # DP Pause/Resume Tests |
| 254 | # ============================================================================= |
| 255 | # When expert_parallel=False: uses non-MoE model (DP replicas as separate engines). |
| 256 | # When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path). |
| 257 | |
| 258 | DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM" |
| 259 | DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b" |
| 260 | DP_PAUSE_PROMPT = "This is a test of data parallel pause" |
| 261 | |
| 262 | |
| 263 | def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs: |
| 264 | """Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama.""" |
| 265 | model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL |
| 266 | return AsyncEngineArgs( |
| 267 | model=model, |
| 268 | enforce_eager=True, |
| 269 | tensor_parallel_size=int(os.getenv("TP_SIZE", 1)), |
| 270 | data_parallel_size=DP_SIZE, |
| 271 | data_parallel_backend="mp", |
| 272 | enable_expert_parallel=expert_parallel, |
| 273 | ) |
| 274 | |
| 275 | |
| 276 | @pytest.mark.asyncio |
| 277 | @pytest.mark.parametrize("expert_parallel", [False, True]) |
| 278 | async def test_dp_pause_resume_basic(expert_parallel: bool): |
| 279 | """Pausing from the client (one call) pauses all DP ranks; resume clears it.""" |
| 280 | with ExitStack() as after: |
| 281 | engine_args = _get_dp_pause_engine_args(expert_parallel) |
| 282 | engine = AsyncLLM.from_engine_args(engine_args) |
| 283 | after.callback(engine.shutdown) |
| 284 | |
| 285 | assert not await engine.is_paused() |
| 286 | await engine.pause_generation(mode="abort") |
| 287 | assert await engine.is_paused() |
| 288 | await engine.resume_generation() |
| 289 | assert not await engine.is_paused() |
| 290 | |
| 291 | # Engine still works after resume |
| 292 | sampling_params = SamplingParams(max_tokens=5) |
| 293 | async for out in engine.generate( |
| 294 | request_id="after-resume", |
| 295 | prompt=DP_PAUSE_PROMPT, |
| 296 | sampling_params=sampling_params, |
| 297 | ): |
| 298 | pass |
| 299 | assert out.finished |
| 300 | |
| 301 | |
| 302 | async def _consume(generator) -> None: |
| 303 | async for _ in generator: |
| 304 | pass |
| 305 | |
| 306 | |
| 307 | async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool: |
| 308 | """Wait for the front-end's view of the DP engines to reach ``want``.""" |
| 309 | deadline = time.monotonic() + timeout |
| 310 | while time.monotonic() < deadline: |
| 311 | if engine.engine_core.dp_engines_running() == want: |
| 312 | return True |
| 313 | await asyncio.sleep(0.05) |
| 314 | return False |
| 315 | |
| 316 | |
| 317 | @pytest.mark.asyncio |
| 318 | async def test_dp_pause_late_request_does_not_block_drain(): |
| 319 | """A request arriving after pause must not leave the coordinator believing |
| 320 | the engines are running. |
| 321 | |
| 322 | Paused engines discard START_DP_WAVE, so nothing can report wave |
| 323 | completion afterwards; if forwarding the wake also marks the engines as |
| 324 | running, drain has no way back and can only end in a timeout. |
| 325 | |
| 326 | MoE only: wave coordination is enabled iff the model is MoE, so a dense |
| 327 | model never reaches the coordinator state this exercises. |
| 328 | """ |
| 329 | with ExitStack() as after: |
| 330 | engine_args = _get_dp_pause_engine_args(expert_parallel=True) |
| 331 | engine = AsyncLLM.from_engine_args(engine_args) |
| 332 | after.callback(engine.shutdown) |
| 333 | |
| 334 | # Run a wave first, so the engines are quiesced by the pause rather |
| 335 | # than by never having started, and drain has something to observe. |
| 336 | long_request = asyncio.create_task( |
| 337 | _consume( |
| 338 | engine.generate( |
| 339 | request_id="warmup", |
| 340 | prompt=DP_PAUSE_PROMPT, |
| 341 | sampling_params=SamplingParams(max_tokens=400, ignore_eos=True), |
| 342 | ) |
| 343 | ) |
| 344 | ) |
| 345 | |
| 346 | # A design that never reports the engines as running would satisfy |
| 347 | # every drain assertion below while destroying the signal, so pin it |
| 348 | # down first. The front-end sets its own copy optimistically when it |
| 349 | # forwards the wake, so sample only after several coordinator |
| 350 | # publishes (every 100ms while stats change) have overwritten it. |
| 351 | await asyncio.sleep(2) |
| 352 | assert not long_request.done(), "the warmup request was too short to sample" |
| 353 | assert engine.engine_core.dp_engines_running(), ( |
| 354 | "the coordinator does not report the engines as running while they are" |
| 355 | ) |
| 356 | |
| 357 | await long_request |
| 358 | assert await _poll_flag(engine, False, timeout=30) |
| 359 | |
| 360 | await engine.pause_generation(mode="abort") |
| 361 | await engine.wait_for_requests_to_drain(drain_timeout=30) |
| 362 | |
| 363 | # Awaiting add_request guarantees the new-request notification has been |
| 364 | # sent to the coordinator - the message that used to latch the flag. |
| 365 | collector = await engine.add_request( |
| 366 | request_id="late", |
| 367 | prompt=DP_PAUSE_PROMPT, |
| 368 | params=SamplingParams(max_tokens=5), |
| 369 | ) |
| 370 | |
| 371 | # The front-end marks the engines running off the back of that |
| 372 | # notification. This is what makes the test non-vacuous: it is the |
| 373 | # path that used to leave the coordinator stuck. |
| 374 | assert await _poll_flag(engine, True, timeout=5), ( |
| 375 | "the late request did not notify the coordinator" |
| 376 | ) |
| 377 | |
| 378 | # It must settle back by itself. Unfixed it never does, because the |
| 379 | # paused engines discard the wake and so never report wave completion. |
| 380 | assert await _poll_flag(engine, False, timeout=60) |
| 381 | |
| 382 | # The late request was held rather than dropped: it completes on resume. |
| 383 | await engine.resume_generation() |
| 384 | while True: |
| 385 | out = await asyncio.wait_for(collector.get(), timeout=60) |
| 386 | if out.finished: |
| 387 | break |
| 388 | |
| 389 | |
| 390 | @pytest.mark.asyncio |
| 391 | async def test_dp_sleep_late_request_does_not_block_drain(): |
| 392 | """The same latch, reached through sleep rather than pause. |
| 393 | |
| 394 | Sleep stops the engines stepping just as pause does, so a request arriving |
| 395 | while they are asleep can leave the coordinator believing they are running |
| 396 | with nothing able to say otherwise. This is worth pinning separately from |
| 397 | the pause case because it is the shape that reaches |
| 398 | `_drain_requests_for_elastic_ep`, which decides from the same signal |
| 399 | whether it is safe to scale. |
| 400 | """ |
| 401 | with ExitStack() as after: |
| 402 | engine_args = _get_dp_pause_engine_args(expert_parallel=True) |
| 403 | engine = AsyncLLM.from_engine_args(engine_args) |
| 404 | after.callback(engine.shutdown) |
| 405 | |
| 406 | async for _ in engine.generate( |
| 407 | request_id="warmup", |
| 408 | prompt=DP_PAUSE_PROMPT, |
| 409 | sampling_params=SamplingParams(max_tokens=5), |
| 410 | ): |
| 411 | pass |
| 412 | assert await _poll_flag(engine, False, timeout=30) |
| 413 | |
| 414 | await engine.sleep(level=1) |
| 415 | assert await engine.is_sleeping() |
| 416 | |
| 417 | collector = await engine.add_request( |
| 418 | request_id="while-asleep", |
| 419 | prompt=DP_PAUSE_PROMPT, |
| 420 | params=SamplingParams(max_tokens=5), |
| 421 | ) |
| 422 | |
| 423 | assert await _poll_flag(engine, True, timeout=5), ( |
| 424 | "the request did not notify the coordinator" |
| 425 | ) |
| 426 | |
| 427 | # Sleeping engines cannot report wave completion, so a coordinator that |
| 428 | # marked them running when it forwarded the wake never hears otherwise. |
| 429 | assert await _poll_flag(engine, False, timeout=60) |
| 430 | |
| 431 | await engine.wake_up() |
| 432 | assert not await engine.is_sleeping() |
| 433 | while True: |
| 434 | out = await asyncio.wait_for(collector.get(), timeout=60) |
| 435 | if out.finished: |
| 436 | break |
| 437 | |
| 438 | |
| 439 | @pytest.mark.asyncio |
| 440 | @pytest.mark.parametrize("expert_parallel", [False, True]) |
| 441 | async def test_dp_pause_abort(expert_parallel: bool): |
| 442 | """Pause with abort from one client aborts in-flight requests on all DP ranks.""" |
| 443 | with ExitStack() as after: |
| 444 | engine_args = _get_dp_pause_engine_args(expert_parallel) |
| 445 | engine = AsyncLLM.from_engine_args(engine_args) |
| 446 | after.callback(engine.shutdown) |
| 447 | |
| 448 | # Start several requests so they are distributed across ranks |
| 449 | sampling_params = SamplingParams(max_tokens=500, ignore_eos=True) |
| 450 | num_requests = 4 |
| 451 | outputs_by_id: dict[str, list[RequestOutput]] = {} |
| 452 | |
| 453 | async def gen(rid: str): |
| 454 | out_list: list[RequestOutput] = [] |
| 455 | outputs_by_id[rid] = out_list |
| 456 | async for out in engine.generate( |
| 457 | request_id=rid, |
| 458 | prompt=DP_PAUSE_PROMPT, |
| 459 | sampling_params=sampling_params, |
| 460 | ): |
| 461 | out_list.append(out) |
| 462 | return out_list[-1] if out_list else None |
| 463 | |
| 464 | tasks = [asyncio.create_task(gen(f"req-{i}")) for i in range(num_requests)] |
| 465 | # Wait for some tokens on at least one request |
| 466 | while not any(len(o) >= 2 for o in outputs_by_id.values()): |
| 467 | await asyncio.sleep(0.02) |
| 468 | |
| 469 | await engine.pause_generation(mode="abort") |
| 470 | |
| 471 | finals = await asyncio.gather(*tasks) |
| 472 | for i, final in enumerate(finals): |
| 473 | assert final is not None, f"req-{i} had no output" |
| 474 | assert final.finished |
| 475 | assert final.outputs[0].finish_reason == "abort" |
| 476 | |
| 477 | assert await engine.is_paused() |
| 478 | await engine.resume_generation() |
| 479 | assert not await engine.is_paused() |
| 480 | |
| 481 | # New request completes after resume |
| 482 | async for out in engine.generate( |
| 483 | request_id="after-abort", |
| 484 | prompt=DP_PAUSE_PROMPT, |
| 485 | sampling_params=SamplingParams(max_tokens=5), |
| 486 | ): |
| 487 | pass |
| 488 | assert out.finished |
| 489 | assert not engine.output_processor.has_unfinished_requests() |
| 490 | |
| 491 | |
| 492 | @pytest.mark.asyncio |
| 493 | @pytest.mark.parametrize("expert_parallel", [False, True]) |
| 494 | async def test_dp_pause_keep_then_resume(expert_parallel: bool): |
| 495 | """Start generation, pause after a few tokens (keep mode), resume; verify gap.""" |
| 496 | pause_duration = 2.0 |
| 497 | min_tokens_before_pause = 3 |
| 498 | |
| 499 | with ExitStack() as after: |
| 500 | engine_args = _get_dp_pause_engine_args(expert_parallel) |
| 501 | engine = AsyncLLM.from_engine_args(engine_args) |
| 502 | after.callback(engine.shutdown) |
| 503 | |
| 504 | sampling_params = SamplingParams(max_tokens=15, ignore_eos=True) |
| 505 | token_times: list[tuple[int, float]] = [] |
| 506 | pause_token_idx = 0 |
| 507 | |
| 508 | async def generator_task(): |
| 509 | nonlocal pause_token_idx |
| 510 | out = None |
| 511 | async for output in engine.generate( |
| 512 | request_id="keep-resume-req", |
| 513 | prompt=DP_PAUSE_PROMPT, |
| 514 | sampling_params=sampling_params, |
| 515 | ): |
| 516 | token_count = len(output.outputs[0].token_ids) |
| 517 | token_times.append((token_count, time.monotonic())) |
| 518 | out = output |
| 519 | return out |
| 520 | |
| 521 | async def controller_task(): |
| 522 | nonlocal pause_token_idx |
| 523 | while len(token_times) < min_tokens_before_pause: |
| 524 | await asyncio.sleep(0.01) |
| 525 | await engine.pause_generation(mode="keep") |
| 526 | await asyncio.sleep(pause_duration) |
| 527 | pause_token_idx = len(token_times) |
| 528 | await engine.resume_generation() |
| 529 | |
| 530 | gen_task = asyncio.create_task(generator_task()) |
| 531 | ctrl_task = asyncio.create_task(controller_task()) |
| 532 | final_output, _ = await asyncio.gather(gen_task, ctrl_task) |
| 533 | |
| 534 | assert final_output is not None and final_output.finished |
| 535 | assert await engine.is_paused() is False |
| 536 | assert pause_token_idx >= min_tokens_before_pause |
| 537 | if pause_token_idx > 0 and pause_token_idx < len(token_times): |
| 538 | pause_gap = ( |
| 539 | token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1] |
| 540 | ) |
| 541 | assert pause_gap >= pause_duration * 0.8, ( |
| 542 | f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s" |
| 543 | ) |
| 544 | |
| 545 | |
| 546 | @pytest.mark.asyncio |
| 547 | async def test_dp_pause_keep_race_staggered_engines(): |
| 548 | """Race: send pause(keep) to engine 0, then add two requests, |
| 549 | then pause(keep) to engine 1. Ensures no deadlock when pause |
| 550 | requests are staggered and requests arrive in between.""" |
| 551 | if DP_SIZE != 2: |
| 552 | pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2") |
| 553 | |
| 554 | with ExitStack() as after: |
| 555 | engine_args = _get_dp_pause_engine_args(expert_parallel=True) |
| 556 | engine = AsyncLLM.from_engine_args(engine_args) |
| 557 | after.callback(engine.shutdown) |
| 558 | |
| 559 | client = engine.engine_core |
| 560 | |
| 561 | original_call_utility = client.call_utility_async |
| 562 | mid_pause_tasks: list[asyncio.Task] = [] |
| 563 | |
| 564 | async def staggered_pause_keep(method: str, *args) -> Any: |
| 565 | if method != "pause_scheduler" or not args or args[0] != "keep": |
| 566 | return await original_call_utility(method, *args) |
| 567 | # Fire pause(keep) to engine 0 (don't await — with DP |
| 568 | # two-phase pause, consensus requires all ranks). |
| 569 | pause_0 = asyncio.create_task( |
| 570 | client._call_utility_async(method, *args, engine=client.core_engines[0]) |
| 571 | ) |
| 572 | # Let the event loop send the message to engine 0. |
| 573 | await asyncio.sleep(0.5) |
| 574 | # In the middle: send two requests (race window) |
| 575 | sp = SamplingParams(max_tokens=5, ignore_eos=True) |
| 576 | |
| 577 | async def consume_gen(req_id: str) -> None: |
| 578 | async for _ in engine.generate( |
| 579 | request_id=req_id, |
| 580 | prompt=DP_PAUSE_PROMPT, |
| 581 | sampling_params=sp, |
| 582 | ): |
| 583 | pass |
| 584 | |
| 585 | t1 = asyncio.create_task(consume_gen("race-1")) |
| 586 | t2 = asyncio.create_task(consume_gen("race-2")) |
| 587 | mid_pause_tasks.extend([t1, t2]) |
| 588 | await asyncio.sleep(3) |
| 589 | # Fire pause(keep) to engine 1, then await both so |
| 590 | # consensus can be reached. |
| 591 | pause_1 = asyncio.create_task( |
| 592 | client._call_utility_async(method, *args, engine=client.core_engines[1]) |
| 593 | ) |
| 594 | results = await asyncio.gather(pause_0, pause_1) |
| 595 | return results[0] |
| 596 | |
| 597 | client.call_utility_async = staggered_pause_keep |
| 598 | |
| 599 | await engine.pause_generation(mode="keep") |
| 600 | assert await engine.is_paused() |
| 601 | await engine.resume_generation() |
| 602 | assert not await engine.is_paused() |
| 603 | # Let the two requests we sent mid-pause complete |
| 604 | await asyncio.gather(*mid_pause_tasks) |
| 605 | |
| 606 | |
| 607 | @pytest.mark.asyncio |
| 608 | async def test_dp_pause_barrier_request_deadlock(): |
| 609 | """Test that start_dp_wave is ignored while paused. |
| 610 | |
| 611 | Sequence: |
| 612 | 1. Pause all engines (PAUSED_ALL). |
| 613 | 2. Send barrier to engine 0 only — blocks in dist.barrier(dp_group). |
| 614 | 3. Send a request routed to engine 1. |
| 615 | 4. Wait for any (buggy) START_DP_WAVE propagation. |
| 616 | 5. Send barrier to engine 1 — completes in fixed code, deadlocks |
| 617 | in buggy code because engine 1 is stuck in EP all-to-all. |
| 618 | """ |
| 619 | if DP_SIZE != 2: |
| 620 | pytest.skip("requires DP_SIZE=2") |
| 621 | |
| 622 | with ExitStack() as after: |
| 623 | engine_args = _get_dp_pause_engine_args(expert_parallel=True) |
| 624 | engine = AsyncLLM.from_engine_args(engine_args) |
| 625 | after.callback(engine.shutdown) |
| 626 | |
| 627 | client = engine.engine_core |
| 628 | |
| 629 | # Cache get_supported_tasks so that generate() won't need to |
| 630 | # send a utility call to all engines (which would hang once |
| 631 | # engine 0 is blocked in the barrier). |
| 632 | await engine.get_supported_tasks() |
| 633 | |
| 634 | # Pause all engines normally — no staggering. |
| 635 | await engine.pause_generation(mode="keep") |
| 636 | assert await engine.is_paused() |
| 637 | |
| 638 | original_call_utility = client.call_utility_async |
| 639 | mid_barrier_tasks: list[asyncio.Task] = [] |
| 640 | |
| 641 | async def staggered_barrier(method: str, *args) -> Any: |
| 642 | if method != "barrier": |
| 643 | return await original_call_utility(method, *args) |
| 644 | |
| 645 | # Send barrier to engine 0 only — it blocks in |
| 646 | # dist.barrier(dp_group) waiting for engine 1. |
| 647 | barrier_0 = asyncio.create_task( |
| 648 | client._call_utility_async(method, *args, engine=client.core_engines[0]) |
| 649 | ) |
| 650 | await asyncio.sleep(1) |
| 651 | |
| 652 | # While engine 0 is blocked, send a request routed |
| 653 | # specifically to engine 1. |
| 654 | sp = SamplingParams(max_tokens=5, ignore_eos=True) |
| 655 | |
| 656 | engine_1 = client.core_engines[1] |
| 657 | original_get_engine = client.get_core_engine_for_request |
| 658 | |
| 659 | def route_to_engine_1(req): |
| 660 | client.reqs_in_flight[req.request_id] = engine_1 |
| 661 | return engine_1 |
| 662 | |
| 663 | client.get_core_engine_for_request = route_to_engine_1 |
| 664 | |
| 665 | async def consume_gen(req_id: str) -> None: |
| 666 | async for _ in engine.generate( |
| 667 | request_id=req_id, |
| 668 | prompt=DP_PAUSE_PROMPT, |
| 669 | sampling_params=sp, |
| 670 | ): |
| 671 | pass |
| 672 | |
| 673 | t1 = asyncio.create_task(consume_gen("race-1")) |
| 674 | mid_barrier_tasks.append(t1) |
| 675 | |
| 676 | # Yield so generate() preprocessing completes and |
| 677 | # add_request_async is called (which, in buggy code, |
| 678 | # would send FIRST_REQ and wake engine 1). |
| 679 | for _ in range(200): |
| 680 | await asyncio.sleep(0) |
| 681 | |
| 682 | client.get_core_engine_for_request = original_get_engine |
| 683 | |
| 684 | # Wait for any START_DP_WAVE to propagate and for |
| 685 | # engine 1 to potentially enter execute_dummy_batch. |
| 686 | await asyncio.sleep(5) |
| 687 | |
| 688 | # Now send barrier to engine 1. In buggy code engine 1 |
| 689 | # is stuck in execute_dummy_batch (EP all-to-all) while |
| 690 | # engine 0 is stuck in dist.barrier(dp_group) — deadlock. |
| 691 | result = await client._call_utility_async( |
| 692 | method, *args, engine=client.core_engines[1] |
| 693 | ) |
| 694 | await barrier_0 |
| 695 | return result |
| 696 | |
| 697 | client.call_utility_async = staggered_barrier |
| 698 | |
| 699 | # Drive the staggered barrier. Old code deadlocks here. |
| 700 | try: |
| 701 | await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30) |
| 702 | except asyncio.TimeoutError: |
| 703 | for t in mid_barrier_tasks: |
| 704 | t.cancel() |
| 705 | pytest.fail( |
| 706 | "Staggered barrier deadlocked — FIRST_REQ sent while " |
| 707 | "paused caused collective-op mismatch between engines" |
| 708 | ) |
| 709 | |
| 710 | await engine.resume_generation() |
| 711 | assert not await engine.is_paused() |
| 712 | # Let the two requests we sent mid-barrier complete. |
| 713 | await asyncio.gather(*mid_barrier_tasks) |
| 714 | |
| 715 | |
| 716 | @pytest.mark.asyncio |
| 717 | async def test_dp_pause_wait_mode_drains_in_flight(): |
| 718 | """mode="wait" through the DP consensus path: the pause lets the |
| 719 | in-flight request run to completion while the consensus is pending, |
| 720 | resolves only once drained, and the engine generates again on resume.""" |
| 721 | with ExitStack() as after: |
| 722 | engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) |
| 723 | after.callback(engine.shutdown) |
| 724 | |
| 725 | collector = await engine.add_request( |
| 726 | request_id="inflight", |
| 727 | prompt=DP_PAUSE_PROMPT, |
| 728 | params=SamplingParams(max_tokens=64), |
| 729 | ) |
| 730 | await engine.pause_generation(mode="wait") |
| 731 | assert await engine.is_paused() |
| 732 | while True: |
| 733 | out = await asyncio.wait_for(collector.get(), timeout=30) |
| 734 | if out.finished: |
| 735 | break |
| 736 | |
| 737 | await engine.resume_generation() |
| 738 | async for out in engine.generate( |
| 739 | request_id="after-wait", |
| 740 | prompt=DP_PAUSE_PROMPT, |
| 741 | sampling_params=SamplingParams(max_tokens=5), |
| 742 | ): |
| 743 | pass |
| 744 | assert out.finished |
| 745 | |
| 746 | |
| 747 | @pytest.mark.asyncio |
| 748 | async def test_dp_pause_while_asleep(): |
| 749 | """Pausing a sleeping DP engine: sleeping ranks skip dummy batches yet |
| 750 | must still reach the pause consensus, and the completion barrier must |
| 751 | be harmless on workers whose memory is unmapped.""" |
| 752 | with ExitStack() as after: |
| 753 | engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) |
| 754 | after.callback(engine.shutdown) |
| 755 | |
| 756 | async for _ in engine.generate( |
| 757 | request_id="warmup", |
| 758 | prompt=DP_PAUSE_PROMPT, |
| 759 | sampling_params=SamplingParams(max_tokens=5), |
| 760 | ): |
| 761 | pass |
| 762 | await engine.sleep(level=1) |
| 763 | assert await engine.is_sleeping() |
| 764 | |
| 765 | await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=30) |
| 766 | assert await engine.is_paused() |
| 767 | |
| 768 | # A full wake makes the memory resident and resumes the scheduler. |
| 769 | await engine.wake_up() |
| 770 | assert not await engine.is_sleeping() |
| 771 | async for out in engine.generate( |
| 772 | request_id="after-wake", |
| 773 | prompt=DP_PAUSE_PROMPT, |
| 774 | sampling_params=SamplingParams(max_tokens=5), |
| 775 | ): |
| 776 | pass |
| 777 | assert out.finished |
| 778 | |
| 779 | |
| 780 | @pytest.mark.asyncio |
| 781 | async def test_dp_pause_completion_implies_device_idle( |
| 782 | monkeypatch: pytest.MonkeyPatch, |
| 783 | ): |
| 784 | """A resolved pause future promises an idle device. Inflate the dummy |
| 785 | batches the pause consensus manufactures so any work the pause fails to |
| 786 | wait on stays visible, then require a quiet worker stream the moment |
| 787 | pause_generation() returns.""" |
| 788 | with ExitStack() as after: |
| 789 | # The probes below ship functions through collective_rpc. |
| 790 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 791 | engine = AsyncLLM.from_engine_args(_get_dp_pause_engine_args(True)) |
| 792 | after.callback(engine.shutdown) |
| 793 | |
| 794 | async for _ in engine.generate( |
| 795 | request_id="warmup", |
| 796 | prompt=DP_PAUSE_PROMPT, |
| 797 | sampling_params=SamplingParams(max_tokens=5), |
| 798 | ): |
| 799 | pass |
| 800 | |
| 801 | def inflate_dummy_batches(worker: Any) -> bool: |
| 802 | import torch |
| 803 | |
| 804 | inner = worker.execute_dummy_batch |
| 805 | |
| 806 | def slow_dummy_batch() -> None: |
| 807 | inner() |
| 808 | # Keep the stream busy after launch, like a large-MoE forward. |
| 809 | torch.cuda._sleep(200_000_000) |
| 810 | |
| 811 | worker.execute_dummy_batch = slow_dummy_batch |
| 812 | return True |
| 813 | |
| 814 | assert all(await engine.engine_core.collective_rpc_async(inflate_dummy_batches)) |
| 815 | |
| 816 | await asyncio.wait_for(engine.pause_generation(mode="abort"), timeout=60) |
| 817 | |
| 818 | def device_idle(worker: Any) -> bool: |
| 819 | import torch |
| 820 | |
| 821 | return bool(torch.cuda.current_stream().query()) |
| 822 | |
| 823 | assert all(await engine.engine_core.collective_rpc_async(device_idle)) |