[Feature] support eplb in api_server (#4782)

* support eplb in api_server

* update code

* add eplb test case

* update eplb

* support tp+dp eplb

* update test cese

* update code

* update code

* fix bug

* update copilot review

* update test case name

fastdeploy/eplb/experts_manager.py

@@ -1,19 +1,33 @@
 """
-redundant expert manger
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 """
 
 import threading
 import time
 from http import HTTPStatus
-from multiprocessing import Pipe, Process, shared_memory
+from multiprocessing import Pipe, Process
 
 import numpy as np
 import requests
 
+from fastdeploy.config import FDConfig
 from fastdeploy.eplb.async_expert_loader import load_model_weights_process
 from fastdeploy.eplb.eplb import rebalance_experts
-from fastdeploy.eplb.utils import RearrangeExpertState, RedundantExpertWorkload
-from fastdeploy.utils import envs, get_logger
+from fastdeploy.eplb.utils import RedundantExpertWorkload
+from fastdeploy.inter_communicator import IPCSignal, RearrangeExpertStatus
+from fastdeploy.utils import get_logger
 
 
 class RedundantExpertManager:
@@ -21,7 +35,13 @@ class RedundantExpertManager:
     RedundantExpertManger
     """
 
-    def __init__(self, rank=0, ep_size=64, fd_config=None):
+    def __init__(
+        self,
+        rank: int = 0,
+        ep_size: int = 32,
+        fd_config: FDConfig = None,
+        ipc_signal_suffix: int = 0,
+    ):
         self.logger = get_logger("eplb_expert_manager", "eplb_{0}.log".format(rank))
 
         self.rank = rank
@@ -30,9 +50,11 @@ class RedundantExpertManager:
         self.eplb_config = fd_config.eplb_config
         self.api_user = self.eplb_config.redundant_expert_api_user
         self.api_passwd = self.eplb_config.redundant_expert_api_password
-        self.num_hidden_layers = self.eplb_config.model_config.num_layers
-        self.num_logical_experts = self.eplb_config.model_config.moe_num_experts
         self.num_redundant_experts = self.eplb_config.redundant_experts_num
+        self.num_hidden_layers = self.fd_config.model_config.num_hidden_layers
+        self.num_logical_experts = self.fd_config.model_config.moe_num_experts
+        self.ipc_signal_suffix = ipc_signal_suffix
+        self.local_rank = self.rank % self.fd_config.parallel_config.tensor_parallel_size
 
         self.num_replicas = self.num_logical_experts + self.num_redundant_experts
         self.num_groups = self.num_logical_experts
@@ -112,9 +134,12 @@ class RedundantExpertManager:
             name=f"eplb::async_load_model_{rank}",
             args=(
                 self.rank,
+                self.fd_config.model_config.model,
                 self.expert_per_rank,
                 self.fd_config.model_config.moe_layer_start_index,
                 self.eplb_config.moe_quant_type,
+                self.ipc_signal_suffix,
+                self.eplb_config,
                 child_data_conn,
                 child_mg_conn,
             ),
@@ -130,9 +155,6 @@ class RedundantExpertManager:
             strategy {self.eplb_config.redundant_expert_eplb_strategy}"
         )
 
-    def get_unique_name(self, name):
-        return f"{envs.get_unique_name(name + '_dprank_' + str(self.rank))}"
-
     def get_ep_rank_to_expert_id_list(self):
         """
         get_ep_rank_to_expert_id_list
@@ -147,66 +169,84 @@ class RedundantExpertManager:
         """
         listen_rearrange_expert_signal
         """
-        if self.rank == 0:
-            rearrange_experts_ips_size = np.zeros([1], dtype=np.int32)
-            shm_rearrange_experts_ips_size = shared_memory.SharedMemory(
+        dp_ipc_signal_suffix = f"{self.ipc_signal_suffix}_dp{self.fd_config.parallel_config.local_data_parallel_id}"
+        if self.local_rank == 0:
+            rearrange_experts_ips_size_array = np.zeros([1], dtype=np.int32)
+            rearrange_experts_ips_size_signal = IPCSignal(
+                name="rearrange_experts_ips_size",
+                array=rearrange_experts_ips_size_array,
+                dtype=np.int32,
+                suffix=dp_ipc_signal_suffix,
                 create=False,
-                size=rearrange_experts_ips_size.nbytes,
-                name=self.get_unique_name("rearrange_experts_ips_size"),
             )
-            rearrange_experts_ips_size_array = np.ndarray(
-                rearrange_experts_ips_size.shape,
-                dtype=rearrange_experts_ips_size.dtype,
-                buffer=shm_rearrange_experts_ips_size.buf,
-            )
-            shm_rearrange_experts_ips_list = shared_memory.SharedMemory(
+
+            shm_rearrange_experts_ips_list = IPCSignal(
+                name="rearrange_experts_ips_list",
+                shm_size=self.eplb_config.redundant_expert_ip_shm_size,
+                suffix=dp_ipc_signal_suffix,
                 create=False,
-                size=1024,
-                name=self.get_unique_name("rearrange_experts_ips_list"),
             )
 
             rearrange_experts_status = np.zeros([1], dtype=np.int32)
-            shm_rearrange_experts_status = shared_memory.SharedMemory(
+            rearrange_experts_signal = IPCSignal(
+                name="rearrange_experts_status",
+                array=rearrange_experts_status,
+                dtype=np.int32,
+                suffix=dp_ipc_signal_suffix,
                 create=False,
-                size=rearrange_experts_status.nbytes,
-                name=self.get_unique_name("rearrange_experts_status"),
-            )
-            rearrange_experts_status_array = np.ndarray(
-                rearrange_experts_status.shape,
-                dtype=rearrange_experts_status.dtype,
-                buffer=shm_rearrange_experts_status.buf,
             )
 
+            signal_update_weight_from_tensor = np.zeros([1], dtype=np.int32)
+            self.signal_update_weight_from_tensor_array = IPCSignal(
+                name="signal_update_weight_from_tensor",
+                array=signal_update_weight_from_tensor,
+                dtype=np.int32,
+                suffix=dp_ipc_signal_suffix,
+                create=False,
+            )
+
+        tp_ipc_signal_suffix = f"{dp_ipc_signal_suffix}_tp{self.local_rank}"
         signal_update_weight_from_disk = np.zeros([1], dtype=np.int32)
-        shm_signal_update_weight_from_disk = shared_memory.SharedMemory(
+        signal_update_weight_from_disk_array = IPCSignal(
+            name="signal_update_weight_from_disk",
+            array=signal_update_weight_from_disk,
+            dtype=np.int32,
+            suffix=tp_ipc_signal_suffix,
             create=False,
-            size=signal_update_weight_from_disk.nbytes,
-            name=self.get_unique_name("signal_update_weight_from_disk"),
-        )
-        signal_update_weight_from_disk_array = np.ndarray(
-            signal_update_weight_from_disk.shape,
-            dtype=signal_update_weight_from_disk.dtype,
-            buffer=shm_signal_update_weight_from_disk.buf,
         )
 
-        experts_token_stats = np.zeros((self.num_hidden_layers, 64), dtype=np.int32)
-        shm_all_experts_token_stats = shared_memory.SharedMemory(
+        experts_token_stats = np.zeros(
+            (self.fd_config.model_config.num_hidden_layers, self.fd_config.model_config.moe_num_experts),
+            dtype=np.int32,
+        )
+        shm_all_experts_token_stats = IPCSignal(
+            name="all_experts_token_stats",
+            array=experts_token_stats,
+            dtype=np.int32,
+            suffix=tp_ipc_signal_suffix,
+            create=False,
+        )
+
+        result_update_weight_from_disk = np.zeros([1], dtype=np.int32)
+        self.update_weight_from_disk_result = IPCSignal(
+            name="result_update_weight_from_disk",
+            array=result_update_weight_from_disk,
+            dtype=np.int32,
+            suffix=tp_ipc_signal_suffix,
             create=False,
-            size=experts_token_stats.nbytes,
-            name=self.get_unique_name("all_experts_token_stats"),
         )
 
         while True:
-            if self.rank == 0:
+            if self.local_rank == 0:
                 now = int(time.time())
-                if rearrange_experts_ips_size_array[0] > 0:
+                if rearrange_experts_ips_size_signal.value[0] > 0:
                     # step 1. all reduce experts token stats
-                    address = bytes(shm_rearrange_experts_ips_list.buf[: rearrange_experts_ips_size_array[0]]).decode(
-                        "utf-8"
-                    )
+                    address = bytes(
+                        shm_rearrange_experts_ips_list.shm.buf[: rearrange_experts_ips_size_signal.value[0]]
+                    ).decode("utf-8")
                     self.logger.info(f"redundant_expert: all rank ips {address}")
-                    rearrange_experts_ips_size_array[0] = 0
-                    rearrange_experts_status_array[0] = RearrangeExpertState.doing.value
+                    rearrange_experts_ips_size_signal.value[0] = 0
+                    rearrange_experts_signal.value[0] = RearrangeExpertStatus.DOING.value
 
                     self.dp_rank_address = address.strip().split(";")
                     if self.allreduce_experts_stat():
@@ -214,30 +254,25 @@ class RedundantExpertManager:
                         self.load_weight_begin_ts = now
                         self.logger.info("redundant_expert: all-reduce experts stats success")
                     else:
-                        rearrange_experts_status_array[0] = RearrangeExpertState.free.value
+                        rearrange_experts_signal.value[0] = RearrangeExpertStatus.FREE.value
                         self.logger.warning("redundant_expert: all-reduce experts stats fail")
                 elif self.need_allgather_load_weight_result and self.allreduce_load_weight_result():
                     # step 3. all reduce the result of load weight from disk
                     self.need_allgather_load_weight_result = False
-                    rearrange_experts_status_array[0] = RearrangeExpertState.load_succ.value
+                    rearrange_experts_signal.value[0] = RearrangeExpertStatus.LOAD_SUCC.value
                     self.rearrange_end_ts = now
-                if rearrange_experts_status_array[0] > 1 and (
+                if rearrange_experts_signal.value[0] > 1 and (
                     now - self.rearrange_end_ts > self.rearrange_reset_interval
                 ):
                     # reset rearrange status
-                    rearrange_experts_status_array[0] = RearrangeExpertState.free.value
+                    rearrange_experts_signal.value[0] = RearrangeExpertStatus.FREE.value
 
-            if signal_update_weight_from_disk_array[0] == 1:
+            if signal_update_weight_from_disk_array.value[0] == 1:
                 # step 2. async load weight: disk -> memory
-                expert_token_stats = np.ndarray(
-                    experts_token_stats.shape,
-                    dtype=experts_token_stats.dtype,
-                    buffer=shm_all_experts_token_stats.buf,
-                )
-                self.model_tokens_per_expert_stats_list[:] = expert_token_stats[:]
+                self.model_tokens_per_expert_stats_list[:] = shm_all_experts_token_stats.value[:]
                 self.caculate_expert_rank_table()
                 self.update_weight_from_disk()
-                signal_update_weight_from_disk_array[0] = 0
+                signal_update_weight_from_disk_array.value[0] = 0
             time.sleep(0.5)
 
     def caculate_expert_rank_table(self, is_init=False):
@@ -274,7 +309,7 @@ class RedundantExpertManager:
         self.model_expert_id_to_ep_rank_array[..., : logical_to_physical_map.shape[-1]] = logical_to_physical_map[:]
         self.model_expert_in_rank_num_list[:] = expert_count[:]
 
-        if self.rank == 0:
+        if self.local_rank == 0:
             workload = RedundantExpertWorkload()
             workload.tokens_per_expert_stats_list = self.model_tokens_per_expert_stats_list.tolist()
             workload.ep_rank_to_expert_id_list = rank_expert_list.tolist()
@@ -287,18 +322,7 @@ class RedundantExpertManager:
         update_weight_from_disk
         """
         begin_time = time.time()
-        result_update_weight_from_disk = np.zeros([1], dtype=np.int32)
-        shm_result_update_weight_from_disk = shared_memory.SharedMemory(
-            create=False,
-            size=result_update_weight_from_disk.nbytes,
-            name=self.get_unique_name("result_update_weight_from_disk"),
-        )
-        result_update_weight_from_disk_array = np.ndarray(
-            result_update_weight_from_disk.shape,
-            dtype=result_update_weight_from_disk.dtype,
-            buffer=shm_result_update_weight_from_disk.buf,
-        )
-        result_update_weight_from_disk_array[0] = 0
+        self.update_weight_from_disk_result.value[0] = 0
 
         self.logger.info(f"redundant_expert: update_weight_from_disk send to async process, rank {self.rank}")
         self.parent_mg_conn.send(
@@ -312,7 +336,7 @@ class RedundantExpertManager:
         self.tensor_infos = response["weights"]
 
         # 更新权重加载结果
-        result_update_weight_from_disk_array[0] = 1 if response["result"] else -1
+        self.update_weight_from_disk_result.value[0] = 1 if response["result"] else -1
         self.logger.info(
             "redundant_expert: update_weight_from_disk end, rank"
             + f" {self.rank} {response['result']}, cost {int(time.time() - begin_time)}s"
@@ -330,8 +354,8 @@ class RedundantExpertManager:
         """
         allgather_expert_token_stats
         """
-        success_count = 0
         expert_token_stats = np.zeros((self.num_hidden_layers, self.num_logical_experts), dtype=np.int32)
+        success_count = 0
         for addr in self.dp_rank_address:
             try:
                 # TODO: 请求失败重试
@@ -347,8 +371,10 @@ class RedundantExpertManager:
                         + f"addr {addr}, res {res.status_code} {res.json()}"
                     )
                     break
+
+                for meta_data in res.json()["data"]:
+                    expert_token_stats += np.array(meta_data, dtype=np.int32)
                 success_count += 1
-                expert_token_stats += np.array(res.json()["data"], dtype=np.int32)
             except Exception as e:
                 self.logger.error(f"redundant_expert: allgather_expert_token_stats fail. addr {addr}, error {e}")
         if success_count == len(self.dp_rank_address):
@@ -426,18 +452,7 @@ class RedundantExpertManager:
                 or not self.eplb_config.redundant_expert_enable_schedule_cordon
             ):
                 self.logger.info("redundant_expert: allreduce_load_weight_result success, notify infer.py")
-                signal_update_weight_from_tensor = np.zeros([1], dtype=np.int32)
-                shm_signal_update_weight_from_tensor = shared_memory.SharedMemory(
-                    create=False,
-                    size=signal_update_weight_from_tensor.nbytes,
-                    name=self.get_unique_name("signal_update_weight_from_tensor"),
-                )
-                signal_update_weight_from_tensor_array = np.ndarray(
-                    signal_update_weight_from_tensor.shape,
-                    dtype=signal_update_weight_from_tensor.dtype,
-                    buffer=shm_signal_update_weight_from_tensor.buf,
-                )
-                signal_update_weight_from_tensor_array[0] = 1
+                self.signal_update_weight_from_tensor_array.value[0] = 1
         return True
 
     def allgather_load_weight_result(self):
@@ -465,140 +480,28 @@ class RedundantExpertManager:
                         + f"addr {addr}, res {res.status_code} {res.json()}"
                     )
                     break
-                result = res.json()["data"]
+                result_list = res.json()["data"]
                 self.logger.info(
-                    f"redundant_expert: allgather_load_weight_result success. addr {addr}, result {result}"
+                    f"redundant_expert: allgather_load_weight_result success. addr {addr}, result_list {result_list}"
                 )
-                if result == 1:
-                    success_count += 1
-                elif result == -1:
-                    fail_count += 1
-                    self.logger.error(
-                        f"redundant_expert: allgather_load_weight_result fail. addr {addr}, result {result}"
-                    )
-                    exist_fail = True
+                for result in result_list:
+                    if result == 1:
+                        success_count += 1
+                    elif result == -1:
+                        fail_count += 1
+                        self.logger.error(
+                            f"redundant_expert: allgather_load_weight_result fail. addr {addr}, result {result}"
+                        )
+                        exist_fail = True
             except Exception as e:
                 self.logger.error(f"redundant_expert: allgather_load_weight_result error. addr {addr}, error {e}")
-        if success_count == len(self.dp_rank_address):
-            self.logger.info("redundant_expert: allgather_load_weight_result all success")
-            all_success = True
-        else:
+
+        if fail_count > 0:
             self.logger.info(
                 "redundant_expert: allgather_load_weight_result not all ready, "
                 + f"succ {success_count} fail {fail_count} total {len(self.dp_rank_address)}"
             )
+        else:
+            self.logger.info("redundant_expert: allgather_load_weight_result all success")
+            all_success = True
         return all_success, exist_fail
-
-
-def init_shared_memory_for_eplb_rank0(rank):
-    rearrange_experts_ips_size = np.zeros([1], dtype=np.int32)
-    shm_rearrange_experts_ips_size = shared_memory.SharedMemory(
-        create=True,
-        size=rearrange_experts_ips_size.nbytes,
-        name=f"{envs.get_unique_name('rearrange_experts_ips_size_dprank' + rank)}",
-    )
-    rearrange_experts_ips_size_array = np.ndarray(
-        rearrange_experts_ips_size.shape,
-        dtype=rearrange_experts_ips_size.dtype,
-        buffer=shm_rearrange_experts_ips_size.buf,
-    )
-    shm_rearrange_experts_ips_list = shared_memory.SharedMemory(
-        create=True,
-        size=envs.FD_REDUNDANT_EXPERT_IP_SHM_SIZE,
-        name=f"{envs.get_unique_name('rearrange_experts_ips_list_dprank' + rank)}",
-    )
-    # 记录专家重排状态
-    rearrange_experts_status = np.zeros([1], dtype=np.int32)
-    shm_rearrange_experts_status = shared_memory.SharedMemory(
-        create=True,
-        size=rearrange_experts_status.nbytes,
-        name=f"{envs.get_unique_name('rearrange_experts_status_dprank' + rank)}",
-    )
-    rearrange_experts_status_array = np.ndarray(
-        rearrange_experts_status.shape, dtype=rearrange_experts_status.dtype, buffer=shm_rearrange_experts_status.buf
-    )
-    # 接收更新权重的信号
-    signal_update_weight_from_tensor = np.zeros([1], dtype=np.int32)
-    shm_signal_update_weight_from_tensor = shared_memory.SharedMemory(
-        create=True,
-        size=signal_update_weight_from_tensor.nbytes,
-        name=f"{envs.get_unique_name('signal_update_weight_from_tensor_dprank' + rank) }",
-    )
-    signal_update_weight_from_tensor_array = np.ndarray(
-        signal_update_weight_from_tensor.shape,
-        dtype=signal_update_weight_from_tensor.dtype,
-        buffer=shm_signal_update_weight_from_tensor.buf,
-    )
-    return (
-        rearrange_experts_ips_size_array,
-        shm_rearrange_experts_ips_list,
-        rearrange_experts_status_array,
-        signal_update_weight_from_tensor_array,
-    )
-
-
-def init_shared_memory_for_eplb_each_rank(fd_config, rank):
-    # 记录专家负载
-    num_layers = fd_config.model_config.num_hidden_layers
-    num_experts = fd_config.model_config.moe_num_experts
-    experts_token_stats = np.zeros((num_layers, num_experts), dtype=np.int32)
-    shm_local_experts_token_stats = shared_memory.SharedMemory(
-        create=True,
-        size=experts_token_stats.nbytes,
-        name=f"{envs.get_unique_name('local_experts_token_stats_dprank' + rank)}",
-    )
-    local_experts_token_stats_array = np.ndarray(
-        experts_token_stats.shape, dtype=experts_token_stats.dtype, buffer=shm_local_experts_token_stats.buf
-    )
-    # TODO: 全局专家负载状态是一样的，节点上的所有DP可以共用一份，但需要避免多个DP同时更新
-    shm_all_experts_token_stats = shared_memory.SharedMemory(
-        create=True,
-        size=experts_token_stats.nbytes,
-        name=f"{envs.get_unique_name('all_experts_token_stats_dprank' + rank)}",
-    )
-    expert_tokens_stats_array = np.ndarray(
-        experts_token_stats.shape, dtype=experts_token_stats.dtype, buffer=shm_all_experts_token_stats.buf
-    )
-    # 接收加载权重的信号
-    signal_update_weight_from_disk = np.zeros([1], dtype=np.int32)
-    shm_signal_update_weight_from_disk = shared_memory.SharedMemory(
-        create=True,
-        size=signal_update_weight_from_disk.nbytes,
-        name=f"{envs.get_unique_name('signal_update_weight_from_disk_dprank' + rank)}",
-    )
-    signal_update_weight_from_disk_array = np.ndarray(
-        signal_update_weight_from_disk.shape,
-        dtype=signal_update_weight_from_disk.dtype,
-        buffer=shm_signal_update_weight_from_disk.buf,
-    )
-    # 记录加载权重的结果
-    result_update_weight_from_disk = np.zeros([1], dtype=np.int32)
-    shm_result_update_weight_from_disk = shared_memory.SharedMemory(
-        create=True,
-        size=result_update_weight_from_disk.nbytes,
-        name=f"{envs.get_unique_name('result_update_weight_from_disk_dprank' + rank)}",
-    )
-    result_update_weight_from_disk_array = np.ndarray(
-        result_update_weight_from_disk.shape,
-        dtype=result_update_weight_from_disk.dtype,
-        buffer=shm_result_update_weight_from_disk.buf,
-    )
-    # 接收清零专家负载的信号
-    signal_clear_experts_token_stats = np.zeros([1], dtype=np.int32)
-    shm_signal_clear_experts_token_stats = shared_memory.SharedMemory(
-        create=True,
-        size=signal_clear_experts_token_stats.nbytes,
-        name=f"{envs.get_unique_name('signal_clear_experts_token_stats_dprank' + rank)}",
-    )
-    signal_clear_experts_token_stats_array = np.ndarray(
-        signal_clear_experts_token_stats.shape,
-        dtype=signal_clear_experts_token_stats.dtype,
-        buffer=shm_signal_clear_experts_token_stats.buf,
-    )
-    return (
-        local_experts_token_stats_array,
-        expert_tokens_stats_array,
-        signal_update_weight_from_disk_array,
-        result_update_weight_from_disk_array,
-        signal_clear_experts_token_stats_array,
-    )