replace mem_queue with shm

Author: WANDY666

lightllm/common/mem_manager.py

@@ -15,8 +15,11 @@
 from lightllm.utils.dist_utils import get_current_device_id
 from lightllm.utils.config_utils import get_num_key_value_heads
 from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io
+from lightllm.utils.shm_utils import create_or_link_shm
+from multiprocessing.reduction import ForkingPickler
 
 logger = init_logger(__name__)
+LIGHTLLM_MEM_MANAGER_SHM_SIZE = int(os.getenv("LIGHTLLM_MEM_MANAGER_SHM_SIZE", 1024 * 1024))
 
 
 class MemoryManager:
@@ -431,6 +434,23 @@ def copy_kv_from_other_dp_ranks(
                 rank_in_dp=rank_in_dp,
             )
 
+    def create_shm(self):
+        obj_bytes = ForkingPickler.dumps(self)
+        shm = create_or_link_shm(
+            f"{get_unique_server_name()}_mem_manager_{get_current_rank_in_node()}", LIGHTLLM_MEM_MANAGER_SHM_SIZE
+        )
+        logger.info(f"create shm {shm.name} size {shm.size} obj size {len(obj_bytes)}")
+        shm.buf[0:4] = len(obj_bytes).to_bytes(4, "little")
+        shm.buf[4 : 4 + len(obj_bytes)] = obj_bytes
+
+    @staticmethod
+    def from_shm(rank_in_node):
+        shm = create_or_link_shm(
+            f"{get_unique_server_name()}_mem_manager_{rank_in_node}", LIGHTLLM_MEM_MANAGER_SHM_SIZE
+        )
+        bytes_len = int.from_bytes(shm.buf[0:4], "little")
+        return ForkingPickler.loads(shm.buf[4 : 4 + bytes_len])
+
 
 class ReadOnlyStaticsMemoryManager:
     """

lightllm/server/router/manager.py

@@ -116,9 +116,6 @@ async def wait_to_model_ready(self):
         self.model_rpc_servers = []
         # 用于 kv move 管理进程 和 推理进程进行task信息的交互。
         self.info_queue: mp.Queue = mp.Queue()
-        self.mem_queues: List[torch.multiprocessing.Queue] = [
-            torch.multiprocessing.Queue() for _ in range(self.node_world_size)
-        ]
         self.rpc_event = multiprocessing.Event()
         self.rpc_finished_event = multiprocessing.Event()
 
@@ -137,9 +134,7 @@ async def wait_to_model_ready(self):
                     rpc_event=self.rpc_event,
                     rpc_finished_event=self.rpc_finished_event,
                     info_queue=self.info_queue,
-                    mem_queue=self.mem_queues[(rank_id % node_world_size)],
                     router_lock=self.router_lock,
-                    mem_queues=self.mem_queues,
                 )
             )
             tasks.append(task)
@@ -206,29 +201,29 @@ async def wait_to_model_ready(self):
                 start_prefill_kv_move_manager_process,
             )
 
-            start_prefill_kv_move_manager_process(self.args, self.info_queue, self.mem_queues)
+            start_prefill_kv_move_manager_process(self.args, self.info_queue)
 
         if self.args.run_mode == "nixl_prefill":
             from lightllm.server.router.model_infer.mode_backend.pd_nixl.prefill_node_impl import (
                 start_prefill_kv_move_manager_process,
             )
 
-            start_prefill_kv_move_manager_process(self.args, self.info_queue, self.mem_queues)
+            start_prefill_kv_move_manager_process(self.args, self.info_queue)
 
         if self.args.run_mode == "decode":
             # 启动 decode kv move 管理进程
             from lightllm.server.router.model_infer.mode_backend.continues_batch.pd_mode.decode_node_impl import (
                 start_decode_kv_move_manager_process,
             )
 
-            start_decode_kv_move_manager_process(self.args, self.info_queue, self.mem_queues)
+            start_decode_kv_move_manager_process(self.args, self.info_queue)
 
         if self.args.run_mode == "nixl_decode":
             from lightllm.server.router.model_infer.mode_backend.pd_nixl.decode_node_impl import (
                 start_decode_kv_move_manager_process,
             )
 
-            start_decode_kv_move_manager_process(self.args, self.info_queue, self.mem_queues)
+            start_decode_kv_move_manager_process(self.args, self.info_queue)
 
         return
 

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

@@ -19,10 +19,9 @@
 
 
 class DecodeNode(ChunkedPrefillBackend):
-    def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
+    def __init__(self, info_queue: mp.Queue) -> None:
         super().__init__()
         self.info_queue: mp.Queue = info_queue
-        self.mem_queue: mp.Queue = mem_queue
         self.classed_req_strict_prefill = False
 
     def init_custom(self):

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

@@ -9,8 +9,8 @@
 
 
 class DPForDecodeNode(DPChunkedPrefillBackend):
-    def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
-        super().__init__(mem_queue=mem_queue)
+    def __init__(self, info_queue: mp.Queue) -> None:
+        super().__init__()
         self.info_queue: mp.Queue = info_queue
         self.classed_req_strict_prefill = False
         return

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

@@ -166,9 +166,9 @@ def exposed_fail_to_realese_forzen_tokens(self, group_req_ids: List[int]):
         release_acquired_lock()
         return
 
-    def exposed_put_mem_manager_to_mem_queue(self):
-        self.backend.mem_queue.put(self.backend.model.mem_manager)
-        logger.info("put mem manager to info_queues ok")
+    def exposed_put_mem_manager_to_shm(self):
+        self.backend.model.mem_manager.create_shm()
+        logger.info("put mem manager to shm ok")
         return
 
     def exposed_unfrozen_time_out_reqs_tokens(self):

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

@@ -33,7 +33,7 @@
 
 
 class DecodeKVMoveManager(rpyc.Service):
-    def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
+    def __init__(self, args, info_queue: mp.Queue):
         super().__init__()
         self.args = args
         # args.dp // args.nnodes 在跨机tp的场景下，可能为0
@@ -44,7 +44,6 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         assert self.dp_world_size <= self.node_world_size
 
         self.info_queue = info_queue
-        self.mem_queues = mem_queues
         self.infer_rpyc_lock = threading.Lock()
         self.infer_rpyc_objs: List[PDDecodeInferRpcServer] = []
 
@@ -87,7 +86,7 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
     # _put_kv_received_to_radix_cache
     # _fail_to_realese_forzen_tokens
     # _unfrozen_time_out_reqs_tokens
-    # _put_mem_manager_to_mem_queue
+    # _put_mem_manager_to_shm
     # 上述接口都是 kv move manager 与推理进程进行交互的接口，主要用于申请锁定kv资源或者释放
     # kv资源的接口
     # ==================================================================================
@@ -155,10 +154,10 @@ def _unfrozen_time_out_reqs_tokens(self) -> None:
             asyncio.run(self.wait_all_future_finish(futures))
         return
 
-    def _put_mem_manager_to_mem_queue(self) -> None:
+    def _put_mem_manager_to_shm(self) -> None:
         with self.infer_rpyc_lock:
             for obj in self.infer_rpyc_objs:
-                obj.put_mem_manager_to_mem_queue()
+                obj.put_mem_manager_to_shm()
         return
 
     # ==================================================================================
@@ -362,14 +361,14 @@ def remove_trans_obj(self, connect_id):
         return
 
 
-def _init_env(args, info_queue: mp.Queue, mem_queues: List[mp.Queue], event: mp.Event):
+def _init_env(args, info_queue: mp.Queue, event: mp.Event):
     import lightllm.utils.rpyc_fix_utils as _
 
     # 注册graceful 退出的处理
     graceful_registry(inspect.currentframe().f_code.co_name)
     setproctitle.setproctitle(f"lightllm::{get_unique_server_name()}::decode_kv_move_manager")
 
-    manager = DecodeKVMoveManager(args, info_queue, mem_queues)
+    manager = DecodeKVMoveManager(args, info_queue)
     t = ThreadedServer(manager, port=args.pd_decode_rpyc_port, protocol_config={"allow_pickle": True})
     threading.Thread(target=lambda: t.start(), daemon=True).start()
 
@@ -381,9 +380,9 @@ def _init_env(args, info_queue: mp.Queue, mem_queues: List[mp.Queue], event: mp.
     return
 
 
-def start_decode_kv_move_manager_process(args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
+def start_decode_kv_move_manager_process(args, info_queue: mp.Queue):
     event = mp.Event()
-    proc = mp.Process(target=_init_env, args=(args, info_queue, mem_queues, event))
+    proc = mp.Process(target=_init_env, args=(args, info_queue, event))
     proc.start()
     event.wait()
     assert proc.is_alive()

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

@@ -279,10 +279,9 @@ def init_all(self, device_id: int, manager: "DecodeKVMoveManager"):
                 device_id,
                 self.task_in_queue,
                 self.task_out_queue,
-                manager.mem_queues,
             )
             assert self.task_out_queue.get(timeout=30) == "proc_start"
-            manager._put_mem_manager_to_mem_queue()
+            manager._put_mem_manager_to_shm()
             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

@@ -91,7 +91,7 @@ def async_connect():
         logger.warning(f"error while connect to prefill node: {e}")
 
 
-def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue, mem_queues: List[mp.Queue]):
+def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.Queue):
     import os
 
     # os.environ["NCCL_DEBUG"] = "INFO"
@@ -111,7 +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_managers: List[MemoryManager] = [mem_queue.get(timeout=60) for mem_queue in mem_queues]
+        # 从共享内存读取所有rank的mem_manager
+        node_world_size = args.tp // args.nnodes
+        mem_managers: List[MemoryManager] = [MemoryManager.from_shm(rank) for rank in range(node_world_size)]
 
         task_out_queue.put("get_mem_managers_ok")
         connect_id_to_comm: Dict[str, PyNcclCommunicator] = {}
@@ -143,9 +145,8 @@ def start_decode_trans_process(
     device_id: int,
     task_in_queue: mp.Queue,
     task_out_queue: mp.Queue,
-    mem_queues: List[mp.Queue],
 ):
-    proc = mp.Process(target=_init_env, args=(args, device_id, task_in_queue, task_out_queue, mem_queues))
+    proc = mp.Process(target=_init_env, args=(args, device_id, task_in_queue, task_out_queue))
     proc.start()
     assert proc.is_alive()
     logger.info(f"decode trans kv process for device: {device_id} start!")

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

@@ -20,11 +20,10 @@
 
 
 class ChunckedPrefillForPrefillNode(ChunkedPrefillBackend):
-    def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue) -> None:
+    def __init__(self, info_queue: mp.Queue) -> None:
         super().__init__()
         self.support_overlap = False
         self.info_queue: mp.Queue = info_queue
-        self.mem_queue: mp.Queue = mem_queue
         self.classed_req_no_decode = True
 
     def init_custom(self):

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

@@ -9,8 +9,8 @@
 
 
 class DPChunkedForPrefillNode(DPChunkedPrefillBackend):
-    def __init__(self, info_queue: mp.Queue, mem_queue: mp.Queue, mem_queues: List[mp.Queue]) -> None:
-        super().__init__(mem_queue=mem_queue, mem_queues=mem_queues)
+    def __init__(self, info_queue: mp.Queue) -> None:
+        super().__init__()
         self.support_overlap = False
         self.info_queue: mp.Queue = info_queue
         self.classed_req_no_decode = True