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 8, 08:45 UTC
[Core] Move pause and resume functions into engine (#34125)
Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Signed-off-by: hao-aaron <ahao@anyscale.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Aaron Hao committed
Feb 13, 2026 at 00:15 UTC
dddbff46242a9292085e2ae3309dc559f242cad6
9 files changed
+621
-136
examples/online_serving/data_parallel_pause_resume.py
new
+135
@@ -0,0 +1,135 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test pause/resume with Data Parallel (DP) via HTTP API.
5
+
6
+This example demonstrates coordinated pause/resume across multiple DP ranks.
7
+The pause synchronizes across all DP engines via all-reduce.
8
+
9
+Prerequisites:
10
+ Start a vLLM server with data parallelism:
11
+
12
+ $ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
13
+ --enforce-eager \
14
+ --data-parallel-size 4 \
15
+ --tensor-parallel-size 1
16
+
17
+ Then run this script:
18
+
19
+ $ python data_parallel_pause_resume.py
20
+
21
+The test verifies pause works by:
22
+1. Starting a streaming generation request
23
+2. Pausing the server mid-generation
24
+3. Sleeping for PAUSE_DURATION seconds
25
+4. Resuming the server
26
+5. Verifying there was a gap in token generation matching the pause duration
27
+"""
28
+
29
+import argparse
30
+import threading
31
+import time
32
+
33
+import requests
34
+from openai import OpenAI
35
+
36
+BASE_URL = "http://localhost:8000"
37
+MODEL_NAME = "facebook/opt-125m"
38
+PAUSE_DURATION = 3.0
39
+
40
+
41
+def pause_generation(base_url: str, mode: str = "keep") -> None:
42
+ """Pause generation via HTTP endpoint."""
43
+ url = f"{base_url}/pause"
44
+ response = requests.post(url, params={"mode": mode}, timeout=60)
45
+ response.raise_for_status()
46
+ print("Server paused")
47
+
48
+
49
+def resume_generation(base_url: str) -> None:
50
+ """Resume generation via HTTP endpoint."""
51
+ url = f"{base_url}/resume"
52
+ response = requests.post(url, timeout=60)
53
+ response.raise_for_status()
54
+ print("Server resumed")
55
+
56
+
57
+def main():
58
+ parser = argparse.ArgumentParser()
59
+ parser.add_argument("--base-url", default=BASE_URL)
60
+ parser.add_argument("--model", default=MODEL_NAME)
61
+ args = parser.parse_args()
62
+
63
+ client = OpenAI(
64
+ base_url=f"{args.base_url}/v1",
65
+ api_key="EMPTY",
66
+ )
67
+
68
+ prompt = "Write a long story about a dragon. Once upon a time"
69
+ token_times: list[float] = []
70
+ pause_token_idx = 0
71
+ pause_triggered = threading.Event()
72
+
73
+ def generator_thread():
74
+ """Stream tokens and record timestamps."""
75
+ stream = client.completions.create(
76
+ model=args.model,
77
+ prompt=prompt,
78
+ max_tokens=50,
79
+ stream=True,
80
+ )
81
+ for chunk in stream:
82
+ if chunk.choices[0].text:
83
+ token_times.append(time.monotonic())
84
+ token_count = len(token_times)
85
+ print(f"Token {token_count}: {chunk.choices[0].text!r}")
86
+
87
+ # Signal controller after some tokens
88
+ if token_count >= 5 and not pause_triggered.is_set():
89
+ pause_triggered.set()
90
+
91
+ def controller_thread():
92
+ """Pause and resume the server."""
93
+ nonlocal pause_token_idx
94
+
95
+ # Wait for some tokens
96
+ pause_triggered.wait()
97
+
98
+ print(f"\nPausing server (keep mode) at token {len(token_times)}...")
99
+ pause_generation(args.base_url, mode="keep")
100
+ pause_token_idx = len(token_times)
101
+ print(f"Sleeping for {PAUSE_DURATION}s...")
102
+
103
+ time.sleep(PAUSE_DURATION)
104
+
105
+ print("Resuming server...")
106
+ resume_generation(args.base_url)
107
+ print("Resumed!\n")
108
+
109
+ # Run both threads
110
+ gen_thread = threading.Thread(target=generator_thread)
111
+ ctrl_thread = threading.Thread(target=controller_thread)
112
+
113
+ gen_thread.start()
114
+ ctrl_thread.start()
115
+
116
+ gen_thread.join()
117
+ ctrl_thread.join()
118
+
119
+ # Check gap at the pause point
120
+ if pause_token_idx < len(token_times):
121
+ pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
122
+ print(
123
+ f"\nGap after pause (token {pause_token_idx} -> "
124
+ f"{pause_token_idx + 1}): {pause_gap:.3f}s"
125
+ )
126
+ if pause_gap >= PAUSE_DURATION * 0.9:
127
+ print("Test passed! Pause synchronized across DP ranks.")
128
+ else:
129
+ print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
130
+ else:
131
+ print("Test failed! No tokens were generated after resuming.")
132
+
133
+
134
+if __name__ == "__main__":
135
+ main()
tests/v1/distributed/test_async_llm_dp.py
+143
@@ -12,6 +12,7 @@ from vllm import SamplingParams
12
from vllm.config import VllmConfig
13
from vllm.engine.arg_utils import AsyncEngineArgs
14
from vllm.inputs import PromptType
15
+from vllm.outputs import RequestOutput
16
from vllm.platforms import current_platform
17
from vllm.sampling_params import RequestOutputKind
18
from vllm.v1.engine.async_llm import AsyncLLM
@@ -181,3 +182,145 @@ async def test_load(
182
assert slogger.finished_req_count > NUM_REQUESTS // (DP_SIZE + 1), (
183
f"requests are imbalanced: {stats_loggers}"
184
)
185
+
186
+
187
+# =============================================================================
188
+# DP Pause/Resume Tests
189
+# =============================================================================
190
+
191
+DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
192
+DP_PAUSE_PROMPT = "This is a test of data parallel pause"
193
+
194
+
195
+@pytest.mark.asyncio
196
+async def test_dp_pause_resume_basic():
197
+ """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")
200
+ 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
+ )
208
+ engine = AsyncLLM.from_engine_args(engine_args)
209
+ after.callback(engine.shutdown)
210
+
211
+ assert not await engine.is_paused()
212
+ await engine.pause_generation(mode="abort")
213
+ assert await engine.is_paused()
214
+ await engine.resume_generation()
215
+ assert not await engine.is_paused()
216
+
217
+ # Engine still works after resume
218
+ sampling_params = SamplingParams(max_tokens=5)
219
+ async for out in engine.generate(
220
+ request_id="after-resume",
221
+ prompt=DP_PAUSE_PROMPT,
222
+ sampling_params=sampling_params,
223
+ ):
224
+ pass
225
+ assert out.finished
226
+
227
+
228
+@pytest.mark.asyncio
229
+async def test_dp_pause_abort():
230
+ """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")
233
+ 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
+ )
241
+ engine = AsyncLLM.from_engine_args(engine_args)
242
+ after.callback(engine.shutdown)
243
+
244
+ # Start several requests so they are distributed across ranks
245
+ sampling_params = SamplingParams(max_tokens=500, ignore_eos=True)
246
+ num_requests = 4
247
+ outputs_by_id: dict[str, list[RequestOutput]] = {}
248
+
249
+ async def gen(rid: str):
250
+ out_list: list[RequestOutput] = []
251
+ outputs_by_id[rid] = out_list
252
+ async for out in engine.generate(
253
+ request_id=rid,
254
+ prompt=DP_PAUSE_PROMPT,
255
+ sampling_params=sampling_params,
256
+ ):
257
+ out_list.append(out)
258
+ return out_list[-1] if out_list else None
259
+
260
+ tasks = [asyncio.create_task(gen(f"req-{i}")) for i in range(num_requests)]
261
+ # Wait for some tokens on at least one request
262
+ while not any(len(o) >= 2 for o in outputs_by_id.values()):
263
+ await asyncio.sleep(0.02)
264
+
265
+ await engine.pause_generation(mode="abort")
266
+
267
+ finals = await asyncio.gather(*tasks)
268
+ for i, final in enumerate(finals):
269
+ assert final is not None, f"req-{i} had no output"
270
+ assert final.finished
271
+ assert final.outputs[0].finish_reason == "abort"
272
+
273
+ assert await engine.is_paused()
274
+ await engine.resume_generation()
275
+ assert not await engine.is_paused()
276
+
277
+ # New request completes after resume
278
+ async for out in engine.generate(
279
+ request_id="after-abort",
280
+ prompt=DP_PAUSE_PROMPT,
281
+ sampling_params=SamplingParams(max_tokens=5),
282
+ ):
283
+ pass
284
+ assert out.finished
285
+ assert not engine.output_processor.has_unfinished_requests()
286
+
287
+
288
+@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
+ 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 = AsyncLLM.from_engine_args(engine_args)
302
+ after.callback(engine.shutdown)
303
+
304
+ await engine.pause_generation(mode="keep")
305
+ assert await engine.is_paused()
306
+
307
+ request_done = asyncio.Event()
308
+
309
+ async def gen():
310
+ async for out in engine.generate(
311
+ request_id="queued-keep",
312
+ prompt=DP_PAUSE_PROMPT,
313
+ sampling_params=SamplingParams(max_tokens=5),
314
+ ):
315
+ pass
316
+ request_done.set()
317
+ return out
318
+
319
+ task = asyncio.create_task(gen())
320
+ await asyncio.sleep(0.2)
321
+ assert not request_done.is_set()
322
+
323
+ await engine.resume_generation()
324
+ final = await asyncio.wait_for(task, timeout=10.0)
325
+ assert final.finished
326
+ assert not await engine.is_paused()
tests/v1/engine/test_async_llm.py
+48
-3
@@ -708,9 +708,7 @@ async def test_pause_resume_basic():
708
# Test all modes with no requests in flight
709
for mode in ("abort", "wait", "keep"):
710
await engine.pause_generation(mode=mode)
711
- # "keep" only freezes the scheduler; it does not set _paused
712
- if mode != "keep":
713
- assert await engine.is_paused()
711
+ assert await engine.is_paused()
712
await engine.resume_generation()
713
assert not await engine.is_paused()
714
@@ -808,6 +806,53 @@ async def test_pause_abort():
806
assert final_output2.finished
807
808
809
+@pytest.mark.asyncio
810
+async def test_pause_then_abort_queued_request():
811
+ """Test that aborting a request that was submitted while paused (in
812
+ _paused_adds_queue) aborts it and notifies the client; the request does
813
+ not run after resume.
814
+ """
815
+ with ExitStack() as after:
816
+ with set_default_torch_num_threads(1):
817
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
818
+ after.callback(engine.shutdown)
819
+
820
+ request_id = "abort-queued-request"
821
+ sampling_params = SamplingParams(max_tokens=20, ignore_eos=True)
822
+ outputs: list[RequestOutput] = []
823
+
824
+ # Pause first so the next add goes to _paused_adds_queue
825
+ await engine.pause_generation(mode="keep")
826
+ assert await engine.is_paused()
827
+
828
+ async def gen():
829
+ async for out in engine.generate(
830
+ request_id=request_id,
831
+ prompt=TEXT_PROMPT,
832
+ sampling_params=sampling_params,
833
+ ):
834
+ outputs.append(out)
835
+ return outputs[-1] if outputs else None
836
+
837
+ gen_task = asyncio.create_task(gen())
838
+
839
+ # Give the request time to reach the engine and sit in _paused_adds_queue
840
+ await asyncio.sleep(0.2)
841
+
842
+ # Abort the queued request
843
+ await engine.abort(request_id, internal=False)
844
+
845
+ # Resume so the engine can process and deliver the abort output
846
+ await engine.resume_generation()
847
+
848
+ final_output = await asyncio.wait_for(gen_task, timeout=10.0)
849
+ assert final_output is not None
850
+ assert final_output.finished
851
+ assert final_output.outputs[0].finish_reason == "abort"
852
+ # Request was never run, so no tokens
853
+ assert len(final_output.outputs[0].token_ids) == 0
854
+
855
+
856
@pytest.mark.asyncio
857
async def test_pause_wait():
858
"""Test that mode='wait' waits for in-flight requests to complete."""
tests/v1/engine/test_engine_core_client.py
+83
@@ -8,6 +8,7 @@ import os
8
import signal
9
import time
10
import uuid
11
+from concurrent.futures import Future
12
from dataclasses import dataclass
13
from threading import Thread
14
from types import SimpleNamespace
@@ -278,6 +279,24 @@ def echo_dc_nested(
279
return structures.get(structure_type, val)
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).
285
+ """
286
+ future: Future = Future()
287
+ remaining = [num_wait_loops]
288
+
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
295
+
296
+ self.per_step_hooks.add(_step)
297
+ return future
298
+
299
+
300
# --- Fixtures for subprocess patching ---
301
# These create sitecustomize.py files that patch EngineCore in spawned
302
# subprocesses. This is necessary because ROCm requires 'spawn' multiprocessing
@@ -383,6 +402,28 @@ def subprocess_echo_dc_nested_patch(monkeypatch, tmp_path):
402
)
403
404
405
+@pytest.fixture
406
+def subprocess_future_echo_patch(monkeypatch, tmp_path):
407
+ """Create sitecustomize.py so spawned subprocesses have future_echo method."""
408
+ sc = tmp_path / "sitecustomize.py"
409
+ sc.write_text(
410
+ "\n".join(
411
+ [
412
+ "from concurrent.futures import Future",
413
+ "from typing import Any",
414
+ "",
415
+ "from vllm.v1.engine.core import EngineCore",
416
+ inspect.getsource(future_echo),
417
+ "EngineCore.future_echo = future_echo",
418
+ ]
419
+ )
420
+ )
421
+ monkeypatch.setenv(
422
+ "PYTHONPATH",
423
+ os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
424
+ )
425
+
426
+
427
@create_new_process_for_each_test()
428
@pytest.mark.parametrize("multiprocessing_mode", [True, False])
429
def test_engine_core_client(
@@ -786,6 +827,48 @@ async def test_engine_core_client_util_method_nested_structures(
827
client.shutdown()
828
829
830
+@pytest.mark.asyncio(loop_scope="function")
831
+async def test_engine_core_client_future_utility_async(
832
+ monkeypatch: pytest.MonkeyPatch,
833
+ subprocess_future_echo_patch,
834
+):
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).
837
+ """
838
+ with monkeypatch.context() as m:
839
+ m.setattr(EngineCore, "future_echo", future_echo, raising=False)
840
+
841
+ engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
842
+ vllm_config = engine_args.create_engine_config(
843
+ usage_context=UsageContext.UNKNOWN_CONTEXT
844
+ )
845
+ executor_class = Executor.get_class(vllm_config)
846
+
847
+ with set_default_torch_num_threads(1):
848
+ client = EngineCoreClient.make_client(
849
+ multiprocess_mode=True,
850
+ asyncio_mode=True,
851
+ vllm_config=vllm_config,
852
+ executor_class=executor_class,
853
+ log_stats=True,
854
+ )
855
+
856
+ try:
857
+ core_client: AsyncMPClient = client
858
+
859
+ # Completes after 2 engine steps (num_wait_loops=2)
860
+ result = await core_client.call_utility_async(
861
+ "future_echo", "future_result", 2
862
+ )
863
+ assert result == "future_result"
864
+
865
+ # None is a valid result (num_wait_loops=0 → completes on first step)
866
+ result = await core_client.call_utility_async("future_echo", None, 0)
867
+ assert result is None
868
+ finally:
869
+ client.shutdown()
870
+
871
+
872
@pytest.mark.parametrize(
873
"multiprocessing_mode,publisher_config",
874
[(True, "tcp"), (False, "inproc")],
vllm/v1/core/sched/interface.py
+33
-4
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import enum
4
from abc import ABC, abstractmethod
5
from collections.abc import Iterable
6
from typing import TYPE_CHECKING
@@ -18,6 +19,20 @@ if TYPE_CHECKING:
19
from vllm.v1.structured_output import StructuredOutputManager
20
21
22
+class PauseState(enum.IntEnum):
23
+ """Scheduler pause state.
24
+
25
+ - UNPAUSED: Normal operation
26
+ - PAUSE_NEW: No new requests are scheduled, requests already in
27
+ running state are scheduled.
28
+ - PAUSE_ALL: No requests are scheduled
29
+ """
30
+
31
+ UNPAUSED = 0
32
+ PAUSED_NEW = 1
33
+ PAUSED_ALL = 2
34
+
35
+
36
class SchedulerInterface(ABC):
37
@abstractmethod
38
def __init__(
@@ -120,11 +135,11 @@ class SchedulerInterface(ABC):
135
@abstractmethod
136
def finish_requests(
137
self,
123
- request_ids: str | Iterable[str],
138
+ request_ids: str | Iterable[str] | None,
139
finished_status: "RequestStatus",
125
- ) -> None:
140
+ ) -> list[tuple[str, int]]:
141
"""Finish the requests in the scheduler's internal queue. If the request
127
- is not in the queue, this method will do nothing.
142
+ is not in the queue, this method will do nothing for that request.
143
144
This method is called in two cases:
145
1. When the request is aborted by the client.
@@ -132,8 +147,12 @@ class SchedulerInterface(ABC):
147
de-tokenizing its generated tokens.
148
149
Args:
135
- request_ids: A single or a list of request IDs.
150
+ request_ids: A single or a list of request IDs, or None to finish all.
151
finished_status: The finished status of the given requests.
152
+
153
+ Returns:
154
+ Tuple of (req_id, client_index) for requests that were aborted. Will not
155
+ include any that were already finished.
156
"""
157
raise NotImplementedError
158
@@ -167,6 +186,16 @@ class SchedulerInterface(ABC):
186
not yet returned in SchedulerOutputs."""
187
return self.has_unfinished_requests() or self.has_finished_requests()
188
189
+ @property
190
+ @abstractmethod
191
+ def pause_state(self) -> PauseState:
192
+ """Current pause state of the scheduler."""
193
+ raise NotImplementedError
194
+
195
+ @abstractmethod
196
+ def set_pause_state(self, pause_state: PauseState) -> None:
197
+ raise NotImplementedError
198
+
199
@abstractmethod
200
def reset_prefix_cache(
201
self, reset_running_requests: bool = False, reset_connector: bool = False
vllm/v1/core/sched/scheduler.py
+40
-12
@@ -38,7 +38,7 @@ from vllm.v1.core.encoder_cache_manager import (
38
)
39
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
40
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
41
-from vllm.v1.core.sched.interface import SchedulerInterface
41
+from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
42
from vllm.v1.core.sched.output import (
43
CachedRequestData,
44
GrammarOutput,
@@ -271,6 +271,8 @@ class Scheduler(SchedulerInterface):
271
vllm_config=self.vllm_config,
272
)
273
274
+ self._pause_state: PauseState = PauseState.UNPAUSED
275
+
276
def _mamba_block_aligned_split(
277
self,
278
request: Request,
@@ -341,6 +343,10 @@ class Scheduler(SchedulerInterface):
343
req_to_new_blocks: dict[str, KVCacheBlocks] = {}
344
num_scheduled_tokens: dict[str, int] = {}
345
token_budget = self.max_num_scheduled_tokens
346
+ if self._pause_state == PauseState.PAUSED_ALL:
347
+ # Do not schedule any requests when paused.
348
+ token_budget = 0
349
+
350
# Encoder-related.
351
scheduled_encoder_inputs: dict[str, list[int]] = {}
352
encoder_compute_budget = self.max_num_encoder_input_tokens
@@ -530,12 +536,12 @@ class Scheduler(SchedulerInterface):
536
)
537
assert len(scheduled_loras) <= self.lora_config.max_loras
538
533
- # Use a temporary RequestQueue to collect requests that need to be
534
- # skipped and put back at the head of the waiting queue later
535
- skipped_waiting_requests = create_request_queue(self.policy)
536
-
539
# Next, schedule the WAITING requests.
538
- if not preempted_reqs:
540
+ if not preempted_reqs and self._pause_state == PauseState.UNPAUSED:
541
+ # Use a temporary RequestQueue to collect requests that need to be
542
+ # skipped and put back at the head of the waiting queue later
543
+ skipped_waiting_requests = create_request_queue(self.policy)
544
+
545
while self.waiting and token_budget > 0:
546
if len(self.running) == self.max_num_running_reqs:
547
break
@@ -802,9 +808,10 @@ class Scheduler(SchedulerInterface):
808
self.encoder_cache_manager.allocate(request, i)
809
if self.ec_connector is not None:
810
self.ec_connector.update_state_after_alloc(request, i)
805
- # Put back any skipped requests at the head of the waiting queue
806
- if skipped_waiting_requests:
807
- self.waiting.prepend_requests(skipped_waiting_requests)
811
+
812
+ # Put back any skipped requests at the head of the waiting queue
813
+ if skipped_waiting_requests:
814
+ self.waiting.prepend_requests(skipped_waiting_requests)
815
816
# Check if the scheduling constraints are satisfied.
817
total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
@@ -1672,18 +1679,26 @@ class Scheduler(SchedulerInterface):
1679
request.record_event(EngineCoreEventType.QUEUED)
1680
1681
def finish_requests(
1675
- self, request_ids: str | Iterable[str], finished_status: RequestStatus
1676
- ) -> None:
1682
+ self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus
1683
+ ) -> list[tuple[str, int]]:
1684
"""Handles the finish signal from outside the scheduler.
1685
1686
For example, the API server can abort a request when the client
1687
disconnects.
1688
+
1689
+ If request_ids is None, all requests will be finished.
1690
+
1691
+ Returns:
1692
+ Tuple of (req_id, client_index) for requests that were aborted. Will not
1693
+ include any that were already finished.
1694
"""
1695
assert RequestStatus.is_finished(finished_status)
1696
if isinstance(request_ids, str):
1697
request_ids = (request_ids,)
1685
- else:
1698
+ elif request_ids is not None:
1699
request_ids = set(request_ids)
1700
+ else:
1701
+ request_ids = self.requests.keys()
1702
1703
running_requests_to_remove = set()
1704
waiting_requests_to_remove = []
@@ -1723,6 +1738,8 @@ class Scheduler(SchedulerInterface):
1738
request.status = finished_status
1739
self._free_request(request, delay_free_blocks=delay_free_blocks)
1740
1741
+ return [(r.request_id, r.client_index) for r in valid_requests]
1742
+
1743
def _free_request(
1744
self, request: Request, delay_free_blocks: bool = False
1745
) -> dict[str, Any] | None:
@@ -1746,7 +1763,18 @@ class Scheduler(SchedulerInterface):
1763
self.kv_cache_manager.free(request)
1764
del self.requests[request.request_id]
1765
1766
+ @property
1767
+ def pause_state(self) -> PauseState:
1768
+ return self._pause_state
1769
+
1770
+ def set_pause_state(self, pause_state: PauseState) -> None:
1771
+ self._pause_state = pause_state
1772
+
1773
def get_num_unfinished_requests(self) -> int:
1774
+ if self._pause_state == PauseState.PAUSED_ALL:
1775
+ return 0
1776
+ if self._pause_state == PauseState.PAUSED_NEW:
1777
+ return len(self.running)
1778
num_waiting = len(self.waiting) - self.num_waiting_for_streaming_input
1779
return num_waiting + len(self.running)
1780
vllm/v1/engine/async_llm.py
+6
-49
@@ -172,9 +172,6 @@ class AsyncLLM(EngineClient):
172
)
173
self.logger_manager.log_engine_initialized()
174
175
- # Pause / resume state for async RL workflows.
176
- self._pause_cond = asyncio.Condition()
177
- self._paused = False
175
self._client_count = client_count
176
177
self.output_handler: asyncio.Task | None = None
@@ -387,10 +384,6 @@ class AsyncLLM(EngineClient):
384
# to handle startup failure gracefully in the OpenAI server.
385
self._run_output_handler()
386
390
- # Respect pause state before accepting new requests.
391
- async with self._pause_cond:
392
- await self._pause_cond.wait_for(lambda: not self._paused)
393
-
387
# Create a new output collector for the request.
388
queue = RequestOutputCollector(params.output_kind, request.request_id)
389
@@ -741,7 +734,9 @@ class AsyncLLM(EngineClient):
734
"""
735
Pause generation to allow model weight updates.
736
744
- New generation/encoding requests are blocked until resume.
737
+ All mode handling (abort / wait / keep) and cache clearing is done
738
+ in the engine. New generation/encoding requests will not be scheduled
739
+ until resume is called.
740
741
Args:
742
mode: How to handle in-flight requests:
@@ -751,11 +746,8 @@ class AsyncLLM(EngineClient):
746
- ``"keep"``: Freeze requests in queue; they resume on
747
:meth:`resume_generation`.
748
wait_for_inflight_requests: DEPRECATED: use mode argument.
754
- Whether to wait for in-flight requests to complete before pausing.
749
clear_cache: Whether to clear KV cache and prefix cache after
750
draining. Set to ``False`` to preserve cache for faster resume.
757
- Default is ``True`` (clear caches).
758
-
751
"""
752
if wait_for_inflight_requests:
753
warnings.warn(
@@ -766,50 +758,15 @@ class AsyncLLM(EngineClient):
758
stacklevel=2,
759
)
760
mode = "wait"
769
-
770
- if mode == "keep":
771
- # Freeze requests in the scheduler - they will resume on
772
- # resume_generation().
773
- await self.engine_core.pause_scheduler_async()
774
- else:
775
- if self._client_count > 1:
776
- raise NotImplementedError(
777
- "pause_generation is not supported with --api-server-count > 1"
778
- " when mode is not 'keep'"
779
- )
780
- async with self._pause_cond:
781
- if not self._paused:
782
- self._paused = True
783
-
784
- if mode == "abort":
785
- request_ids = list(self.output_processor.request_states.keys())
786
- if request_ids:
787
- await self.abort(request_ids, internal=True)
788
- elif mode == "wait":
789
- if self.output_processor.has_unfinished_requests():
790
- await self.output_processor.wait_for_requests_to_drain()
791
- else:
792
- raise ValueError(f"Invalid mode: {mode}")
793
-
794
- # Clear cache
795
- if clear_cache:
796
- await self.reset_prefix_cache(reset_running_requests=True)
797
- await self.reset_mm_cache()
798
- await self.reset_encoder_cache()
761
+ await self.engine_core.pause_scheduler_async(mode=mode, clear_cache=clear_cache)
762
763
async def resume_generation(self) -> None:
764
"""Resume generation after :meth:`pause_generation`."""
802
-
803
- async with self._pause_cond:
804
- await self.engine_core.resume_scheduler_async()
805
- self._paused = False
806
- self._pause_cond.notify_all() # Wake up all waiting requests
765
+ await self.engine_core.resume_scheduler_async()
766
767
async def is_paused(self) -> bool:
768
"""Return whether the engine is currently paused."""
810
-
811
- async with self._pause_cond:
812
- return self._paused
769
+ return await self.engine_core.is_scheduler_paused_async()
770
771
async def encode(
772
self,
vllm/v1/engine/core.py
+125
-50
@@ -5,7 +5,7 @@ import queue
5
import signal
6
import threading
7
import time
8
-from collections import deque
8
+from collections import defaultdict, deque
9
from collections.abc import Callable, Generator
10
from concurrent.futures import Future
11
from contextlib import ExitStack, contextmanager
@@ -40,7 +40,7 @@ from vllm.v1.core.kv_cache_utils import (
40
get_request_block_hasher,
41
init_none_hash,
42
)
43
-from vllm.v1.core.sched.interface import SchedulerInterface
43
+from vllm.v1.core.sched.interface import PauseState, SchedulerInterface
44
from vllm.v1.core.sched.output import SchedulerOutput
45
from vllm.v1.engine import (
46
EngineCoreOutput,
@@ -48,6 +48,7 @@ from vllm.v1.engine import (
48
EngineCoreRequest,
49
EngineCoreRequestType,
50
FinishReason,
51
+ PauseMode,
52
ReconfigureDistributedRequest,
53
ReconfigureRankType,
54
UtilityOutput,
@@ -210,8 +211,7 @@ class EngineCore:
211
212
self.aborts_queue = queue.Queue[list[str]]()
213
213
- # Pause state for "keep" mode - freezes requests in queue.
214
- self._scheduler_paused = False
214
+ self.per_step_hooks: set[Callable] = set()
215
216
# Mark the startup heap as static so that it's ignored by GC.
217
# Reduces pause times of oldest generation collections.
@@ -326,20 +326,6 @@ class EngineCore:
326
# (i.e. client-aborted vs stop criteria met).
327
self.scheduler.finish_requests(request_ids, RequestStatus.FINISHED_ABORTED)
328
329
- def pause_scheduler(self) -> None:
330
- """Pause the scheduler, keeping requests frozen in queue.
331
-
332
- Requests are kept frozen in queue and can be resumed later.
333
- """
334
- self._scheduler_paused = True
335
-
336
- def resume_scheduler(self) -> None:
337
- """Resume the scheduler after a pause.
338
-
339
- Resumes processing of frozen requests in the queue.
340
- """
341
- self._scheduler_paused = False
342
-
329
@contextmanager
330
def log_error_detail(self, scheduler_output: SchedulerOutput):
331
"""Execute the model and log detailed info on failure."""
@@ -393,10 +379,6 @@ class EngineCore:
379
was executed.
380
"""
381
396
- # If paused, don't schedule any work.
397
- if self._scheduler_paused:
398
- return {}, False
399
-
382
# Check for any requests remaining in the scheduler - unfinished,
383
# or finished and not yet removed from the batch.
384
if not self.scheduler.has_requests():
@@ -447,9 +429,6 @@ class EngineCore:
429
batch in the job queue is finished.
430
3. Update the scheduler from the output.
431
"""
450
- # If paused, don't schedule any work.
451
- if self._scheduler_paused:
452
- return {}, False
432
433
batch_queue = self.batch_queue
434
assert batch_queue is not None
@@ -613,6 +592,20 @@ class EngineCore:
592
# Reset the GPU model runner's encoder cache (physical storage)
593
self.model_executor.reset_encoder_cache()
594
595
+ def pause_scheduler(
596
+ self, mode: PauseMode = "abort", clear_cache: bool = True
597
+ ) -> Future[Any] | None:
598
+ """Pause scheduling. No-op in base EngineCore; overridden in EngineCoreProc."""
599
+ return None
600
+
601
+ def resume_scheduler(self) -> None:
602
+ """Resume scheduling. No-op in base EngineCore; overridden in EngineCoreProc."""
603
+
604
+ 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
608
+
609
def sleep(self, level: int = 1):
610
"""Put the engine to sleep at the specified level.
611
@@ -650,7 +643,7 @@ class EngineCore:
643
644
def is_sleeping(self) -> bool:
645
"""Check if engine is sleeping at any level."""
653
- return self._scheduler_paused or self.model_executor.is_sleeping
646
+ return self.is_scheduler_paused() or self.model_executor.is_sleeping
647
648
def execute_dummy_batch(self):
649
self.model_executor.execute_dummy_batch()
@@ -1053,13 +1046,9 @@ class EngineCoreProc(EngineCore):
1046
# 1) Poll the input queue until there is work to do.
1047
self._process_input_queue()
1048
# 2) Step the engine core and return the outputs.
1056
- # Skip if scheduling is paused (level 0 sleep)
1057
- if not self._scheduler_paused:
1058
- self._process_engine_step()
1059
- else:
1060
- # When scheduling is paused, still need to check for wake up
1061
- # by processing any utility requests that might resume scheduling
1062
- pass
1049
+ self._process_engine_step()
1050
+ # 3) Run any per-step hooks.
1051
+ self._process_per_step_hooks()
1052
1053
def _process_input_queue(self):
1054
"""Exits when an engine step needs to be performed."""
@@ -1067,9 +1056,9 @@ class EngineCoreProc(EngineCore):
1056
waited = False
1057
while (
1058
not self.engines_running
1070
- and (not self.scheduler.has_requests() or self._scheduler_paused)
1059
+ and not self.scheduler.has_requests()
1060
and not self.batch_queue
1072
- and not self._scheduler_paused
1061
+ and not self.per_step_hooks
1062
):
1063
if self.input_queue.empty():
1064
# Drain aborts queue; all aborts are also processed via input_queue.
@@ -1109,6 +1098,13 @@ class EngineCoreProc(EngineCore):
1098
1099
return model_executed
1100
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)
1107
+
1108
def _handle_client_request(
1109
self, request_type: EngineCoreRequestType, request: Any
1110
) -> None:
@@ -1122,18 +1118,14 @@ class EngineCoreProc(EngineCore):
1118
elif request_type == EngineCoreRequestType.UTILITY:
1119
client_idx, call_id, method_name, args = request
1120
output = UtilityOutput(call_id)
1125
- try:
1126
- method = getattr(self, method_name)
1127
- result = method(*self._convert_msgspec_args(method, args))
1128
- output.result = UtilityResult(result)
1129
- except BaseException as e:
1130
- logger.exception("Invocation of %s method failed", method_name)
1131
- output.failure_message = (
1132
- f"Call to {method_name} method failed: {str(e)}"
1133
- )
1134
- self.output_queue.put_nowait(
1135
- (client_idx, EngineCoreOutputs(utility_output=output))
1121
+ # Lazily look-up utility method so that failure will be handled/returned.
1122
+ get_result = lambda: (method := getattr(self, method_name)) and method(
1123
+ *self._convert_msgspec_args(method, args)
1124
)
1125
+ enqueue_output = lambda out: self.output_queue.put_nowait(
1126
+ (client_idx, EngineCoreOutputs(utility_output=out))
1127
+ )
1128
+ self._invoke_utility_method(method_name, get_result, output, enqueue_output)
1129
elif request_type == EngineCoreRequestType.EXECUTOR_FAILED:
1130
raise RuntimeError("Executor failed.")
1131
else:
@@ -1141,6 +1133,25 @@ class EngineCoreProc(EngineCore):
1133
"Unrecognized input request type encountered: %s", request_type
1134
)
1135
1136
+ @staticmethod
1137
+ def _invoke_utility_method(
1138
+ name: str, get_result: Callable, output: UtilityOutput, enqueue_output: Callable
1139
+ ):
1140
+ try:
1141
+ result = get_result()
1142
+ if isinstance(result, Future):
1143
+ # Defer utility output handling until future completion.
1144
+ callback = lambda future: EngineCoreProc._invoke_utility_method(
1145
+ name, future.result, output, enqueue_output
1146
+ )
1147
+ result.add_done_callback(callback)
1148
+ return
1149
+ output.result = UtilityResult(result)
1150
+ except Exception as e:
1151
+ logger.exception("Invocation of %s method failed", name)
1152
+ output.failure_message = f"Call to {name} method failed: {str(e)}"
1153
+ enqueue_output(output)
1154
+
1155
@staticmethod
1156
def _convert_msgspec_args(method, args):
1157
"""If a provided arg type doesn't match corresponding target method
@@ -1347,6 +1358,74 @@ class EngineCoreProc(EngineCore):
1358
)
1359
)
1360
1361
+ def pause_scheduler(
1362
+ self, mode: PauseMode = "abort", clear_cache: bool = True
1363
+ ) -> Future | None:
1364
+ """Pause generation; behavior depends on mode.
1365
+
1366
+ All pause states queue new adds. PAUSE_ABORT and PAUSE_KEEP skip step();
1367
+ PAUSE_WAIT allows step() so in-flight requests can drain.
1368
+
1369
+ - ``abort``: Set PAUSE_ABORT, abort all requests, wait for abort
1370
+ outputs to be sent (when running with output_queue), clear caches,
1371
+ then complete the returned Future.
1372
+ - ``wait``: Set PAUSE_WAIT (queue adds, keep stepping); when drained,
1373
+ set PAUSE_KEEP, clear caches, complete the returned Future.
1374
+ - ``keep``: Set PAUSE_KEEP; return a Future that completes when the
1375
+ output queue is empty.
1376
+ """
1377
+ if mode not in ("keep", "abort", "wait"):
1378
+ raise ValueError(f"Invalid pause mode: {mode}")
1379
+
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
1387
+ if clear_cache:
1388
+ engine.reset_prefix_cache(reset_running_requests=True)
1389
+ engine.reset_mm_cache()
1390
+ engine.reset_encoder_cache()
1391
+ future.set_result(None)
1392
+ return True
1393
+
1394
+ if mode == "abort":
1395
+ aborted_reqs = self.scheduler.finish_requests(
1396
+ None, RequestStatus.FINISHED_ABORTED
1397
+ )
1398
+ self._send_abort_outputs(aborted_reqs)
1399
+
1400
+ pause_state = PauseState.PAUSED_ALL if mode == "keep" else PauseState.PAUSED_NEW
1401
+ 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
1406
+
1407
+ def _send_abort_outputs(self, aborted_reqs: list[tuple[str, int]]) -> None:
1408
+ if aborted_reqs:
1409
+ # Map client_index to list of request_ids that belong to that client.
1410
+ by_client = defaultdict[int, set[str]](set)
1411
+ for req_id, client_index in aborted_reqs:
1412
+ by_client[client_index].add(req_id)
1413
+ for client_index, req_ids in by_client.items():
1414
+ outputs = [
1415
+ EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)
1416
+ for req_id in req_ids
1417
+ ]
1418
+ eco = EngineCoreOutputs(finished_requests=req_ids, outputs=outputs)
1419
+ self.output_queue.put_nowait((client_index, eco))
1420
+
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
+
1429
1430
class DPEngineCoreProc(EngineCoreProc):
1431
"""ZMQ-wrapper for running EngineCore in background process
@@ -1450,10 +1529,6 @@ class DPEngineCoreProc(EngineCoreProc):
1529
# 1) Poll the input queue until there is work to do.
1530
self._process_input_queue()
1531
1453
- # Skip processing if scheduling is paused (level 0 sleep)
1454
- if self._scheduler_paused:
1455
- continue
1456
-
1532
# 2) Step the engine core.
1533
executed = self._process_engine_step()
1534
self._maybe_publish_request_counts()
vllm/v1/engine/core_client.py
+8
-18
@@ -36,6 +36,7 @@ from vllm.v1.engine import (
36
EngineCoreOutputs,
37
EngineCoreRequest,
38
EngineCoreRequestType,
39
+ PauseMode,
40
ReconfigureDistributedRequest,
41
ReconfigureRankType,
42
UtilityOutput,
@@ -979,16 +980,17 @@ class AsyncMPClient(MPClient):
980
if request_ids and not self.resources.engine_dead:
981
await self._send_input(EngineCoreRequestType.ABORT, request_ids)
982
982
- async def pause_scheduler_async(self) -> None:
983
- """Pause the scheduler, keeping requests frozen in queue.
984
- Blocks until the EngineCore acknowledges the pause.
985
- """
986
- await self.call_utility_async("pause_scheduler")
983
+ async def pause_scheduler_async(
984
+ self, mode: PauseMode = "abort", clear_cache: bool = True
985
+ ) -> None:
986
+ await self.call_utility_async("pause_scheduler", mode, clear_cache)
987
988
async def resume_scheduler_async(self) -> None:
989
- """Resume the scheduler after a pause."""
989
await self.call_utility_async("resume_scheduler")
990
991
+ async def is_scheduler_paused_async(self) -> bool:
992
+ return await self.call_utility_async("is_scheduler_paused")
993
+
994
async def profile_async(
995
self, is_start: bool = True, profile_prefix: str | None = None
996
) -> None:
@@ -1203,18 +1205,6 @@ class DPAsyncMPClient(AsyncMPClient):
1205
def get_core_engine_for_request(self, request: EngineCoreRequest):
1206
return self.core_engine
1207
1206
- async def pause_scheduler_async(self) -> None:
1207
- """Pause the scheduler, keeping requests frozen in queue."""
1208
- raise NotImplementedError(
1209
- "pause_scheduler_async is not yet supported for data parallel"
1210
- )
1211
-
1212
- async def resume_scheduler_async(self) -> None:
1213
- """Resume the scheduler after a pause."""
1214
- raise NotImplementedError(
1215
- "resume_scheduler_async is not yet supported for data parallel"
1216
- )
1217
-
1208
1209
class DPLBAsyncMPClient(DPAsyncMPClient):
1210
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)