fix(pd): align handshake deadlines with resource wait budgets

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

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",
             )

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

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())

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)