fix pre-commit

Signed-off-by: 白永斌 <[email protected]>

Author: 白永斌

vllm/distributed/eplb/eplb_expert_mapper.py

@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
 import networkx as nx
 import numpy as np
 import torch
 
 
-from typing import Dict, List
-
 class ComposeExpertUpdate:
     def __init__(self, updated_expert_maps, current_expert_maps):
         self.updated_org = updated_expert_maps
@@ -39,8 +38,8 @@ def generate(self):
             updated_layer = self.updated[layer_id]
             current_layer = self.current[layer_id]
 
-            expert_send_info_this_layer: Dict[int, List[int]] = {}
-            expert_recv_info_this_layer: Dict[int, List[int]] = {}
+            expert_send_info_this_layer: dict[int, list[int]] = {}
+            expert_recv_info_this_layer: dict[int, list[int]] = {}
 
             # Guard Clause: if there is no expert weight update,
             # avoid subsequent processing.
@@ -49,7 +48,7 @@ def generate(self):
                     expert_send_info_this_layer,
                     expert_recv_info_this_layer,
                     self._map_to_yield(layer_id),
-                    layer_id
+                    layer_id,
                 )
 
             # Main planning
@@ -66,7 +65,7 @@ def generate(self):
                 expert_send_info_this_layer,
                 expert_recv_info_this_layer,
                 self._map_to_yield(layer_id),
-                layer_id
+                layer_id,
             )
 
 
