fixup.

Author: Weichao Luo

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

@@ -29,7 +29,7 @@
 thread_local_data = threading.local()
 
 KV_MOVE_MAX_NUM = 16
-KV_MOVE_MAX_RESTART_CNT = 3
+KV_MOVE_MAX_START_CNT = 3
 
 
 @dataclass
@@ -288,16 +288,12 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
         # 需要每个卡有一个锁来规划每次只能有一个tran obj 操作对应显卡上的传输任务。
         self.device_locks = [threading.Lock() for _ in range(self.node_world_size)]
 
-        self.kv_trans_processes = []
-        self.kv_trans_task_in_queues = []
-        self.kv_trans_task_out_queues = []
-        self.kv_trans_process_restart_cnt = []
+        self.kv_trans_processes = [None] * self.node_world_size
+        self.kv_trans_task_in_queues = [None] * self.node_world_size
+        self.kv_trans_task_out_queues = [None] * self.node_world_size
+        self.kv_trans_process_start_cnt = [0] * self.node_world_size
 
         for device_id in range(self.node_world_size):
-            self.kv_trans_task_in_queues.append(mp.Queue())
-            self.kv_trans_task_out_queues.append(mp.Queue())
-            self.kv_trans_process_restart_cnt.append(0)
-            self.kv_trans_processes.append(None)
             assert self.start_trans_process(device_id)
 
         return
@@ -525,9 +521,9 @@ def remove_trans_obj_by_deviceid(self, device_id):
                 self.remove_dead_trans_obj(node_id)
 
     def start_trans_process(self, device_id: int):
-        task_in_queue = self.kv_trans_task_in_queues[device_id]
-        task_out_queue = self.kv_trans_task_out_queues[device_id]
-        self.kv_trans_process_restart_cnt[device_id] += 1
+        task_in_queue = mp.Queue()
+        task_out_queue = mp.Queue()
+        self.kv_trans_process_start_cnt[device_id] += 1
 
         if self.kv_trans_processes[device_id]:
             # force kill
@@ -554,14 +550,16 @@ def start_trans_process(self, device_id: int):
             assert task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             self.kv_trans_processes[device_id] = kv_trans_process
+            self.kv_trans_task_in_queues[device_id] = task_in_queue
+            self.kv_trans_task_out_queues[device_id] = task_out_queue
 
             return True
         except Exception as e:
             logger.warning(f"Failed start kv trans process for device {device_id}: {e}")
             return False
 
     def is_kv_trans_process_alive(self, device_id):
-        return self.kv_trans_process_restart_cnt[device_id] <= KV_MOVE_MAX_RESTART_CNT
+        return self.kv_trans_process_start_cnt[device_id] <= KV_MOVE_MAX_START_CNT
 
     def check_trans_process(self, raise_exception=True):
         at_least_one_alive = False

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

@@ -29,7 +29,7 @@
 from lightllm.utils.envs_utils import get_unique_server_name
 
 KV_MOVE_MAX_NUM = 16
-KV_MOVE_MAX_RESTART_CNT = 3
+KV_MOVE_MAX_START_CNT = 3
 
 logger = init_logger(__name__)
 
@@ -348,20 +348,13 @@ def __init__(self, args, info_queue: mp.Queue, mem_queues: List[mp.Queue]):
 
         from .prefill_trans_process import start_prefill_trans_process
 
-        self.kv_trans_ports = []
-        self.kv_trans_processes = []
-        self.kv_trans_task_in_queues = []
-        self.kv_trans_task_out_queues = []
-        self.kv_trans_process_restart_cnt = []
+        self.kv_trans_ports = [None] * self.node_world_size
+        self.kv_trans_processes = [None] * self.node_world_size
+        self.kv_trans_task_in_queues = [None] * self.node_world_size
+        self.kv_trans_task_out_queues = [None] * self.node_world_size
+        self.kv_trans_process_start_cnt = [0] * self.node_world_size
 
         for device_id in range(self.node_world_size):
-            self.kv_trans_task_in_queues.append(mp.Queue())
-            self.kv_trans_task_out_queues.append(mp.Queue())
-            self.kv_trans_ports.append(
-                find_available_port(self.args.pd_p_allowed_port_min, self.args.pd_p_allowed_port_max)
-            )
-            self.kv_trans_process_restart_cnt.append(0)
-            self.kv_trans_processes.append(None)
             assert self.start_trans_process(device_id)
 
         return
@@ -385,10 +378,10 @@ def handle_release_task_loop(self):
         return
 
     def start_trans_process(self, device_id: int):
-        task_in_queue = self.kv_trans_task_in_queues[device_id]
-        task_out_queue = self.kv_trans_task_out_queues[device_id]
-        kv_trans_port = self.kv_trans_ports[device_id]
-        self.kv_trans_process_restart_cnt[device_id] += 1
+        task_in_queue = mp.Queue()
+        task_out_queue = mp.Queue()
+        kv_trans_port = find_available_port(self.args.pd_p_allowed_port_min, self.args.pd_p_allowed_port_max)
+        self.kv_trans_process_start_cnt[device_id] += 1
 
         if self.kv_trans_processes[device_id]:
             # force kill
@@ -417,6 +410,9 @@ def start_trans_process(self, device_id: int):
             assert task_out_queue.get(timeout=60) == "get_mem_managers_ok"
 
             self.kv_trans_processes[device_id] = kv_trans_process
+            self.kv_trans_task_in_queues[device_id] = task_in_queue
+            self.kv_trans_task_out_queues[device_id] = task_out_queue
+            self.kv_trans_ports[device_id] = kv_trans_port
 
             return True
         except Exception as e:
@@ -454,7 +450,7 @@ def check_trans_process_loop(self):
             raise e
 
     def is_kv_trans_process_alive(self, device_id):
-        return self.kv_trans_process_restart_cnt[device_id] <= KV_MOVE_MAX_RESTART_CNT
+        return self.kv_trans_process_start_cnt[device_id] <= KV_MOVE_MAX_START_CNT
 
     def get_next_device_index(self):