fix(pd): harden production cutover lifecycles

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
 
 

lightllm/server/httpserver/async_queue.py

@@ -7,9 +7,9 @@ def __init__(self):
         self.event = asyncio.Event()
         self.lock = asyncio.Lock()
 
-    async def wait_to_ready(self):
+    async def wait_to_ready(self, timeout=3):
         try:
-            await asyncio.wait_for(self.event.wait(), timeout=3)
+            await asyncio.wait_for(self.event.wait(), timeout=timeout)
         except asyncio.TimeoutError:
             pass
 
@@ -26,7 +26,7 @@ async def put(self, obj):
             self.event.set()
         return
 
-    async def wait_to_get_all_data(self):
-        await self.wait_to_ready()
+    async def wait_to_get_all_data(self, timeout=3):
+        await self.wait_to_ready(timeout=timeout)
         handle_list = await self.get_all_data()
         return handle_list

lightllm/server/httpserver/manager.py

@@ -311,6 +311,40 @@ def alloc_req_id(self, sampling_params):
             assert False, "dead code path"
         return group_request_id
 
+    async def _alloc_req_objs(self, group_request_id, prompt_ids, sampling_params) -> List[Req]:
+        """Allocate and initialize request slots without leaking partial allocations."""
+        alloced_req_indexes = []
+        req_objs = []
+        try:
+            while len(alloced_req_indexes) < sampling_params.n:
+                alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
+                sleep_time = 0.1
+                while alloc_req_index is None:
+                    await asyncio.sleep(sleep_time)
+                    sleep_time *= 1.1
+                    sleep_time = min(1, sleep_time)
+
+                    alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
+                alloced_req_indexes.append(alloc_req_index)
+
+            for i, req_index in enumerate(alloced_req_indexes):
+                req_obj = await self.shm_req_manager.async_get_req_obj_by_index(req_index)
+                req_objs.append(req_obj)
+                req_obj.init(
+                    group_request_id + i,
+                    prompt_ids,
+                    sampling_params,
+                    self.tokenizer,
+                    chunked_prefill_size=self.args.chunked_prefill_size,
+                )
+            return req_objs
+        except BaseException:
+            for req_obj in req_objs:
+                await self.shm_req_manager.async_put_back_req_obj(req_obj)
+            for req_index in alloced_req_indexes:
+                await self.shm_req_manager.async_release_req_index(req_index)
+            raise
+
     async def generate(
         self,
         prompt: Union[str, List[int]],
@@ -414,28 +448,7 @@ async def generate(
                     raise PDPrefillNodeStopGenToken(group_request_id=group_request_id)
 
             # 申请资源并存储
-            alloced_req_indexes = []
-            while len(alloced_req_indexes) < sampling_params.n:
-                alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
-                sleep_time = 0.1
-                while alloc_req_index is None:
-                    await asyncio.sleep(sleep_time)
-                    sleep_time *= 1.1
-                    sleep_time = min(1, sleep_time)
-
-                    alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
-                alloced_req_indexes.append(alloc_req_index)
-            req_objs: List[Req] = []
-            for i, req_index in enumerate(alloced_req_indexes):
-                req_obj = await self.shm_req_manager.async_get_req_obj_by_index(req_index)
-                req_obj.init(
-                    group_request_id + i,
-                    prompt_ids,
-                    sampling_params,
-                    self.tokenizer,
-                    chunked_prefill_size=self.args.chunked_prefill_size,
-                )
-                req_objs.append(req_obj)
+            req_objs = await self._alloc_req_objs(group_request_id, prompt_ids, sampling_params)
             self._log_stage_timing(
                 group_request_id,
                 start_time,

lightllm/server/httpserver/pd_loop.py

@@ -34,6 +34,48 @@ async def timer_log(manager: HttpServerManager):
     return
 
 
+async def _abort_pd_request(
+    manager: HttpServerManager,
+    group_req_id: int,
+    generation_tasks: Dict[int, asyncio.Task],
+) -> bool:
+    generation_task = generation_tasks.get(group_req_id)
+    if generation_task is not None and not generation_task.done() and not generation_task.cancelling():
+        # The request may still be waiting for a shared-memory slot and therefore
+        # not yet be visible to HttpServerManager.abort(). Cancelling its owning
+        # task makes the abort effective at every point in the admission path.
+        generation_task.cancel()
+
+    return (await manager.abort(group_req_id)) or generation_task is not None
+
+
+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 +126,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 +157,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,18 +183,22 @@ 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 remove_generation_task(done_task, request_id=group_req_id):
+                            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 connection_failure.done():
+                                    connection_failure.set_exception(error)
+
+                        generation_task.add_done_callback(remove_generation_task)
                     elif obj[0] == ObjType.ABORT:
                         group_req_id = obj[1]
                         logger.warning(f"recv cmd aborted req id {group_req_id}")
-                        if not (await manager.abort(group_req_id)):
-
-                            async def delayed_abort_task(group_req_id, retry_count):
-                                for _ in range(retry_count):
-                                    await asyncio.sleep(5.0)
-                                    if await manager.abort(group_req_id):
-                                        break
-
-                            asyncio.create_task(delayed_abort_task(group_req_id=group_req_id, retry_count=4))
+                        group_req_id_to_event.pop(group_req_id, None)
+                        await _abort_pd_request(manager, group_req_id, generation_tasks)
 
                     elif obj[0] == ObjType.PD_REQ_DECODE_NODE_INFO:
                         _, group_req_id, decode_node_info = obj
@@ -169,7 +221,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 +294,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