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 11, 13:45 UTC
[CI] Fix mypy for vllm/v1/executor (#30517)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 12, 2025 at 13:05 UTC
02a58803948e6b493a9bde6d38b69423a638ae49
5 files changed
+20
-13
tools/pre_commit/mypy.py
+1
-1
@@ -43,6 +43,7 @@ FILES = [
43
"vllm/worker",
44
"vllm/v1/core",
45
"vllm/v1/engine",
46
+ "vllm/v1/executor",
47
"vllm/v1/metrics",
48
"vllm/v1/pool",
49
"vllm/v1/sample",
@@ -60,7 +61,6 @@ SEPARATE_GROUPS = [
61
"vllm/model_executor",
62
# v1 related
63
"vllm/v1/attention",
63
- "vllm/v1/executor",
64
"vllm/v1/kv_offload",
65
"vllm/v1/spec_decode",
66
"vllm/v1/structured_output",
vllm/v1/executor/abstract.py
+1
-1
@@ -219,7 +219,7 @@ class Executor(ABC):
219
220
def sample_tokens(
221
self, grammar_output: GrammarOutput | None, non_block: bool = False
222
- ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
222
+ ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
223
output = self.collective_rpc( # type: ignore[call-overload]
224
"sample_tokens", args=(grammar_output,), non_block=non_block
225
)
vllm/v1/executor/multiproc_executor.py
+7
-3
@@ -294,8 +294,8 @@ class MultiprocExecutor(Executor):
294
kwargs: dict | None = None,
295
non_block: bool = False,
296
unique_reply_rank: int | None = None,
297
- kv_output_aggregator: KVOutputAggregator = None,
298
- ) -> Any | list[Any] | Future[Any | list[Any]]:
297
+ kv_output_aggregator: KVOutputAggregator | None = None,
298
+ ) -> Any:
299
"""Returns single result if unique_reply_rank and/or kv_output_aggregator
300
is provided, otherwise list."""
301
assert self.rpc_broadcast_mq is not None, (
@@ -476,6 +476,8 @@ class WorkerProc:
476
"""Wrapper that runs one Worker in a separate process."""
477
478
READY_STR = "READY"
479
+ rpc_broadcast_mq: MessageQueue | None
480
+ worker_response_mq: MessageQueue | None
481
482
def _init_message_queues(
483
self, input_shm_handle: Handle, vllm_config: VllmConfig
@@ -487,7 +489,7 @@ class WorkerProc:
489
)
490
491
# Initializes a message queue for sending the model output
490
- self.worker_response_mq: MessageQueue = MessageQueue(1, 1)
492
+ self.worker_response_mq = MessageQueue(1, 1)
493
self.peer_response_handles = []
494
else:
495
# Initialize remote MessageQueue for receiving SchedulerOutput across nodes
@@ -720,6 +722,7 @@ class WorkerProc:
722
try:
723
reader.close()
724
worker = WorkerProc(*args, **kwargs)
725
+ assert worker.worker_response_mq is not None
726
727
# Send READY once we know everything is loaded
728
ready_writer.send(
@@ -804,6 +807,7 @@ class WorkerProc:
807
808
def worker_busy_loop(self, cancel: threading.Event | None = None):
809
"""Main busy loop for Multiprocessing Workers"""
810
+ assert self.rpc_broadcast_mq is not None
811
while True:
812
method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue(
813
cancel=cancel, indefinite=True
vllm/v1/executor/ray_executor.py
+3
-3
@@ -413,7 +413,7 @@ class RayDistributedExecutor(Executor):
413
self,
414
grammar_output: "GrammarOutput | None",
415
non_block: bool = False,
416
- ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
416
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
417
"""Execute the model on the Ray workers.
418
419
The scheduler output to use should have been provided in
@@ -428,7 +428,7 @@ class RayDistributedExecutor(Executor):
428
"""
429
scheduler_output = self.scheduler_output
430
if scheduler_output is None:
431
- return COMPLETED_NONE_FUTURE if non_block else None # noqa
431
+ return COMPLETED_NONE_FUTURE if non_block else None
432
433
self.scheduler_output = None
434
@@ -439,7 +439,7 @@ class RayDistributedExecutor(Executor):
439
scheduler_output: SchedulerOutput,
440
grammar_output: "GrammarOutput | None",
441
non_block: bool = False,
442
- ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
442
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
443
# Build the compiled DAG for the first time.
444
if self.forward_dag is None: # type: ignore
445
self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
vllm/v1/executor/uniproc_executor.py
+8
-5
@@ -67,7 +67,7 @@ class UniProcExecutor(Executor):
67
kwargs: dict | None = None,
68
non_block: bool = False,
69
single_value: bool = False,
70
- ) -> Any | list[Any] | Future[Any | list[Any]]:
70
+ ) -> Any:
71
if kwargs is None:
72
kwargs = {}
73
@@ -79,10 +79,13 @@ class UniProcExecutor(Executor):
79
result = run_method(self.driver_worker, method, args, kwargs)
80
if isinstance(result, AsyncModelRunnerOutput):
81
if (async_thread := self.async_output_thread) is not None:
82
- get_output = result.get_output
83
- if not single_value:
84
- get_output = lambda go=result.get_output: [go()]
85
- return async_thread.submit(get_output)
82
+ if single_value:
83
+ return async_thread.submit(result.get_output)
84
+
85
+ def get_output_list() -> list[Any]:
86
+ return [result.get_output()]
87
+
88
+ return async_thread.submit(get_output_list)
89
result = result.get_output()
90
future = Future[Any]()
91
future.set_result(result if single_value else [result])