diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index 1d68f81a9e..e35c507f91 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -170,6 +170,7 @@ class PDChunckedTransTask: start_trans_time: float = None # 用于标记传输开始的时间。同时标记是否正在传输中 error_info: Optional[str] = None + transfer_quiesced: bool = False transfer_time_out_secs: int = 66 page_kind: str = "kv" # Only valid for the local task owner; remote notify copies may carry the sender-local req_idx. diff --git a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py index 036c6f162b..86e2c8258f 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/decode_node_impl/decode_trans_process.py @@ -138,6 +138,7 @@ def __init__( self.recv_task_group_queue = queue.Queue() self.waiting_dict_lock = threading.Lock() self.waiting_dict: Dict[str, PDChunckedTransTask] = {} + self.quarantined_dict: Dict[str, PDChunckedTransTask] = {} self.request_page_task_queue = queue.Queue() self.ready_page_task_queue = queue.Queue() self.success_queue = queue.Queue() @@ -198,7 +199,7 @@ def _abort(self, request_id: int, error_info: str = "aborted req"): for trans_task in aborted_tasks: trans_task.error_info = error_info - self.failed_queue.put(trans_task) + self._queue_failed_task(trans_task) return @log_exception @@ -263,15 +264,20 @@ def accept_peer_task_loop( # 请求有错误 if notify_obj.error_info is not None: - # 直接清理掉所有的相关请求。 + quarantined_task = None with self.waiting_dict_lock: local_trans_task = self.waiting_dict.pop(notify_obj.get_key(), None) + if notify_obj.transfer_quiesced: + quarantined_task = self.quarantined_dict.pop(notify_obj.get_key(), None) if local_trans_task is not None: local_trans_task.error_info = notify_obj.error_info - # 软性的调整超时时间，防止一些特殊情况，过快的释放task - # 占用的page 页面，导致多p 复写引起脏内容的问题。 - local_trans_task.transfer_time_out_secs = 12 - self.failed_queue.put(local_trans_task) + local_trans_task.transfer_quiesced = notify_obj.transfer_quiesced + + if local_trans_task is not None: + self._queue_failed_task(local_trans_task) + + if quarantined_task is not None: + self._recycle_quarantined_page(quarantined_task) self._abort( request_id=notify_obj.request_id, @@ -311,13 +317,18 @@ def accept_peer_task_loop( # prefill 写完数据到了 done 阶段 if remote_trans_task.write_stage == "done": + quarantined_task = None with self.waiting_dict_lock: local_trans_task = self.waiting_dict.pop(remote_trans_task.get_key(), None) + if local_trans_task is None: + quarantined_task = self.quarantined_dict.pop(remote_trans_task.get_key(), None) if local_trans_task is not None: local_trans_task.first_gen_token_id = remote_trans_task.first_gen_token_id local_trans_task.first_gen_token_logprob = remote_trans_task.first_gen_token_logprob self.ready_page_task_queue.put(local_trans_task) logger.info(f"recv WRITE done from prefill: {remote_trans_task.to_str()}") + elif quarantined_task is not None: + self._recycle_quarantined_page(quarantined_task) else: # Same race as the WRITE request stage: decode may have cleaned the # waiting task because the request was aborted, then a late done notify @@ -346,7 +357,7 @@ def _check_tasks_time_out(self): for trans_task in timeout_tasks: trans_task.error_info = "time out in accept_peer_task_loop" - self.failed_queue.put(trans_task) + self._queue_failed_task(trans_task) return @log_exception @@ -369,7 +380,7 @@ def request_page_loop(self): logger.exception(str(e)) self.transporter.remove_remote_agent(peer_name=trans_task.prefill_agent_name) trans_task.error_info = f"send write ready task to prefill node failed: {str(e)}" - self.failed_queue.put(trans_task) + self._queue_failed_task(trans_task) continue return @@ -432,9 +443,10 @@ def fail_loop(self): while True: trans_task: PDChunckedTransTask = self.failed_queue.get() - # 回收页面 if trans_task.dst_page_index is not None: - self.page_index_queue.put(trans_task.dst_page_index) + if trans_task.transfer_quiesced: + self.page_index_queue.put(trans_task.dst_page_index) + trans_task.dst_page_index = None if trans_task.xfer_handle is not None: self.transporter.release_xfer_handle(trans_task.xfer_handle) @@ -449,4 +461,18 @@ def fail_loop(self): request_id=trans_task.request_id, error_info=trans_task.error_info, ) - self.transporter.send_error_info_to_prefill_node(trans_task=trans_task) + if not trans_task.transfer_quiesced: + self.transporter.send_error_info_to_prefill_node(trans_task=trans_task) + + def _recycle_quarantined_page(self, trans_task: PDChunckedTransTask): + if trans_task.dst_page_index is not None: + self.page_index_queue.put(trans_task.dst_page_index) + trans_task.dst_page_index = None + trans_task.transfer_quiesced = True + logger.info(f"recycle quiesced decode page for failed task: {trans_task.to_str()}") + + def _queue_failed_task(self, trans_task: PDChunckedTransTask): + if trans_task.dst_page_index is not None and not trans_task.transfer_quiesced: + with self.waiting_dict_lock: + self.quarantined_dict[trans_task.get_key()] = trans_task + self.failed_queue.put(trans_task) diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py index 2ed0335ca5..163304bd53 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nccl_kv_transporter.py @@ -183,7 +183,7 @@ def send_error_info_to_decode_node(self, trans_task: PDChunckedTransTask): new_trans_task.prefill_num_pages = self.num_pages new_trans_task.prefill_page_reg_desc = self.local_page_mem_desc self._send_task_notif(trans_task.decode_agent_name, new_trans_task) - return + return True def write_blocks_paged(self, trans_task: PDChunckedTransTask) -> "_NcclXferHandle": assert trans_task.src_page_index is not None and trans_task.dst_page_index is not None diff --git a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py index bd5e11f05d..7faac18659 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/nixl_kv_transporter.py @@ -229,11 +229,12 @@ def send_error_info_to_decode_node(self, trans_task: PDChunckedTransTask): remote_agent_name=decode_agent_name, notif_msg=pickle.dumps(new_trans_task), ) + return True except BaseException as e: logger.error(f"send error info to decode node failed: {trans_task.to_str()}") logger.exception(str(e)) self.remove_remote_agent(peer_name=decode_agent_name) - return + return False def write_blocks_paged( self, diff --git a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py index e286a10f96..3b61f0b192 100644 --- a/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py +++ b/lightllm/server/router/model_infer/mode_backend/pd/prefill_node_impl/prefill_trans_process.py @@ -116,6 +116,7 @@ def __init__( self.write_peer_kv_queue = queue.Queue() self.success_queue = queue.Queue() self.failed_queue = queue.Queue() + self.draining_queue = queue.Queue() self.page_index_queue = queue.Queue() for page_index in range(self.args.pd_kv_page_num): @@ -133,6 +134,7 @@ def __init__( self.update_task_status_loop, self.success_loop, self.fail_loop, + self.drain_failed_xfer_loop, ]: threading.Thread(target=func, daemon=True).start() return @@ -377,11 +379,12 @@ def success_loop(self): torch.cuda.set_device(self.device_id) while True: trans_task: PDChunckedTransTask = self.success_queue.get() - # 写回后，回收页面 - if trans_task.src_page_index is not None: - self.page_index_queue.put(trans_task.src_page_index) if trans_task.xfer_handle is not None: self.transporter.release_xfer_handle(trans_task.xfer_handle) + trans_task.xfer_handle = None + if trans_task.src_page_index is not None: + self.page_index_queue.put(trans_task.src_page_index) + trans_task.src_page_index = None ret = trans_task.createRetObj() ret.first_gen_token_id = None @@ -399,11 +402,13 @@ def fail_loop(self): while True: trans_task: PDChunckedTransTask = self.failed_queue.get() - # 回收页面 - if trans_task.src_page_index is not None: - self.page_index_queue.put(trans_task.src_page_index) - if trans_task.xfer_handle is not None: - self.transporter.release_xfer_handle(trans_task.xfer_handle) + if trans_task.xfer_handle is None: + if trans_task.src_page_index is not None: + self.page_index_queue.put(trans_task.src_page_index) + trans_task.src_page_index = None + trans_task.transfer_quiesced = True + if trans_task.error_info is not None: + self.draining_queue.put(trans_task) ret = trans_task.createRetObj() self.task_out_queue.put(ret) @@ -411,4 +416,55 @@ def fail_loop(self): if trans_task.error_info is not None: self._abort(request_id=trans_task.request_id, error_info=trans_task.error_info) - self.transporter.send_error_info_to_decode_node(trans_task=trans_task) + + @log_exception + def drain_failed_xfer_loop(self): + torch.cuda.set_device(self.device_id) + draining_tasks: Dict[str, PDChunckedTransTask] = {} + while True: + try: + trans_task = self.draining_queue.get(timeout=0.001) + draining_tasks[trans_task.get_key()] = trans_task + except queue.Empty: + pass + + for key, trans_task in list(draining_tasks.items()): + try: + status = self._try_finish_failed_xfer_drain(trans_task) + if status is None: + continue + except BaseException as e: + if trans_task.transfer_quiesced: + logger.error( + f"failed to send quiesced ack, decode page remains quarantined: {trans_task.to_str()}" + ) + else: + logger.error(f"failed to drain xfer task, keeping page quarantined: {trans_task.to_str()}") + logger.exception(str(e)) + continue + + draining_tasks.pop(key, None) + logger.info(f"failed xfer task reached terminal state {status}: {trans_task.to_str()}") + + def _try_finish_failed_xfer_drain(self, trans_task: PDChunckedTransTask) -> Optional[str]: + if trans_task.transfer_quiesced: + sent = self.transporter.send_error_info_to_decode_node(trans_task=trans_task) + if sent is False: + raise RuntimeError("failed to send quiesced transfer acknowledgement") + return "QUIESCED" + + status = self.transporter.check_task_status(trans_task=trans_task) + if status not in ["DONE", "ERR"]: + return None + + self.transporter.release_xfer_handle(trans_task.xfer_handle) + trans_task.xfer_handle = None + if trans_task.src_page_index is not None: + self.page_index_queue.put(trans_task.src_page_index) + trans_task.src_page_index = None + + trans_task.transfer_quiesced = True + sent = self.transporter.send_error_info_to_decode_node(trans_task=trans_task) + if sent is False: + raise RuntimeError("failed to send quiesced transfer acknowledgement") + return status diff --git a/test/unit/test_pd_transfer_lifecycle.py b/test/unit/test_pd_transfer_lifecycle.py new file mode 100644 index 0000000000..47e7f6d31f --- /dev/null +++ b/test/unit/test_pd_transfer_lifecycle.py @@ -0,0 +1,165 @@ +import queue +import threading +from types import SimpleNamespace + +from lightllm.server.router.model_infer.mode_backend.pd.decode_node_impl.decode_trans_process import ( + _DecodeTransModule, +) +from lightllm.server.router.model_infer.mode_backend.pd.prefill_node_impl.prefill_trans_process import ( + _PrefillTransModule, +) + + +class _RecordingQueue: + def __init__(self, events): + self.events = events + + def put(self, value): + self.events.append(("recycle", value)) + + +class _FakeTransporter: + def __init__(self, status, events, release_error=None, notify_errors=None): + self.status = status + self.events = events + self.release_error = release_error + self.notify_errors = list(notify_errors or []) + + def check_task_status(self, trans_task): + self.events.append(("check", trans_task.xfer_handle)) + return self.status + + def release_xfer_handle(self, handle): + self.events.append(("release", handle)) + if self.release_error is not None: + raise self.release_error + + def send_error_info_to_decode_node(self, trans_task): + self.events.append(("notify", trans_task.transfer_quiesced)) + if self.notify_errors: + raise self.notify_errors.pop(0) + + +def _make_task(): + return SimpleNamespace( + dst_page_index=7, + error_info="timeout", + src_page_index=3, + transfer_quiesced=False, + xfer_handle=11, + get_key=lambda: "task-key", + to_str=lambda: "task-key", + ) + + +def test_prefill_keeps_source_page_while_transfer_is_in_progress(): + events = [] + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.transporter = _FakeTransporter("PROC", events) + module.page_index_queue = _RecordingQueue(events) + task = _make_task() + + status = module._try_finish_failed_xfer_drain(task) + + assert status is None + assert events == [("check", 11)] + assert task.src_page_index == 3 + assert task.xfer_handle == 11 + assert not task.transfer_quiesced + + +def test_prefill_recycles_source_page_only_after_terminal_status_and_release(): + events = [] + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.transporter = _FakeTransporter("DONE", events) + module.page_index_queue = _RecordingQueue(events) + task = _make_task() + + status = module._try_finish_failed_xfer_drain(task) + + assert status == "DONE" + assert events == [ + ("check", 11), + ("release", 11), + ("recycle", 3), + ("notify", True), + ] + assert task.src_page_index is None + assert task.xfer_handle is None + assert task.transfer_quiesced + + +def test_prefill_keeps_page_quarantined_when_handle_release_fails(): + events = [] + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.transporter = _FakeTransporter("ERR", events, release_error=RuntimeError("still active")) + module.page_index_queue = _RecordingQueue(events) + task = _make_task() + + try: + module._try_finish_failed_xfer_drain(task) + except RuntimeError as exc: + assert str(exc) == "still active" + else: + raise AssertionError("release failure must be propagated to the drain loop") + + assert events == [("check", 11), ("release", 11)] + assert task.src_page_index == 3 + assert task.xfer_handle == 11 + assert not task.transfer_quiesced + + +def test_prefill_retries_quiesced_ack_without_rechecking_released_handle(): + events = [] + module = _PrefillTransModule.__new__(_PrefillTransModule) + module.transporter = _FakeTransporter( + "DONE", + events, + notify_errors=[RuntimeError("peer unavailable")], + ) + module.page_index_queue = _RecordingQueue(events) + task = _make_task() + + try: + module._try_finish_failed_xfer_drain(task) + except RuntimeError as exc: + assert str(exc) == "peer unavailable" + else: + raise AssertionError("notify failure must be propagated to the drain loop") + + assert task.transfer_quiesced + assert task.xfer_handle is None + assert task.src_page_index is None + + status = module._try_finish_failed_xfer_drain(task) + + assert status == "QUIESCED" + assert events == [ + ("check", 11), + ("release", 11), + ("recycle", 3), + ("notify", True), + ("notify", True), + ] + + +def test_decode_quarantines_destination_page_until_quiesced_ack(): + module = _DecodeTransModule.__new__(_DecodeTransModule) + module.waiting_dict_lock = threading.Lock() + module.quarantined_dict = {} + module.failed_queue = queue.Queue() + module.page_index_queue = queue.Queue() + task = _make_task() + + module._queue_failed_task(task) + + assert module.quarantined_dict == {"task-key": task} + assert module.failed_queue.get_nowait() is task + assert module.page_index_queue.empty() + + quarantined_task = module.quarantined_dict.pop(task.get_key()) + module._recycle_quarantined_page(quarantined_task) + + assert module.page_index_queue.get_nowait() == 7 + assert task.dst_page_index is None + assert task.transfer_quiesced