diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 43b06f02ef..49b6134509 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -33,6 +33,7 @@ from lightllm.server.core.objs.atomic_array_lock import AtomicShmArrayLock, AsyncLock, AtomicLockItem from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from lightllm.utils.log_utils import init_logger +from lightllm.utils.pd_timeout_utils import get_pd_handshake_timeout_seconds from lightllm.server.metrics.manager import MetricClient from .rl_controller import HttpRlController from .manager_ext import HttpRlManagerHelper @@ -422,10 +423,15 @@ async def _generate( await pd_upload_websocket.send( pickle.dumps((ObjType.PD_UPLOAD_PREFILL_PROMPT_IDS, group_request_id, array("q", prompt_ids))) ) + handshake_timeout = get_pd_handshake_timeout_seconds( + resource_wait_seconds=sampling_params.pd_node_resource_wait_timeout_seconds, prefill=True + ) try: - await asyncio.wait_for(pd_event.wait(), timeout=180) + await asyncio.wait_for(pd_event.wait(), timeout=handshake_timeout) except asyncio.TimeoutError: - logger.error(f"pd prefill node wait pd_event 180s time out, group_req_id {group_request_id}") + logger.error( + f"pd prefill node wait pd_event {handshake_timeout}s time out, group_req_id {group_request_id}" + ) raise Exception(f"group_req_id {group_request_id} wait pd_event time out") decode_node_info: PDDecodeNodeInfo = pd_event.decode_node_info diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index a9a438e11d..5d5daf0556 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -34,6 +34,7 @@ get_pd_node_resource_wait_timeout_seconds, ) from lightllm.utils.shm_port_args import get_shm_port_args +from lightllm.utils.pd_timeout_utils import get_pd_handshake_timeout_seconds from .pd_selector import PDSelectionExtraInfo, create_selector logger = init_logger(__name__) @@ -524,12 +525,12 @@ async def _wait_for_event_or_disconnect( self, event: asyncio.Event, request: Request, - timeout: float, + timeout: Optional[float], group_request_id: int, stage: str, ) -> None: """Wait for an asyncio.Event but abort early if the HTTP client disconnects.""" - deadline = time.time() + timeout + deadline = time.monotonic() + timeout if timeout is not None else None disconnect_reason = f"fetch_pd_stream {stage} period check network disconnected" async def raise_if_disconnected() -> None: @@ -541,7 +542,7 @@ async def raise_if_disconnected() -> None: ) while True: - remaining = deadline - time.time() + remaining = deadline - time.monotonic() if deadline is not None else 1.0 if remaining <= 0: raise ServerBusyError() await raise_if_disconnected() @@ -610,7 +611,9 @@ async def fetch_pd_stream( await self._wait_for_event_or_disconnect( up_status_event, request, - timeout=180, + timeout=get_pd_handshake_timeout_seconds( + resource_wait_seconds=sampling_params.pd_node_resource_wait_timeout_seconds + ), group_request_id=group_request_id, stage="decode", ) diff --git a/lightllm/utils/pd_timeout_utils.py b/lightllm/utils/pd_timeout_utils.py new file mode 100644 index 0000000000..dee29456aa --- /dev/null +++ b/lightllm/utils/pd_timeout_utils.py @@ -0,0 +1,12 @@ +from typing import Optional + + +def get_pd_handshake_timeout_seconds(resource_wait_seconds: int, prefill: bool = False) -> Optional[int]: + """Let node resource timeouts expire before their enclosing PD handshake.""" + if resource_wait_seconds < 0: + return None + # Decode can wait for a shared request slot, then independently in the router. + # Allow its five-second polling loop and error forwarding to finish first. + timeout = max(180, 2 * resource_wait_seconds + 15) + # Prefill waits before Master sends to Decode; leave time for Master's abort. + return timeout + 15 if prefill else timeout diff --git a/unit_tests/server/httpserver/test_pd_handshake_budget.py b/unit_tests/server/httpserver/test_pd_handshake_budget.py new file mode 100644 index 0000000000..cc2fbb87bf --- /dev/null +++ b/unit_tests/server/httpserver/test_pd_handshake_budget.py @@ -0,0 +1,114 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from lightllm.server.core.objs import SamplingParams +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, ReqStatus +from lightllm.server.multimodal_params import MultimodalParams +from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError +from lightllm.utils.pd_timeout_utils import get_pd_handshake_timeout_seconds + + +@pytest.mark.parametrize( + "resource_wait,master,prefill", [(-1, None, None), (0, 180, 195), (60, 180, 195), (200, 415, 430), (450, 915, 930)] +) +def test_handshake_encloses_both_resource_waits(resource_wait: int, master: int | None, prefill: int | None) -> None: + assert get_pd_handshake_timeout_seconds(resource_wait_seconds=resource_wait) == master + assert get_pd_handshake_timeout_seconds(resource_wait_seconds=resource_wait, prefill=True) == prefill + + +@pytest.mark.parametrize("resource_wait,expected", [(60, 180), (200, 415), (-1, None)]) +def test_master_uses_the_current_segments_resource_budget(resource_wait: int, expected: int | None) -> None: + async def run() -> None: + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.args = SimpleNamespace(pd_node_id=1) + manager.req_id_to_out_inf = {} + params = SamplingParams() + params.group_request_id, params.max_new_tokens = 123, 4 + params.pd_node_resource_wait_timeout_seconds = resource_wait + + async def handshake( + event: asyncio.Event, request: object, timeout: float | None, group_request_id: int, stage: str + ) -> None: + if stage == "prefill": + assert timeout == 60 + event.prompt_ids = [1, 2] + else: + assert timeout == expected + raise RuntimeError("decode budget checked") + + manager._wait_for_event_or_disconnect = handshake + node = SimpleNamespace(websocket=SimpleNamespace(send_bytes=AsyncMock())) + stream = manager.fetch_pd_stream( + p_node=node, + d_node=node, + prompt="prompt", + sampling_params=params, + multimodal_params=MultimodalParams(), + request=SimpleNamespace(is_disconnected=AsyncMock(return_value=False)), + ) + try: + with pytest.raises(expected_exception=RuntimeError, match="decode budget checked"): + await anext(stream) + finally: + await stream.aclose() + + asyncio.run(main=run()) + + +@pytest.mark.parametrize("timeout", [None, 415]) +def test_node_error_wakes_handshake_without_retry(timeout: float | None) -> None: + async def run() -> None: + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + status = ReqStatus(req_id=123, p_node=None, d_node=None) + request = SimpleNamespace(is_disconnected=AsyncMock(return_value=False)) + task = asyncio.create_task( + coro=manager._wait_for_event_or_disconnect( + event=status.up_status_event, request=request, timeout=timeout, group_request_id=123, stage="decode" + ) + ) + await status.set_error(error_info="allocation failed") + await asyncio.wait_for(fut=task, timeout=1) + with pytest.raises(expected_exception=RuntimeError, match="allocation failed"): + status.raise_if_error() + + asyncio.run(main=run()) + + +def test_unbounded_handshake_still_checks_disconnect_and_cancellation() -> None: + async def run() -> None: + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + request = SimpleNamespace(is_disconnected=AsyncMock(return_value=True)) + with pytest.raises(expected_exception=ClientDisconnected): + await manager._wait_for_event_or_disconnect( + event=asyncio.Event(), request=request, timeout=None, group_request_id=123, stage="decode" + ) + request.is_disconnected.return_value = False + task = asyncio.create_task( + coro=manager._wait_for_event_or_disconnect( + event=asyncio.Event(), request=request, timeout=None, group_request_id=123, stage="decode" + ) + ) + await asyncio.sleep(delay=0) + task.cancel() + with pytest.raises(expected_exception=asyncio.CancelledError): + await task + + asyncio.run(main=run()) + + +def test_finite_handshake_still_expires() -> None: + async def run() -> None: + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + with pytest.raises(expected_exception=ServerBusyError): + await manager._wait_for_event_or_disconnect( + event=asyncio.Event(), + request=SimpleNamespace(is_disconnected=AsyncMock(return_value=False)), + timeout=0.001, + group_request_id=123, + stage="decode", + ) + + asyncio.run(main=run()) diff --git a/unit_tests/server/httpserver/test_running_request_lifecycle.py b/unit_tests/server/httpserver/test_running_request_lifecycle.py index 0487e45281..e999fb4610 100644 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -271,6 +271,30 @@ async def run(): asyncio.run(run()) +@pytest.mark.parametrize("resource_wait,expected", [(60, 195), (200, 430), (-1, None)]) +def test_prefill_assignment_wait_encloses_master_budget(resource_wait: int, expected: int | None) -> None: + async def run() -> None: + manager = _make_manager(mode=NodeRole.P) + params = _sampling_params() + params.pd_node_resource_wait_timeout_seconds = resource_wait + event = asyncio.Event() + event.decode_node_info = SimpleNamespace(ready_kv_len=2) + event.set() + with patch(target="lightllm.server.httpserver.manager.asyncio.wait_for", wraps=asyncio.wait_for) as wait: + with pytest.raises(expected_exception=PDPrefillNodeStopGenToken): + await _drain_generate( + manager=manager, + sampling_params=params, + multimodal_params=_multimodal_params(), + websocket=AsyncMock(), + pd_event=event, + ) + assert wait.call_args.kwargs["timeout"] == expected + manager._register_running_request.assert_not_awaited() + + asyncio.run(main=run()) + + def test_prefill_upload_preserves_int64_multimodal_prompt_ids(): async def run(): manager = _make_manager(NodeRole.P)