@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7ddb50788d

[Bugfix][DP] Don't assume the engines started when forwarding a wake (#51481)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

aoshen02 committed Aug 19, 2026 at 04:49 UTC 7ddb50788ddf0846518feff972028534e39bcb86

3 files changed +155 -2

tests/v1/distributed/test_async_llm_dp.py

+137

@@ -299,6 +299,143 @@ async def test_dp_pause_resume_basic(expert_parallel: bool):

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

vllm/v1/engine/coordinator.py

+4 -2

@@ -358,8 +358,10 @@ class DPCoordinatorProc:

358 # is handled by all the engines.

359 engine_to_exclude = None

360

361 - engines_running = True

362 - wave_state_changed = True

361 + # engines_running is only set from the engines'

362 + # own notifications; a paused engine discards

363 + # START_DP_WAVE, so sending it is not evidence

364 + # that the engines are running.

365 self._send_start_wave(

366 publish_back, current_wave, engine_to_exclude

367 )

vllm/v1/engine/core.py

+14

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

2156 # Loop until process is sent a SIGINT or SIGTERM

2157 while self._handle_shutdown():

2158 # 1) Poll the input queue until there is work to do.

2159 + was_running = self.engines_running

2160 self._process_input_queue()

2161 # Publish request counts before and after GPU step to ensure freshness.

2162 self._maybe_publish_request_counts()

@@ -2216,6 +2217,19 @@ class DPEngineCoreProc(EngineCoreProc):

2217 # Increment wave count and reset step counter.

2218 self.current_wave += 1

2219 self.step_counter = 0

2220 + elif (

2221 + not was_running

2222 + and self.has_coordinator

2223 + and self.dp_rank == 0

2224 + and not self.pending_pause

2225 + ):

2226 + # Mirror of the wave_complete notification above: the

2227 + # coordinator must observe this edge too rather than assume

2228 + # that a START_DP_WAVE it sent was acted upon, since a paused

2229 + # engine discards it.

2230 + self.output_queue.put_nowait(

2231 + (-1, EngineCoreOutputs(start_wave=self.current_wave))

2232 + )

2233

2234 raise SystemExit

2235