[None][fix] Disaggregated client: bound streamed responses by silence, not total time

tensorrt_llm/commands/serve.py

@@ -2364,7 +2364,9 @@ def serve_embedding(
                   "--request_timeout",
                   type=int,
                   default=180,
-                  help="Request timeout",
+                  help="Request timeout in seconds. A streamed response "
+                  "from a worker may run longer: it times out only after "
+                  "this many seconds without data.",
                   status="beta")
 @stability_option("-l",
                   '--log_level',

tensorrt_llm/serve/openai_client.py

@@ -19,7 +19,8 @@
 import os
 import traceback
 from abc import ABC, abstractmethod
-from contextlib import aclosing
+from contextlib import AbstractAsyncContextManager, aclosing
+from types import TracebackType
 from typing import (
     Any,
     AsyncGenerator,
@@ -244,6 +245,34 @@ def _server_netloc(server_or_url: str) -> str:
     return urlsplit(_server_base_url(server_or_url)).netloc
 
 
+class _ResponseStartBound:
+    """Wrap a worker request so that its response must start within a time limit.
+
+    Entering the request covers connecting, sending the request and receiving
+    the response headers, and only entering is bounded; exiting is delegated
+    unchanged.
+    """
+
+    def __init__(
+        self,
+        request: AbstractAsyncContextManager[aiohttp.ClientResponse],
+        timeout_secs: Optional[float],
+    ) -> None:
+        self._request = request
+        self._timeout_secs = timeout_secs
+
+    async def __aenter__(self) -> aiohttp.ClientResponse:
+        return await asyncio.wait_for(self._request.__aenter__(), self._timeout_secs)
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> Optional[bool]:
+        return await self._request.__aexit__(exc_type, exc, tb)
+
+
 class OpenAIHttpClient(OpenAIClient):
     def __init__(
         self,
@@ -262,6 +291,20 @@ def __init__(
         self._router = router
         self._role = role
         self._metrics_collector = ClientMetricsCollector(role)
+        # A non-streamed response arrives whole, so timeout_secs bounds the whole
+        # request. A streamed response lasts as long as the worker keeps
+        # generating, so timeout_secs bounds its silence instead: the wait for the
+        # response to start, which includes connecting and sending the request,
+        # and every wait for data after that. A total bound would cut a long
+        # generation that is still producing tokens.
+        self._non_stream_timeout = aiohttp.ClientTimeout(total=timeout_secs)
+        # Non-positive values disable every stream bound. A negative sock_read
+        # value would otherwise schedule an immediate timeout in aiohttp.
+        stream_timeout_secs = timeout_secs if timeout_secs > 0 else None
+        self._stream_timeout = aiohttp.ClientTimeout(
+            total=None, connect=stream_timeout_secs, sock_read=stream_timeout_secs
+        )
+        self._response_start_timeout_secs = stream_timeout_secs
         self._session = session or aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(
                 limit=0,
@@ -380,6 +423,24 @@ async def get_json(
                 )
             return response_type(**await response.json())
 
+    def _post(
+        self,
+        url: str,
+        data: bytes,
+        headers: dict[str, str],
+        timeout: aiohttp.ClientTimeout,
+    ) -> _ResponseStartBound:
+        """POST to a worker, failing if its response has not started within timeout_secs.
+
+        aiohttp starts its sock_read bound only once the request is fully written,
+        so for a streamed response, which has no total bound, this is the bound
+        on a worker that stops reading a large request.
+        """
+        return _ResponseStartBound(
+            self._session.post(url, data=data, headers=headers, timeout=timeout),
+            self._response_start_timeout_secs,
+        )
+
     async def _send_request(
         self,
         endpoint: str,
@@ -555,10 +616,11 @@ async def _post_with_retry(
             try:
                 lines_yielded = 0
                 start_time = get_steady_clock_now_in_seconds()
-                async with self._session.post(
+                async with self._post(
                     url,
                     data=body,
                     headers=headers,
+                    timeout=self._stream_timeout if is_stream else self._non_stream_timeout,
                 ) as http_response:
                     response_started = True
                     response_clock = AdjustedSteadyClock()

tests/unittest/disaggregated/test_disagg_openai_client.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import contextlib
 import errno
 import json
 import socket
@@ -25,6 +26,8 @@
 import aiohttp
 import msgspec
 import pytest
+from aiohttp import web
+from aiohttp.test_utils import TestServer
 
 from tensorrt_llm._utils import AdjustedSteadyClock
 from tensorrt_llm.disaggregated_params import DisaggScheduleStyle
@@ -1342,6 +1345,233 @@ async def test_transient_tcp_capped_at_5_when_max_retries_smaller(self, error):
         assert session.post.call_count == 6
 
 
+class TestRequestTimeout:
+    """What timeout_secs (--request_timeout) bounds on a worker's response.
+
+    A streamed response is bounded by silence, so a long generation completes
+    while data keeps arriving. A non-streamed response is bounded by its total
+    time. The bounds are enforced by aiohttp's transport, so these tests run a
+    real worker server and a client that builds its own session, as the
+    disaggregated server does.
+    """
+
+    _TIMEOUT_SECS = 1
+    # Eight chunks a quarter of the timeout apart: twice the timeout in total,
+    # never silent for more than a quarter of it.
+    _CHUNKS = 8
+    _INTERVAL_SECS = 0.25
+
+    @staticmethod
+    def _request(stream: bool, prompt: str = "hi") -> CompletionRequest:
+        return CompletionRequest(
+            model="m",
+            prompt=prompt,
+            stream=stream,
+            disaggregated_params=DisaggregatedParams(
+                request_type="generation_only", first_gen_tokens=[7], ctx_request_id=1
+            ),
+        )
+
+    @staticmethod
+    async def _open_response(request: web.Request, content_type: str) -> web.StreamResponse:
+        response = web.StreamResponse(headers={"Content-Type": content_type})
+        await response.prepare(request)
+        return response
+
+    @contextlib.asynccontextmanager
+    async def _worker(self, handler, timeout_secs=None):
+        """Serve handler as a worker's /v1/completions; yield a client of it and its router."""
+        app = web.Application()
+        app.router.add_post("/v1/completions", handler)
+        async with TestServer(app) as server:
+            _reset_prometheus_registry()
+            router = AsyncMock(spec=Router)
+            router.get_next_server = AsyncMock(return_value=(f"{server.host}:{server.port}", None))
+            client = OpenAIHttpClient(
+                router=router,
+                role=ServerRole.GENERATION,
+                timeout_secs=self._TIMEOUT_SECS if timeout_secs is None else timeout_secs,
+                max_retries=0,
+            )
+            try:
+                yield client, router
+            finally:
+                await client.shutdown()
+
+    @pytest.mark.asyncio
+    async def test_a_stream_longer_than_the_timeout_completes_while_data_keeps_arriving(self):
+        expected = b"".join(f"data: {index}\n\n".encode() for index in range(self._CHUNKS))
+
+        async def generate(request):
+            response = await self._open_response(request, "text/event-stream")
+            for index in range(self._CHUNKS):
+                await asyncio.sleep(self._INTERVAL_SECS)
+                await response.write(f"data: {index}\n\n".encode())
+            return response
+
+        loop = asyncio.get_running_loop()
+        async with self._worker(generate) as (client, router):
+            started = loop.time()
+            response = await client.send_request(self._request(stream=True))
+            received = b"".join([chunk async for chunk in response])
+            elapsed = loop.time() - started
+
+        assert elapsed > self._TIMEOUT_SECS, "the stream must outlast the timeout to test it"
+        assert received == expected
+        router.finish_request.assert_awaited_once()
+        assert router.finish_request.await_args.kwargs["success"] is True
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout_secs", [0, -1])
+    async def test_nonpositive_timeout_allows_delayed_headers_and_stream_data(self, timeout_secs):
+        async def generate(request):
+            await asyncio.sleep(0.05)
+            response = await self._open_response(request, "text/event-stream")
+            await asyncio.sleep(0.05)
+            await response.write(b"data: hello\n\n")
+            await asyncio.sleep(0.05)
+            await response.write(b"data: [DONE]\n\n")
+            return response
+
+        async with self._worker(generate, timeout_secs=timeout_secs) as (client, router):
+            with patch.object(client._session, "post", wraps=client._session.post) as post:
+                response = await client.send_request(self._request(stream=True))
+                received = b"".join([chunk async for chunk in response])
+
+            post.assert_called_once()
+            timeout = post.call_args.kwargs["timeout"]
+            assert timeout.total is None
+            assert timeout.connect is None
+            assert timeout.sock_read is None
+
+        assert received == b"data: hello\n\ndata: [DONE]\n\n"
+        router.finish_request.assert_awaited_once()
+        assert router.finish_request.await_args.kwargs["success"] is True
+
+    @pytest.mark.asyncio
+    async def test_a_stream_silent_for_longer_than_the_timeout_fails(self):
+        async def stall(request):
+            response = await self._open_response(request, "text/event-stream")
+            await response.write(b'data: "first"\n\n')
+            # Far longer than the timeout; the stream would end cleanly here if
+            # nothing bounded the silence.
+            await asyncio.sleep(self._TIMEOUT_SECS * 10)
+            await response.write(b"data: [DONE]\n\n")
+            return response
+
+        received = []
+        async with self._worker(stall) as (client, router):
+            response = await client.send_request(self._request(stream=True))
+            with pytest.raises(asyncio.TimeoutError):
+                async for chunk in response:
+                    received.append(chunk)
+
+        assert b"".join(received) == b'data: "first"\n\n'
+        router.finish_request.assert_awaited_once()
+        assert router.finish_request.await_args.kwargs["success"] is False
+
+    @pytest.mark.asyncio
+    async def test_a_stream_whose_request_the_worker_does_not_read_fails(self):
+        client_gave_up = asyncio.Event()
+
+        async def stall_before_reading(request):
+            # Far longer than the timeout unless the client gives up first; the
+            # stream would end cleanly here if nothing bounded sending the request.
+            with contextlib.suppress(asyncio.TimeoutError):
+                await asyncio.wait_for(client_gave_up.wait(), self._TIMEOUT_SECS * 10)
+            response = await self._open_response(request, "text/event-stream")
+            await response.write(b"data: [DONE]\n\n")
+            return response
+
+        # Far more than the socket buffers between client and worker hold, so
+        # sending it waits for the worker to read it.
+        request = self._request(stream=True, prompt="x" * (32 << 20))
+        async with self._worker(stall_before_reading) as (client, router):
+            response = await client.send_request(request)
+            with pytest.raises(asyncio.TimeoutError):
+                _ = [chunk async for chunk in response]
+            # A worker that is not reading cannot see the client leave, so
+            # release it rather than have the server's shutdown wait out the stall.
+            client_gave_up.set()
+
+        router.finish_request.assert_awaited_once()
+        assert router.finish_request.await_args.kwargs["success"] is False
+
+    @pytest.mark.asyncio
+    async def test_a_non_streamed_response_is_bounded_by_its_total_time(self):
+        body = json.dumps(TestOpenAIHttpClient().dummy_response().model_dump()).encode()
+        piece = -(-len(body) // self._CHUNKS)
+
+        async def trickle(request):
+            response = await self._open_response(request, "application/json")
+            for start in range(0, len(body), piece):
+                await asyncio.sleep(self._INTERVAL_SECS)
+                await response.write(body[start : start + piece])
+            return response
+
+        async with self._worker(trickle) as (client, router):
+            with pytest.raises(asyncio.TimeoutError):
+                await client.send_request(self._request(stream=False))
+
+        router.finish_request.assert_awaited_once()
+        assert router.finish_request.await_args.kwargs["success"] is False
+
+    @pytest.mark.asyncio
+    async def test_a_streamed_request_still_bounds_connecting(
+        self, openai_client, streaming_completion_request, mock_session
+    ):
+        http_response = AsyncMock()
+        http_response.status = 200
+        http_response.headers = {"Content-Type": "text/event-stream"}
+
+        async def chunks():
+            yield b"data: [DONE]\n\n"
+
+        http_response.content.iter_any = chunks
+        http_response.__aenter__.return_value = http_response
+        http_response.__aexit__.return_value = False
+        mock_session.post.return_value = http_response
+
+        response = await openai_client.send_request(streaming_completion_request)
+        _ = [chunk async for chunk in response]
+
+        timeout = mock_session.post.call_args.kwargs["timeout"]
+        assert (timeout.total, timeout.connect, timeout.sock_read) == (None, 180, 180)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("timeout_secs", [0, -1])
+    async def test_a_timeout_that_is_not_positive_does_not_bound_the_response_start(
+        self, mock_router, mock_session, completion_request, monkeypatch, timeout_secs
+    ):
+        bounds = []
+        wait_for = asyncio.wait_for
+
+        async def recording_wait_for(awaitable, timeout):
+            bounds.append(timeout)
+            return await wait_for(awaitable, timeout)
+
+        monkeypatch.setattr(asyncio, "wait_for", recording_wait_for)
+        _reset_prometheus_registry()
+        client = OpenAIHttpClient(
+            router=mock_router,
+            role=ServerRole.GENERATION,
+            timeout_secs=timeout_secs,
+            session=mock_session,
+        )
+        http_response = AsyncMock()
+        http_response.status = 200
+        http_response.headers = {"Content-Type": "application/json"}
+        http_response.json.return_value = TestOpenAIHttpClient().dummy_response().model_dump()
+        http_response.__aenter__.return_value = http_response
+        http_response.__aexit__.return_value = False
+        mock_session.post.return_value = http_response
+
+        response = await client.send_request(completion_request)
+
+        assert isinstance(response, CompletionResponse)
+        assert bounds == [None]
+
+
 def test_a_forwarded_request_keeps_the_field_names_it_arrived_with():
     """`schema` must not reach a worker spelled `schema_`, its pydantic name."""
     request = ResponsesRequest.model_validate(