fix

Author: wangzaijun

lightllm/server/router/model_infer/mode_backend/pd_nixl/decode_node_impl/decode_impl.py

@@ -90,38 +90,37 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]:
         ans_list : List[InferReq] = []
         for request_id in req_ids:
             req_obj: InferReq = g_infer_context.requests_mapping[request_id]
-            if req_obj.infer_aborted:
-                if req_obj.nixl_pd_task_num == (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
-                    ans_list.append(req_obj)
+            if req_obj.nixl_pd_task_num != (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
                 continue
-            
-            if req_obj.nixl_pd_task_num == (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
-                if req_obj.nixl_pd_task_failed_num > 0:
-                    if not req_obj.finish_status.is_finished():
-                        # 强制停止
-                        req_obj.cur_output_len += 1
-                        req_obj.set_next_gen_token_id(0, 0.0, 1)
-                        req_obj.finish_status.set_status(FinishStatus.FINISHED_STOP)
-
-                        if self.is_master_in_dp:
-                            req_obj.shm_req.shm_cur_output_len = req_obj.cur_output_len
-                            req_obj.shm_req.finish_token_index = req_obj.get_cur_total_len() - 1
-                            req_obj.shm_req.finish_status.set_status(FinishStatus.FINISHED_STOP)
-                            req_obj.shm_req.candetoken_out_len = req_obj.cur_output_len
-
-                            logger.error(f"req_id: {req_obj.req_id} forced to finished, it exits kv transfer error")
-                                         
-                    # 提前释放有问题的 mem_index
-                    old_prefix_len = 0 if req_obj.shared_kv_node is None else req_obj.shared_kv_node.node_prefix_total_len
-                    error_mem_len = req_obj.cur_kv_len - old_prefix_len
+
+            if req_obj.nixl_pd_task_failed_num > 0:
+                # 强制停止
+                if not req_obj.finish_status.is_finished():
+                    req_obj.cur_output_len += 1
+                    req_obj.set_next_gen_token_id(0, 0.0, 1)
+                    req_obj.finish_status.set_status(FinishStatus.FINISHED_STOP)
+
+                    if self.is_master_in_dp:
+                        req_obj.shm_req.shm_cur_output_len = req_obj.cur_output_len
+                        req_obj.shm_req.finish_token_index = req_obj.get_cur_total_len() - 1
+                        req_obj.shm_req.finish_status.set_status(FinishStatus.FINISHED_STOP)
+                        req_obj.shm_req.candetoken_out_len = req_obj.cur_output_len
+
+                        logger.error(f"req_id: {req_obj.req_id} forced to finished, it exits kv transfer error")
+                                        
+                # 提前释放有问题的 mem_index
+                old_prefix_len = 0 if req_obj.shared_kv_node is None else req_obj.shared_kv_node.node_prefix_total_len
+                error_mem_len = req_obj.cur_kv_len - old_prefix_len
+                if error_mem_len > 0:
                     req_obj.cur_kv_len -= error_mem_len
 
                     mem_indexes = self.model.req_manager.req_to_token_indexs[req_obj.req_idx, req_obj.cur_kv_len:(req_obj.cur_kv_len + error_mem_len)].detach().cpu()
                     self.model.mem_manager.free(mem_indexes)
                     if self.is_master_in_dp:
                         req_obj.shm_req.shm_cur_kv_len = req_obj.cur_kv_len
-                else:
-                    ans_list.append(req_obj)
+                
+            ans_list.append(req_obj)
+
         g_infer_state_lock.release()
         return ans_list
 

lightllm/server/router/model_infer/mode_backend/pd_nixl/decode_node_impl/decode_trans_process.py

@@ -152,6 +152,7 @@ def accept_peer_task_loop(
             # notify update
             notifies_dict = self.transporter.get_new_notifs()
             if not notifies_dict:
+                self._check_tasks_time_out()
                 time.sleep(0.005)
                 continue
 
@@ -181,7 +182,7 @@ def accept_peer_task_loop(
                             local_trans_task.prefill_page_reg_desc = remote_trans_task.prefill_page_reg_desc
 
                             self.read_peer_kv_queue.put(local_trans_task)
-            
+                            
             self._check_tasks_time_out()
 
 

lightllm/server/router/model_infer/mode_backend/pd_nixl/prefill_node_impl/prefill_impl.py

@@ -42,18 +42,19 @@ def _filter_not_ready_reqs(self, req_ids: List[int]) -> List[InferReq]:
         ans_list : List[InferReq] = []
         for request_id in req_ids:
             req_obj: InferReq = g_infer_context.requests_mapping[request_id]
-            if req_obj.infer_aborted:
-                if req_obj.nixl_pd_task_num == (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
-                    ans_list.append(req_obj)
-                continue
-
             prefill_finished = req_obj.shm_req.input_len <= req_obj.cur_kv_len
             if prefill_finished:
-                # 所有传输任务都已经完成。
+                # 等待所有传输任务都已经完成。
                 if req_obj.nixl_pd_task_num == (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
                     ans_list.append(req_obj)
             else:
-                ans_list.append(req_obj)
+                if req_obj.infer_aborted:
+                    if req_obj.nixl_pd_task_num == (req_obj.nixl_pd_task_failed_num + req_obj.nixl_pd_task_sunccess_num):
+                        ans_list.append(req_obj)
+                    else:
+                        continue
+                else:
+                    ans_list.append(req_obj)
         return ans_list
 
     def _prefill_chuncked_handle_func(self, req_obj: InferReq, next_token_id: int, next_token_prob: float, output_len: int):