Update Optional[x] -> x | None and Union[x, y] to x | y (#26633)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

vllm/model_executor/model_loader/__init__.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from typing import Literal, Optional
+from typing import Literal
 
 from torch import nn
 
@@ -122,7 +122,7 @@ def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
 
 
 def get_model(
-    *, vllm_config: VllmConfig, model_config: Optional[ModelConfig] = None
+    *, vllm_config: VllmConfig, model_config: ModelConfig | None = None
 ) -> nn.Module:
     loader = get_model_loader(vllm_config.load_config)
     if model_config is None:

vllm/model_executor/model_loader/bitsandbytes_loader.py

@@ -6,8 +6,8 @@ import glob
 import itertools
 import math
 import os
-from collections.abc import Generator
-from typing import Any, Callable, Optional
+from collections.abc import Callable, Generator
+from typing import Any
 
 import numpy as np
 import torch
@@ -88,7 +88,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
         self,
         model_name_or_path: str,
         allowed_patterns: list[str],
-        revision: Optional[str] = None,
+        revision: str | None = None,
     ) -> tuple[str, list[str], str]:
         """Retrieve weight files. Download the files if necessary.
 
@@ -122,7 +122,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
         raise RuntimeError(f"No model weights found in: `{model_name_or_path}`")
 
     def _prepare_weights(
-        self, model_name_or_path: str, revision: Optional[str]
+        self, model_name_or_path: str, revision: str | None
     ) -> tuple[list[str], bool]:
         """Prepare weight files for the model."""
 
@@ -196,7 +196,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
     def _get_quantized_weights_iterator(
         self,
         model_name_or_path: str,
-        revision: Optional[str],
+        revision: str | None,
     ) -> tuple[Generator[tuple[str, torch.Tensor], None, None], dict[str, Any]]:
         """Get an iterator to the model weights with bitsandbytes quantization,
         as well as the quantization state dictionary."""

vllm/model_executor/model_loader/default_loader.py

@@ -5,7 +5,7 @@ import glob
 import os
 import time
 from collections.abc import Generator, Iterable
-from typing import Optional, cast
+from typing import cast
 
 import torch
 from torch import nn
@@ -47,7 +47,7 @@ class DefaultModelLoader(BaseModelLoader):
         model_or_path: str
         """The model ID or path."""
 
-        revision: Optional[str]
+        revision: str | None
         """The optional model revision."""
 
         prefix: str = ""
@@ -56,7 +56,7 @@ class DefaultModelLoader(BaseModelLoader):
         fall_back_to_pt: bool = True
         """Whether .pt weights can be used."""
 
-        allow_patterns_overrides: Optional[list[str]] = None
+        allow_patterns_overrides: list[str] | None = None
         """If defined, weights will load exclusively using these patterns."""
 
     counter_before_loading_weights: float = 0.0
@@ -79,9 +79,9 @@ class DefaultModelLoader(BaseModelLoader):
     def _prepare_weights(
         self,
         model_name_or_path: str,
-        revision: Optional[str],
+        revision: str | None,
         fall_back_to_pt: bool,
-        allow_patterns_overrides: Optional[list[str]],
+        allow_patterns_overrides: list[str] | None,
     ) -> tuple[str, list[str], bool]:
         """Prepare weights for the model.
 

vllm/model_executor/model_loader/runai_streamer_loader.py

@@ -3,7 +3,6 @@
 # ruff: noqa: SIM117
 import os
 from collections.abc import Generator
-from typing import Optional
 
 import torch
 from torch import nn
@@ -51,7 +50,7 @@ class RunaiModelStreamerLoader(BaseModelLoader):
                 os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url
 
     def _prepare_weights(
-        self, model_name_or_path: str, revision: Optional[str]
+        self, model_name_or_path: str, revision: str | None
     ) -> list[str]:
         """Prepare weights for the model.
 

vllm/model_executor/model_loader/sharded_state_loader.py

@@ -5,7 +5,7 @@ import collections
 import glob
 import os
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 import torch
 from torch import nn
@@ -89,7 +89,7 @@ class ShardedStateLoader(BaseModelLoader):
                     result[k] = t
         return result
 
-    def _prepare_weights(self, model_name_or_path: str, revision: Optional[str]):
+    def _prepare_weights(self, model_name_or_path: str, revision: str | None):
         if is_s3(model_name_or_path) or os.path.isdir(model_name_or_path):
             return model_name_or_path
         else:
@@ -171,8 +171,8 @@ class ShardedStateLoader(BaseModelLoader):
     def save_model(
         model: torch.nn.Module,
         path: str,
-        pattern: Optional[str] = None,
-        max_size: Optional[int] = None,
+        pattern: str | None = None,
+        max_size: int | None = None,
     ) -> None:
         from safetensors.torch import save_file
 

vllm/model_executor/model_loader/tensorizer.py

@@ -12,7 +12,7 @@ import threading
 import time
 from collections.abc import Generator, MutableMapping
 from dataclasses import asdict, dataclass, field, fields
-from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
 
 import regex as re
 import torch
@@ -67,7 +67,7 @@ __all__ = [
 logger = init_logger(__name__)
 
 
-def is_valid_deserialization_uri(uri: Optional[str]) -> bool:
+def is_valid_deserialization_uri(uri: str | None) -> bool:
     if uri:
         scheme = uri.lower().split("://")[0]
         return scheme in {"s3", "http", "https"} or os.path.exists(uri)
@@ -156,25 +156,23 @@ class _NoInitOrTensorImpl:
 
 @dataclass
 class TensorizerConfig(MutableMapping):
-    tensorizer_uri: Optional[str] = None
-    tensorizer_dir: Optional[str] = None
-    vllm_tensorized: Optional[bool] = None
-    verify_hash: Optional[bool] = None
-    num_readers: Optional[int] = None
-    encryption_keyfile: Optional[str] = None
-    s3_access_key_id: Optional[str] = None
-    s3_secret_access_key: Optional[str] = None
-    s3_endpoint: Optional[str] = None
-    lora_dir: Optional[str] = None
-    stream_kwargs: Optional[dict[str, Any]] = None
-    serialization_kwargs: Optional[dict[str, Any]] = None
-    deserialization_kwargs: Optional[dict[str, Any]] = None
-    _extra_serialization_attrs: Optional[dict[str, Any]] = field(
-        init=False, default=None
-    )
-    model_class: Optional[type[torch.nn.Module]] = field(init=False, default=None)
-    hf_config: Optional[PretrainedConfig] = field(init=False, default=None)
-    dtype: Optional[Union[str, torch.dtype]] = field(init=False, default=None)
+    tensorizer_uri: str | None = None
+    tensorizer_dir: str | None = None
+    vllm_tensorized: bool | None = None
+    verify_hash: bool | None = None
+    num_readers: int | None = None
+    encryption_keyfile: str | None = None
+    s3_access_key_id: str | None = None
+    s3_secret_access_key: str | None = None
+    s3_endpoint: str | None = None
+    lora_dir: str | None = None
+    stream_kwargs: dict[str, Any] | None = None
+    serialization_kwargs: dict[str, Any] | None = None
+    deserialization_kwargs: dict[str, Any] | None = None
+    _extra_serialization_attrs: dict[str, Any] | None = field(init=False, default=None)
+    model_class: type[torch.nn.Module] | None = field(init=False, default=None)
+    hf_config: PretrainedConfig | None = field(init=False, default=None)
+    dtype: str | torch.dtype | None = field(init=False, default=None)
     _is_sharded: bool = field(init=False, default=False)
     _fields: ClassVar[tuple[str, ...]]
     _keys: ClassVar[frozenset[str]]
@@ -362,9 +360,9 @@ TensorizerConfig._keys = frozenset(TensorizerConfig._fields)
 
 @dataclass
 class TensorizerArgs:
-    tensorizer_uri: Optional[str] = None
-    tensorizer_dir: Optional[str] = None
-    encryption_keyfile: Optional[str] = None
+    tensorizer_uri: str | None = None
+    tensorizer_dir: str | None = None
+    encryption_keyfile: str | None = None
 
     def __init__(self, tensorizer_config: TensorizerConfig):
         for k, v in tensorizer_config.items():
@@ -621,7 +619,7 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
 
 
 def serialize_extra_artifacts(
-    tensorizer_args: TensorizerArgs, served_model_name: Union[str, list[str], None]
+    tensorizer_args: TensorizerArgs, served_model_name: str | list[str] | None
 ) -> None:
     if not isinstance(served_model_name, str):
         raise ValueError(

vllm/model_executor/model_loader/tensorizer_loader.py

@@ -3,7 +3,6 @@
 # ruff: noqa: SIM117
 import copy
 from collections.abc import Generator
-from typing import Union
 
 import torch
 from torch import nn
@@ -140,7 +139,7 @@ class TensorizerLoader(BaseModelLoader):
     @staticmethod
     def save_model(
         model: torch.nn.Module,
-        tensorizer_config: Union[TensorizerConfig, dict],
+        tensorizer_config: TensorizerConfig | dict,
         model_config: ModelConfig,
     ) -> None:
         if isinstance(tensorizer_config, dict):

vllm/model_executor/model_loader/tpu.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import time
-from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -30,7 +29,7 @@ class TPUModelLoader(DefaultModelLoader):
         self,
         vllm_config: VllmConfig,
         model_config: ModelConfig,
-        mesh: Optional[xs.Mesh] = None,
+        mesh: xs.Mesh | None = None,
     ) -> nn.Module:
         # Initialize model and load weights on CPU. Then, during SPMD partition,
         # weights are sharded and transferred to TPUs.
@@ -90,7 +89,7 @@ class TPUModelLoader(DefaultModelLoader):
             )
         return model
 
-    def _check_model_is_loaded(self, mesh: Optional[xs.Mesh], model: nn.Module) -> None:
+    def _check_model_is_loaded(self, mesh: xs.Mesh | None, model: nn.Module) -> None:
         """
         Ensure the model is properly loaded.
         1. All model parameters and buffers are on XLA device.

vllm/model_executor/model_loader/utils.py

