update format

Author: lt

vllm_ascend/eplb/eplb_updator.py

@@ -1,3 +1,4 @@
+#
 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,6 +23,7 @@
 from vllm_ascend.eplb.core.worker.eplb_worker import EplbProcess
 from vllm_ascend.eplb.core.loader.device_transfer_loader import D2DExpertWeightLoader
 
+
 class EplbUpdator:
 
     def __init__(self, expert_map_path):
@@ -42,7 +44,7 @@ def init_eplb(self, expert_map_path):
             if not envs.VLLM_ALLOW_EXPERT_LOAD_COLLECTING:
                 self.num_expert_load_gather = self.num_iterations
         except Exception as e:
-                self.num_expert_load_gather = self.num_iterations
+            self.num_expert_load_gather = self.num_iterations
 
         self.weight_update_counter = 0
         self.expert_map_initialized = False
@@ -72,19 +74,18 @@ def init_eplb(self, expert_map_path):
         })
 
         self.eplb = EplbProcess(
-            shared_dict = self.shared_dict,
-            planner_q = self.planner_block_queue,
-            block_update_q = self.block_update_queue,
-            redundant_enable = self.redundant_enable,
-            policy_type = 6,
-            enable_d2d = True
+            shared_dict=self.shared_dict,
+            planner_q=self.planner_block_queue,
+            block_update_q=self.block_update_queue,
+            redundant_enable=self.redundant_enable,
+            policy_type=6,
+            enable_d2d=True
         )
 
         self.eplb_process = self.eplb._launch_process()
 
         logger.info(f"[ModelRunner] Launched EPLB process (pid={self.eplb_process.pid})")
 
-
     def get_update_iteration(self):
         self.cur_iterations = self.cur_iterations + 1
         load_gather_iteration = self.cur_iterations % self.num_expert_load_gather == 0 if not self.gate_eplb else self.cur_iterations == self.num_iterations
@@ -94,7 +95,8 @@ def get_update_iteration(self):
     def get_init_expert_map(self):
         try:
             if not self.expert_map_initialized:
-                self.shared_dict["expert_maps"] = self.adaptor.get_init_expert_map_from_file(self.num_moe_layers, self.expert_map_path)
+                self.shared_dict["expert_maps"] = self.adaptor.get_init_expert_map_from_file(self.num_moe_layers,
+                                                                                             self.expert_map_path)
                 self.expert_map_initialized = True
         except Exception as e:
             logger.warning(f"[ModelRunner] Failed to wake EPLB process: {e}", exc_info=True)
@@ -103,6 +105,7 @@ def wakeup_eplb_worker(self):
         self.planner_block_queue.put(1)
 
     def forward_before(self):
+
         # Batch after eplb process being triggered, get update info provided by eplb process
         if self.update_in_flight and self.weight_update_counter == 0 and self.wait_worker_iterations == self.num_wait_worker_iterations:
             self.wait_worker_iterations = 0
@@ -111,14 +114,16 @@ def forward_before(self):
             self.weight_loading = True
 
         if self.update_in_flight and self.weight_loading and self.weight_update_counter < self.num_moe_layers:
-            (expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = self.update_info_all.pop(0)
+            (expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = self.update_info_all.pop(
+                0)
             rank_id = torch.distributed.get_rank()
             self.eplb_loader.set_log2phy_map(log2phy_map)
             expert_send_info_this_rank = expert_send_info[rank_id] if rank_id in expert_send_info else []
             expert_recv_info_this_rank = expert_recv_info[rank_id] if rank_id in expert_recv_info else []
-            #logger.info(f"check update info, layer = {layer_id}, send = {expert_send_info_this_rank}, recv = {expert_recv_info_this_rank}")
+            # logger.info(f"check update info, layer = {layer_id}, send = {expert_send_info_this_rank}, recv = {expert_recv_info_this_rank}")
             self.eplb_loader.generate_expert_d2d_transfer_task(expert_send_info_this_rank,
-                expert_recv_info_this_rank, updated_expert_map, layer_id + 3)
+                                                               expert_recv_info_this_rank, updated_expert_map,
+                                                               layer_id + 3)
             self.weight_update_counter += 1
             if self.weight_update_counter == self.num_moe_layers:
                 self.weight_update_counter = 0
@@ -129,8 +134,8 @@ def forward_before(self):
         self.reqs = []
         self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
 
-     def forward_end(self,dummy_run=False):
-        self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
+    def forward_end(self, dummy_run=False):
+        self.adaptor.get_rank_expert_workload(self.num_moe_layers, dummy_run)
         if not self.update_in_flight:
             load_gather_iteration, update_iteration = self.get_update_iteration()
             if load_gather_iteration:
@@ -146,8 +151,8 @@ def forward_end(self,dummy_run=False):
 
         self.eplb_loader.update_expert_map_and_weight(self.reqs, self.redundant_enable)
 
-    def compute_and_set_moe_load(self,dummy_run=False):
-        local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
+    def compute_and_set_moe_load(self, dummy_run=False):
+        local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers, dummy_run)
         self._gather_buffer = None
         if dist.is_initialized():
             self.world_size = dist.get_world_size()
@@ -192,7 +197,7 @@ def warm_up_eplb(self):
                 continue
             comm_op_list.append(
                 dist.P2POp(dist.irecv, src_tensor, src_rank)
-        )
+            )
         if comm_op_list:
             reqs = dist.batch_isend_irecv(comm_op_list)
 
@@ -205,7 +210,7 @@ def unpack_update_batch(self, packed_update_info):
         """
         send_all, recv_all, stacked_maps, stacked_log2phy, layer_id_tensor = packed_update_info
 
-        maps     = stacked_maps.unbind(0)
+        maps = stacked_maps.unbind(0)
         layer_ids = layer_id_tensor.tolist()
 
         if self.redundant_enable:
@@ -217,7 +222,7 @@ def unpack_update_batch(self, packed_update_info):
         _send = send_all
         _recv = recv_all
         _maps = maps
-        _l2p  = log2phy_list
+        _l2p = log2phy_list
         _lids = layer_ids
 
         recovered = [
@@ -249,7 +254,6 @@ def update_expert_load_statistical_period(self, num_expert_load_gather: int, num
         self.num_iterations = num_iterations
         logger.info(f" update {self.num_expert_load_gather=}, {self.num_iterations} success...")
 
-
     def shutdown(self):
         """
         Clean up the EPLB process.