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 15, 19:45 UTC
[core][V1] pluggable scheduler (#14466)
Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
Joe Runde committed
Mar 11, 2025 at 19:15 UTC
47532cd9f4bb751955d10989eda2078966deb0aa
3 files changed
+58
-11
tests/plugins_tests/test_scheduler_plugins.py
+40
-9
@@ -1,27 +1,35 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import pytest
4
+
5
from vllm.core.scheduler import Scheduler
6
+from vllm.engine.arg_utils import EngineArgs
7
+from vllm.engine.llm_engine import LLMEngine
8
+from vllm.sampling_params import SamplingParams
9
+from vllm.v1.core.scheduler import Scheduler as V1Scheduler
10
+from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
11
12
6
-class DummyScheduler(Scheduler):
13
+class DummyV0Scheduler(Scheduler):
14
15
def schedule(self):
9
- raise Exception("Exception raised by DummyScheduler")
16
+ raise Exception("Exception raised by DummyV0Scheduler")
17
+
18
19
+class DummyV1Scheduler(V1Scheduler):
20
12
-def test_scheduler_plugins():
13
- import pytest
21
+ def schedule(self):
22
+ raise Exception("Exception raised by DummyV1Scheduler")
23
15
- from vllm.engine.arg_utils import EngineArgs
16
- from vllm.engine.llm_engine import LLMEngine
17
- from vllm.sampling_params import SamplingParams
24
25
+def test_scheduler_plugins_v0(monkeypatch):
26
+ monkeypatch.setenv("VLLM_USE_V1", "0")
27
with pytest.raises(Exception) as exception_info:
28
29
engine_args = EngineArgs(
30
model="facebook/opt-125m",
31
enforce_eager=True, # reduce test time
24
- scheduler_cls=DummyScheduler,
32
+ scheduler_cls=DummyV0Scheduler,
33
)
34
35
engine = LLMEngine.from_engine_args(engine_args=engine_args)
@@ -30,4 +38,27 @@ def test_scheduler_plugins():
38
engine.add_request("0", "foo", sampling_params)
39
engine.step()
40
33
- assert str(exception_info.value) == "Exception raised by DummyScheduler"
41
+ assert str(exception_info.value) == "Exception raised by DummyV0Scheduler"
42
+
43
+
44
+def test_scheduler_plugins_v1(monkeypatch):
45
+ monkeypatch.setenv("VLLM_USE_V1", "1")
46
+ # Explicitly turn off engine multiprocessing so that the scheduler runs in
47
+ # this process
48
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
49
+
50
+ with pytest.raises(Exception) as exception_info:
51
+
52
+ engine_args = EngineArgs(
53
+ model="facebook/opt-125m",
54
+ enforce_eager=True, # reduce test time
55
+ scheduler_cls=DummyV1Scheduler,
56
+ )
57
+
58
+ engine = V1LLMEngine.from_engine_args(engine_args=engine_args)
59
+
60
+ sampling_params = SamplingParams(max_tokens=1)
61
+ engine.add_request("0", "foo", sampling_params)
62
+ engine.step()
63
+
64
+ assert str(exception_info.value) == "Exception raised by DummyV1Scheduler"
vllm/engine/arg_utils.py
+5
@@ -1437,6 +1437,11 @@ class EngineArgs:
1437
1438
# V1 always uses chunked prefills.
1439
self.enable_chunked_prefill = True
1440
+ # V1 should use the new scheduler by default.
1441
+ # Swap it only if this arg is set to the original V0 default
1442
+ if self.scheduler_cls == EngineArgs.scheduler_cls:
1443
+ self.scheduler_cls = "vllm.v1.core.scheduler.Scheduler"
1444
+
1445
# When no user override, set the default values based on the usage
1446
# context.
1447
# Use different default values for different hardware.
vllm/v1/engine/core.py
+13
-2
@@ -19,9 +19,10 @@ from vllm.logger import init_logger
19
from vllm.lora.request import LoRARequest
20
from vllm.transformers_utils.config import (
21
maybe_register_config_serialize_by_value)
22
-from vllm.utils import get_exception_traceback, zmq_socket_ctx
22
+from vllm.utils import (get_exception_traceback, resolve_obj_by_qualname,
23
+ zmq_socket_ctx)
24
from vllm.v1.core.kv_cache_utils import get_kv_cache_configs
24
-from vllm.v1.core.scheduler import Scheduler, SchedulerOutput
25
+from vllm.v1.core.scheduler import SchedulerOutput
26
from vllm.v1.engine import (EngineCoreOutputs, EngineCoreRequest,
27
EngineCoreRequestType, UtilityOutput)
28
from vllm.v1.engine.mm_input_cache import MMInputCacheServer
@@ -65,6 +66,16 @@ class EngineCore:
66
self.structured_output_manager = StructuredOutputManager(vllm_config)
67
68
# Setup scheduler.
69
+ if isinstance(vllm_config.scheduler_config.scheduler_cls, str):
70
+ logger.warning(
71
+ "Using configured V1 scheduler class %s. "
72
+ "This scheduler interface is not public and "
73
+ "compatibility may not be maintained.",
74
+ vllm_config.scheduler_config.scheduler_cls)
75
+ Scheduler = resolve_obj_by_qualname(
76
+ vllm_config.scheduler_config.scheduler_cls)
77
+ else:
78
+ Scheduler = vllm_config.scheduler_config.scheduler_cls
79
self.scheduler = Scheduler(
80
scheduler_config=vllm_config.scheduler_config,
81
model_config=vllm_config.model_config,