fix(pd): release partial request allocations on abort

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]],
@@ -322,7 +356,6 @@ async def generate(
         # 用于等待 pd_master 下发的交换信息
         pd_event: asyncio.Event = None,
     ) -> AsyncGenerator[Tuple[int, str, dict, FinishStatus], None]:
-
         start_time = time.time()
         request_headers = request.headers if request is not None else {}
         group_request_id = self.alloc_req_id(sampling_params)
@@ -414,28 +447,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,
@@ -663,7 +675,6 @@ async def transfer_to_next_module(
         self,
         group_req_objs: Optional[GroupReqObjs] = None,
     ):
-
         if self.pd_mode.is_P_or_NORMAL():
             if not self.args.disable_vision:
                 self.send_to_visual.send_pyobj(group_req_objs.to_group_req_index(), protocol=pickle.HIGHEST_PROTOCOL)
@@ -706,7 +717,6 @@ async def _wait_to_token_package(
         req_status: "ReqStatus",
         request: Request,
     ):
-
         event = req_status.event
         unfinished_count = sampling_params.best_of
         out_token_counter = 0

lightllm/server/httpserver/pd_loop.py

@@ -34,6 +34,21 @@ 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 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")
@@ -94,7 +109,6 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O
                 # 下方应用层心跳已负责存活检测，禁用协议层 keepalive，避免繁忙连接被误断。
                 ping_interval=None,
             ) as websocket:
-
                 sock = websocket.transport.get_extra_info("socket")
                 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
@@ -146,18 +160,8 @@ def remove_generation_task(task: asyncio.Task, request_id: int = group_req_id):
                     elif obj[0] == ObjType.ABORT:
                         group_req_id = obj[1]
                         logger.warning(f"recv cmd aborted req id {group_req_id}")
-                        generation_task = generation_tasks.get(group_req_id)
-                        if generation_task is not None and not generation_task.done():
-                            generation_task.cancel()
-                        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
@@ -182,7 +186,8 @@ async def delayed_abort_task(group_req_id, retry_count):
             child_tasks = [task for task in (forwarding_tokens_task, heartbeat_task) if task is not None]
             child_tasks.extend(generation_tasks.values())
             for task in child_tasks:
-                task.cancel()
+                if not task.done() and not task.cancelling():
+                    task.cancel()
             if child_tasks:
                 await asyncio.gather(*child_tasks, return_exceptions=True)
 
@@ -247,6 +252,9 @@ 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 asyncio.CancelledError:
+        logger.info(f"pd request task cancelled for group_request_id {sampling_params.group_request_id}")
+        raise
     except BaseException as e:
         logger.error(str(e))
 
@@ -270,7 +278,6 @@ async def _send_heartbeat_to_pd_master(websocket: ClientConnection):
 
 # 获取节点负载信息
 def _get_load_info() -> dict:
-
     from lightllm.server.api_http import g_objs
 
     assert g_objs.shared_token_load is not None, "shared_token_load is not initialized"

unit_tests/server/httpserver/test_pd_pending_abort.py

@@ -0,0 +1,70 @@
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from lightllm.server.httpserver.manager import HttpServerManager
+from lightllm.server.httpserver.pd_loop import _abort_pd_request
+
+
+def test_pd_abort_cancels_request_before_http_manager_registration():
+    async def run():
+        manager = MagicMock()
+        manager.abort = AsyncMock(return_value=False)
+        request_started = asyncio.Event()
+
+        async def wait_for_request_slot():
+            request_started.set()
+            await asyncio.Future()
+
+        generation_task = asyncio.create_task(wait_for_request_slot())
+        generation_tasks = {123: generation_task}
+        await request_started.wait()
+
+        assert await _abort_pd_request(manager, 123, generation_tasks)
+        assert generation_task.cancelling() == 1
+        assert await _abort_pd_request(manager, 123, generation_tasks)
+        assert generation_task.cancelling() == 1
+        with pytest.raises(asyncio.CancelledError):
+            await generation_task
+
+        assert manager.abort.await_count == 2
+        manager.abort.assert_awaited_with(123)
+        assert generation_task.cancelled()
+
+    asyncio.run(run())
+
+
+def test_cancelled_slot_allocation_releases_partially_allocated_indexes():
+    async def run():
+        manager = HttpServerManager.__new__(HttpServerManager)
+        manager.shm_req_manager = MagicMock()
+        manager.shm_req_manager.async_release_req_index = AsyncMock()
+        manager.shm_req_manager.async_put_back_req_obj = AsyncMock()
+        manager.tokenizer = MagicMock()
+        manager.args = MagicMock(chunked_prefill_size=16)
+
+        waiting_for_second_slot = asyncio.Event()
+        allocation_count = 0
+
+        async def alloc_req_index():
+            nonlocal allocation_count
+            allocation_count += 1
+            if allocation_count == 1:
+                return 7
+            waiting_for_second_slot.set()
+            return None
+
+        manager.shm_req_manager.async_alloc_req_index = alloc_req_index
+        sampling_params = MagicMock(n=2)
+        allocation_task = asyncio.create_task(manager._alloc_req_objs(123, [1, 2], sampling_params))
+        await waiting_for_second_slot.wait()
+        allocation_task.cancel()
+
+        with pytest.raises(asyncio.CancelledError):
+            await allocation_task
+
+        manager.shm_req_manager.async_release_req_index.assert_awaited_once_with(7)
+        manager.shm_req_manager.async_put_back_req_obj.assert_not_awaited()
+
+    asyncio.run(run())