diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 195ef17e4..6ee841009 100644 --- a/lightllm/server/httpserver/manager.py +++ b/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 diff --git a/lightllm/server/httpserver/pd_loop.py b/lightllm/server/httpserver/pd_loop.py index ba843db0b..6bb24c7c0 100644 --- a/lightllm/server/httpserver/pd_loop.py +++ b/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" diff --git a/unit_tests/server/httpserver/test_pd_pending_abort.py b/unit_tests/server/httpserver/test_pd_pending_abort.py new file mode 100644 index 000000000..0ece4d4d8 --- /dev/null +++ b/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())