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, 17:45 UTC
[Feat][RL] Pause and Resume with keep requests for single engine (#32351)
Signed-off-by: ahao-anyscale <ahao@anyscale.com> Signed-off-by: Aaron Hao <ahao@anyscale.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Aaron Hao committed
Feb 6, 2026 at 16:08 UTC
89a385d79fb511f439f8ebe9d38d6811b5b24a91
8 files changed
+537
-31
examples/offline_inference/pause_resume.py
new
+108
@@ -0,0 +1,108 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Test for pause/resume with keep mode.
4
+
5
+This test uses concurrent tasks to verify the engine truly stops generating
6
+during pause:
7
+1. Generator task: continuously generates and logs time between tokens
8
+2. Controller task: sends pause/resume commands
9
+
10
+If the engine properly pauses, we should see a gap in token timestamps
11
+matching the pause duration.
12
+"""
13
+
14
+import asyncio
15
+import time
16
+
17
+from vllm import SamplingParams
18
+from vllm.engine.arg_utils import AsyncEngineArgs
19
+from vllm.v1.engine.async_llm import AsyncLLM
20
+
21
+PAUSE_DURATION = 3.0 # seconds
22
+
23
+
24
+async def main():
25
+ # Create engine with a small model
26
+ engine_args = AsyncEngineArgs(
27
+ model="facebook/opt-125m",
28
+ enforce_eager=True,
29
+ )
30
+ engine = AsyncLLM.from_engine_args(engine_args)
31
+
32
+ prompt = "Write a story about a dragon. Once upon a time"
33
+ sampling_params = SamplingParams(max_tokens=30, ignore_eos=True)
34
+
35
+ # Track token arrival times
36
+ token_times: list[tuple[int, float]] = [] # (token_count, timestamp)
37
+ pause_time: float = 0
38
+ resume_time: float = 0
39
+ pause_token_idx: int = 0 # Index in token_times when pause occurred
40
+
41
+ async def generator_task():
42
+ """Generate tokens and record timestamps."""
43
+ async for output in engine.generate(
44
+ request_id="test-req",
45
+ prompt=prompt,
46
+ sampling_params=sampling_params,
47
+ ):
48
+ token_count = len(output.outputs[0].token_ids)
49
+ token_times.append((token_count, time.monotonic()))
50
+ print(
51
+ f"Token {token_count} arrived:"
52
+ f"T={token_times[-1][1] - token_times[0][1]:.3f}s"
53
+ )
54
+ return output
55
+
56
+ async def controller_task():
57
+ """Pause and resume the engine after some tokens generated."""
58
+ nonlocal pause_time, resume_time, pause_token_idx
59
+
60
+ # Wait for some tokens to be generated
61
+ while len(token_times) < 5:
62
+ await asyncio.sleep(0.01)
63
+
64
+ print(f"\nPausing engine (keep mode) at token {len(token_times)}")
65
+ pause_time = time.monotonic()
66
+ await engine.pause_generation(mode="keep")
67
+ pause_token_idx = len(token_times)
68
+ print(f"Paused! Sleeping for {PAUSE_DURATION}s...")
69
+
70
+ # Sleep while paused - no tokens should be generated during this time
71
+ await asyncio.sleep(PAUSE_DURATION)
72
+
73
+ print("Resuming engine...")
74
+ await engine.resume_generation()
75
+ resume_time = time.monotonic()
76
+ print("Resumed!\n")
77
+
78
+ # Run both tasks concurrently
79
+ gen_task = asyncio.create_task(generator_task())
80
+ ctrl_task = asyncio.create_task(controller_task())
81
+
82
+ final_output, _ = await asyncio.gather(gen_task, ctrl_task)
83
+
84
+ # Verify the pause actually stopped generation.
85
+ # The gap after the pause token should be approximately the sleep duration.
86
+ pause_gap = token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
87
+ print(
88
+ f"\nGap after pause (token {pause_token_idx - 1} -> {pause_token_idx}): "
89
+ f"{pause_gap:.3f}s"
90
+ )
91
+ if pause_gap >= PAUSE_DURATION * 0.9:
92
+ print(f"✓ Test passed! Engine paused for ~{pause_gap:.1f}s")
93
+ else:
94
+ print(
95
+ f"✗ Test failed! Expected ~{PAUSE_DURATION}s gap after pause, "
96
+ f"got {pause_gap:.3f}s"
97
+ )
98
+ raise AssertionError("Engine did not properly pause")
99
+
100
+ # Verify request completed
101
+ assert final_output.finished, "Request should have finished"
102
+ assert len(final_output.outputs[0].token_ids) == 30, "Should have all tokens"
103
+
104
+ engine.shutdown()
105
+
106
+
107
+if __name__ == "__main__":
108
+ asyncio.run(main())
tests/v1/engine/test_async_llm.py
+299
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import time
6
from contextlib import ExitStack
7
from unittest.mock import MagicMock
8
@@ -661,3 +662,301 @@ async def collect_outputs(
662
outputs_list.append(output)
663
final_output = output
664
return final_output
665
+
666
+
667
+# =============================================================================
668
+# Pause/Resume Tests
669
+# =============================================================================
670
+
671
+
672
+@pytest.mark.asyncio
673
+async def test_pause_resume_basic():
674
+ """Test basic pause/resume flag behavior and idempotency.
675
+
676
+ Tests:
677
+ - pause_generation sets the paused flag
678
+ - resume_generation clears the paused flag
679
+ - calling pause when already paused is a no-op
680
+ - calling resume when not paused is safe
681
+ - all pause modes work with no requests in flight
682
+ - rapid pause/resume cycles don't break the engine
683
+ """
684
+ with ExitStack() as after:
685
+ with set_default_torch_num_threads(1):
686
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
687
+ after.callback(engine.shutdown)
688
+
689
+ # Initially not paused
690
+ assert not await engine.is_paused()
691
+
692
+ # Resume when not paused should be safe
693
+ await engine.resume_generation()
694
+ assert not await engine.is_paused()
695
+
696
+ # Pause sets flag
697
+ await engine.pause_generation(mode="abort")
698
+ assert await engine.is_paused()
699
+
700
+ # Pause when already paused is a no-op
701
+ await engine.pause_generation(mode="abort")
702
+ assert await engine.is_paused()
703
+
704
+ # Resume clears flag
705
+ await engine.resume_generation()
706
+ assert not await engine.is_paused()
707
+
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()
714
+ await engine.resume_generation()
715
+ assert not await engine.is_paused()
716
+
717
+ # Concurrent pause/resume race conditions - should not deadlock or raise
718
+ await asyncio.gather(
719
+ engine.pause_generation(mode="abort"),
720
+ engine.resume_generation(),
721
+ engine.pause_generation(mode="abort"),
722
+ engine.resume_generation(),
723
+ )
724
+
725
+ # Ensure we end in a known state
726
+ await engine.resume_generation()
727
+ assert not await engine.is_paused()
728
+
729
+ # Engine should still work after all cycles
730
+ sampling_params = SamplingParams(max_tokens=5)
731
+ async for out in engine.generate(
732
+ request_id="post-cycles",
733
+ prompt=TEXT_PROMPT,
734
+ sampling_params=sampling_params,
735
+ ):
736
+ pass
737
+ assert out.finished
738
+
739
+
740
+@pytest.mark.asyncio
741
+async def test_pause_abort():
742
+ """Test that mode='abort' aborts in-flight requests immediately."""
743
+ with ExitStack() as after:
744
+ with set_default_torch_num_threads(1):
745
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
746
+ after.callback(engine.shutdown)
747
+
748
+ # Start a long-running request
749
+ sampling_params = SamplingParams(max_tokens=1000, ignore_eos=True)
750
+ outputs: list[RequestOutput] = []
751
+
752
+ async def gen():
753
+ async for out in engine.generate(
754
+ request_id="test-abort-pause",
755
+ prompt=TEXT_PROMPT,
756
+ sampling_params=sampling_params,
757
+ ):
758
+ outputs.append(out)
759
+ return outputs[-1] if outputs else None
760
+
761
+ # Start generation task
762
+ gen_task = asyncio.create_task(gen())
763
+
764
+ # Wait for some tokens to be generated
765
+ while len(outputs) < 3:
766
+ await asyncio.sleep(0.01)
767
+
768
+ # Pause with abort mode
769
+ await engine.pause_generation(mode="abort")
770
+
771
+ # Wait for task to complete (should be aborted)
772
+ final_output = await gen_task
773
+
774
+ # Request should be finished (aborted)
775
+ assert final_output is not None
776
+ assert final_output.finished
777
+ assert final_output.outputs[0].finish_reason == "abort"
778
+
779
+ # Also test that new requests are blocked while paused, then resume
780
+ assert await engine.is_paused()
781
+
782
+ request_completed = False
783
+
784
+ async def gen_blocked():
785
+ nonlocal request_completed
786
+ async for out in engine.generate(
787
+ request_id="test-blocked",
788
+ prompt=TEXT_PROMPT,
789
+ sampling_params=SamplingParams(max_tokens=5),
790
+ ):
791
+ pass
792
+ request_completed = True
793
+ return out
794
+
795
+ # Start a request (should block)
796
+ gen_task2 = asyncio.create_task(gen_blocked())
797
+
798
+ # Wait a bit - request should not have completed
799
+ await asyncio.sleep(0.3)
800
+ assert not request_completed, "Request should be blocked while paused"
801
+
802
+ # Resume
803
+ await engine.resume_generation()
804
+
805
+ # Now request should complete
806
+ final_output2 = await asyncio.wait_for(gen_task2, timeout=10.0)
807
+ assert request_completed
808
+ assert final_output2.finished
809
+
810
+
811
+@pytest.mark.asyncio
812
+async def test_pause_wait():
813
+ """Test that mode='wait' waits for in-flight requests to complete."""
814
+ with ExitStack() as after:
815
+ with set_default_torch_num_threads(1):
816
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
817
+ after.callback(engine.shutdown)
818
+
819
+ # Start a request - use fewer tokens since wait mode waits for completion
820
+ sampling_params = SamplingParams(max_tokens=10, ignore_eos=True)
821
+ got_first_token = asyncio.Event()
822
+ request_completed = False
823
+
824
+ async def gen():
825
+ nonlocal request_completed
826
+ async for out in engine.generate(
827
+ request_id="test-wait",
828
+ prompt=TEXT_PROMPT,
829
+ sampling_params=sampling_params,
830
+ ):
831
+ got_first_token.set()
832
+ request_completed = True
833
+ return out
834
+
835
+ # Start generation
836
+ gen_task = asyncio.create_task(gen())
837
+
838
+ # Wait for generation to start (event-driven)
839
+ await asyncio.wait_for(got_first_token.wait(), timeout=30.0)
840
+
841
+ # Pause with wait mode - should wait for request to finish
842
+ await engine.pause_generation(mode="wait")
843
+
844
+ # By now the request should be done (wait mode waits for completion)
845
+ assert request_completed, "Request should have completed during wait"
846
+
847
+ final_output = gen_task.result()
848
+ assert final_output.finished
849
+ # Should complete normally, not aborted
850
+ assert final_output.outputs[0].finish_reason != "eos"
851
+
852
+
853
+@pytest.mark.asyncio
854
+async def test_pause_keep_single_request():
855
+ """Test that mode='keep' freezes a single request and resumes with timing gap."""
856
+ with ExitStack() as after:
857
+ with set_default_torch_num_threads(1):
858
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
859
+ after.callback(engine.shutdown)
860
+
861
+ sampling_params = SamplingParams(max_tokens=30, ignore_eos=True)
862
+ token_times: list[tuple[int, float]] = []
863
+ pause_duration = 5.0
864
+ pause_token_idx = 0
865
+
866
+ async def generator_task():
867
+ """Generate tokens and record timestamps."""
868
+ async for output in engine.generate(
869
+ request_id="test-keep-single",
870
+ prompt=TEXT_PROMPT,
871
+ sampling_params=sampling_params,
872
+ ):
873
+ token_count = len(output.outputs[0].token_ids)
874
+ token_times.append((token_count, time.monotonic()))
875
+ return output
876
+
877
+ async def controller_task():
878
+ """Pause and resume the engine."""
879
+ nonlocal pause_token_idx
880
+ # Wait for some tokens (event-driven, handles slow token generation)
881
+ while len(token_times) < 5:
882
+ await asyncio.sleep(0.01)
883
+
884
+ # Pause with keep mode
885
+ await engine.pause_generation(mode="keep")
886
+ pause_token_idx = len(token_times)
887
+
888
+ # Sleep while paused
889
+ await asyncio.sleep(pause_duration)
890
+
891
+ # Resume
892
+ await engine.resume_generation()
893
+
894
+ # Run both tasks with timeout for slow generation
895
+ gen_task = asyncio.create_task(generator_task())
896
+ ctrl_task = asyncio.create_task(controller_task())
897
+
898
+ final_output, _ = await asyncio.wait_for(
899
+ asyncio.gather(gen_task, ctrl_task), timeout=60.0
900
+ )
901
+
902
+ # Request should complete with all tokens
903
+ assert final_output.finished
904
+ assert len(final_output.outputs[0].token_ids) == 30
905
+
906
+ # Check the gap at the recorded pause index matches the pause duration
907
+ pause_gap = (
908
+ token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
909
+ )
910
+ assert pause_gap >= pause_duration * 0.8, (
911
+ f"Expected gap of ~{pause_duration}s after pause, got {pause_gap:.3f}s"
912
+ )
913
+
914
+
915
+@pytest.mark.asyncio
916
+async def test_pause_keep_multi_request():
917
+ """Test that mode='keep' freezes multiple concurrent requests and all resume."""
918
+ with ExitStack() as after:
919
+ with set_default_torch_num_threads(1):
920
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
921
+ after.callback(engine.shutdown)
922
+
923
+ num_requests = 3
924
+ sampling_params = SamplingParams(max_tokens=10, ignore_eos=True)
925
+ completed_requests: list[str] = []
926
+ any_token_generated = asyncio.Event()
927
+
928
+ async def gen_multi(request_id: str):
929
+ async for out in engine.generate(
930
+ request_id=request_id,
931
+ prompt=TEXT_PROMPT,
932
+ sampling_params=sampling_params,
933
+ ):
934
+ any_token_generated.set()
935
+ completed_requests.append(request_id)
936
+ return out
937
+
938
+ # Start multiple requests
939
+ tasks = [
940
+ asyncio.create_task(gen_multi(f"req-multi-{i}"))
941
+ for i in range(num_requests)
942
+ ]
943
+
944
+ # Wait for at least one token across any request (event-driven)
945
+ await asyncio.wait_for(any_token_generated.wait(), timeout=30.0)
946
+
947
+ # Pause with keep mode
948
+ await engine.pause_generation(mode="keep")
949
+
950
+ # Wait while paused
951
+ await asyncio.sleep(0.5)
952
+
953
+ # Resume
954
+ await engine.resume_generation()
955
+
956
+ # All requests should complete
957
+ results = await asyncio.wait_for(asyncio.gather(*tasks), timeout=60.0)
958
+
959
+ assert len(completed_requests) == num_requests
960
+ for result in results:
961
+ assert result.finished
962
+ assert len(result.outputs[0].token_ids) == 10
vllm/engine/protocol.py
+14
-5
@@ -3,7 +3,7 @@
3
4
from abc import ABC, abstractmethod
5
from collections.abc import AsyncGenerator, Iterable, Mapping
6
-from typing import Any
6
+from typing import TYPE_CHECKING, Any
7
8
from vllm.config import ModelConfig, VllmConfig
9
from vllm.distributed.weight_transfer.base import (
@@ -22,6 +22,9 @@ from vllm.tasks import SupportedTask
22
from vllm.v1.engine import EngineCoreRequest
23
from vllm.v1.engine.input_processor import InputProcessor
24
25
+if TYPE_CHECKING:
26
+ from vllm.v1.engine import PauseMode
27
+
28
29
class EngineClient(ABC):
30
"""Protocol class for Clients to Engine"""
@@ -158,16 +161,22 @@ class EngineClient(ABC):
161
async def pause_generation(
162
self,
163
*,
164
+ mode: "PauseMode" = "abort",
165
wait_for_inflight_requests: bool = False,
166
clear_cache: bool = True,
167
) -> None:
168
"""Pause new generation/encoding requests.
169
170
Args:
167
- wait_for_inflight_requests: When ``True`` waits for in-flight requests
168
- to finish before pausing. When ``False`` (default), aborts in-flight
169
- requests immediately.
170
- clear_cache: Whether to clear KV and prefix caches after draining.
171
+ mode: How to handle in-flight requests:
172
+ - ``"abort"``: Abort all in-flight requests immediately
173
+ and return partial results with "abort" reason (default).
174
+ - ``"wait"``: Wait for in-flight requests to complete.
175
+ - ``"keep"``: Freeze requests in queue; they resume on
176
+ :meth:`resume_generation`.
177
+ wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead.
178
+ clear_cache: DEPRECATED. Whether to clear KV and prefix caches
179
+ after draining.
180
"""
181
...
182
vllm/entrypoints/serve/rlhf/api_router.py
+13
-6
@@ -3,6 +3,7 @@
3
4
import json
5
from http import HTTPStatus
6
+from typing import Annotated
7
8
from fastapi import APIRouter, FastAPI, HTTPException, Query, Request
9
from fastapi.responses import JSONResponse
@@ -14,6 +15,7 @@ from vllm.distributed.weight_transfer.base import (
15
)
16
from vllm.engine.protocol import EngineClient
17
from vllm.logger import init_logger
18
+from vllm.v1.engine import PauseMode
19
20
logger = init_logger(__name__)
21
@@ -28,24 +30,29 @@ router = APIRouter()
30
@router.post("/pause")
31
async def pause_generation(
32
raw_request: Request,
33
+ mode: Annotated[PauseMode, Query()] = "abort",
34
wait_for_inflight_requests: bool = Query(False),
32
- clear_cache: bool = Query(True),
35
+ clear_cache: Annotated[bool, Query()] = True,
36
) -> JSONResponse:
37
"""Pause generation requests to allow weight updates.
38
39
Args:
37
- wait_for_inflight_requests: When ``True`` waits for in-flight
38
- requests to finish before pausing. When ``False`` (default),
39
- aborts any in-flight requests immediately.
40
- clear_cache: Whether to clear KV/prefix caches after draining.
40
+ mode: How to handle in-flight requests:
41
+ - ``"abort"``: Abort all in-flight requests immediately (default).
42
+ - ``"wait"``: Wait for in-flight requests to complete.
43
+ - ``"keep"``: Freeze requests in queue; they resume on /resume.
44
+ wait_for_inflight_requests: DEPRECATED. Use ``mode="wait"`` instead.
45
+ clear_cache: DEPRECATED. Whether to clear KV/prefix caches after
46
+ draining. Ignored when mode="keep".
47
"""
48
49
engine = engine_client(raw_request)
50
51
try:
52
await engine.pause_generation(
47
- wait_for_inflight_requests=wait_for_inflight_requests,
53
+ mode=mode,
54
clear_cache=clear_cache,
55
+ wait_for_inflight_requests=wait_for_inflight_requests,
56
)
57
return JSONResponse(
58
content={"status": "paused"},
vllm/v1/engine/__init__.py
+7
-1
@@ -4,7 +4,7 @@
4
import enum
5
import time
6
from collections.abc import Mapping
7
-from typing import Any
7
+from typing import Any, Literal
8
9
import msgspec
10
import numpy as np
@@ -18,6 +18,12 @@ from vllm.v1.metrics.stats import SchedulerStats
18
from vllm.v1.outputs import LogprobsLists, LogprobsTensors
19
from vllm.v1.serial_utils import UtilityResult
20
21
+# Type for pause_generation mode parameter.
22
+# - "abort": Abort all in-flight requests immediately (default).
23
+# - "wait": Wait for in-flight requests to complete before pausing.
24
+# - "keep": Freeze requests in queue; they resume on resume_generation().
25
+PauseMode = Literal["abort", "wait", "keep"]
26
+
27
# These are possible values of RequestOutput.finish_reason,
28
# so form part of the external API.
29
FINISH_REASON_STRINGS = ("stop", "length", "abort", "error")
vllm/v1/engine/async_llm.py
+46
-18
@@ -38,7 +38,7 @@ from vllm.transformers_utils.config import maybe_register_config_serialize_by_va
38
from vllm.usage.usage_lib import UsageContext
39
from vllm.utils.async_utils import cancel_task_threadsafe
40
from vllm.utils.collection_utils import as_list
41
-from vllm.v1.engine import EngineCoreRequest
41
+from vllm.v1.engine import EngineCoreRequest, PauseMode
42
from vllm.v1.engine.core_client import EngineCoreClient
43
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
44
from vllm.v1.engine.input_processor import InputProcessor
@@ -170,6 +170,7 @@ class AsyncLLM(EngineClient):
170
# Pause / resume state for async RL workflows.
171
self._pause_cond = asyncio.Condition()
172
self._paused = False
173
+ self._client_count = client_count
174
175
self.output_handler: asyncio.Task | None = None
176
try:
@@ -728,7 +729,8 @@ class AsyncLLM(EngineClient):
729
async def pause_generation(
730
self,
731
*,
731
- wait_for_inflight_requests: bool = False,
732
+ mode: PauseMode = "abort",
733
+ wait_for_inflight_requests: bool | None = None,
734
clear_cache: bool = True,
735
) -> None:
736
"""
@@ -737,27 +739,52 @@ class AsyncLLM(EngineClient):
739
New generation/encoding requests are blocked until resume.
740
741
Args:
740
- wait_for_inflight_requests: When ``True`` waits for in-flight
741
- requests to finish before pausing. When ``False`` (default),
742
- immediately aborts any in-flight requests.
742
+ mode: How to handle in-flight requests:
743
+ - ``"abort"``: Abort all in-flight requests immediately
744
+ (default).
745
+ - ``"wait"``: Wait for in-flight requests to complete.
746
+ - ``"keep"``: Freeze requests in queue; they resume on
747
+ :meth:`resume_generation`.
748
+ wait_for_inflight_requests: DEPRECATED: use mode argument.
749
+ Whether to wait for in-flight requests to complete before pausing.
750
clear_cache: Whether to clear KV cache and prefix cache after
751
draining. Set to ``False`` to preserve cache for faster resume.
752
Default is ``True`` (clear caches).
746
- """
747
-
748
- async with self._pause_cond:
749
- if self._paused:
750
- return
751
- self._paused = True
753
753
- if not wait_for_inflight_requests:
754
- request_ids = list(self.output_processor.request_states.keys())
755
- if request_ids:
756
- await self.abort(request_ids, internal=True)
754
+ """
755
+ if wait_for_inflight_requests:
756
+ warnings.warn(
757
+ "The `wait_for_inflight_requests` parameter in "
758
+ "`AsyncLLM.pause_generation()` is deprecated. "
759
+ "Please use `mode` argument instead.",
760
+ DeprecationWarning,
761
+ stacklevel=2,
762
+ )
763
+ mode = "wait"
764
758
- # Wait for running requests to drain before clearing cache.
759
- if self.output_processor.has_unfinished_requests():
760
- await self.output_processor.wait_for_requests_to_drain()
765
+ if mode == "keep":
766
+ # Freeze requests in the scheduler - they will resume on
767
+ # resume_generation().
768
+ await self.engine_core.pause_scheduler_async()
769
+ else:
770
+ if self._client_count > 1:
771
+ raise NotImplementedError(
772
+ "pause_generation is not supported with --api-server-count > 1"
773
+ " when mode is not 'keep'"
774
+ )
775
+ async with self._pause_cond:
776
+ if not self._paused:
777
+ self._paused = True
778
+
779
+ if mode == "abort":
780
+ request_ids = list(self.output_processor.request_states.keys())
781
+ if request_ids:
782
+ await self.abort(request_ids, internal=True)
783
+ elif mode == "wait":
784
+ if self.output_processor.has_unfinished_requests():
785
+ await self.output_processor.wait_for_requests_to_drain()
786
+ else:
787
+ raise ValueError(f"Invalid mode: {mode}")
788
789
# Clear cache
790
if clear_cache:
@@ -769,6 +796,7 @@ class AsyncLLM(EngineClient):
796
"""Resume generation after :meth:`pause_generation`."""
797
798
async with self._pause_cond:
799
+ await self.engine_core.resume_scheduler_async()
800
self._paused = False
801
self._pause_cond.notify_all() # Wake up all waiting requests
802
vllm/v1/engine/core.py
+27
@@ -209,6 +209,10 @@ class EngineCore:
209
self.async_scheduling = vllm_config.scheduler_config.async_scheduling
210
211
self.aborts_queue = queue.Queue[list[str]]()
212
+
213
+ # Pause state for "keep" mode - freezes requests in queue.
214
+ self._scheduler_paused = False
215
+
216
# Mark the startup heap as static so that it's ignored by GC.
217
# Reduces pause times of oldest generation collections.
218
freeze_gc_heap()
@@ -322,6 +326,20 @@ 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
+
343
@contextmanager
344
def log_error_detail(self, scheduler_output: SchedulerOutput):
345
"""Execute the model and log detailed info on failure."""
@@ -375,6 +393,10 @@ class EngineCore:
393
was executed.
394
"""
395
396
+ # If paused, don't schedule any work.
397
+ if self._scheduler_paused:
398
+ return {}, False
399
+
400
# Check for any requests remaining in the scheduler - unfinished,
401
# or finished and not yet removed from the batch.
402
if not self.scheduler.has_requests():
@@ -425,6 +447,10 @@ class EngineCore:
447
batch in the job queue is finished.
448
3. Update the scheduler from the output.
449
"""
450
+ # If paused, don't schedule any work.
451
+ if self._scheduler_paused:
452
+ return {}, False
453
+
454
batch_queue = self.batch_queue
455
assert batch_queue is not None
456
@@ -1007,6 +1033,7 @@ class EngineCoreProc(EngineCore):
1033
not self.engines_running
1034
and not self.scheduler.has_requests()
1035
and not self.batch_queue
1036
+ and not self._scheduler_paused
1037
):
1038
if self.input_queue.empty():
1039
# Drain aborts queue; all aborts are also processed via input_queue.
vllm/v1/engine/core_client.py
+23
-1
@@ -105,7 +105,7 @@ class EngineCoreClient(ABC):
105
client_addresses: dict[str, str] | None = None,
106
client_count: int = 1,
107
client_index: int = 0,
108
- ) -> "MPClient":
108
+ ) -> "AsyncMPClient":
109
parallel_config = vllm_config.parallel_config
110
client_args = (
111
vllm_config,
@@ -976,6 +976,16 @@ class AsyncMPClient(MPClient):
976
if request_ids and not self.resources.engine_dead:
977
await self._send_input(EngineCoreRequestType.ABORT, request_ids)
978
979
+ async def pause_scheduler_async(self) -> None:
980
+ """Pause the scheduler, keeping requests frozen in queue.
981
+ Blocks until the EngineCore acknowledges the pause.
982
+ """
983
+ await self.call_utility_async("pause_scheduler")
984
+
985
+ async def resume_scheduler_async(self) -> None:
986
+ """Resume the scheduler after a pause."""
987
+ await self.call_utility_async("resume_scheduler")
988
+
989
async def profile_async(self, is_start: bool = True) -> None:
990
await self.call_utility_async("profile", is_start)
991
@@ -1188,6 +1198,18 @@ class DPAsyncMPClient(AsyncMPClient):
1198
def get_core_engine_for_request(self, request: EngineCoreRequest):
1199
return self.core_engine
1200
1201
+ async def pause_scheduler_async(self) -> None:
1202
+ """Pause the scheduler, keeping requests frozen in queue."""
1203
+ raise NotImplementedError(
1204
+ "pause_scheduler_async is not yet supported for data parallel"
1205
+ )
1206
+
1207
+ async def resume_scheduler_async(self) -> None:
1208
+ """Resume the scheduler after a pause."""
1209
+ raise NotImplementedError(
1210
+ "resume_scheduler_async is not yet supported for data parallel"
1211
+ )
1212
+
1213
1214
class DPLBAsyncMPClient(DPAsyncMPClient):
1215
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)