mypy for vllm/utils

Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>

Author: wwl2755

tools/pre_commit/mypy.py

@@ -36,6 +36,7 @@
     "vllm/transformers_utils",
     "vllm/triton_utils",
     "vllm/usage",
+    "vllm/utils",
 ]
 
 # After fixing errors resulting from changing follow_imports

vllm/engine/arg_utils.py

@@ -1240,10 +1240,12 @@ def create_engine_config(
         self.model = model_config.model
         self.tokenizer = model_config.tokenizer
 
+        # After ModelConfig init, tokenizer must be resolved (never None).
+        assert self.tokenizer is not None
         (self.model, self.tokenizer, self.speculative_config) = (
             maybe_override_with_speculators(
                 model=self.model,
-                tokenizer=self.tokenizer if self.tokenizer is not None else self.model,
+                tokenizer=self.tokenizer,
                 revision=self.revision,
                 trust_remote_code=self.trust_remote_code,
                 vllm_speculative_config=self.speculative_config,

vllm/multimodal/parse.py

@@ -13,7 +13,6 @@
     Optional,
     TypeVar,
     Union,
-    cast,
 )
 
 import numpy as np
@@ -367,8 +366,7 @@ def _is_embeddings(
         if isinstance(data, torch.Tensor):
             return data.ndim == 3
         if is_list_of(data, torch.Tensor):
-            tensors = cast(list[torch.Tensor], data)
-            return tensors[0].ndim == 2
+            return data[0].ndim == 2
 
         return False
 
@@ -426,6 +424,8 @@ def _parse_audio_data(
         if self._is_embeddings(data):
             return AudioEmbeddingItems(data)
 
+        # Normalize into a list of audio items
+        data_items: list[AudioItem]
         if (
             is_list_of(data, float)
             or isinstance(data, (np.ndarray, torch.Tensor))
@@ -436,7 +436,7 @@ def _parse_audio_data(
         elif isinstance(data, (np.ndarray, torch.Tensor)):
             data_items = [elem for elem in data]
         else:
-            data_items = data  # type: ignore[assignment]
+            data_items = data
 
         new_audios = list[np.ndarray]()
         for data_item in data_items:

vllm/utils/__init__.py

@@ -81,7 +81,7 @@
 import setproctitle
 import torch
 import torch.types
-import yaml
+import yaml  # type: ignore[import-untyped]
 import zmq
 import zmq.asyncio
 from packaging import version
@@ -486,7 +486,10 @@ async def merge_async_iterators(
 
     loop = asyncio.get_running_loop()
 
-    awaits = {loop.create_task(anext(it)): (i, it) for i, it in enumerate(iterators)}
+    awaits: dict[asyncio.Task[T], tuple[int, AsyncGenerator[T, None]]] = {
+        loop.create_task(anext(it)): (i, it)  # type: ignore[arg-type]
+        for i, it in enumerate(iterators)
+    }
     try:
         while awaits:
             done, _ = await asyncio.wait(awaits.keys(), return_when=FIRST_COMPLETED)
@@ -495,7 +498,7 @@ async def merge_async_iterators(
                 try:
                     item = await d
                     i, it = pair
-                    awaits[loop.create_task(anext(it))] = pair
+                    awaits[loop.create_task(anext(it))] = pair  # type: ignore[arg-type]
                     yield i, item
                 except StopAsyncIteration:
                     pass
@@ -1163,11 +1166,13 @@ def find_nccl_include_paths() -> list[str] | None:
         import importlib.util
 
         spec = importlib.util.find_spec("nvidia.nccl")
-        if spec and getattr(spec, "submodule_search_locations", None):
-            for loc in spec.submodule_search_locations:
-                inc_dir = os.path.join(loc, "include")
-                if os.path.exists(os.path.join(inc_dir, "nccl.h")):
-                    paths.append(inc_dir)
+        if spec:
+            locations = getattr(spec, "submodule_search_locations", None)
+            if locations:
+                for loc in locations:
+                    inc_dir = os.path.join(loc, "include")
+                    if os.path.exists(os.path.join(inc_dir, "nccl.h")):
+                        paths.append(inc_dir)
     except Exception:
         pass
 

vllm/utils/jsontree.py

@@ -4,7 +4,7 @@
 
 from collections.abc import Iterable
 from functools import reduce
-from typing import TYPE_CHECKING, Callable, TypeVar, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union, overload
 
 if TYPE_CHECKING:
     import torch
@@ -94,7 +94,7 @@ def json_map_leaves(
             for k, v in value.items()
         }
     elif isinstance(value, list):
-        return [json_map_leaves(func, v) for v in value]
+        return [json_map_leaves(func, v) for v in value]  # type: ignore[return-value]
     elif isinstance(value, tuple):
         return tuple(json_map_leaves(func, v) for v in value)
     else:
@@ -143,11 +143,11 @@ def json_reduce_leaves(
 
 
 def json_reduce_leaves(
-    func: Callable[..., Union[_T, _U]],
-    value: _JSONTree[_T],
-    initial: _U = cast(_U, ...),  # noqa: B008
+    func: Callable[..., Any],
+    value: _JSONTree[Any],
+    initial: Any = ...,  # noqa: B008
     /,
-) -> Union[_T, _U]:
+) -> Any:
     """
     Apply a function of two arguments cumulatively to each leaf in a
     nested JSON structure, from left to right, so as to reduce the