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, 22:15 UTC
[Feature] Pluggable platform-specific scheduler (#13161)
Signed-off-by: Yannick Schnider <yannick.schnider1@ibm.com> Signed-off-by: Yannick Schnider <Yannick.Schnider1@ibm.com>
Yannick Schnider committed
Feb 19, 2025 at 10:16 UTC
423330263ba2a460bd93443b06e1ebcd4b87efe2
5 files changed
+56
-3
.buildkite/test-pipeline.yaml
+1
@@ -531,6 +531,7 @@ steps:
531
- pip uninstall vllm_add_dummy_platform -y
532
# end platform plugin tests
533
# other tests continue here:
534
+ - pytest -v -s plugins_tests/test_scheduler_plugins.py
535
- pip install -e ./plugins/vllm_add_dummy_model
536
- pytest -v -s distributed/test_distributed_oot.py
537
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
tests/plugins_tests/test_scheduler_plugins.py
new
+33
@@ -0,0 +1,33 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from vllm.core.scheduler import Scheduler
4
+
5
+
6
+class DummyScheduler(Scheduler):
7
+
8
+ def schedule(self):
9
+ raise Exception("Exception raised by DummyScheduler")
10
+
11
+
12
+def test_scheduler_plugins():
13
+ import pytest
14
+
15
+ from vllm.engine.arg_utils import EngineArgs
16
+ from vllm.engine.llm_engine import LLMEngine
17
+ from vllm.sampling_params import SamplingParams
18
+
19
+ with pytest.raises(Exception) as exception_info:
20
+
21
+ engine_args = EngineArgs(
22
+ model="facebook/opt-125m",
23
+ enforce_eager=True, # reduce test time
24
+ scheduler_cls=DummyScheduler,
25
+ )
26
+
27
+ engine = LLMEngine.from_engine_args(engine_args=engine_args)
28
+
29
+ sampling_params = SamplingParams(max_tokens=1)
30
+ engine.add_request("0", "foo", sampling_params)
31
+ engine.step()
32
+
33
+ assert str(exception_info.value) == "Exception raised by DummyScheduler"
vllm/config.py
+4
@@ -1495,6 +1495,10 @@ class SchedulerConfig:
1495
1496
chunked_prefill_enabled: bool = field(init=False)
1497
1498
+ # scheduler class or path. "vllm.core.scheduler.Scheduler" (default)
1499
+ # or "mod.custom_class".
1500
+ scheduler_cls: Union[str, Type[object]] = "vllm.core.scheduler.Scheduler"
1501
+
1502
def compute_hash(self) -> str:
1503
"""
1504
WARNING: Whenever a new field is added to this config,
vllm/engine/arg_utils.py
+10
@@ -192,6 +192,7 @@ class EngineArgs:
192
collect_detailed_traces: Optional[str] = None
193
disable_async_output_proc: bool = False
194
scheduling_policy: Literal["fcfs", "priority"] = "fcfs"
195
+ scheduler_cls: Union[str, Type[object]] = "vllm.core.scheduler.Scheduler"
196
197
override_neuron_config: Optional[Dict[str, Any]] = None
198
override_pooler_config: Optional[PoolerConfig] = None
@@ -938,6 +939,13 @@ class EngineArgs:
939
'priority (lower value means earlier handling) and time of '
940
'arrival deciding any ties).')
941
942
+ parser.add_argument(
943
+ '--scheduler-cls',
944
+ default=EngineArgs.scheduler_cls,
945
+ help='The scheduler class to use. "vllm.core.scheduler.Scheduler" '
946
+ 'is the default scheduler. Can be a class directly or the path to '
947
+ 'a class of form "mod.custom_class".')
948
+
949
parser.add_argument(
950
'--override-neuron-config',
951
type=json.loads,
@@ -1273,10 +1281,12 @@ class EngineArgs:
1281
send_delta_data=(envs.VLLM_USE_RAY_SPMD_WORKER
1282
and parallel_config.use_ray),
1283
policy=self.scheduling_policy,
1284
+ scheduler_cls=self.scheduler_cls,
1285
max_num_partial_prefills=self.max_num_partial_prefills,
1286
max_long_partial_prefills=self.max_long_partial_prefills,
1287
long_prefill_token_threshold=self.long_prefill_token_threshold,
1288
)
1289
+
1290
lora_config = LoRAConfig(
1291
bias_enabled=self.enable_lora_bias,
1292
max_lora_rank=self.max_lora_rank,
vllm/engine/llm_engine.py
+8
-3
@@ -19,8 +19,7 @@ import vllm.envs as envs
19
from vllm.config import (DecodingConfig, LoRAConfig, ModelConfig,
20
ObservabilityConfig, ParallelConfig, SchedulerConfig,
21
VllmConfig)
22
-from vllm.core.scheduler import (ScheduledSequenceGroup, Scheduler,
23
- SchedulerOutputs)
22
+from vllm.core.scheduler import ScheduledSequenceGroup, SchedulerOutputs
23
from vllm.engine.arg_utils import EngineArgs
24
from vllm.engine.metrics_types import StatLoggerBase, Stats
25
from vllm.engine.output_processor.interfaces import (
@@ -58,7 +57,8 @@ from vllm.transformers_utils.tokenizer_group import (
57
BaseTokenizerGroup, init_tokenizer_from_configs)
58
from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,
59
usage_message)
61
-from vllm.utils import Counter, Device, deprecate_kwargs, weak_bind
60
+from vllm.utils import (Counter, Device, deprecate_kwargs,
61
+ resolve_obj_by_qualname, weak_bind)
62
from vllm.version import __version__ as VLLM_VERSION
63
64
logger = init_logger(__name__)
@@ -346,6 +346,11 @@ class LLMEngine:
346
# Create the scheduler.
347
# NOTE: the cache_config here have been updated with the numbers of
348
# GPU and CPU blocks, which are profiled in the distributed executor.
349
+ if isinstance(self.vllm_config.scheduler_config.scheduler_cls, str):
350
+ Scheduler = resolve_obj_by_qualname(
351
+ self.vllm_config.scheduler_config.scheduler_cls)
352
+ else:
353
+ Scheduler = self.vllm_config.scheduler_config.scheduler_cls
354
self.scheduler = [
355
Scheduler(
356
self.scheduler_config, self.cache_config, self.lora_config,