@@ -97,9 +96,9 @@ def _plan_transfers(
         for idx in range(len(dst_rank_indices)):
             expert_id = experts_to_recv[idx].item()
             if expert_id not in src_ranks_set:
-                src_ranks_set[expert_id] = np.where(
-                    current_layer[:, expert_id] != -1
-                )[0]
+                src_ranks_set[expert_id] = np.where(current_layer[:, expert_id] != -1)[
+                    0
+                ]
 
         # Loop until all experts are scheduled
         while len(dst_rank_indices) > 0:
@@ -129,7 +128,7 @@ def _plan_transfers(
             for src_rank, dst_rank in all_matches.items():
                 dst_rank = int(dst_rank)
                 assert src_rank != dst_rank
-                if graph_expert_update.nodes[src_rank].get('bipartite') == 0:
+                if graph_expert_update.nodes[src_rank].get("bipartite") == 0:
                     # currently not scheduled experts in rank dst_rank
                     experts_v = experts_to_recv[np.where(dst_rank_indices == dst_rank)]
                     # src: src_rank, dest: dst_rank, expert: expert_id
@@ -161,6 +160,7 @@ def _plan_transfers(
 
 class GreedyExpertUpdate(ComposeExpertUpdate):
     """Greedy version."""
+
     def _prepare_internal(self, updated, current):
         # align devices
         if not torch.is_tensor(updated):
@@ -209,4 +209,3 @@ def _plan_transfers(
 
             send_dict[src_rank_id].append((dst_rank_id, expert_id))
             recv_dict[dst_rank_id].append((src_rank_id, expert_id))
-

vllm/distributed/eplb/eplb_process.py

@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import multiprocessing as mp
 import random
+from collections.abc import Callable
 from contextlib import suppress
 from multiprocessing import Queue
-from typing import Callable
 from queue import Empty
 
 import torch
@@ -13,7 +13,7 @@
 from vllm.logger import init_logger
 
 from .eplb_expert_mapper import BipartiteExpertUpdate, GreedyExpertUpdate
-from .eplb_state import RebalanceTaskArgs, ExpertMapperArgs
+from .eplb_state import ExpertMapperArgs, RebalanceTaskArgs
 
 logger = init_logger(__name__)
 
@@ -44,7 +44,6 @@ def __init__(self, target_func: Callable, num_wait_worker_iterations: int):
         self._exception_queue: Queue | None = None
         self._step_counter = 0
         self._result: tuple | None = None
-        self._args: tuple | None = None
         self._is_running = False
         self._has_pending_task = False
         self._is_post_processing = False
@@ -65,13 +64,13 @@ def _initialize_process(self) -> None:
             self._process = mp.Process(
                 target=self._worker_loop,
                 name="EPLBProcess",
-                args=(self._input_queue, self._result_queue, self._exception_queue)
+                args=(self._input_queue, self._result_queue, self._exception_queue),
             )
             self._process.start()
             self._is_running = True
             logger.debug("EPLB background process started")
 
-        except Exception as e:
+        except Exception:
             self.cleanup()
             raise
 
@@ -126,9 +125,8 @@ def generate_log2phy_map(self, expert_map):
         num_ranks, num_global_expert = log2phy_map.shape
 
         row_indices = (
-            torch.arange(num_ranks).view(-1, 1).expand(
-                num_ranks, num_global_expert
-            ) * num_local_experts
+            torch.arange(num_ranks).view(-1, 1).expand(num_ranks, num_global_expert)
+            * num_local_experts
         )
         log2phy_map[log2phy_map != -1] += row_indices[log2phy_map != -1]
 
@@ -182,7 +180,8 @@ def _worker_loop(
                             expert_mapper_args.num_moe_layers,
                             args.num_gpus,
                             -1,
-                    ))
+                        )
+                    )
                     if policy_type == "bipartite":
                         update_info = BipartiteExpertUpdate(
                             new_deployment, old_deployment
@@ -212,9 +211,7 @@ def _worker_loop(
             logger.debug("EPLB worker process exiting")
 
     def submit_task(
-        self,
-        args: RebalanceTaskArgs,
-        expert_mapper_args: ExpertMapperArgs
+        self, args: RebalanceTaskArgs, expert_mapper_args: ExpertMapperArgs
     ) -> bool:
         """
         Submit a task to the asynchronous process
@@ -242,7 +239,6 @@ def submit_task(
             # Put arguments to the input queue
             combined_args = (args, expert_mapper_args)
             self._input_queue.put(combined_args)
-            self._args = args
             self._has_pending_task = True
             self._step_counter = 0
             self._result = None
@@ -302,7 +298,7 @@ def cleanup(self) -> None:
         # Send sentinel value to stop the process
         if self._input_queue:
             with suppress(Exception):
-                self._input_queue.put(None, None)
+                self._input_queue.put(None)
 
         if self._process:
             if self._process.is_alive():

vllm/distributed/eplb/eplb_state.py

@@ -51,6 +51,7 @@
 
 logger = init_logger(__name__)
 
+
 @dataclass
 class RebalanceTaskArgs:
     global_expert_load_window: torch.Tensor
@@ -62,7 +63,7 @@ class RebalanceTaskArgs:
 @dataclass
 class ExpertMapperArgs:
     num_moe_layers: int
-    policy_type: Literal["greedy","bipartite"]
+    policy_type: Literal["greedy", "bipartite"]
     phyhsical_to_logical_map: torch.Tensor
 
 @dataclass
@@ -211,7 +212,7 @@ class EplbState:
     """
     Records the current moe layer being precessed for expert weight transfer.
     """
-    
+
     @staticmethod
     def build_initial_global_physical_to_logical_map(
         num_routed_experts: int,
@@ -381,9 +382,11 @@ def build(
                 rank_mapping,
             )
             expert_rearrangement_step = 0
-        expert_mapper_args = ExpertMapperArgs()
-        expert_mapper_args.num_moe_layers = model.num_moe_layers
-        expert_mapper_args.policy_type = parallel_config.eplb_config.expert_mapper_policy_type
+        expert_mapper_args = ExpertMapperArgs(
+            model.num_moe_layers,
+            parallel_config.eplb_config.expert_mapper_policy_type,
+            None
+        )
         return cls(
             physical_to_logical_map,
             logical_to_physical_map,
@@ -393,9 +396,11 @@ def build(
             expert_load_window_size=expert_load_window_size,
             expert_rearrangement_step=expert_rearrangement_step,
             expert_rearrangement_step_interval=eplb_step_interval,
-            num_wait_worker_iterations=parallel_config.eplb_config.num_wait_worker_iterations,
+            num_wait_worker_iterations=(
+                parallel_config.eplb_config.num_wait_worker_iterations
+            ),
             enable_async=parallel_config.eplb_config.enable_async,
-            expert_mapper_args=expert_mapper_args
+            expert_mapper_args=expert_mapper_args,
         )
 
     def __post_init__(self):
@@ -518,7 +523,13 @@ def step(
                 )
                 input_args = self.rebalance_task_args
 
-                self.expert_mapper_args.phyhsical_to_logical_map = self.physical_to_logical_map.cpu()
+                assert(
+                    self.expert_mapper_args is not None,
+                    "expert_mapper_args is not initialized",
+                )
+                self.expert_mapper_args.phyhsical_to_logical_map = (
+                    self.physical_to_logical_map.cpu()
+                )
                 expert_mapper_args = self.expert_mapper_args
 
                 self.rebalance_task(input_args, expert_mapper_args)
@@ -529,7 +540,7 @@ def step(
                 + self.num_wait_worker_iterations
                 + model.num_moe_layers
             ):
-                self.expert_rearrangement_step = 0 
+                self.expert_rearrangement_step = 0
 
     def rearrange(
         self,
@@ -747,9 +758,10 @@ def get_at_index(self, model, result, layer_id) -> list[Any]:
         size = len(result)
         # check if queue length matches the of layers
         if size != model.num_moe_layers:
-            logger.info(f"size={size}, num_moe_layers={model.num_moe_layers}")
+            logger.info("size=%s, num_moe_layers=%s", size, model.num_moe_layers)
             raise ValueError(
-                f"Queue length {size} does not match the number of moe layers in the model"
+                f"Queue length {size} does not match "
+                "the number of moe layers in the model"
             )
         if layer_id < 0 or layer_id >= size:
             raise ValueError(f"Index {layer_id} out of range for queue of size {size}")
@@ -884,6 +896,7 @@ def should_trigger_rebalance(self):
         return self.expert_rearrangement_step == (
             self.expert_rearrangement_step_interval - 1
         )
+
     def compute_and_set_moe_load(self):
         """
         Computes the MoE load across all ranks and sets it in the shared dictionary.
@@ -992,7 +1005,6 @@ def rebalance_task(self, input_args, expert_mapper_args):
             logger.error("Failed to submit rebalance task to async process")
         return None
 
-
     def __del__(self):
         """Clean up async process resources"""
         if self._async_processor:

vllm/v1/worker/gpu_model_runner.py

@@ -2410,9 +2410,8 @@ def execute_model(
     ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
         with record_function_or_nullcontext("Preprocess"):
             with self.synchronize_input_prep():
-                if self.parallel_config.eplb_config.enable_async:
-                    if self.eplb_state is not None:
-                        self.eplb_state.step_before_forward(self.get_model())
+                if self.parallel_config.eplb_config.enable_async and self.eplb_state:
+                    self.eplb_state.step_before_forward(self.get_model())
                 # Update persistent batch states.
                 self._update_states(scheduler_output)