fix process manager

Author: wangzaijun

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

@@ -23,6 +23,8 @@
 
 logger = init_logger(__name__)
 
+thread_local_data = threading.local()
+
 
 @dataclass
 class TransProcessObj:
@@ -68,8 +70,8 @@ def check_trans_process(self):
     def __del__(self):
         # 强制关闭连接和杀掉传输进程
         if self.process is not None:
-            os.kill(self.process.pid, signal.SIGKILL)
             logger.warning(f"trans kv process {self.process.pid} is killed")
+            os.kill(self.process.pid, signal.SIGKILL)
         pass
 
 
@@ -104,8 +106,24 @@ async def wait_all_future_finish(self, futures: List[AsyncResult]):
         await asyncio.gather(*[asyncio.to_thread(future.wait) for future in futures])
         return
 
+    def on_connect(self, conn):
+        thread_local_data.prefill_node_id = None
+        pass
+
+    def on_disconnect(self, conn):
+        if thread_local_data.prefill_node_id is not None:
+            self.node_id_to_trans_obj.pop(thread_local_data.prefill_node_id, None)
+            logger.info(f"prefill node id {thread_local_data.prefill_node_id} disconnect")
+        pass
+
+    def exposed_check_alive(self):
+        # 用于 prefill node check 通信连接的状态。
+        return
+
     def exposed_build_trans_process(self, prefill_node_id, nccl_ip, nccl_port):
         prefill_node_id, nccl_ip, nccl_port = list(map(obtain, [prefill_node_id, nccl_ip, nccl_port]))
+        thread_local_data.prefill_node_id = prefill_node_id
+
         logger.info(f"build trans infos {prefill_node_id} {nccl_ip} {nccl_port}")
         if prefill_node_id in self.node_id_to_trans_obj:
             self.node_id_to_trans_obj.pop(prefill_node_id, None)
@@ -194,6 +212,8 @@ def handle_loop(self, task_queue: queue.Queue):
                         asyncio.run(self.wait_all_future_finish(futures))
                     logger.error(f"decode kv move task {task.to_decode_log_info()} has error, remove the trans_obj")
                     self.node_id_to_trans_obj.pop(task.prefill_node_id, None)
+                finally:
+                    # 去除引用否则进程无法自动退出
                     trans_obj = None
         except BaseException as e:
             logger.exception(str(e))

lightllm/server/router/model_infer/mode_backend/continues_batch/prefill_node_impl/prefill_kv_move_manager.py

@@ -2,6 +2,7 @@
 import rpyc
 import sys
 import os
+import gc
 import signal
 import copy
 import numpy as np
@@ -84,8 +85,8 @@ def check_trans_process(self):
     def __del__(self):
         # 强制关闭连接和杀掉传输进程
         if self.process is not None:
-            os.kill(self.process.pid, signal.SIGKILL)
             logger.warning(f"prefill trans process {self.process.pid} is killed")
+            os.kill(self.process.pid, signal.SIGKILL)
         pass
 
 
@@ -114,11 +115,28 @@ def get_next_device_index(self):
 
     def get_trans_obj(self, task: KVMoveTask):
         if task.decode_node.node_id not in self.node_id_to_trans_obj:
+            # 先遍历删除老的不能用的连接
+            self.remove_dead_trans_obj()
             trans_obj = TransProcessObj()
             trans_obj.create(task.decode_node.node_id, task.decode_node.ip, task.decode_node.rpyc_port, self)
             self.node_id_to_trans_obj[task.decode_node.node_id] = trans_obj
         return self.node_id_to_trans_obj[task.decode_node.node_id]
 
+    def remove_dead_trans_obj(self):
+        del_node_ids = []
+        for node_id, t_obj in self.node_id_to_trans_obj.items():
+            try:
+                t_obj.rpyc_conn.root.check_alive()
+            except BaseException as e:
+                logger.error(f"check error {str(e)}")
+                del_node_ids.append(node_id)
+        for node_id in del_node_ids:
+            self.node_id_to_trans_obj.pop(node_id)
+
+        if len(del_node_ids) != 0:
+            gc.collect()
+        return
+
     def handle_loop(self):
         try:
             while True:
@@ -154,9 +172,10 @@ def handle_loop(self):
                     logger.exception(str(e))
                     logger.error(f"kv move task {move_task.to_prefill_log_info()} has error, remove the trans_obj")
                     self.node_id_to_trans_obj.pop(move_task.decode_node.node_id, None)
-                    trans_obj = None
 
                 finally:
+                    # 去引用否则进程无法杀掉
+                    trans_obj = None
                     # 解除对prefill token的占用状态。
                     futures: List[AsyncResult] = []
                     for infer_rpyc in self.infer_rpyc_objs: