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 12, 09:45 UTC
[mypy] Pass type checking for `vllm/utils` and `vllm/v1/pool` (#29666)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Nov 28, 2025 at 20:40 UTC
953d9c820b0d8f9a55ee455814287536a498498c
9 files changed
+37
-43
tools/pre_commit/mypy.py
+2
-1
@@ -36,8 +36,10 @@ FILES = [
36
"vllm/transformers_utils",
37
"vllm/triton_utils",
38
"vllm/usage",
39
+ "vllm/utils",
40
"vllm/v1/core",
41
"vllm/v1/engine",
42
+ "vllm/v1/pool",
43
"vllm/v1/worker",
44
]
45
@@ -59,7 +61,6 @@ SEPARATE_GROUPS = [
61
"vllm/v1/executor",
62
"vllm/v1/kv_offload",
63
"vllm/v1/metrics",
62
- "vllm/v1/pool",
64
"vllm/v1/sample",
65
"vllm/v1/spec_decode",
66
"vllm/v1/structured_output",
vllm/utils/async_utils.py
+8
-1
@@ -12,7 +12,7 @@ from asyncio import FIRST_COMPLETED, AbstractEventLoop, Future, Task
12
from collections.abc import AsyncGenerator, Awaitable, Callable
13
from concurrent.futures import Executor, ThreadPoolExecutor
14
from functools import partial
15
-from typing import TypeVar
15
+from typing import TYPE_CHECKING, TypeVar
16
17
from transformers.tokenization_utils_base import BatchEncoding
18
from typing_extensions import ParamSpec
@@ -257,6 +257,13 @@ def in_loop(event_loop: AbstractEventLoop) -> bool:
257
return False
258
259
260
+# A hack to pass mypy
261
+if TYPE_CHECKING:
262
+
263
+ def anext(it: AsyncGenerator[T, None]):
264
+ return it.__anext__()
265
+
266
+
267
async def merge_async_iterators(
268
*iterators: AsyncGenerator[T, None],
269
) -> AsyncGenerator[tuple[int, T], None]:
vllm/utils/jsontree.py
+8
-15
@@ -4,7 +4,7 @@
4
5
from collections.abc import Callable, Iterable
6
from functools import reduce
7
-from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast, overload
7
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload
8
9
if TYPE_CHECKING:
10
import torch
@@ -82,16 +82,13 @@ def json_map_leaves(
82
83
def json_map_leaves(
84
func: Callable[[_T], _U],
85
- value: "BatchedTensorInputs" | _JSONTree[_T],
85
+ value: Any,
86
) -> "BatchedTensorInputs" | _JSONTree[_U]:
87
"""Apply a function to each leaf in a nested JSON structure."""
88
if isinstance(value, dict):
89
- return {
90
- k: json_map_leaves(func, v) # type: ignore[arg-type]
91
- for k, v in value.items()
92
- }
89
+ return {k: json_map_leaves(func, v) for k, v in value.items()} # type: ignore
90
elif isinstance(value, list):
94
- return [json_map_leaves(func, v) for v in value]
91
+ return [json_map_leaves(func, v) for v in value] # type: ignore
92
elif isinstance(value, tuple):
93
return tuple(json_map_leaves(func, v) for v in value)
94
else:
@@ -140,9 +137,9 @@ def json_reduce_leaves(
137
138
139
def json_reduce_leaves(
143
- func: Callable[..., _T | _U],
140
+ func: Callable[[_T, _T], _T] | Callable[[_U, _T], _U],
141
value: _JSONTree[_T],
145
- initial: _U = cast(_U, ...), # noqa: B008
142
+ initial: _U = ..., # type: ignore[assignment]
143
/,
144
) -> _T | _U:
145
"""
@@ -151,13 +148,9 @@ def json_reduce_leaves(
148
sequence to a single value.
149
"""
150
if initial is ...:
154
- return reduce(func, json_iter_leaves(value)) # type: ignore[arg-type]
151
+ return reduce(func, json_iter_leaves(value)) # type: ignore
152
156
- return reduce(
157
- func, # type: ignore[arg-type]
158
- json_iter_leaves(value),
159
- initial,
160
- )
153
+ return reduce(func, json_iter_leaves(value), initial) # type: ignore
154
155
156
def json_count_leaves(value: JSONTree[_T]) -> int:
vllm/utils/mem_utils.py
+2
-2
@@ -68,11 +68,11 @@ class MemorySnapshot:
68
timestamp: float = 0.0
69
auto_measure: bool = True
70
71
- def __post_init__(self):
71
+ def __post_init__(self) -> None:
72
if self.auto_measure:
73
self.measure()
74
75
- def measure(self):
75
+ def measure(self) -> None:
76
from vllm.platforms import current_platform
77
78
# we measure the torch peak memory usage via allocated_bytes,
vllm/utils/nccl.py
+3
-3
@@ -3,7 +3,7 @@
3
4
from __future__ import annotations
5
6
-import importlib
6
+import importlib.util
7
import os
8
9
import torch
@@ -47,8 +47,8 @@ def find_nccl_include_paths() -> list[str] | None:
47
48
try:
49
spec = importlib.util.find_spec("nvidia.nccl")
50
- if spec and getattr(spec, "submodule_search_locations", None):
51
- for loc in spec.submodule_search_locations:
50
+ if spec and (locs := getattr(spec, "submodule_search_locations", None)):
51
+ for loc in locs:
52
inc_dir = os.path.join(loc, "include")
53
if os.path.exists(os.path.join(inc_dir, "nccl.h")):
54
paths.append(inc_dir)
vllm/utils/network_utils.py
+1
-1
@@ -72,7 +72,7 @@ def get_ip() -> str:
72
return "0.0.0.0"
73
74
75
-def test_loopback_bind(address, family):
75
+def test_loopback_bind(address: str, family: int) -> bool:
76
try:
77
s = socket.socket(family, socket.SOCK_DGRAM)
78
s.bind((address, 0)) # Port 0 = auto assign
vllm/utils/registry.py
+4
-2
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any
3
+from typing import Any, TypeVar
4
+
5
+_T = TypeVar("_T", bound=type)
6
7
8
class ExtensionManager:
@@ -34,7 +36,7 @@ class ExtensionManager:
36
Decorator to register a class with the given name.
37
"""
38
37
- def wrap(cls_to_register):
39
+ def wrap(cls_to_register: _T) -> _T:
40
self.name2class[name] = cls_to_register
41
return cls_to_register
42
vllm/utils/torch_utils.py
+6
-15
@@ -13,7 +13,7 @@ import numpy.typing as npt
13
import torch
14
from packaging import version
15
from packaging.version import Version
16
-from torch.library import Library
16
+from torch.library import Library, infer_schema
17
18
import vllm.envs as envs
19
@@ -78,7 +78,6 @@ def guard_cuda_initialization():
78
yield
79
return
80
81
- had_key = "CUDA_VISIBLE_DEVICES" in os.environ
81
old_value = os.environ.get("CUDA_VISIBLE_DEVICES")
82
os.environ["CUDA_VISIBLE_DEVICES"] = ""
83
try:
@@ -90,10 +89,10 @@ def guard_cuda_initialization():
89
err_msg = str(e)
90
raise RuntimeError(err_msg) from e
91
finally:
93
- if had_key:
94
- os.environ["CUDA_VISIBLE_DEVICES"] = old_value
92
+ if old_value is None:
93
+ del os.environ["CUDA_VISIBLE_DEVICES"]
94
else:
96
- os.environ.pop("CUDA_VISIBLE_DEVICES")
95
+ os.environ["CUDA_VISIBLE_DEVICES"] = old_value
96
97
98
def get_dtype_size(dtype: torch.dtype) -> int:
@@ -525,8 +524,7 @@ def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
524
525
# Helper function used in testing.
526
def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:
528
- torch_version = version.parse(torch_version)
529
- return torch_version >= version.parse(target)
527
+ return version.parse(torch_version) >= version.parse(target)
528
529
530
def is_torch_equal_or_newer(target: str) -> bool:
@@ -640,15 +638,8 @@ def direct_register_custom_op(
638
639
dispatch_key = current_platform.dispatch_key
640
643
- import torch.library
644
-
645
- if hasattr(torch.library, "infer_schema"):
646
- schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
647
- else:
648
- # for pytorch 2.4
649
- import torch._custom_op.impl
641
+ schema_str = infer_schema(op_func, mutates_args=mutates_args)
642
651
- schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
643
my_lib = target_lib or vllm_lib
644
my_lib.define(op_name + schema_str, tags=tags)
645
my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
vllm/v1/pool/metadata.py
+3
-3
@@ -67,16 +67,16 @@ def build_pooling_cursor(
67
68
n_seq = len(num_scheduled_tokens)
69
index = list(range(n_seq))
70
- num_scheduled_tokens = torch.tensor(num_scheduled_tokens, device="cpu")
70
+ num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, device="cpu")
71
cumsum = torch.zeros(
72
n_seq + 1, dtype=torch.int64, pin_memory=pin_memory, device="cpu"
73
)
74
- torch.cumsum(num_scheduled_tokens, dim=0, out=cumsum[1:])
74
+ torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
75
cumsum = cumsum.to(device, non_blocking=True)
76
return PoolingCursor(
77
index=index,
78
first_token_indices_gpu=cumsum[:n_seq],
79
last_token_indices_gpu=cumsum[1:] - 1,
80
prompt_lens_cpu=prompt_lens,
81
- num_scheduled_tokens_cpu=num_scheduled_tokens,
81
+ num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
82
)