[None][fix] Disaggregated Responses: open the stream and keep it alive while the context phase runs

docs/source/features/disagg-serving.md

@@ -236,7 +236,7 @@ curl http://localhost:8000/v1/completions \
 
 #### Responses API
 
-The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events unchanged. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history:
+The disaggregated server also serves `/v1/responses`, with both `context_first` and `generation_first` scheduling, and `/v1/models`, which it answers from a context worker. A streamed response carries the generation worker's events. If the context phase is still running after a second, the disaggregated server starts the stream itself with `response.created` and `response.in_progress`, repeats that `response.in_progress` whenever the stream is silent, and relays the generation worker's other events, numbered to continue the stream; a failure after that point ends the stream with `error` and `response.failed` instead of an HTTP error status. Unlike a single `trtllm-serve` instance, the disaggregated server keeps no response history:
 
 - A request with `previous_response_id` is rejected with a 400, because a stored response lives in one worker's memory and a later request may reach another worker. Send the earlier turns in `input` instead.
 - `store` is ignored: nothing is persisted, and the server logs a warning once when a client asks for it.

tensorrt_llm/serve/openai_disagg_server.py

@@ -21,7 +21,7 @@
 import traceback
 import uuid
 from contextlib import asynccontextmanager
-from typing import Callable, Optional
+from typing import AsyncIterator, Callable, Optional
 
 import aiohttp
 import uvicorn
@@ -67,10 +67,11 @@
                                              PerfMetricsJsonlWriter,
                                              PerfMetricsMiddleware,
                                              combine_disagg_metrics)
-from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware,
-                                                get_steady_clock_now_in_seconds,
-                                                guard_responses_stream,
-                                                stream_error_event)
+from tensorrt_llm.serve.responses_utils import (
+    STREAM_TERMINATION_UPSTREAM_ERROR, ClosingStreamingResponse,
+    ResponsesStreamingProcessor, ServerArrivalTimeMiddleware,
+    get_steady_clock_now_in_seconds, guard_responses_stream,
+    keep_responses_stream_alive, relay_after_opening, stream_error_event)
 from tensorrt_llm.serve.router import Router
 from tensorrt_llm.usage import TerminalOutcome, record_termination_observation
 from tensorrt_llm.version import __version__ as VERSION
@@ -81,6 +82,17 @@
     for code in (*range(32), 127)
 }
 
+# Seconds a streamed /v1/responses request waits for its pipeline before the
+# response starts. The pipeline has a stream to return only once its context
+# phase is done, and until the response starts the client receives nothing,
+# not even the status line, which clients and proxies take for a dead
+# connection after a while. A pipeline that answers within the wait is relayed
+# or, if it failed, answered with its error status. For one still running, the
+# stream starts at once with response.created and response.in_progress, and
+# response.in_progress is sent again whenever the stream is silent, as an
+# aggregated server does.
+_RESPONSES_STREAM_START_WAIT_S = 1.0
+
 # Most round trips to spend estimating one ctx/gen server's steady-clock
 # offset. Only the least-delayed sample is kept; see `_sync_server_clock`.
 # Each probe costs ~0.2 s because the `/steady_clock_offset` handler sleeps
@@ -184,6 +196,30 @@ def on_resp_done(
             ))
 
 
