From 929b39ed42f050353b6a720df3940231a9279243 Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 20 Aug 2026 23:11:20 +0800 Subject: [PATCH] fix(pd): fail requests when assigned nodes disconnect Track registrations by their owning WebSocket so a stale disconnect cannot evict a replacement connection. Wake requests assigned to a replaced or disconnected node and surface the connection failure at every PD wait stage. --- lightllm/server/api_http_pd.py | 2 +- .../httpserver_for_pd_master/manager.py | 53 ++++++++++++++++--- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/lightllm/server/api_http_pd.py b/lightllm/server/api_http_pd.py index 1d8b2112f..c23e5e7a7 100644 --- a/lightllm/server/api_http_pd.py +++ b/lightllm/server/api_http_pd.py @@ -57,7 +57,7 @@ async def register_and_keep_alive(websocket: WebSocket): logger.exception(str(e)) finally: logger.error(f"client {regist_json} removed") - await g_objs.httpserver_manager.remove_pd(regist_json) + await g_objs.httpserver_manager.remove_pd(regist_json, websocket) return diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index 8d307abdc..7084b7683 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -74,11 +74,26 @@ def is_healthy(self): return False async def register_pd(self, pd_info_json, websocket): - self.pd_manager.register_pd(pd_info_json, websocket) + replaced_node = self.pd_manager.register_pd(pd_info_json, websocket) + if replaced_node is not None: + await self._fail_requests_for_node(replaced_node, "PD node connection was replaced") return - async def remove_pd(self, pd_info_json): - self.pd_manager.remove_pd(pd_info_json) + async def remove_pd(self, pd_info_json, websocket): + removed_node = self.pd_manager.remove_pd(pd_info_json, websocket) + if removed_node is not None: + await self._fail_requests_for_node(removed_node, "PD node disconnected") + return + + async def _fail_requests_for_node(self, node: PD_Client_Obj, reason: str): + error_message = f"{reason}: {node.mode} node {node.client_ip_port}" + affected_requests = [ + req_status + for req_status in list(self.req_id_to_out_inf.values()) + if req_status.p_node is node or req_status.d_node is node + ] + for req_status in affected_requests: + await req_status.fail(RuntimeError(error_message)) return async def update_req_status(self, upkv_status: PDUpKVStatus): @@ -395,6 +410,7 @@ async def fetch_pd_stream( group_request_id=group_request_id, stage="prefill", ) + req_status.raise_if_failed() except ServerBusyError: logger.warning(f"group_request_id: {group_request_id} wait prefill prompt ids time out") raise @@ -415,6 +431,7 @@ async def fetch_pd_stream( group_request_id=group_request_id, stage="decode", ) + req_status.raise_if_failed() except ServerBusyError: logger.warning(f"group_request_id: {group_request_id} wait decode stage time out err, server is busy now.") raise @@ -430,6 +447,7 @@ async def fetch_pd_stream( first_token_gen = False while True: await req_status.wait_to_ready() + req_status.raise_if_failed() if await request.is_disconnected(): raise ClientDisconnected( group_request_id=group_request_id, @@ -636,6 +654,21 @@ def __init__(self, req_id, p_node, d_node) -> None: self.out_token_info_list: List[Tuple[int, str, dict, FinishStatus]] = [] self.p_node: PD_Client_Obj = p_node self.d_node: PD_Client_Obj = d_node + self.error: Optional[BaseException] = None + + async def fail(self, error: BaseException): + async with self.lock: + if self.error is None: + self.error = error + self.event.set() + self.up_status_event.set() + self.prefill_prompt_ids_event.set() + return + + def raise_if_failed(self): + if self.error is not None: + raise self.error + return async def wait_to_ready(self): try: @@ -646,6 +679,7 @@ async def wait_to_ready(self): async def can_read(self, req_id_to_out_inf): async with self.lock: self.event.clear() + self.raise_if_failed() assert self.req_id in req_id_to_out_inf, f"error state req_id {self.req_id}" if len(self.out_token_info_list) == 0: return False @@ -720,6 +754,7 @@ async def check_pd_nodes_health(self): def register_pd(self, pd_info_json, websocket): pd_client = PD_Client_Obj(**pd_info_json) + replaced_node = self.url_to_pd_nodes.get(pd_client.client_ip_port) client_max_req_total_len = pd_client.start_args["max_req_total_len"] if client_max_req_total_len != self.args.max_req_total_len: logger.error( @@ -756,10 +791,16 @@ def register_pd(self, pd_info_json, websocket): self.selector.update_nodes(self.prefill_nodes, self.decode_nodes) logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} registed") - return + if replaced_node is not None and replaced_node.websocket is not websocket: + return replaced_node + return None - def remove_pd(self, pd_info_json): + def remove_pd(self, pd_info_json, websocket): pd_client = PD_Client_Obj(**pd_info_json) + registered_node = self.url_to_pd_nodes.get(pd_client.client_ip_port) + if registered_node is None or registered_node.websocket is not websocket: + logger.info(f"ignore stale disconnect for {pd_client.mode} node {pd_client.client_ip_port}") + return None self.url_to_pd_nodes.pop(pd_client.client_ip_port, None) self.prefill_nodes = [e for e in self.prefill_nodes if e.client_ip_port != pd_client.client_ip_port] @@ -768,7 +809,7 @@ def remove_pd(self, pd_info_json): self.selector.update_nodes(self.prefill_nodes, self.decode_nodes) logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} removed") - return + return registered_node def update_node_load_info(self, load_info: Optional[dict]): """更新节点负载信息