diff --git a/lightllm/server/httpserver/pd_loop.py b/lightllm/server/httpserver/pd_loop.py index a69364c95..13dc545a5 100644 --- a/lightllm/server/httpserver/pd_loop.py +++ b/lightllm/server/httpserver/pd_loop.py @@ -34,6 +34,33 @@ async def timer_log(manager: HttpServerManager): return +async def _recv_or_raise_on_background_failure(websocket, background_tasks): + recv_task = asyncio.create_task(websocket.recv()) + try: + done, _ = await asyncio.wait( + [recv_task, *background_tasks], + return_when=asyncio.FIRST_COMPLETED, + ) + for task in background_tasks: + if task not in done: + continue + + recv_task.cancel() + await asyncio.gather(recv_task, return_exceptions=True) + if task.cancelled(): + raise asyncio.CancelledError + error = task.exception() + if error is not None: + raise error + raise RuntimeError("PD connection background task exited unexpectedly") + + return recv_task.result() + finally: + if not recv_task.done(): + recv_task.cancel() + await asyncio.gather(recv_task, return_exceptions=True) + + async def pd_handle_loop(manager: HttpServerManager): if manager.args.host in ["127.0.0.1", "localhost"]: logger.error("pd mode must specify host ip, not use 127.0.0.1 or localhost") @@ -84,6 +111,8 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O while True: forwarding_tokens_task = None heartbeat_task = None + connection_failure = None + generation_tasks: Dict[int, asyncio.Task] = {} try: uri = f"ws://{pd_master_obj.host_ip_port}/pd_register" async with websockets.connect( @@ -113,18 +142,22 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O # 转发任务 forwarding_tokens_task = asyncio.create_task(_up_tokens_to_pd_master(forwarding_queue, websocket)) heartbeat_task = asyncio.create_task(_send_heartbeat_to_pd_master(websocket)) + connection_failure = asyncio.get_running_loop().create_future() group_req_id_to_event: Dict[int, asyncio.Event] = weakref.WeakValueDictionary() # 接收 pd master 发来的请求，并推理后，将生成的token转发回pd master。 while True: - recv_bytes = await websocket.recv() + recv_bytes = await _recv_or_raise_on_background_failure( + websocket, + (forwarding_tokens_task, heartbeat_task, connection_failure), + ) obj = pickle.loads(recv_bytes) if obj[0] == ObjType.REQ: prompt, sampling_params, multimodal_params = obj[1] group_req_id = sampling_params.group_request_id pd_event = asyncio.Event() group_req_id_to_event[group_req_id] = pd_event - asyncio.create_task( + generation_task = asyncio.create_task( _pd_process_generate( manager=manager, prompt=prompt, @@ -135,6 +168,21 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O pd_event=pd_event, ) ) + generation_tasks[group_req_id] = generation_task + + def report_generation_failure( + done_task, + request_id=group_req_id, + failure_future=connection_failure, + ): + if generation_tasks.get(request_id) is done_task: + generation_tasks.pop(request_id, None) + if not done_task.cancelled(): + error = done_task.exception() + if error is not None and not failure_future.done(): + failure_future.set_exception(error) + + generation_task.add_done_callback(report_generation_failure) elif obj[0] == ObjType.ABORT: group_req_id = obj[1] logger.warning(f"recv cmd aborted req id {group_req_id}") @@ -169,7 +217,14 @@ async def delayed_abort_task(group_req_id, retry_count): logger.exception(str(e)) finally: child_tasks = [task for task in (forwarding_tokens_task, heartbeat_task) if task is not None] + child_tasks.extend(generation_tasks.values()) + if connection_failure is not None: + child_tasks.append(connection_failure) for task in child_tasks: + if task.done(): + continue + if isinstance(task, asyncio.Task) and task.cancelling(): + continue task.cancel() if child_tasks: await asyncio.gather(*child_tasks, return_exceptions=True) @@ -235,8 +290,12 @@ async def _pd_process_generate( await forwarding_queue.put((sub_req_id, request_output, metadata, finish_status)) except PDPrefillNodeStopGenToken as e: logger.info(f"pd prefill node stop gen token for group_request_id {e.group_request_id}") - except BaseException as e: - logger.error(str(e)) + except asyncio.CancelledError: + logger.info(f"pd request task cancelled for group_request_id {sampling_params.group_request_id}") + raise + except BaseException: + logger.exception(f"PD request generation failed for group_request_id {sampling_params.group_request_id}") + raise # 转发token的task diff --git a/unit_tests/server/httpserver/test_pd_connection_tasks.py b/unit_tests/server/httpserver/test_pd_connection_tasks.py new file mode 100644 index 000000000..eaf2a8e08 --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_connection_tasks.py @@ -0,0 +1,61 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from lightllm.server.httpserver.async_queue import AsyncQueue +from lightllm.server.httpserver.pd_loop import ( + _pd_process_generate, + _recv_or_raise_on_background_failure, +) + + +def test_background_failure_interrupts_blocked_pd_receive(): + async def run(): + recv_started = asyncio.Event() + recv_cancelled = asyncio.Event() + + class BlockingWebsocket: + async def recv(self): + recv_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + recv_cancelled.set() + raise + + websocket = BlockingWebsocket() + failure = asyncio.get_running_loop().create_future() + receive_task = asyncio.create_task(_recv_or_raise_on_background_failure(websocket, (failure,))) + await recv_started.wait() + + failure.set_exception(RuntimeError("generation failed")) + with pytest.raises(RuntimeError, match="generation failed"): + await receive_task + assert recv_cancelled.is_set() + + asyncio.run(run()) + + +def test_pd_generation_failure_is_not_swallowed(): + async def run(): + class FailingManager: + args = SimpleNamespace(run_mode="prefill") + + async def generate(self, **_kwargs): + yield 1, "token", {}, None + raise RuntimeError("generation failed") + + sampling_params = SimpleNamespace(group_request_id=123) + with pytest.raises(RuntimeError, match="generation failed"): + await _pd_process_generate( + manager=FailingManager(), + prompt="prompt", + sampling_params=sampling_params, + multimodal_params={}, + forwarding_queue=AsyncQueue(), + pd_upload_websocket=object(), + pd_event=asyncio.Event(), + ) + + asyncio.run(run())