+class _ResponsesStreamingResponse(ClosingStreamingResponse):
+    """Cancel the pipeline even when a disconnect precedes body iteration."""
+
+    def __init__(self, pipeline: asyncio.Future,
+                 content: AsyncIterator[str]) -> None:
+        super().__init__(content=content, media_type="text/event-stream")
+        self._pipeline = pipeline
+
+    async def __call__(self, scope, receive, send) -> None:
+        try:
+            await super().__call__(scope, receive, send)
+        finally:
+            self._pipeline.cancel()
+
+
+def _error_response_message(response: Response) -> str:
+    """The message of an error response, for a stream to report it."""
+    try:
+        payload = json.loads(response.body)
+        return str(payload.get("message") or payload.get("detail") or payload)
+    except (json.JSONDecodeError, AttributeError, TypeError):
+        return "Internal server error"
+
+
 def _upstream_error_message(error: aiohttp.ClientResponseError) -> str:
     """Recover the worker's own error text without leaking the worker.
 
@@ -348,7 +384,7 @@ def register_routes(self):
         self.app.add_api_route("/v1/chat/completions", self._wrap_entry_point(self._service.openai_chat_completion, ChatCompletionRequest), methods=["POST"])
         self.app.add_api_route("/v1/messages", self.anthropic_messages, methods=["POST"])
         self.app.add_api_route("/v1/messages/count_tokens", self.anthropic_count_tokens, methods=["POST"])
-        self.app.add_api_route("/v1/responses", self._wrap_entry_point(self._service.openai_responses, ResponsesRequest), methods=["POST"])
+        self.app.add_api_route("/v1/responses", self.openai_responses, methods=["POST"])
         self.app.add_api_route("/v1/models", self.get_model, methods=["GET"])
         self.app.add_api_route("/health", self.health, methods=["GET"])
         self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"])
@@ -460,6 +496,87 @@ async def wrapper(req: request_type, raw_req: Request) -> Response:
                 return self._handle_exception(e)
         return wrapper
 
+    async def openai_responses(self, request: ResponsesRequest,
+                               raw_request: Request) -> Response:
+        """Serve Responses through the disaggregated pipeline.
+
+        A streamed request whose pipeline is still running after
+        _RESPONSES_STREAM_START_WAIT_S is answered from here at once; see
+        _responses_stream_when_done.
+        """
+        responses = self._wrap_entry_point(self._service.openai_responses)
+        if not request.stream:
+            return await responses(request, raw_request)
+        # The workers name their response after request_id, which is forwarded
+        # only if set. Assigning it marks it set, so that the workers answer
+        # under the id that an opening sent from here names.
+        request.request_id = request.request_id
+        pipeline = asyncio.ensure_future(responses(request, raw_request))
+        try:
+            await asyncio.wait((pipeline, ),
+                               timeout=_RESPONSES_STREAM_START_WAIT_S)
+        except asyncio.CancelledError:
+            # Cancelling the request cancels its pipeline.
+            pipeline.cancel()
+            raise
+        if pipeline.done():
+            return pipeline.result()
+        try:
+            processor = ResponsesStreamingProcessor(
+                request=request,
+                sampling_params=request.to_sampling_params(),
+                model_name=request.model)
+            opening = processor.get_initial_responses()
+        except ValueError:
+            # The opening is built from the request alone; a request it cannot
+            # be built from is answered when its pipeline answers.
+            return await pipeline
+        # keep_responses_stream_alive repeats the opening's response.in_progress
+        # whenever the stream is silent, as it is while the pipeline runs.
+        return _ResponsesStreamingResponse(
+            pipeline=pipeline,
+            content=keep_responses_stream_alive(
+                self._responses_stream_when_done(processor, opening,
+                                                 pipeline)))
+
+    async def _responses_stream_when_done(
+            self, processor: ResponsesStreamingProcessor, opening: list[str],
+            pipeline: asyncio.Future) -> AsyncIterator[str]:
+        """The stream of a Responses request whose pipeline outlasted the start wait.
+
+        ``opening`` is the request's response.created and response.in_progress,
+        and the pipeline's stream follows it without the worker's own
+        response.created and response.in_progress events (relay_after_opening),
+        so that every keepalive the client gets repeats ``opening``'s
+        response.in_progress. The status line is already out, so a pipeline
+        that fails ends the stream with error and response.failed, which carry
+        the message of its error response. A stream abandoned before the
+        pipeline answers, as when the client disconnects, cancels the pipeline.
+        """
+        try:
+            for frame in opening:
+                yield frame
+            try:
+                answer = await pipeline
+            except HTTPException as e:
+                failure = str(e.detail)
+            else:
+                failure = (None if isinstance(answer, StreamingResponse) else
+                           _error_response_message(answer))
+        finally:
+            pipeline.cancel()
+        if failure is not None:
+            for frame in processor.get_stream_failed_events(
+                    STREAM_TERMINATION_UPSTREAM_ERROR, failure, len(opening)):
+                yield frame
+            return
+        frames = relay_after_opening(answer.body_iterator, len(opening))
+        try:
+            async for frame in frames:
+                yield frame
+        finally:
+            await frames.aclose()
+
     async def anthropic_messages(self, request: AnthropicMessagesRequest,
                                  raw_request: Request) -> Response:
         """Serve Anthropic Messages through the disaggregated chat pipeline."""

tensorrt_llm/serve/openai_disagg_service.py

@@ -333,7 +333,7 @@ async def _send_disagg_request_ctx_first(
                     if hooks:
                         hooks.on_disagg_request_id(disagg_request_id)
                 gen_req = self._get_gen_request(request, ctx_response, disagg_request_id)
-            except Exception:
+            except (Exception, asyncio.CancelledError):
                 if gen_server:
                     await self._gen_router.finish_request(
                         request, success=False, req_id=gen_reservation_id

tensorrt_llm/serve/responses_utils.py

@@ -3475,6 +3475,47 @@ def until_keepalive() -> Optional[float]:
         await items.aclose()
 
 
+_OPENING_FRAME_STARTS = ("event: response.created\n", _IN_PROGRESS_FRAME_START)
+
+
+async def relay_after_opening(stream: AsyncGenerator[Union[str, bytes], None],
+                              events_sent: int) -> AsyncGenerator[str, None]:
+    """Relay a worker's Responses stream to a client that already has its opening.
+
+    The client holds ``events_sent`` events, ``response.created`` and
+    ``response.in_progress`` among them, and is kept alive with that
+    ``response.in_progress`` by the stream this relay feeds
+    (``keep_responses_stream_alive``). So every ``response.created`` and
+    ``response.in_progress`` of the worker's is dropped, wherever it comes:
+    the worker's opening and its keepalives, which repeat that opening, carry
+    the worker's response object, not the one the client holds, and a second
+    source of keepalives would only add to the relaying stream's own. Every
+    other event is numbered to follow on, and is yielded whole, one per item
+    as ``keep_responses_stream_alive`` takes them, wherever ``stream`` splits
+    it. ``stream`` is closed when the relay ends.
+    """
+    delimiter = _SSE_EVENT_DELIMITER.encode()
+    pending = b""
+    try:
+        async for chunk in stream:
+            if isinstance(chunk, str):
+                chunk = chunk.encode()
+            *events, pending = (pending + chunk).split(delimiter)
+            for event in events:
+                frame = event.decode(errors="replace") + _SSE_EVENT_DELIMITER
+                if not event or frame.startswith(_OPENING_FRAME_STARTS):
+                    continue
+                yield _stamp_sequence_number(frame, events_sent)
+                events_sent += 1
+        if pending:
+            # The stream stopped inside an event.
+            yield pending.decode(errors="replace")
+    finally:
+        aclose = getattr(stream, "aclose", None)
+        if aclose is not None:
+            await aclose()
+
+
 class ServerArrivalTimeMiddleware:
     """
     Custom ASGI middleware to track server arrival time.