fix pd mem_manager get failed

Author: WANDY666

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

@@ -281,8 +281,11 @@ def init_all(self, device_id: int, manager: "DecodeKVMoveManager"):
                 self.task_out_queue,
             )
             assert self.task_out_queue.get(timeout=30) == "proc_start"
+            # 确保在子进程读取共享内存之前，主进程已经将 mem_manager 写入共享内存
             if self.device_id == 0:
                 manager._put_mem_manager_to_shm()
+            # 通知子进程可以从共享内存读取 mem_manager
+            self.task_in_queue.put("mem_managers_ready")
             assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             return True

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

@@ -111,6 +111,9 @@ def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.
         graceful_registry(inspect.currentframe().f_code.co_name)
         task_out_queue.put("proc_start")
 
+        # 等待主进程将 mem_manager 写入共享内存后的信号
+        assert task_in_queue.get(timeout=60) == "mem_managers_ready"
+
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [MemoryManager.from_shm(rank, device_id) for rank in range(node_world_size)]
@@ -136,7 +139,7 @@ def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.
                 logger.warning(f"unexpected task type: {task}")
 
     except Exception as e:
-        logger.error(f"Fatal error happened in kv trans process: {e}")
+        logger.error(f"Fatal error happened in kv trans process: {e} in device {device_id}")
         raise
 
 

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

@@ -357,6 +357,7 @@ def init_all(self, device_id: int, manager: "PrefillKVMoveManager"):
             assert self.task_out_queue.get(timeout=30) == "proc_start"
             if self.device_id == 0:
                 manager._put_mem_manager_to_shm()
+            self.task_in_queue.put("mem_managers_ready")
             assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             return True

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

@@ -116,6 +116,9 @@ def _init_env(
         )
         task_out_queue.put("proc_start")
 
+        # 等待主进程将 mem_manager 写入共享内存后的信号
+        assert task_in_queue.get(timeout=60) == "mem_managers_ready"
+
         # 从共享内存读取所有rank的mem_manager
         node_world_size = args.tp // args.nnodes
         mem_managers: List[MemoryManager] = [MemoryManager.from_shm(rank, device_id) for rank in range(node_world_size)]