fix

Author: wangzaijun

lightllm/common/kv_cache_mem_manager/mem_manager.py

@@ -19,6 +19,7 @@
 from lightllm.utils.device_utils import kv_trans_use_p2p
 from lightllm.utils.shm_utils import create_or_link_shm
 from multiprocessing.reduction import ForkingPickler
+from filelock import FileLock
 
 logger = init_logger(__name__)
 
@@ -450,25 +451,39 @@ def write_to_shm(self, req_manager):
         # 避免过多无用的数据复制和传输开销。
         self.req_to_token_indexs: torch.Tensor = req_manager.req_to_token_indexs
 
-        shm_name = f"{get_unique_server_name()}_mem_manager_{get_current_rank_in_node()}"
-        for rank_in_node in range(0, get_node_world_size() * 2):
-            obj_bytes = ForkingPickler.dumps(self).tobytes()
+        lock = FileLock(f"/tmp/{get_unique_server_name()}_mem_manager_lock")
+        with lock:
+            node_world_size = get_node_world_size()
+            shm_name = f"{get_unique_server_name()}_mem_manager_{get_current_rank_in_node()}"
+            obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)]
+            obj_size = len(obj_bytes_array[0])
             shm = create_or_link_shm(
-                name=f"{shm_name}_{rank_in_node}", expected_size=len(obj_bytes) + 4, force_mode="create"
+                name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create"
             )
             logger.info(f"create shm {shm.name} size {shm.size} for mem manger shared buffer")
-            shm.buf[0:4] = len(obj_bytes).to_bytes(4, "little")
-            shm.buf[4 : 4 + len(obj_bytes)] = obj_bytes
+            shm.buf[0:4] = (node_world_size * 2).to_bytes(4, "little")
+            shm.buf[4:8] = obj_size.to_bytes(4, "little")
+            start_index = 8
+            for obj_bytes in obj_bytes_array:
+                shm.buf[start_index : start_index + obj_size] = obj_bytes
+                start_index += obj_size
 
     @staticmethod
-    def loads_from_shm(rank_in_node: int, current_rank_in_node: int) -> "MemoryManager":
-        shm_name = f"{get_unique_server_name()}_mem_manager_{rank_in_node}_{current_rank_in_node}"
+    def loads_from_shm(rank_in_node: int) -> "MemoryManager":
+        shm_name = f"{get_unique_server_name()}_mem_manager_{rank_in_node}"
+        lock = FileLock(f"/tmp/{get_unique_server_name()}_mem_manager_lock")
         logger.info(f"get memmanager from shm {shm_name}")
-        shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link")
-        bytes_len = int.from_bytes(shm.buf[0:4], "little")
-        obj_bytes = shm.buf[4 : 4 + bytes_len].tobytes()
-        shm.close()
-        return ForkingPickler.loads(obj_bytes)
+        with lock:
+            shm = create_or_link_shm(name=shm_name, expected_size=-1, force_mode="link")
+            left_num = int.from_bytes(shm.buf[0:4], "little")
+            obj_size = int.from_bytes(shm.buf[4:8], "little")
+            assert left_num > 0
+            end_index = 8 + left_num * obj_size
+            start_index = 8 + (left_num - 1) * obj_size
+            obj_bytes = shm.buf[start_index:end_index].tobytes()
+            shm.buf[0:4] = (left_num - 1).to_bytes(4, byteorder="little")
+            shm.close()
+            return ForkingPickler.loads(obj_bytes)
 
 
 class ReadOnlyStaticsMemoryManager:

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -245,7 +245,7 @@ def init_custom(self):
 
     def init_dp_kv_shared(self):
         from lightllm.server.router.model_infer.mode_backend.dp_backend.dp_shared_kv_trans import DPKVSharedMoudle
-        from lightllm.common.mem_manager import MemoryManager
+        from lightllm.common.kv_cache_mem_manager import MemoryManager
 
         torch.cuda.set_device(get_current_device_id())
 
@@ -260,7 +260,7 @@ def init_dp_kv_shared(self):
         self.mem_managers = []
         for rank_idx in range(self.node_world_size):
             if rank_idx != self.rank_in_node:
-                self.mem_managers.append(MemoryManager.loads_from_shm(rank_idx, self.rank_in_node))
+                self.mem_managers.append(MemoryManager.loads_from_shm(rank_idx))
             else:
                 self.mem_managers.append(self.model.mem_manager)
         return

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_trans_process.py

@@ -114,8 +114,7 @@ def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [
-            MemoryManager.loads_from_shm(rank_in_node=rank, current_rank_in_node=device_id + node_world_size)
-            for rank in range(node_world_size)
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
         ]
 
         task_out_queue.put("get_mem_managers_ok")

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/prefill_node_impl/prefill_trans_process.py

@@ -119,8 +119,7 @@ def _init_env(
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [
-            MemoryManager.loads_from_shm(rank_in_node=rank, current_rank_in_node=device_id + node_world_size)
-            for rank in range(node_world_size)
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
         ]
         task_out_queue.put("get_mem_managers_ok")
         connect_id_to_comm: Dict[str, PyNcclCommunicator] = {}

lightllm/server/router/model_infer/mode_backend/pd_nixl/decode_node_impl/decode_trans_process.py

@@ -58,8 +58,7 @@ def _init_env(
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [
-            MemoryManager.loads_from_shm(rank_in_node=rank, current_rank_in_node=device_id + node_world_size)
-            for rank in range(node_world_size)
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
         ]
 
         task_out_queue.put("get_mem_managers_ok")

lightllm/server/router/model_infer/mode_backend/pd_nixl/prefill_node_impl/prefill_trans_process.py

@@ -50,8 +50,7 @@ def _init_env(
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [
-            MemoryManager.loads_from_shm(rank_in_node=rank, current_rank_in_node=device_id + node_world_size)
-            for rank in range(node_world_size)
+            MemoryManager.loads_from_shm(rank_in_node=rank) for rank in range(node_world_size)
         ]
 
         task_out_queue.put("get_mem_managers_ok")