format

Author: hiworldwzj

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

@@ -30,6 +30,7 @@
 
 KV_MOVE_MAX_NUM = 16
 
+
 class DecodeKVMoveManager(rpyc.Service):
     def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         super().__init__()
@@ -45,7 +46,7 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         self.mem_queues = mem_queues
         self.infer_rpyc_lock = threading.Lock()
         self.infer_rpyc_objs: List[PDDecodeInferRpcServer] = []
-        
+
         from .decode_trans_obj import KVTransConnectObj
 
         self.connect_id_to_trans_obj: Dict[str, KVTransConnectObj] = {}
@@ -70,16 +71,16 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
 
         # 在不使用p2p 复制kv 的方案时，需要全局的传输锁进行控制。这个时候kv传输的效率会下降。
         self.kv_trans_lock = threading.Lock()
-        
+
         from .decode_trans_obj import KVTransProcess
-        
+
         self.kv_trans_processes: List[KVTransProcess] = [None] * self.node_world_size
         for device_id in range(self.node_world_size):
             self.kv_trans_processes[device_id] = KVTransProcess()
             assert self.kv_trans_processes[device_id].init_all(device_id, self)
 
         return
-    
+
     # ==================================================================================
     # _dp_alloc_to_frozen_some_tokens
     # _put_kv_received_to_radix_cache
@@ -158,13 +159,13 @@ def _put_mem_manager_to_mem_queue(self) -> None:
             for obj in self.infer_rpyc_objs:
                 obj.put_mem_manager_to_mem_queue()
         return
-    
+
     # ==================================================================================
     # put_to_fail_release_task_queue 将因为一些原因失败，需要释放锁定的kv资源的请求放入到
     # 对应的处理队列中，handle_fail_release_task_loop 是一个循环的线程，专门处理这些失败的请求
     # 通过调用与推理进程交互的接口，释放掉申请锁定的 kv 资源。
     # ==================================================================================
-    
+
     def put_to_fail_release_task_queue(self, task: Union[KVMoveTask, List[KVMoveTask]]):
         if isinstance(task, KVMoveTask):
             self.fail_to_release_queue.put(task)
@@ -182,9 +183,9 @@ def handle_fail_release_task_loop(self):
             else:
                 self._fail_to_realese_forzen_tokens(handle_list)
         return
-    
+
     # ==================================================================================
-    # on_connect 
+    # on_connect
     # on_disconnect
     # exposed_check_alive
     # exposed_build_trans_process
@@ -278,12 +279,14 @@ def exposed_request_data_transfer(self, tasks: List[KVMoveTask]) -> List[Optiona
             self.remove_trans_obj(tasks[0].connect_id)
             logger.exception(str(e))
             raise e
-        
+
         if alloc_tokened_tasks:
-            trans_obj.ready_to_move_queue.put(alloc_tokened_tasks, error_handle_func=self.put_to_fail_release_task_queue)
+            trans_obj.ready_to_move_queue.put(
+                alloc_tokened_tasks, error_handle_func=self.put_to_fail_release_task_queue
+            )
 
         return ans_list
-    
+
     # ==================================================================================
     # 定时检测kv 传输成功，但是长时间没有pd master来触发推理的请求，
     # 释放这些超时请求占用的kv资源
@@ -308,24 +311,24 @@ def check_trans_process_loop(self):
                 for device_id in range(self.node_world_size):
                     if not self.kv_trans_processes[device_id].is_trans_process_health():
                         raise Exception(f"device_id {device_id} kv process is unhealth")
-                    
+
                 time.sleep(10.0)
         except (BaseException, RuntimeError) as e:
             logger.exception(str(e))
-            
+
             for device_id in range(self.node_world_size):
                 self.kv_trans_processes[device_id].killself()
 
             # 杀掉当前进程的父进程（router), 触发全局崩溃
             os.kill(os.getppid(), signal.SIGKILL)
             os.kill(os.getpid(), signal.SIGKILL)
             raise e
-        
+
     # ==================================================================================
     # 常用辅助功能函数
     # ==================================================================================
     def get_next_device_index(self):
-        counts = [0  for _ in range(self.node_world_size)]
+        counts = [0 for _ in range(self.node_world_size)]
         for obj in self.connect_id_to_trans_obj.values():
             counts[obj.device_index] += 1
         device_index = int(np.argmin(counts))

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

@@ -16,11 +16,12 @@
 
 KV_MOVE_MAX_NUM = 16
 
+
 @dataclass
 class KVTransConnectObj:
     connect_id: str = None
     prefill_node_id: int = None
-    kv_trans_process: 'KVTransProcess' = None
+    kv_trans_process: "KVTransProcess" = None
     pd_prefill_nccl_ip: str = None
     pd_prefill_nccl_port: int = None
     device_index: int = None
@@ -33,8 +34,13 @@ class KVTransConnectObj:
     timer_checker: TimeChecker = None
 
     def create(
-        self, connect_id: str, prefill_node_id: str, pd_prefill_nccl_ip: str, pd_prefill_nccl_port: int, manager: "DecodeKVMoveManager"
-    ): 
+        self,
+        connect_id: str,
+        prefill_node_id: str,
+        pd_prefill_nccl_ip: str,
+        pd_prefill_nccl_port: int,
+        manager: "DecodeKVMoveManager",
+    ):
         self.connect_id = connect_id
         self.device_index = manager.get_next_device_index()
         self.kv_trans_process = manager.kv_trans_processes[self.device_index]
