diff --git a/vllm/executor/ray_gpu_executor.py b/vllm/executor/ray_gpu_executor.py index bc5d7f2102a6b..62594d72ae266 100644 --- a/vllm/executor/ray_gpu_executor.py +++ b/vllm/executor/ray_gpu_executor.py @@ -2,7 +2,7 @@ import os from collections import defaultdict from itertools import islice, repeat -from typing import TYPE_CHECKING, Any, DefaultDict, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union import msgspec @@ -53,8 +53,9 @@ def _init_executor(self) -> None: "VLLM_USE_RAY_COMPILED_DAG=1") assert self.uses_ray - placement_group: Optional["PlacementGroup"] = self.parallel_config.placement_group - + placement_group: Optional[PlacementGroup] = ( + self.parallel_config.placement_group + ) # Disable Ray usage stats collection. ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0") if ray_usage != "1": @@ -205,9 +206,10 @@ def sort_by_driver_then_worker_ip(worker): worker_node_and_gpu_ids.append( ray.get(worker.get_node_and_gpu_ids.remote()) \ ) # type: ignore - - node_workers = defaultdict[str, list[int]](list) # node id -> list of worker ranks - node_gpus = defaultdict[str, list[Union[int, str]]](list) # node id -> list of gpu ids + # node id -> list of worker ranks + node_workers = defaultdict[str, list[int]](list) + # node id -> list of gpu ids + node_gpus = defaultdict[str, list[Union[int, str]]](list) for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids): node_workers[node_id].append(i)