@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 480bdf5a7b
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 14, 05:30 UTC

[Core] Support custom executor qualname (#23314)

Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>

22quinn committed Aug 21, 2025 at 18:40 UTC 480bdf5a7b23ae9b9c06e56b3e9fab8e28f1f7c8
6 files changed +133 -8
.buildkite/test-pipeline.yaml
+1
@@ -244,6 +244,7 @@ steps:
244 - pytest -v -s v1/core
245 - pytest -v -s v1/engine
246 - pytest -v -s v1/entrypoints
247 + - pytest -v -s v1/executor
248 - pytest -v -s v1/sample
249 - pytest -v -s v1/logits_processors
250 - pytest -v -s v1/worker
tests/v1/executor/__init__.py
tests/v1/executor/test_executor.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import asyncio
5 +import os
6 +from typing import Any, Callable, Optional, Union
7 +
8 +import pytest
9 +
10 +from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
11 +from vllm.sampling_params import SamplingParams
12 +from vllm.v1.engine.async_llm import AsyncLLM
13 +from vllm.v1.engine.llm_engine import LLMEngine
14 +from vllm.v1.executor.multiproc_executor import MultiprocExecutor
15 +
16 +
17 +class Mock:
18 + ...
19 +
20 +
21 +class CustomMultiprocExecutor(MultiprocExecutor):
22 +
23 + def collective_rpc(self,
24 + method: Union[str, Callable],
25 + timeout: Optional[float] = None,
26 + args: tuple = (),
27 + kwargs: Optional[dict] = None,
28 + non_block: bool = False,
29 + unique_reply_rank: Optional[int] = None) -> list[Any]:
30 + # Drop marker to show that this was ran
31 + with open(".marker", "w"):
32 + ...
33 + return super().collective_rpc(method, timeout, args, kwargs)
34 +
35 +
36 +CustomMultiprocExecutorAsync = CustomMultiprocExecutor
37 +MODEL = "Qwen/Qwen3-0.6B"
38 +
39 +
40 +def test_custom_executor_type_checking():
41 + with pytest.raises(ValueError):
42 + engine_args = EngineArgs(
43 + model=MODEL,
44 + gpu_memory_utilization=0.2,
45 + max_model_len=8192,
46 + distributed_executor_backend=Mock,
47 + )
48 + LLMEngine.from_engine_args(engine_args)
49 + with pytest.raises(ValueError):
50 + engine_args = AsyncEngineArgs(model=MODEL,
51 + gpu_memory_utilization=0.2,
52 + max_model_len=8192,
53 + distributed_executor_backend=Mock)
54 + AsyncLLM.from_engine_args(engine_args)
55 +
56 +
57 +@pytest.mark.parametrize("distributed_executor_backend", [
58 + CustomMultiprocExecutor,
59 + "tests.v1.executor.test_executor.CustomMultiprocExecutor"
60 +])
61 +def test_custom_executor(distributed_executor_backend, tmp_path):
62 + cwd = os.path.abspath(".")
63 + os.chdir(tmp_path)
64 + try:
65 + assert not os.path.exists(".marker")
66 +
67 + engine_args = EngineArgs(
68 + model=MODEL,
69 + gpu_memory_utilization=0.2,
70 + max_model_len=8192,
71 + distributed_executor_backend=distributed_executor_backend,
72 + enforce_eager=True, # reduce test time
73 + )
74 + engine = LLMEngine.from_engine_args(engine_args)
75 + sampling_params = SamplingParams(max_tokens=1)
76 +
77 + engine.add_request("0", "foo", sampling_params)
78 + engine.step()
79 +
80 + assert os.path.exists(".marker")
81 + finally:
82 + os.chdir(cwd)
83 +
84 +
85 +@pytest.mark.parametrize("distributed_executor_backend", [
86 + CustomMultiprocExecutorAsync,
87 + "tests.v1.executor.test_executor.CustomMultiprocExecutorAsync"
88 +])
89 +def test_custom_executor_async(distributed_executor_backend, tmp_path):
90 + cwd = os.path.abspath(".")
91 + os.chdir(tmp_path)
92 + try:
93 + assert not os.path.exists(".marker")
94 +
95 + engine_args = AsyncEngineArgs(
96 + model=MODEL,
97 + gpu_memory_utilization=0.2,
98 + max_model_len=8192,
99 + distributed_executor_backend=distributed_executor_backend,
100 + enforce_eager=True, # reduce test time
101 + )
102 + engine = AsyncLLM.from_engine_args(engine_args)
103 + sampling_params = SamplingParams(max_tokens=1)
104 +
105 + async def t():
106 + stream = engine.generate(request_id="0",
107 + prompt="foo",
108 + sampling_params=sampling_params)
109 + async for x in stream:
110 + ...
111 +
112 + asyncio.run(t())
113 +
114 + assert os.path.exists(".marker")
115 + finally:
116 + os.chdir(cwd)
vllm/config/parallel.py
+7 -7
@@ -143,7 +143,8 @@ class ParallelConfig:
143 placement_group: Optional[PlacementGroup] = None
144 """ray distributed model workers placement group."""
145
146 - distributed_executor_backend: Optional[Union[DistributedExecutorBackend,
146 + distributed_executor_backend: Optional[Union[str,
147 + DistributedExecutorBackend,
148 type[ExecutorBase]]] = None
149 """Backend to use for distributed model
150 workers, either "ray" or "mp" (multiprocessing). If the product
@@ -416,23 +417,22 @@ class ParallelConfig:
417 def use_ray(self) -> bool:
418 return self.distributed_executor_backend == "ray" or (
419 isinstance(self.distributed_executor_backend, type)
419 - and self.distributed_executor_backend.uses_ray)
420 + and getattr(self.distributed_executor_backend, "uses_ray", False))
421
422 @model_validator(mode='after')
423 def _verify_args(self) -> Self:
424 # Lazy import to avoid circular import
425 from vllm.executor.executor_base import ExecutorBase
426 from vllm.platforms import current_platform
426 - if self.distributed_executor_backend not in (
427 - "ray", "mp", "uni",
428 - "external_launcher", None) and not (isinstance(
427 + if self.distributed_executor_backend is not None and not isinstance(
428 + self.distributed_executor_backend, str) and not (isinstance(
429 self.distributed_executor_backend, type) and issubclass(
430 self.distributed_executor_backend, ExecutorBase)):
431 raise ValueError(
432 "Unrecognized distributed executor backend "
433 f"{self.distributed_executor_backend}. Supported "
434 - "values are 'ray', 'mp' 'uni', 'external_launcher' or"
435 - " custom ExecutorBase subclass.")
434 + "values are 'ray', 'mp' 'uni', 'external_launcher', "
435 + " custom ExecutorBase subclass or its import path.")
436 if self.use_ray:
437 from vllm.executor import ray_utils
438 ray_utils.assert_ray_available()
vllm/engine/arg_utils.py
+1 -1
@@ -290,7 +290,7 @@ class EngineArgs:
290 # is intended for expert use only. The API may change without
291 # notice.
292 distributed_executor_backend: Optional[Union[
293 - DistributedExecutorBackend,
293 + str, DistributedExecutorBackend,
294 Type[ExecutorBase]]] = ParallelConfig.distributed_executor_backend
295 # number of P/D disaggregation (or other disaggregation) workers
296 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
vllm/v1/executor/abstract.py
+8
@@ -13,6 +13,7 @@ from vllm.executor.uniproc_executor import ( # noqa
13 ExecutorWithExternalLauncher as ExecutorWithExternalLauncherV0)
14 from vllm.executor.uniproc_executor import ( # noqa
15 UniProcExecutor as UniProcExecutorV0)
16 +from vllm.utils import resolve_obj_by_qualname
17 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
18 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
19
@@ -50,6 +51,13 @@ class Executor(ExecutorBase):
51 # TODO: make v1 scheduling deterministic
52 # to support external launcher
53 executor_class = ExecutorWithExternalLauncher
54 + elif isinstance(distributed_executor_backend, str):
55 + executor_class = resolve_obj_by_qualname(
56 + distributed_executor_backend)
57 + if not issubclass(executor_class, ExecutorBase):
58 + raise TypeError(
59 + "distributed_executor_backend must be a subclass of "
60 + f"ExecutorBase. Got {executor_class}.")
61 else:
62 raise ValueError("Unknown distributed executor backend: "
63 f"{distributed_executor_backend}")