@@ -49,7 +55,7 @@ def create(
                     pd_prefill_nccl_port=pd_prefill_nccl_port,
                     decode_id=decode_node_id,
                     decode_device_id=self.device_index,
-                    connect_id=self.connect_id
+                    connect_id=self.connect_id,
                 )
             )
             assert self.kv_trans_process.task_out_queue.get(timeout=60) == "nccl_ok"
@@ -74,7 +80,7 @@ def create(
         self.put_to_radix_thread = threading.Thread(target=self.put_to_radix_loop, daemon=True)
         self.put_to_radix_thread.start()
         return
-    
+
     # ==================================================================================
     # 处理接受所有进行 kv 传输的请求，完成后，将请求放入到 move_finished_queue 中
     # ==================================================================================
@@ -106,7 +112,7 @@ def kv_move_loop(self):
                 logger.error(f"error get need 1, but get {len(move_tasks)}")
                 assert False
 
-            move_tasks:List[KVMoveTask] = move_tasks[0]
+            move_tasks: List[KVMoveTask] = move_tasks[0]
             for task in move_tasks:
                 logger.info(f"{func_name} get task {task.to_decode_log_info()}")
 
@@ -128,7 +134,7 @@ def kv_move_loop(self):
 
         logger.error(f"{func_name}  thread quit")
         return
-    
+
     # ==================================================================================
     # 将传输完成的请求，放入到 radix cache 中进行管理。
     # ==================================================================================
@@ -168,11 +174,11 @@ def put_to_radix_loop(self):
 
         logger.error(f"{func_name} thread quit, info: {self.to_log_info()}")
         return
-    
+
     # ==================================================================================
     # 错误处理检测操作的一些通用函数
     # ==================================================================================
-    
+
     def timer_to_check_status(self, raise_exception=True):
         if self.timer_checker.has_exceeded():
             try:
@@ -203,10 +209,10 @@ def set_has_error(self):
 
         if self.ready_to_move_queue is not None:
             self.ready_to_move_queue.has_error = True
-        
+
         if self.move_finished_queue is not None:
             self.move_finished_queue.has_error = True
-        
+
         if self.manager is not None:
             self.manager.remove_trans_obj(self.connect_id)
         return
@@ -219,15 +225,19 @@ def __del__(self):
 
             join_if_alive(self.kv_move_thread)
             join_if_alive(self.put_to_radix_thread)
-            
+
             if self.connect_id is not None and self.kv_trans_process is not None:
-                self.kv_trans_process.task_in_queue.put(PDTransLeaveInfo(decode_id=self.decode_node_id, prefill_id=self.prefill_node_id, connect_id=self.connect_id))
+                self.kv_trans_process.task_in_queue.put(
+                    PDTransLeaveInfo(
+                        decode_id=self.decode_node_id, prefill_id=self.prefill_node_id, connect_id=self.connect_id
+                    )
+                )
 
             if self.ready_to_move_queue is not None:
                 self.ready_to_move_queue.clear_tasks()
             if self.move_finished_queue is not None:
                 self.move_finished_queue.clear_tasks()
-            
+
         except BaseException as e:
             logger.exception(str(e))
 
@@ -240,6 +250,7 @@ def to_log_info(self):
         log += f"device_index: {self.device_index} "
         return log
 
+
 @dataclass
 class KVTransProcess:
     process: mp.Process = None
@@ -249,7 +260,6 @@ class KVTransProcess:
     task_out_queue: mp.Queue = None
     device_id: int = None
 
-
     def init_all(self, device_id: int, manager: "DecodeKVMoveManager"):
         self.device_lock = threading.Lock()
         self.device_id = device_id
@@ -271,12 +281,12 @@ def init_all(self, device_id: int, manager: "DecodeKVMoveManager"):
             assert self.task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             return True
-        
+
         except Exception as e:
             logger.warning(f"Failed start kv trans process for device {device_id}: {e}")
             logger.exception(str(e))
             return False
-    
+
     def is_trans_process_health(self):
         try:
             process = psutil.Process(self.process.pid)
@@ -287,6 +297,6 @@ def is_trans_process_health(self):
                 return True
         except:
             return False
-        
+
     def killself(self):
         self.process.kill()

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

@@ -83,7 +83,9 @@ def _init_env(args, device_id: int, task_in_queue: mp.Queue, task_out_queue: mp.
         while True:
             task: Union[KVMoveTaskGroup, PDTransJoinInfo, PDTransLeaveInfo] = task_in_queue.get()
             if isinstance(task, KVMoveTaskGroup):
-                _handle_kvmove_task(task, task_out_queue, mem_managers, connect_id_to_comm, task.connect_id, dp_size_in_node)
+                _handle_kvmove_task(
+                    task, task_out_queue, mem_managers, connect_id_to_comm, task.connect_id, dp_size_in_node
+                )
             elif isinstance(task, PDTransJoinInfo):
                 _handle_prefill_join(task, task_out_queue, connect_id_to_comm)
             elif isinstance(task, PDTransLeaveInfo):

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

@@ -66,14 +66,14 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         self.release_tasks_thread.start()
 
         from .prefill_trans_obj import KVTransProcess
-        
+
         self.kv_trans_processes: List[KVTransProcess] = [None] * self.node_world_size
         for device_id in range(self.node_world_size):
             self.kv_trans_processes[device_id] = KVTransProcess()
             assert self.kv_trans_processes[device_id].init_all(device_id, self)
 
         return
-    
+
     # ==================================================================================
     # 主任务循环，接收需要进行kv传输的请求进行处理
     # ==================================================================================
@@ -95,7 +95,7 @@ def task_dispatcher_loop(self):
         except (BaseException, RuntimeError) as e:
             logger.exception(str(e))
             raise e
-    
+
     # ==================================================================================
     # 请求出错或者完成kv传输后的处理队列和线程loop
     # ==================================================================================
@@ -117,7 +117,7 @@ def handle_release_task_loop(self):
             else:
                 self._remove_req_refs_from_prompt_cache(handle_list)
         return
-    
+
     # ==================================================================================
     # 定时检测传输进程的健康状态，出现问题拉崩整个系统触发重启
     # ==================================================================================
@@ -128,21 +128,21 @@ def check_trans_process_loop(self):
                 for device_id in range(self.node_world_size):
                     if not self.kv_trans_processes[device_id].is_trans_process_health():
                         raise Exception(f"device_id {device_id} kv process is unhealth")
-                    
+
                 time.sleep(10.0)
         except (BaseException, RuntimeError) as e:
             logger.exception(str(e))
-            
+
             for device_id in range(self.node_world_size):
                 self.kv_trans_processes[device_id].killself()
 
             # 杀掉当前进程的父进程（router), 触发全局崩溃
             os.kill(os.getppid(), signal.SIGKILL)
             os.kill(os.getpid(), signal.SIGKILL)
             raise e
-    
+
     # ==================================================================================
-    # 与推理进程交互接口,  _remove_req_refs_from_prompt_cache 和 
+    # 与推理进程交互接口,  _remove_req_refs_from_prompt_cache 和
     # _put_mem_manager_to_mem_queue 都是通过 rpyc 与推理进程进行交互的接口
     # ==================================================================================
 
@@ -172,7 +172,7 @@ def _put_mem_manager_to_mem_queue(self):
     async def wait_all_future_finish(self, futures: List[AsyncResult]):
         await asyncio.gather(*[asyncio.to_thread(future.wait) for future in futures])
         return
-    
+
     # ==================================================================================
     # 辅助功能接口
     # ==================================================================================
@@ -191,18 +191,18 @@ def remove_trans_obj(self, connect_id):
                 trans_obj.set_has_error()
                 logger.error(f"remove tran obj id {trans_obj.decode_node_id}")
         return
-    
+
     def __get_trans_obj(self, task: KVMoveTask):
         self.__remove_dead_trans_obj()
         # 如果已经存在连接对象，直接返回
         for obj in self.connect_id_to_trans_obj.values():
             if obj.decode_node_id == task.decode_node.node_id:
                 return obj
-        
+
         # 如果不存在连接对象，创建新的连接对象
         gc.collect()
         from .prefill_trans_obj import KVTransConnectObj
-        
+
         trans_obj = KVTransConnectObj()
         trans_obj.create(task.decode_node.node_id, task.decode_node.ip, task.decode_node.rpyc_port, self)
         self.connect_id_to_trans_obj[trans_obj.connect_id] = trans_obj
@@ -221,6 +221,7 @@ def __remove_dead_trans_obj(self):
             gc.collect()
         return
 
+
 def _init_env(args, info_queue: mp.Queue, mem_queues: List[mp.Queue], event: mp.Event):
     import lightllm.utils.rpyc_fix_utils as _