fix(pd): recover NIXL transfer task failures

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

@@ -21,7 +21,7 @@
 from ..kv_transporter import create_kv_transporter
 from lightllm.utils.error_utils import log_exception
 from lightllm.utils.envs_utils import get_unique_server_name
-from lightllm.utils.process_check import start_parent_check_thread
+from lightllm.utils.process_check import install_fatal_thread_excepthook, start_parent_check_thread
 
 logger = init_logger(__name__)
 
@@ -47,6 +47,7 @@ def _init_env(
     task_out_queue: mp.Queue,
     up_status_in_queue: Optional[mp.SimpleQueue],
 ):
+    install_fatal_thread_excepthook()
     start_parent_check_thread()
     import lightllm.utils.rpyc_fix_utils as _
 

lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py

@@ -15,7 +15,7 @@
 from ..kv_transporter import create_kv_transporter
 from lightllm.utils.error_utils import log_exception
 from lightllm.utils.envs_utils import get_unique_server_name
-from lightllm.utils.process_check import start_parent_check_thread
+from lightllm.utils.process_check import install_fatal_thread_excepthook, start_parent_check_thread
 
 
 logger = init_logger(__name__)
@@ -41,6 +41,7 @@ def _init_env(
     task_in_queue: mp.Queue,
     task_out_queue: mp.Queue,
 ):
+    install_fatal_thread_excepthook()
     start_parent_check_thread()
     import lightllm.utils.rpyc_fix_utils as _
 
@@ -331,48 +332,76 @@ def update_task_status_loop(
                 time.sleep(0.001)
                 continue
 
-            with self.waiting_dict_lock:
-                tasks = list(self.waiting_dict.values())
-                for trans_task in tasks:
-                    if trans_task.xfer_handle is None:
-                        continue
+            self._update_task_status_once()
+            time.sleep(0.001)
 
-                    # 传输任务状态检查
+    def _update_task_status_once(self):
+        with self.waiting_dict_lock:
+            tasks = list(self.waiting_dict.values())
+            for trans_task in tasks:
+                if trans_task.xfer_handle is None:
+                    continue
+
+                # A native NIXL error must fail this task without terminating the
+                # status thread and stranding every transfer queued behind it.
+                try:
                     ret = self.transporter.check_task_status(trans_task=trans_task)
-                    if ret == "DONE":
-                        trans_task = self.waiting_dict.pop(trans_task.get_key(), None)
-                        if self.transporter.capture_telemetry:
-                            telem = self.transporter.nixl_agent.get_xfer_telemetry(trans_task.xfer_handle)
+                except BaseException as e:
+                    failed_task = self.waiting_dict.pop(trans_task.get_key(), None)
+                    if failed_task is not None:
+                        logger.exception(f"check transfer status failed: {failed_task.to_str()}")
+                        failed_task.error_info = f"check transfer status failed: {str(e)}"
+                        self.failed_queue.put(failed_task)
+                    continue
+
+                if ret == "DONE":
+                    completed_task = self.waiting_dict.pop(trans_task.get_key(), None)
+                    if completed_task is None:
+                        continue
+                    if self.transporter.capture_telemetry:
+                        try:
+                            telem = self.transporter.nixl_agent.get_xfer_telemetry(completed_task.xfer_handle)
                             total_us = telem.xferDuration
                             post_us = telem.postDuration
                             backend_us = telem.xferDuration - telem.postDuration
-                            nixl_backend = self.transporter.nixl_agent.query_xfer_backend(trans_task.xfer_handle)
+                            nixl_backend = self.transporter.nixl_agent.query_xfer_backend(completed_task.xfer_handle)
                             logger.info(
-                                f"write trans task request_id={trans_task.request_id} "
-                                f"kv=[{trans_task.start_kv_index},{trans_task.end_kv_index}) "
-                                f"src_page={trans_task.src_page_index} dst_page={trans_task.dst_page_index} "
+                                f"write trans task request_id={completed_task.request_id} "
+                                f"kv=[{completed_task.start_kv_index},{completed_task.end_kv_index}) "
+                                f"src_page={completed_task.src_page_index} "
+                                f"dst_page={completed_task.dst_page_index} "
                                 f"xfer time: {total_us:.3f} us, "
                                 f"post time: {post_us:.3f} us, backend time: {backend_us:.3f} us, "
                                 f"nixl_backend: {nixl_backend}, total_bytes: {telem.totalBytes}"
                             )
-                        self.transporter.send_write_done_task_to_decode_node(trans_task)
-                        logger.info(
-                            f"send WRITE done nixl notify "
-                            f"request_id={trans_task.request_id} "
-                            f"kv=[{trans_task.start_kv_index},{trans_task.end_kv_index}) "
-                            f"src_page={trans_task.src_page_index} dst_page={trans_task.dst_page_index}"
-                        )
-                        self.success_queue.put(trans_task)
-                    elif ret == "ERR":
-                        trans_task = self.waiting_dict.pop(trans_task.get_key(), None)
-                        trans_task.error_info = "xfer error"
-                        self.failed_queue.put(trans_task)
-                    elif trans_task.time_out():
-                        trans_task = self.waiting_dict.pop(trans_task.get_key(), None)
-                        trans_task.error_info = "time out in update_task_status_loop"
-                        self.failed_queue.put(trans_task)
+                        except BaseException:
+                            logger.exception(f"get transfer telemetry failed: {completed_task.to_str()}")
+
+                    try:
+                        self.transporter.send_write_done_task_to_decode_node(completed_task)
+                    except BaseException as e:
+                        logger.exception(f"send WRITE done nixl notify failed: {completed_task.to_str()}")
+                        completed_task.error_info = f"send WRITE done nixl notify failed: {str(e)}"
+                        self.failed_queue.put(completed_task)
+                        continue
 
-            time.sleep(0.001)
+                    logger.info(
+                        f"send WRITE done nixl notify "
+                        f"request_id={completed_task.request_id} "
+                        f"kv=[{completed_task.start_kv_index},{completed_task.end_kv_index}) "
+                        f"src_page={completed_task.src_page_index} dst_page={completed_task.dst_page_index}"
+                    )
+                    self.success_queue.put(completed_task)
+                elif ret == "ERR":
+                    failed_task = self.waiting_dict.pop(trans_task.get_key(), None)
+                    if failed_task is not None:
+                        failed_task.error_info = "xfer error"
+                        self.failed_queue.put(failed_task)
+                elif trans_task.time_out():
+                    failed_task = self.waiting_dict.pop(trans_task.get_key(), None)
+                    if failed_task is not None:
+                        failed_task.error_info = "time out in update_task_status_loop"
+                        self.failed_queue.put(failed_task)
 
     @log_exception
     def success_loop(self):

lightllm/utils/process_check.py

@@ -8,6 +8,21 @@
 logger = init_logger(__name__)
 
 
+def install_fatal_thread_excepthook():
+    """Terminate the process when an unexpected daemon-thread exception escapes."""
+
+    def exit_process(args):
+        try:
+            logger.error(
+                f"fatal background thread failure in {getattr(args.thread, 'name', 'unknown')}",
+                exc_info=(args.exc_type, args.exc_value, args.exc_traceback),
+            )
+        finally:
+            os._exit(1)
+
+    threading.excepthook = exit_process
+
+
 def is_process_active(pid):
     try:
         process = psutil.Process(pid)