@@ -7,7 +7,6 @@ import inspect
 import warnings
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Optional
 
 import torch
 from torch import nn
@@ -46,8 +45,8 @@ def initialize_model(
     vllm_config: VllmConfig,
     *,
     prefix: str = "",
-    model_class: Optional[type[nn.Module]] = None,
-    model_config: Optional[ModelConfig] = None,
+    model_class: type[nn.Module] | None = None,
+    model_config: ModelConfig | None = None,
 ) -> nn.Module:
     """Initialize a model with the given configurations."""
     if model_config is None:
@@ -268,7 +267,7 @@ class ParamMapping:
                     index,
                 )
 
-    def get_sub_modules(self, module_name: str) -> Optional[tuple[str, list[str]]]:
+    def get_sub_modules(self, module_name: str) -> tuple[str, list[str]] | None:
         for key, value in self.packed_mapping.items():
             if module_name.endswith(key):
                 return key, value

vllm/model_executor/model_loader/weight_utils.py

@@ -11,10 +11,10 @@ import os
 import tempfile
 import time
 from collections import defaultdict
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import IO, Any, Callable, Optional, Union
+from typing import IO, Any
 
 import filelock
 import huggingface_hub.constants
@@ -85,7 +85,7 @@ class DisabledTqdm(tqdm):
         super().__init__(*args, **kwargs, disable=True)
 
 
-def get_lock(model_name_or_path: Union[str, Path], cache_dir: Optional[str] = None):
+def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
     lock_dir = cache_dir or temp_dir
     model_name_or_path = str(model_name_or_path)
     os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
@@ -100,7 +100,7 @@ def get_lock(model_name_or_path: Union[str, Path], cache_dir: Optional[str] = No
 
 @contextmanager
 def atomic_writer(
-    filepath: Union[str, Path], mode: str = "w", encoding: Optional[str] = None
+    filepath: str | Path, mode: str = "w", encoding: str | None = None
 ) -> Generator[IO]:
     """
     Context manager that provides an atomic file writing routine.
@@ -143,11 +143,11 @@ def atomic_writer(
 
 def maybe_download_from_modelscope(
     model: str,
-    revision: Optional[str] = None,
-    download_dir: Optional[str] = None,
-    ignore_patterns: Optional[Union[str, list[str]]] = None,
-    allow_patterns: Optional[Union[list[str], str]] = None,
-) -> Optional[str]:
+    revision: str | None = None,
+    download_dir: str | None = None,
+    ignore_patterns: str | list[str] | None = None,
+    allow_patterns: list[str] | str | None = None,
+) -> str | None:
     """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.
 
     Returns the path to the downloaded model, or None if the model is not
@@ -370,10 +370,10 @@ def get_sparse_attention_config(
 
 def download_weights_from_hf(
     model_name_or_path: str,
-    cache_dir: Optional[str],
+    cache_dir: str | None,
     allow_patterns: list[str],
-    revision: Optional[str] = None,
-    ignore_patterns: Optional[Union[str, list[str]]] = None,
+    revision: str | None = None,
+    ignore_patterns: str | list[str] | None = None,
 ) -> str:
     """Download model weights from Hugging Face Hub.
 
@@ -448,8 +448,8 @@ def download_weights_from_hf(
 def download_safetensors_index_file_from_hf(
     model_name_or_path: str,
     index_file: str,
-    cache_dir: Optional[str],
-    revision: Optional[str] = None,
+    cache_dir: str | None,
+    revision: str | None = None,
 ) -> None:
     """Download hf safetensors index file from Hugging Face Hub.
 
@@ -540,7 +540,7 @@ def enable_tqdm(use_tqdm_on_load: bool):
 
 def np_cache_weights_iterator(
     model_name_or_path: str,
-    cache_dir: Optional[str],
+    cache_dir: str | None,
     hf_folder: str,
     hf_weights_files: list[str],
     use_tqdm_on_load: bool,
@@ -746,7 +746,7 @@ def fastsafetensors_weights_iterator(
 def pt_weights_iterator(
     hf_weights_files: list[str],
     use_tqdm_on_load: bool,
-    pt_load_map_location: Union[str, dict[str, str]] = "cpu",
+    pt_load_map_location: str | dict[str, str] = "cpu",
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model bin/pt files."""
     for bin_file in tqdm(
@@ -765,7 +765,7 @@ def pt_weights_iterator(
 def multi_thread_pt_weights_iterator(
     hf_weights_files: list[str],
     use_tqdm_on_load: bool,
-    pt_load_map_location: Union[str, dict[str, str]] = "cpu",
+    pt_load_map_location: str | dict[str, str] = "cpu",
     max_workers: int = 4,
 ) -> Generator[tuple[str, torch.Tensor], None, None]:
     """Multi-Thread iterate over the weights in the model bin/pt files."""
@@ -985,7 +985,7 @@ def initialize_dummy_weights(
                 param.uniform_(low, high, generator=generator)
 
 
-def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
+def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
     """Remap the name of FP8 k/v_scale parameters.
 
     This function handles the remapping of FP8 k/v_scale parameter names.