@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dbf0da817a
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 7, 19:00 UTC

[Core] Cleanup engine pause/sleep logic (#34528)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Feb 24, 2026 at 19:33 UTC dbf0da817adf06edf9906d666de7be788dfb66d6
10 files changed +302 -198
tests/v1/distributed/test_async_llm_dp.py
+120 -46
@@ -3,8 +3,10 @@
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
@@ -187,24 +189,33 @@ async def test_load(
189 # =============================================================================
190 # DP Pause/Resume Tests
191 # =============================================================================
192 +# When expert_parallel=False: uses non-MoE model (DP replicas as separate engines).
193 +# When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path).
194
195 DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
196 +DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b"
197 DP_PAUSE_PROMPT = "This is a test of data parallel pause"
198
199
200 +def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs:
201 + """Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama."""
202 + model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL
203 + return AsyncEngineArgs(
204 + model=model,
205 + enforce_eager=True,
206 + tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
207 + data_parallel_size=DP_SIZE,
208 + data_parallel_backend="mp",
209 + enable_expert_parallel=expert_parallel,
210 + )
211 +
212 +
213 @pytest.mark.asyncio
196 -async def test_dp_pause_resume_basic():
214 +@pytest.mark.parametrize("expert_parallel", [False, True])
215 +async def test_dp_pause_resume_basic(expert_parallel: bool):
216 """Pausing from the client (one call) pauses all DP ranks; resume clears it."""
198 - if current_platform.is_rocm():
199 - pytest.skip("DP pause tests use mp backend only")
217 with ExitStack() as after:
201 - engine_args = AsyncEngineArgs(
202 - model=DP_PAUSE_MODEL,
203 - enforce_eager=True,
204 - tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
205 - data_parallel_size=DP_SIZE,
206 - data_parallel_backend="mp",
207 - )
218 + engine_args = _get_dp_pause_engine_args(expert_parallel)
219 engine = AsyncLLM.from_engine_args(engine_args)
220 after.callback(engine.shutdown)
221
@@ -226,18 +237,11 @@ async def test_dp_pause_resume_basic():
237
238
239 @pytest.mark.asyncio
229 -async def test_dp_pause_abort():
240 +@pytest.mark.parametrize("expert_parallel", [False, True])
241 +async def test_dp_pause_abort(expert_parallel: bool):
242 """Pause with abort from one client aborts in-flight requests on all DP ranks."""
231 - if current_platform.is_rocm():
232 - pytest.skip("DP pause tests use mp backend only")
243 with ExitStack() as after:
234 - engine_args = AsyncEngineArgs(
235 - model=DP_PAUSE_MODEL,
236 - enforce_eager=True,
237 - tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
238 - data_parallel_size=DP_SIZE,
239 - data_parallel_backend="mp",
240 - )
244 + engine_args = _get_dp_pause_engine_args(expert_parallel)
245 engine = AsyncLLM.from_engine_args(engine_args)
246 after.callback(engine.shutdown)
247
@@ -286,41 +290,111 @@ async def test_dp_pause_abort():
290
291
292 @pytest.mark.asyncio
289 -async def test_dp_pause_keep_then_resume():
290 - """Pause with keep queues new requests; resume allows them to run."""
291 - if current_platform.is_rocm():
292 - pytest.skip("DP pause tests use mp backend only")
293 +@pytest.mark.parametrize("expert_parallel", [False, True])
294 +async def test_dp_pause_keep_then_resume(expert_parallel: bool):
295 + """Start generation, pause after a few tokens (keep mode), resume; verify gap."""
296 +
297 + pause_duration = 2.0
298 + min_tokens_before_pause = 3
299 +
300 with ExitStack() as after:
294 - engine_args = AsyncEngineArgs(
295 - model=DP_PAUSE_MODEL,
296 - enforce_eager=True,
297 - tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
298 - data_parallel_size=DP_SIZE,
299 - data_parallel_backend="mp",
300 - )
301 + engine_args = _get_dp_pause_engine_args(expert_parallel)
302 engine = AsyncLLM.from_engine_args(engine_args)
303 after.callback(engine.shutdown)
304
304 - await engine.pause_generation(mode="keep")
305 - assert await engine.is_paused()
306 -
307 - request_done = asyncio.Event()
305 + sampling_params = SamplingParams(max_tokens=15, ignore_eos=True)
306 + token_times: list[tuple[int, float]] = []
307 + pause_token_idx = 0
308
309 - async def gen():
310 - async for out in engine.generate(
311 - request_id="queued-keep",
309 + async def generator_task():
310 + nonlocal pause_token_idx
311 + out = None
312 + async for output in engine.generate(
313 + request_id="keep-resume-req",
314 prompt=DP_PAUSE_PROMPT,
313 - sampling_params=SamplingParams(max_tokens=5),
315 + sampling_params=sampling_params,
316 ):
315 - pass
316 - request_done.set()
317 + token_count = len(output.outputs[0].token_ids)
318 + token_times.append((token_count, time.monotonic()))
319 + out = output
320 return out
321
319 - task = asyncio.create_task(gen())
320 - await asyncio.sleep(0.2)
321 - assert not request_done.is_set()
322 + async def controller_task():
323 + nonlocal pause_token_idx
324 + while len(token_times) < min_tokens_before_pause:
325 + await asyncio.sleep(0.01)
326 + await engine.pause_generation(mode="keep")
327 + await asyncio.sleep(pause_duration)
328 + pause_token_idx = len(token_times)
329 + await engine.resume_generation()
330 +
331 + gen_task = asyncio.create_task(generator_task())
332 + ctrl_task = asyncio.create_task(controller_task())
333 + final_output, _ = await asyncio.gather(gen_task, ctrl_task)
334 +
335 + assert final_output is not None and final_output.finished
336 + assert await engine.is_paused() is False
337 + assert pause_token_idx >= min_tokens_before_pause
338 + if pause_token_idx > 0 and pause_token_idx < len(token_times):
339 + pause_gap = (
340 + token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
341 + )
342 + assert pause_gap >= pause_duration * 0.8, (
343 + f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s"
344 + )
345
346 +
347 +@pytest.mark.asyncio
348 +async def test_dp_pause_keep_race_staggered_engines():
349 + """Race: send pause(keep) to engine 0, then add two requests,
350 + then pause(keep) to engine 1. Ensures no deadlock when pause
351 + requests are staggered and requests arrive in between."""
352 + if DP_SIZE != 2:
353 + pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2")
354 +
355 + with ExitStack() as after:
356 + engine_args = _get_dp_pause_engine_args(expert_parallel=True)
357 + engine = AsyncLLM.from_engine_args(engine_args)
358 + after.callback(engine.shutdown)
359 +
360 + client = engine.engine_core
361 +
362 + original_call_utility = client.call_utility_async
363 + mid_pause_tasks: list[asyncio.Task] = []
364 +
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]
371 + )
372 + # In the middle: send two requests (race window)
373 + sp = SamplingParams(max_tokens=5, ignore_eos=True)
374 +
375 + async def consume_gen(req_id: str) -> None:
376 + async for _ in engine.generate(
377 + request_id=req_id,
378 + prompt=DP_PAUSE_PROMPT,
379 + sampling_params=sp,
380 + ):
381 + pass
382 +
383 + t1 = asyncio.create_task(consume_gen("race-1"))
384 + t2 = asyncio.create_task(consume_gen("race-2"))
385 + mid_pause_tasks.extend([t1, t2])
386 + 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 + )
391 + return result
392 +
393 + client.call_utility_async = staggered_pause_keep
394 +
395 + await engine.pause_generation(mode="keep")
396 + assert await engine.is_paused()
397 await engine.resume_generation()
324 - final = await asyncio.wait_for(task, timeout=10.0)
325 - assert final.finished
398 assert not await engine.is_paused()
399 + # Let the two requests we sent mid-pause complete
400 + await asyncio.gather(*mid_pause_tasks)
tests/v1/engine/test_engine_core_client.py
+7 -12
@@ -280,20 +280,15 @@ def echo_dc_nested(
280
281
282 def future_echo(self, value: Any, num_wait_loops: int = 2) -> Future:
283 - """Utility that returns a Future completed by a per_step_hook after
284 - num_wait_loops engine steps (tests deferred utility path).
283 + """Utility that returns a Future completed once the engine is idle
284 + (tests deferred utility path).
285 """
286 future: Future = Future()
287 - remaining = [num_wait_loops]
287
289 - def _step(engine: EngineCore) -> bool:
290 - remaining[0] -= 1
291 - if remaining[0] <= 0:
292 - future.set_result(value)
293 - return True # remove hook
294 - return False
288 + def idle(engine: EngineCore):
289 + future.set_result(value)
290
296 - self.per_step_hooks.add(_step)
291 + self._idle_state_callbacks.append(idle)
292 return future
293
294
@@ -832,8 +827,8 @@ async def test_engine_core_client_future_utility_async(
827 monkeypatch: pytest.MonkeyPatch,
828 subprocess_future_echo_patch,
829 ):
835 - """Test that a utility returning a Future (completed by a per_step_hook
836 - after N steps) completes when the future is done (engine uses add_done_callback).
830 + """Test that a utility returning a Future completes when the future is done
831 + (engine uses add_done_callback).
832 """
833 with monkeypatch.context() as m:
834 m.setattr(EngineCore, "future_echo", future_echo, raising=False)
vllm/engine/protocol.py
+1 -1
@@ -148,7 +148,7 @@ class EngineClient(ABC):
148 ...
149
150 @abstractmethod
151 - async def sleep(self, level: int = 1) -> None:
151 + async def sleep(self, level: int = 1, mode: "PauseMode" = "abort") -> None:
152 """Sleep the engine"""
153 ...
154
vllm/entrypoints/llm.py
+41 -47
@@ -87,6 +87,7 @@ from vllm.usage.usage_lib import UsageContext
87 from vllm.utils.counter import Counter
88 from vllm.utils.mistral import is_mistral_tokenizer
89 from vllm.utils.tqdm_utils import maybe_tqdm
90 +from vllm.v1.engine import PauseMode
91 from vllm.v1.engine.llm_engine import LLMEngine
92 from vllm.v1.sample.logits_processor import LogitsProcessor
93
@@ -441,8 +442,7 @@ class LLM:
442 A list of `RequestOutput` objects containing the
443 generated completions in the same order as the input prompts.
444 """
444 - model_config = self.model_config
445 - runner_type = model_config.runner_type
445 + runner_type = self.model_config.runner_type
446 if runner_type != "generate":
447 raise ValueError(
448 "LLM.generate() is only supported for generative models. "
@@ -489,46 +489,22 @@ class LLM:
489 Returns:
490 A list of request IDs for the enqueued requests.
491 """
492 - model_config = self.model_config
493 - runner_type = model_config.runner_type
492 + runner_type = self.model_config.runner_type
493 if runner_type != "generate":
494 raise ValueError("LLM.enqueue() is only supported for generative models.")
495
496 if sampling_params is None:
497 sampling_params = self.get_default_sampling_params()
498
500 - # Use the same preprocessing as _run_completion
501 - seq_prompts = prompt_to_seq(prompts)
502 - seq_params = self._params_to_seq(sampling_params, len(seq_prompts))
503 - seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_prompts))
504 - seq_tok_kwargs = [
505 - merge_kwargs(
506 - tokenization_kwargs,
507 - dict(truncate_prompt_tokens=param.truncate_prompt_tokens),
508 - )
509 - for param in seq_params
510 - ]
511 - seq_priority = self._priority_to_seq(priority, len(prompts))
512 -
513 - request_ids = self._render_and_add_requests(
514 - prompts=(
515 - self._preprocess_cmpl_one(prompt, tok_kwargs)
516 - for prompt, tok_kwargs in zip(
517 - maybe_tqdm(
518 - seq_prompts,
519 - use_tqdm=use_tqdm,
520 - desc="Rendering prompts",
521 - ),
522 - seq_tok_kwargs,
523 - )
524 - ),
525 - params=seq_params,
526 - lora_requests=seq_lora_requests,
527 - priorities=seq_priority,
499 + return self._add_completion_requests(
500 + prompts=prompts,
501 + params=sampling_params,
502 + use_tqdm=use_tqdm,
503 + lora_request=lora_request,
504 + priority=priority,
505 + tokenization_kwargs=tokenization_kwargs,
506 )
507
530 - return request_ids
531 -
508 @overload
509 def wait_for_completion(
510 self,
@@ -1659,7 +1635,7 @@ class LLM:
1635 reset_running_requests, reset_connector
1636 )
1637
1662 - def sleep(self, level: int = 1):
1638 + def sleep(self, level: int = 1, mode: PauseMode = "abort"):
1639 """
1640 Put the engine to sleep. The engine should not process any requests.
1641 The caller should guarantee that no requests are being processed
@@ -1679,10 +1655,10 @@ class LLM:
1655 a different model or update the model, where
1656 previous model weights are not needed. It reduces
1657 CPU memory pressure.
1658 + mode: How to handle any existing requests, can be "abort", "wait",
1659 + or "keep".
1660 """
1683 - if level > 0:
1684 - self.reset_prefix_cache()
1685 - self.llm_engine.sleep(level=level)
1661 + self.llm_engine.sleep(level=level, mode=mode)
1662
1663 def wake_up(self, tags: list[str] | None = None):
1664 """
@@ -1759,19 +1735,18 @@ class LLM:
1735
1736 return [0] * num_requests
1737
1762 - def _run_completion(
1738 + def _add_completion_requests(
1739 self,
1740 prompts: PromptType | Sequence[PromptType],
1741 params: SamplingParams
1742 | PoolingParams
1743 | Sequence[SamplingParams | PoolingParams],
1768 - output_type: type[_O],
1744 *,
1745 use_tqdm: bool | Callable[..., tqdm] = True,
1746 lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
1747 priority: list[int] | None = None,
1748 tokenization_kwargs: dict[str, Any] | None = None,
1774 - ):
1749 + ) -> list[str]:
1750 seq_prompts = prompt_to_seq(prompts)
1751 seq_params = self._params_to_seq(params, len(seq_prompts))
1752 seq_lora_requests = self._lora_request_to_seq(lora_request, len(seq_prompts))
@@ -1784,25 +1759,44 @@ class LLM:
1759 ]
1760 seq_priority = self._priority_to_seq(priority, len(prompts))
1761
1787 - return self._render_and_run_requests(
1762 + return self._render_and_add_requests(
1763 prompts=(
1764 self._preprocess_cmpl_one(prompt, tok_kwargs)
1765 for prompt, tok_kwargs in zip(
1766 maybe_tqdm(
1792 - seq_prompts,
1793 - use_tqdm=use_tqdm,
1794 - desc="Rendering prompts",
1767 + seq_prompts, use_tqdm=use_tqdm, desc="Rendering prompts"
1768 ),
1769 seq_tok_kwargs,
1770 )
1771 ),
1772 params=seq_params,
1800 - output_type=output_type,
1801 - use_tqdm=use_tqdm,
1773 lora_requests=seq_lora_requests,
1774 priorities=seq_priority,
1775 )
1776
1777 + def _run_completion(
1778 + self,
1779 + prompts: PromptType | Sequence[PromptType],
1780 + params: SamplingParams
1781 + | PoolingParams
1782 + | Sequence[SamplingParams | PoolingParams],
1783 + output_type: type[_O],
1784 + *,
1785 + use_tqdm: bool | Callable[..., tqdm] = True,
1786 + lora_request: Sequence[LoRARequest] | LoRARequest | None = None,
1787 + priority: list[int] | None = None,
1788 + tokenization_kwargs: dict[str, Any] | None = None,
1789 + ):
1790 + self._add_completion_requests(
1791 + prompts=prompts,
1792 + params=params,
1793 + use_tqdm=use_tqdm,
1794 + lora_request=lora_request,
1795 + priority=priority,
1796 + tokenization_kwargs=tokenization_kwargs,
1797 + )
1798 + return self._run_engine(use_tqdm=use_tqdm, output_type=output_type)
1799 +
1800 def _run_chat(
1801 self,
1802 messages: list[ChatCompletionMessageParam]
vllm/entrypoints/serve/sleep/api_router.py
+2 -1
@@ -23,7 +23,8 @@ router = APIRouter()
23 async def sleep(raw_request: Request):
24 # get POST params
25 level = raw_request.query_params.get("level", "1")
26 - await engine_client(raw_request).sleep(int(level))
26 + mode = raw_request.query_params.get("mode", "abort")
27 + await engine_client(raw_request).sleep(int(level), mode)
28 # FIXME: in v0 with frontend multiprocessing, the sleep command
29 # is sent but does not finish yet when we return a response.
30 return Response(status_code=200)
vllm/v1/engine/async_llm.py
+9 -4
@@ -753,6 +753,13 @@ class AsyncLLM(EngineClient):
753 )
754 mode = "wait"
755 await self.engine_core.pause_scheduler_async(mode=mode, clear_cache=clear_cache)
756 + # Small sleep to help ensure that final outputs from any in-flight requests are
757 + # returned prior to this method returning. These outputs come out of the engine
758 + # prior to the wait-for-idle completion event, but involve additional async
759 + # tasks in output processing.
760 + # Note that this is not required for correctness, just more intuitive ordering
761 + # of events from caller's pov.
762 + await asyncio.sleep(0.02)
763
764 async def resume_generation(self) -> None:
765 """Resume generation after :meth:`pause_generation`."""
@@ -890,10 +897,8 @@ class AsyncLLM(EngineClient):
897 async def reset_encoder_cache(self) -> None:
898 await self.engine_core.reset_encoder_cache_async()
899
893 - async def sleep(self, level: int = 1) -> None:
894 - if level > 0:
895 - await self.reset_prefix_cache()
896 - await self.engine_core.sleep_async(level)
900 + async def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
901 + await self.engine_core.sleep_async(level, mode)
902
903 if self.logger_manager is not None:
904 self.logger_manager.record_sleep_state(1, level)
vllm/v1/engine/core.py
+108 -62
@@ -9,6 +9,7 @@ from collections import defaultdict, deque
9 from collections.abc import Callable, Generator
10 from concurrent.futures import Future
11 from contextlib import ExitStack, contextmanager
12 +from functools import partial
13 from inspect import isclass, signature
14 from logging import DEBUG
15 from typing import Any, TypeVar, cast
@@ -211,7 +212,7 @@ class EngineCore:
212
213 self.aborts_queue = queue.Queue[list[str]]()
214
214 - self.per_step_hooks: set[Callable] = set()
215 + self._idle_state_callbacks: list[Callable] = []
216
217 # Mark the startup heap as static so that it's ignored by GC.
218 # Reduces pause times of oldest generation collections.
@@ -592,21 +593,51 @@ class EngineCore:
593 # Reset the GPU model runner's encoder cache (physical storage)
594 self.model_executor.reset_encoder_cache()
595
596 + def _reset_caches(self, reset_running_requests=True) -> None:
597 + self.reset_prefix_cache(reset_running_requests=reset_running_requests)
598 + self.reset_mm_cache()
599 + self.reset_encoder_cache()
600 +
601 def pause_scheduler(
602 self, mode: PauseMode = "abort", clear_cache: bool = True
597 - ) -> Future[Any] | None:
598 - """Pause scheduling. No-op in base EngineCore; overridden in EngineCoreProc."""
603 + ) -> Future | None:
604 + """Pause generation; behavior depends on mode.
605 +
606 + All pause modes queue new adds -- "abort" and "keep" skip step();
607 + "wait" allows step() so in-flight requests can drain.
608 +
609 + - ``abort``: Set PAUSED_NEW, abort all requests, wait for abort
610 + outputs to be sent (when running with output_queue), optionally
611 + clear caches, then complete the returned Future.
612 + - ``wait``: Set PAUSED_NEW (queue adds, keep stepping); when drained,
613 + optionally clear caches, then complete the returned Future.
614 + - ``keep``: Set PAUSED_ALL; return a Future that completes when the
615 + output queue is empty.
616 + """
617 + if mode not in ("keep", "abort", "wait"):
618 + raise ValueError(f"Invalid pause mode: {mode}")
619 + if mode == "wait":
620 + raise ValueError("'wait' mode can't be used in inproc-engine mode")
621 +
622 + if mode == "abort":
623 + self.scheduler.finish_requests(None, RequestStatus.FINISHED_ABORTED)
624 +
625 + pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
626 + self.scheduler.set_pause_state(pause_state)
627 + if clear_cache:
628 + self._reset_caches()
629 +
630 return None
631
632 def resume_scheduler(self) -> None:
602 - """Resume scheduling. No-op in base EngineCore; overridden in EngineCoreProc."""
633 + """Resume the scheduler and flush any requests queued while paused."""
634 + self.scheduler.set_pause_state(PauseState.UNPAUSED)
635
636 def is_scheduler_paused(self) -> bool:
605 - """Return whether the scheduler is in any pause state. False in base EngineCore
606 - and overridden in EngineCoreProc."""
607 - return False
637 + """Return whether the scheduler is in any pause state."""
638 + return self.scheduler.pause_state != PauseState.UNPAUSED
639
609 - def sleep(self, level: int = 1):
640 + def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None | Future:
641 """Put the engine to sleep at the specified level.
642
643 Args:
@@ -615,13 +646,34 @@ class EngineCore:
646 but not processed. No GPU memory changes.
647 - Level 1: Offload model weights to CPU, discard KV cache.
648 - Level 2: Discard all GPU memory.
649 + mode: Pause mode - how to deal with any existing requests, see
650 + documentation of pause_scheduler method.
651 """
619 - if level == 0:
620 - # Level 0: Just pause scheduling, don't touch GPU
621 - self.pause_scheduler()
622 - else:
623 - # Level 1+: Delegate to executor for GPU memory management
624 - self.model_executor.sleep(level)
652 +
653 + # Pause scheduler before sleeping.
654 + clear_prefix_cache = level >= 1
655 + pause_future = self.pause_scheduler(mode=mode, clear_cache=clear_prefix_cache)
656 + if level < 1:
657 + return pause_future
658 +
659 + # Level 1+: Delegate to executor for GPU memory management
660 + model_executor = self.model_executor
661 + if pause_future is None:
662 + model_executor.sleep(level)
663 + return None
664 +
665 + future = Future[Any]()
666 +
667 + def pause_complete(f: Future):
668 + try:
669 + f.result() # propagate any exception
670 + future.set_result(model_executor.sleep(level))
671 + except Exception as e:
672 + future.set_exception(e)
673 +
674 + logger.info("Waiting for in-flight requests to complete before sleeping...")
675 + pause_future.add_done_callback(pause_complete)
676 + return future
677
678 def wake_up(self, tags: list[str] | None = None):
679 """Wake up the engine from sleep.
@@ -630,17 +682,15 @@ class EngineCore:
682 tags: Tags to wake up. Use ["scheduling"] for level 0 wake up.
683 """
684 if tags is not None and "scheduling" in tags:
633 - # Level 0 wake up: Resume scheduling
634 - self.resume_scheduler()
635 - # Remove "scheduling" from tags if there are other tags to process
636 - remaining_tags = [t for t in tags if t != "scheduling"]
637 - if remaining_tags:
638 - self.model_executor.wake_up(remaining_tags)
639 - else:
640 - # Full wake up
641 - self.resume_scheduler()
685 + # Remove "scheduling" from tags if there are other tags to process.
686 + tags = [t for t in tags if t != "scheduling"]
687 +
688 + if tags is None or tags:
689 self.model_executor.wake_up(tags)
690
691 + # Resume scheduling (applies to all levels)
692 + self.resume_scheduler()
693 +
694 def is_sleeping(self) -> bool:
695 """Check if engine is sleeping at any level."""
696 return self.is_scheduler_paused() or self.model_executor.is_sleeping
@@ -1038,6 +1088,14 @@ class EngineCoreProc(EngineCore):
1088 def _init_data_parallel(self, vllm_config: VllmConfig):
1089 pass
1090
1091 + def has_work(self) -> bool:
1092 + """Returns true if the engine should be stepped."""
1093 + return (
1094 + self.engines_running
1095 + or self.scheduler.has_requests()
1096 + or bool(self.batch_queue)
1097 + )
1098 +
1099 def run_busy_loop(self):
1100 """Core busy loop of the EngineCore."""
1101
@@ -1047,19 +1105,14 @@ class EngineCoreProc(EngineCore):
1105 self._process_input_queue()
1106 # 2) Step the engine core and return the outputs.
1107 self._process_engine_step()
1050 - # 3) Run any per-step hooks.
1051 - self._process_per_step_hooks()
1108
1109 def _process_input_queue(self):
1110 """Exits when an engine step needs to be performed."""
1111
1112 waited = False
1057 - while (
1058 - not self.engines_running
1059 - and not self.scheduler.has_requests()
1060 - and not self.batch_queue
1061 - and not self.per_step_hooks
1062 - ):
1113 + while not self.has_work():
1114 + # Notify callbacks waiting for engine to become idle.
1115 + self._notify_idle_state_callbacks()
1116 if self.input_queue.empty():
1117 # Drain aborts queue; all aborts are also processed via input_queue.
1118 with self.aborts_queue.mutex:
@@ -1098,12 +1151,10 @@ class EngineCoreProc(EngineCore):
1151
1152 return model_executed
1153
1101 - def _process_per_step_hooks(self) -> None:
1102 - if self.per_step_hooks:
1103 - for hook in list(self.per_step_hooks):
1104 - finished = hook(self)
1105 - if finished:
1106 - self.per_step_hooks.discard(hook)
1154 + def _notify_idle_state_callbacks(self) -> None:
1155 + while self._idle_state_callbacks:
1156 + callback = self._idle_state_callbacks.pop()
1157 + callback(self)
1158
1159 def _handle_client_request(
1160 self, request_type: EngineCoreRequestType, request: Any
@@ -1377,19 +1428,10 @@ class EngineCoreProc(EngineCore):
1428 if mode not in ("keep", "abort", "wait"):
1429 raise ValueError(f"Invalid pause mode: {mode}")
1430
1380 - future: Future[Any] = Future()
1381 -
1382 - def wait_until_idle(engine: "EngineCoreProc") -> bool:
1383 - scheduler = engine.scheduler
1384 - out_queue = engine.output_queue
1385 - if scheduler.has_requests() or engine.batch_queue or not out_queue.empty():
1386 - return False
1431 + def engine_idle_callback(engine: "EngineCoreProc", future: Future[Any]) -> None:
1432 if clear_cache:
1388 - engine.reset_prefix_cache(reset_running_requests=True)
1389 - engine.reset_mm_cache()
1390 - engine.reset_encoder_cache()
1433 + engine._reset_caches()
1434 future.set_result(None)
1392 - return True
1435
1436 if mode == "abort":
1437 aborted_reqs = self.scheduler.finish_requests(
@@ -1399,12 +1441,17 @@ class EngineCoreProc(EngineCore):
1441
1442 pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
1443 self.scheduler.set_pause_state(pause_state)
1402 - if not wait_until_idle(self):
1403 - self.per_step_hooks.add(wait_until_idle)
1404 - return future
1405 - return None
1444 + if not self.has_work():
1445 + if clear_cache:
1446 + self._reset_caches()
1447 + return None
1448 +
1449 + future = Future[Any]()
1450 + self._idle_state_callbacks.append(partial(engine_idle_callback, future=future))
1451 + return future
1452
1453 def _send_abort_outputs(self, aborted_reqs: list[tuple[str, int]]) -> None:
1454 + # TODO(nick) this will be moved inside the scheduler
1455 if aborted_reqs:
1456 # Map client_index to list of request_ids that belong to that client.
1457 by_client = defaultdict[int, set[str]](set)
@@ -1418,14 +1465,6 @@ class EngineCoreProc(EngineCore):
1465 eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1466 self.output_queue.put_nowait((client_index, eco))
1467
1421 - def resume_scheduler(self) -> None:
1422 - """Resume the scheduler and flush any requests queued while paused."""
1423 - self.scheduler.set_pause_state(PauseState.UNPAUSED)
1424 -
1425 - def is_scheduler_paused(self) -> bool:
1426 - """Return whether the scheduler is in any pause state."""
1427 - return self.scheduler.pause_state != PauseState.UNPAUSED
1428 -
1468
1469 class DPEngineCoreProc(EngineCoreProc):
1470 """ZMQ-wrapper for running EngineCore in background process
@@ -1481,6 +1520,7 @@ class DPEngineCoreProc(EngineCoreProc):
1520 stateless_destroy_torch_distributed_process_group(dp_group)
1521
1522 def add_request(self, request: Request, request_wave: int = 0):
1523 + super().add_request(request, request_wave)
1524 if self.has_coordinator and request_wave != self.current_wave:
1525 if request_wave > self.current_wave:
1526 self.current_wave = request_wave
@@ -1491,7 +1531,13 @@ class DPEngineCoreProc(EngineCoreProc):
1531 (-1, EngineCoreOutputs(start_wave=self.current_wave))
1532 )
1533
1494 - super().add_request(request, request_wave)
1534 + def resume_scheduler(self):
1535 + super().resume_scheduler()
1536 + if not self.engines_running and self.scheduler.has_unfinished_requests():
1537 + # Wake up other DP engines.
1538 + self.output_queue.put_nowait(
1539 + (-1, EngineCoreOutputs(start_wave=self.current_wave))
1540 + )
1541
1542 def _handle_client_request(
1543 self, request_type: EngineCoreRequestType, request: Any
@@ -1532,8 +1578,8 @@ class DPEngineCoreProc(EngineCoreProc):
1578 # 2) Step the engine core.
1579 executed = self._process_engine_step()
1580 self._maybe_publish_request_counts()
1535 - local_unfinished_reqs = self.scheduler.has_unfinished_requests()
1581
1582 + local_unfinished_reqs = self.scheduler.has_unfinished_requests()
1583 if not executed:
1584 if not local_unfinished_reqs and not self.engines_running:
1585 # All engines are idle.
vllm/v1/engine/core_client.py
+11 -8
@@ -150,7 +150,7 @@ class EngineCoreClient(ABC):
150 def reset_encoder_cache(self) -> None:
151 raise NotImplementedError
152
153 - def sleep(self, level: int = 1) -> None:
153 + def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
154 raise NotImplementedError
155
156 def wake_up(self, tags: list[str] | None = None) -> None:
@@ -227,7 +227,7 @@ class EngineCoreClient(ABC):
227 async def reset_encoder_cache_async(self) -> None:
228 raise NotImplementedError
229
230 - async def sleep_async(self, level: int = 1) -> None:
230 + async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
231 raise NotImplementedError
232
233 async def wake_up_async(self, tags: list[str] | None = None) -> None:
@@ -314,8 +314,11 @@ class InprocClient(EngineCoreClient):
314 def reset_encoder_cache(self) -> None:
315 self.engine_core.reset_encoder_cache()
316
317 - def sleep(self, level: int = 1) -> None:
318 - self.engine_core.sleep(level)
317 + def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
318 + if mode == "wait":
319 + raise ValueError("'wait' pause mode is not supported in inproc-engine mode")
320 + result = self.engine_core.sleep(level, mode)
321 + assert result is None
322
323 def wake_up(self, tags: list[str] | None = None) -> None:
324 self.engine_core.wake_up(tags)
@@ -796,8 +799,8 @@ class SyncMPClient(MPClient):
799 def pin_lora(self, lora_id: int) -> bool:
800 return self.call_utility("pin_lora", lora_id)
801
799 - def sleep(self, level: int = 1) -> None:
800 - self.call_utility("sleep", level)
802 + def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:
803 + self.call_utility("sleep", level, mode)
804
805 def wake_up(self, tags: list[str] | None = None) -> None:
806 self.call_utility("wake_up", tags)
@@ -1009,8 +1012,8 @@ class AsyncMPClient(MPClient):
1012 async def reset_encoder_cache_async(self) -> None:
1013 await self.call_utility_async("reset_encoder_cache")
1014
1012 - async def sleep_async(self, level: int = 1) -> None:
1013 - await self.call_utility_async("sleep", level)
1015 + async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:
1016 + await self.call_utility_async("sleep", level, mode)
1017
1018 async def wake_up_async(self, tags: list[str] | None = None) -> None:
1019 await self.call_utility_async("wake_up", tags)
vllm/v1/engine/llm_engine.py
+3 -3
@@ -28,7 +28,7 @@ from vllm.tasks import SupportedTask
28 from vllm.tokenizers import TokenizerLike
29 from vllm.tracing import init_tracer
30 from vllm.usage.usage_lib import UsageContext
31 -from vllm.v1.engine import EngineCoreRequest
31 +from vllm.v1.engine import EngineCoreRequest, PauseMode
32 from vllm.v1.engine.core_client import EngineCoreClient
33 from vllm.v1.engine.input_processor import InputProcessor
34 from vllm.v1.engine.output_processor import OutputProcessor
@@ -355,8 +355,8 @@ class LLMEngine:
355 """
356 self.engine_core.reset_encoder_cache()
357
358 - def sleep(self, level: int = 1):
359 - self.engine_core.sleep(level)
358 + def sleep(self, level: int = 1, mode: PauseMode = "abort"):
359 + self.engine_core.sleep(level, mode)
360
361 if self.logger_manager is not None:
362 self.logger_manager.record_sleep_state(1, level)
vllm/v1/engine/output_processor.py
-14
@@ -429,8 +429,6 @@ class OutputProcessor:
429 self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list)
430 self.lora_states = LoRARequestStates(log_stats)
431 self.tracing_enabled = tracing_enabled
432 - self._requests_drained = asyncio.Event()
433 - self._requests_drained.set()
432
433 def get_num_unfinished_requests(self):
434 return len(self.request_states)
@@ -438,11 +436,6 @@ class OutputProcessor:
436 def has_unfinished_requests(self) -> bool:
437 return len(self.request_states) > 0
438
441 - async def wait_for_requests_to_drain(self) -> None:
442 - if not self.request_states:
443 - return
444 - await self._requests_drained.wait()
445 -
439 def propagate_error(self, e: Exception):
440 """Propagate error to all generate() tasks."""
441
@@ -510,8 +503,6 @@ class OutputProcessor:
503 child_reqs = self.abort_requests(child_reqs, internal=True)
504 request_ids_to_abort.extend(child_reqs)
505 self.parent_requests.pop(request_id, None)
513 - if not self.request_states:
514 - self._requests_drained.set()
506 return request_ids_to_abort
507
508 def add_request(
@@ -538,8 +529,6 @@ class OutputProcessor:
529 log_stats=self.log_stats,
530 stream_interval=self.stream_interval,
531 )
541 - if self._requests_drained.is_set():
542 - self._requests_drained.clear()
532 self.request_states[request_id] = req_state
533 if parent_req:
534 self.parent_requests[parent_req.request_id] = parent_req
@@ -706,9 +695,6 @@ class OutputProcessor:
695 if parent_req and not parent_req.child_requests:
696 self.parent_requests.pop(parent_req.request_id, None)
697
709 - if not self.request_states:
710 - self._requests_drained.set()
711 -
698 def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None):
699 self.lora_states.update_scheduler_stats(scheduler_stats)
700