From 69143f7e922d654371344101b90000405bf86b02 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:44:28 +0000 Subject: [PATCH 1/4] [None][fix] Disaggregated Anthropic Messages: start the stream and ping while the context phase runs On the disaggregated server, a streamed /v1/messages request receives no bytes at all, not even the HTTP status line, until the context phase is done: the handler awaits the chat pipeline before it creates the streaming response, and the pipeline returns its stream only once the context worker has answered, in both schedule orders. Queueing plus prefill of a long agent prompt takes tens of seconds to minutes under load, and clients, or the proxies in front of them, drop a stream that has sent nothing for about 15-35 s. The keepalive pings of #21030 cannot help there, because the stream starts only after the context phase. A streamed request now waits for its pipeline for at most _MESSAGES_STREAM_START_WAIT_S (1 s). A pipeline that answers within the wait is answered as before: its stream is relayed, and a failure, such as a context worker's fast over-length rejection, keeps its HTTP status and Anthropic error body. A pipeline still running after the wait has its stream started at once: the reframer sends its pings while the pipeline runs and then relays the generation stream, with message_start from its first chunk. A failure after that ends the stream with an error event that carries the message of the error response the request would have received. A client that disconnects while the context phase runs cancels the wait, which stops the proxy's context request as for any abandoned request. The mapping from the pipeline's outcome to the Anthropic response moves into _messages_answer so that both paths share it. Non-streamed requests and the other routes are unchanged. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 81 +++- .../test_openai_disagg_server.py | 381 +++++++++++++++++- 2 files changed, 451 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 92a298d4260f..dd3b1e8da141 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/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, Awaitable, Callable, Optional, Union import aiohttp import uvicorn @@ -108,6 +108,15 @@ # strictly better than applying it. _CLOCK_SYNC_MAX_DELAY_SECONDS = 0.010 +# Seconds a streamed /v1/messages request waits for its chat 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; one still running has its +# stream started and kept alive with pings until it answers. +_MESSAGES_STREAM_START_WAIT_S = 1.0 + class RawRequestResponseHooks(ResponseHooks): def __init__(self, raw_req: Request, queue_latency_metric, collect_perf_metrics: bool): @@ -469,9 +478,43 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, return anthropic_error_response(str(e), "invalid_request_error", 400) + chat = self._wrap_entry_point(self._service.openai_chat_completion)( + chat_request, raw_request) + if request.stream: + chat = asyncio.ensure_future(chat) + try: + await asyncio.wait((chat, ), + timeout=_MESSAGES_STREAM_START_WAIT_S) + except asyncio.CancelledError: + # Cancelling the request cancels its pipeline. + chat.cancel() + raise + if not chat.done(): + # The reframer pings while it waits for the pipeline. + return StreamingResponse( + content=reframe_openai_stream( + self._openai_stream_when_done(chat), + model=request.model), + media_type="text/event-stream", + ) + + answer = await self._messages_answer(chat) + if isinstance(answer, StreamingResponse): + return StreamingResponse( + content=reframe_openai_stream(answer.body_iterator, + model=request.model), + media_type="text/event-stream", + ) + return answer + + async def _messages_answer(self, chat: Awaitable[Response]) -> Response: + """Await a /v1/messages request's chat pipeline and answer from it. + + A stream is returned as the pipeline's OpenAI stream, anything else as + the Anthropic response. + """ try: - openai_response = await self._wrap_entry_point( - self._service.openai_chat_completion)(chat_request, raw_request) + openai_response = await chat except HTTPException as e: error_type = ("invalid_request_error" if 400 <= e.status_code < 500 else "api_error") @@ -479,11 +522,7 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, e.status_code) if isinstance(openai_response, StreamingResponse): - return StreamingResponse( - content=reframe_openai_stream(openai_response.body_iterator, - model=request.model), - media_type="text/event-stream", - ) + return openai_response status = getattr(openai_response, "status_code", 500) if status != 200: @@ -510,6 +549,32 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, return JSONResponse(content=anthropic_response.model_dump( exclude_none=True)) + async def _openai_stream_when_done( + self, chat: asyncio.Future) -> AsyncIterator[Union[str, bytes]]: + """The OpenAI stream of a chat pipeline that outlasted the start wait. + + The status line is already out, so a pipeline that fails raises + AnthropicResponseError with the message of its error response, and + the reframer ends the stream with an ``error`` event. A stream + abandoned before the pipeline answers, as when the client disconnects, + cancels the pipeline. + """ + try: + answer = await self._messages_answer(chat) + finally: + chat.cancel() + if not isinstance(answer, StreamingResponse): + raise AnthropicResponseError( + json.loads(answer.body)["error"]["message"]) + payloads = answer.body_iterator + try: + async for payload in payloads: + yield payload + finally: + aclose = getattr(payloads, "aclose", None) + if aclose is not None: + await aclose() + async def get_model(self) -> Response: """List the served model, as a context worker reports it.""" try: diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 4984a1ad749e..6576c1510e9d 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -11,21 +11,36 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import asyncio +import contextlib +import json +import re from types import SimpleNamespace from unittest.mock import AsyncMock, Mock, patch +import aiohttp import pytest -from fastapi import Request +from fastapi import FastAPI, Request from starlette.datastructures import Headers -from tensorrt_llm.llmapi.disagg_utils import ServerRole, extract_disagg_cfg -from tensorrt_llm.serve import openai_disagg_server +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.llmapi.disagg_utils import DisaggServerConfig, ServerRole, extract_disagg_cfg +from tensorrt_llm.serve import anthropic_adapter, openai_disagg_server +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer +from tensorrt_llm.serve.openai_disagg_service import OpenAIDisaggregatedService from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionStreamResponse, + ChatMessage, CompletionRequest, ConversationParams, DisaggregatedParams, + UsageInfo, ) +from tensorrt_llm.serve.responses_utils import ServerArrivalTimeMiddleware +from tensorrt_llm.serve.router import Router pytestmark = pytest.mark.cpu_only @@ -324,3 +339,363 @@ def test_disagg_config_rejects_non_bool_request_chat_template_opt_in(value): generation_servers={"num_instances": 0}, allow_request_chat_template=value, ) + + +# --------------------------------------------------------------------------- +# Streamed /v1/messages: a slow context phase must not leave the client silent +# --------------------------------------------------------------------------- + +_MODEL = "test-model" +_SCHEDULE_STYLES = ["context_first", "generation_first"] +# The events of the message the generation worker streams (_generation_frames). +_MESSAGE_EVENTS = [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", +] +_OVER_LENGTH = context_length_exceeded_message(8, 9) + + +def _generation_frames() -> list[bytes]: + """A generation worker's chat stream answering "hello".""" + chunks = [ + ChatCompletionStreamResponse( + model=_MODEL, choices=[{"index": 0, "delta": {"role": "assistant"}}] + ), + ChatCompletionStreamResponse( + model=_MODEL, choices=[{"index": 0, "delta": {"content": "hello"}}] + ), + ChatCompletionStreamResponse( + model=_MODEL, + choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], + usage=UsageInfo(prompt_tokens=3, completion_tokens=1, total_tokens=4), + ), + ] + frames = [f"data: {chunk.model_dump_json()}\n\n".encode() for chunk in chunks] + return frames + [b"data: [DONE]\n\n"] + + +def _worker_rejection(status: int, message: str) -> aiohttp.ClientResponseError: + """A worker's error response, as the HTTP client raises it.""" + body = json.dumps({"object": "error", "message": message, "code": status}) + return aiohttp.ClientResponseError( + Mock(real_url="http://ctx:8001/v1/chat/completions"), + (), + status=status, + message=f"Error: {body}", + ) + + +class _Workers: + """A context and a generation worker behind the disaggregated service, without a network. + + The context phase runs ``context_phase`` and then answers with the KV-cache + handoff. The generation worker streams ``_generation_frames`` once the + context phase has succeeded, as it does once it has the KV cache. + """ + + def __init__(self, context_phase): + self.context_phase = context_phase + self.context_done = asyncio.Event() + self.context_cancelled = asyncio.Event() + + async def send_context_request(self, request, server=None, hooks=None, req_id=None): + try: + await self.context_phase() + except asyncio.CancelledError: + self.context_cancelled.set() + raise + self.context_done.set() + disagg_request_id = request.disaggregated_params.disagg_request_id + return ChatCompletionResponse( + model=_MODEL, + usage=UsageInfo(prompt_tokens=3, completion_tokens=1), + prompt_token_ids=[1, 2, 3], + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content=""), + finish_reason="length", + disaggregated_params=DisaggregatedParams( + request_type="context_only", + disagg_request_id=disagg_request_id, + ctx_request_id=disagg_request_id, + ), + ) + ], + ) + + async def send_generation_request(self, request, server=None, hooks=None, req_id=None): + async def stream(): + await self.context_done.wait() + for frame in _generation_frames(): + yield frame + + return stream() + + +def _messages_app(schedule_style: str, workers: _Workers) -> FastAPI: + """The /v1/messages route of a disaggregated server in front of ``workers``.""" + config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) + service = OpenAIDisaggregatedService( + config, DisaggCoordinatorService(config, client_factory=Mock()), client_factory=Mock() + ) + service.is_ready = AsyncMock(return_value=True) + service._ctx_router = AsyncMock(spec=Router) + service._ctx_router.get_next_server.return_value = ("ctx:8001", {}) + service._gen_router = AsyncMock(spec=Router) + service._gen_router.get_next_server.return_value = ("gen:8002", {}) + service._ctx_client = SimpleNamespace(send_request=workers.send_context_request) + service._gen_client = SimpleNamespace(send_request=workers.send_generation_request) + + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._config = config + server._service = service + server._perf_metrics_collector = Mock() + server._allow_request_chat_template = False + server._collect_perf_metrics = False + app = FastAPI() + app.add_middleware(ServerArrivalTimeMiddleware) + app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) + return app + + +async def _post_messages(app: FastAPI, *, stream=True, on_send=None, leave=None) -> list[dict]: + """Send one /v1/messages request to ``app`` as a client and return the messages sent back. + + ``on_send`` sees each message as it is sent. The client disconnects once + ``leave`` is set. + """ + body = { + "model": _MODEL, + "max_tokens": 16, + "stream": stream, + "messages": [{"role": "user", "content": "hi"}], + } + requested = False + + async def receive(): + nonlocal requested + if not requested: + requested = True + return {"type": "http.request", "body": json.dumps(body).encode(), "more_body": False} + await (leave or asyncio.Event()).wait() + return {"type": "http.disconnect"} + + sent = [] + + async def send(message): + sent.append(message) + if on_send is not None: + on_send(message) + + scope = { + "type": "http", + # As uvicorn declares it: a disconnect cancels a streaming response. + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "POST", + "scheme": "http", + "path": "/v1/messages", + "raw_path": b"/v1/messages", + "query_string": b"", + "root_path": "", + "headers": [(b"content-type", b"application/json")], + "client": ("127.0.0.1", 40000), + "server": ("127.0.0.1", 8000), + } + await asyncio.wait_for(app(scope, receive, send), timeout=10) + return sent + + +def _status(sent: list[dict]) -> int: + return next(message["status"] for message in sent if message["type"] == "http.response.start") + + +def _body(sent: list[dict]) -> str: + return b"".join( + message.get("body", b"") for message in sent if message["type"] == "http.response.body" + ).decode() + + +def _sse_events(sent: list[dict]) -> list[tuple[str, dict]]: + """The (event, data) pairs of a streamed response.""" + events = [] + for frame in _body(sent).split("\n\n"): + if frame: + fields = dict(line.split(": ", 1) for line in frame.splitlines()) + events.append((fields["event"], json.loads(fields["data"]))) + return events + + +def _short_waits(monkeypatch): + monkeypatch.setattr(openai_disagg_server, "_MESSAGES_STREAM_START_WAIT_S", 0.05) + monkeypatch.setattr(anthropic_adapter, "STREAM_PING_INTERVAL_S", 0.05) + + +def _waits_for(event: asyncio.Event): + """A context phase that lasts until ``event`` is set, or five seconds without it.""" + + async def context_phase(): + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(event.wait(), 5) + + return context_phase + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_messages_stream_starts_and_pings_while_the_context_phase_runs( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + pinged = asyncio.Event() + workers = _Workers(_waits_for(pinged)) + started_during_context_phase = [] + + def on_send(message): + if message["type"] == "http.response.start": + started_during_context_phase.append(not workers.context_done.is_set()) + elif message.get("body", b"").startswith(b"event: ping\n"): + pinged.set() + + sent = await _post_messages(_messages_app(schedule_style, workers), on_send=on_send) + + assert _status(sent) == 200 + assert started_during_context_phase == [True] + events = [event for event, _ in _sse_events(sent)] + message_start = events.index("message_start") + assert message_start > 0 + assert set(events[:message_start]) == {"ping"} + assert [event for event in events[message_start:] if event != "ping"] == _MESSAGE_EVENTS + delta = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "hello"}, + } + assert ("content_block_delta", delta) in _sse_events(sent) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_messages_stream_after_a_fast_context_phase_is_the_reframed_generation( + schedule_style, +): + async def context_phase(): + pass + + sent = await _post_messages(_messages_app(schedule_style, _Workers(context_phase))) + + async def generation(): + for frame in _generation_frames(): + yield frame + + reframed = anthropic_adapter.reframe_openai_stream(generation(), model=_MODEL) + expected = "".join([frame async for frame in reframed]) + + def without_message_id(text): + return re.sub(r'"id":"msg_[0-9a-f]+"', '"id":"msg"', text) + + assert _status(sent) == 200 + assert "event: ping" not in _body(sent) + assert without_message_id(_body(sent)) == without_message_id(expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize( + "status, worker_message, error_type", + [ + pytest.param(400, _OVER_LENGTH, "invalid_request_error", id="over_length_400"), + pytest.param(503, "no capacity left", "api_error", id="503"), + ], +) +async def test_messages_stream_with_a_fast_context_failure_is_an_http_error( + schedule_style, status, worker_message, error_type +): + async def context_phase(): + raise _worker_rejection(status, worker_message) + + streamed = await _post_messages(_messages_app(schedule_style, _Workers(context_phase))) + not_streamed = await _post_messages( + _messages_app(schedule_style, _Workers(context_phase)), stream=False + ) + + assert _status(streamed) == status + error = json.loads(_body(streamed)) + assert error["type"] == "error" + assert error["error"]["type"] == error_type + # The error response of the same failure without a stream. + assert _status(not_streamed) == status + assert error == json.loads(_body(not_streamed)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize( + "status, worker_message", + [ + pytest.param(400, _OVER_LENGTH, id="over_length_400"), + pytest.param(503, "no capacity left", id="503"), + ], +) +async def test_messages_stream_with_a_late_context_failure_ends_with_an_error_event( + monkeypatch, schedule_style, status, worker_message +): + _short_waits(monkeypatch) + pinged = asyncio.Event() + wait = _waits_for(pinged) + + async def context_phase(): + await wait() + raise _worker_rejection(status, worker_message) + + def on_send(message): + if message.get("body", b"").startswith(b"event: ping\n"): + pinged.set() + + streamed = await _post_messages( + _messages_app(schedule_style, _Workers(context_phase)), on_send=on_send + ) + pinged.set() + not_streamed = await _post_messages( + _messages_app(schedule_style, _Workers(context_phase)), stream=False + ) + + assert _status(streamed) == 200 + events = _sse_events(streamed) + assert {event for event, _ in events[:-1]} == {"ping"} + assert events[-1][0] == "error" + # An error event as the stream reports any failure, with the message of + # the error response of the same failure without a stream. + assert events[-1][1] == { + "type": "error", + "error": { + "type": "api_error", + "message": json.loads(_body(not_streamed))["error"]["message"], + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("leave", ["before_the_response_starts", "once_the_response_starts"]) +async def test_client_leaving_during_the_context_phase_cancels_it( + monkeypatch, schedule_style, leave +): + _short_waits(monkeypatch) + workers = _Workers(_waits_for(asyncio.Event())) + left = asyncio.Event() + if leave == "before_the_response_starts": + left.set() + + def on_send(message): + if message["type"] == "http.response.start": + left.set() + + await _post_messages(_messages_app(schedule_style, workers), on_send=on_send, leave=left) + + await asyncio.wait_for(workers.context_cancelled.wait(), 2) From ecb824aa09bc53808cbdddbfe7cd8c5e63130328 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 13:58:33 +0000 Subject: [PATCH 2/4] test: cover cancellation during the messages handler wait Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../test_openai_disagg_server.py | 59 ++++++++++++++++--- 1 file changed, 51 insertions(+), 8 deletions(-) diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 6576c1510e9d..65486c670b78 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -26,6 +26,7 @@ from tensorrt_llm.executor.utils import context_length_exceeded_message from tensorrt_llm.llmapi.disagg_utils import DisaggServerConfig, ServerRole, extract_disagg_cfg from tensorrt_llm.serve import anthropic_adapter, openai_disagg_server +from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_disagg_service import OpenAIDisaggregatedService @@ -437,8 +438,8 @@ async def stream(): return stream() -def _messages_app(schedule_style: str, workers: _Workers) -> FastAPI: - """The /v1/messages route of a disaggregated server in front of ``workers``.""" +def _messages_server(schedule_style: str, workers: _Workers) -> OpenAIDisaggServer: + """A disaggregated server in front of ``workers``.""" config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) service = OpenAIDisaggregatedService( config, DisaggCoordinatorService(config, client_factory=Mock()), client_factory=Mock() @@ -457,6 +458,12 @@ def _messages_app(schedule_style: str, workers: _Workers) -> FastAPI: server._perf_metrics_collector = Mock() server._allow_request_chat_template = False server._collect_perf_metrics = False + return server + + +def _messages_app(schedule_style: str, workers: _Workers) -> FastAPI: + """The /v1/messages route of a disaggregated server in front of ``workers``.""" + server = _messages_server(schedule_style, workers) app = FastAPI() app.add_middleware(ServerArrivalTimeMiddleware) app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) @@ -682,15 +689,12 @@ def on_send(message): @pytest.mark.asyncio @pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) -@pytest.mark.parametrize("leave", ["before_the_response_starts", "once_the_response_starts"]) -async def test_client_leaving_during_the_context_phase_cancels_it( - monkeypatch, schedule_style, leave -): +async def test_client_leaving_after_the_response_starts_cancels_the_context_phase( + monkeypatch, schedule_style +) -> None: _short_waits(monkeypatch) workers = _Workers(_waits_for(asyncio.Event())) left = asyncio.Event() - if leave == "before_the_response_starts": - left.set() def on_send(message): if message["type"] == "http.response.start": @@ -699,3 +703,42 @@ def on_send(message): await _post_messages(_messages_app(schedule_style, workers), on_send=on_send, leave=left) await asyncio.wait_for(workers.context_cancelled.wait(), 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_cancelling_the_messages_handler_during_its_wait_cancels_the_context_phase( + monkeypatch, schedule_style +) -> None: + monkeypatch.setattr(openai_disagg_server, "_MESSAGES_STREAM_START_WAIT_S", 60) + context_started = asyncio.Event() + release_context = asyncio.Event() + + async def context_phase() -> None: + context_started.set() + await release_context.wait() + + workers = _Workers(context_phase) + server = _messages_server(schedule_style, workers) + request = AnthropicMessagesRequest( + model=_MODEL, + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "hi"}], + ) + raw_request = Request({"type": "http", "headers": []}) + raw_request.state.server_arrival_time = 0 + handler = asyncio.create_task(server.anthropic_messages(request, raw_request)) + try: + await asyncio.wait_for(context_started.wait(), 2) + assert not handler.done() + handler.cancel() + with pytest.raises(asyncio.CancelledError): + await handler + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + assert not workers.context_done.is_set() + finally: + release_context.set() + handler.cancel() + with contextlib.suppress(asyncio.CancelledError): + await handler From 396fe59aa1d12a31bb752efe83969d298b9abf33 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:21:58 +0000 Subject: [PATCH 3/4] fix: cancel pending Messages chat when the response disconnects Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 21 ++++++- .../test_openai_disagg_server.py | 63 ++++++++++++++++++- 2 files changed, 78 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index dd3b1e8da141..741bc871f2f0 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -67,7 +67,8 @@ PerfMetricsJsonlWriter, PerfMetricsMiddleware, combine_disagg_metrics) -from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, +from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse, + ServerArrivalTimeMiddleware, get_steady_clock_now_in_seconds, guard_responses_stream, stream_error_event) @@ -220,6 +221,20 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: return f"context worker rejected the request with status {error.status}" +class _MessagesStreamingResponse(ClosingStreamingResponse): + """Cancel the chat pipeline even when a disconnect precedes body iteration.""" + + def __init__(self, chat: asyncio.Future, content: AsyncIterator[str]) -> None: + super().__init__(content=content, media_type="text/event-stream") + self._chat = chat + + async def __call__(self, scope, receive, send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + self._chat.cancel() + + class OpenAIDisaggServer: def __init__(self, config: DisaggServerConfig, @@ -491,11 +506,11 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, raise if not chat.done(): # The reframer pings while it waits for the pipeline. - return StreamingResponse( + return _MessagesStreamingResponse( + chat=chat, content=reframe_openai_stream( self._openai_stream_when_done(chat), model=request.model), - media_type="text/event-stream", ) answer = await self._messages_answer(chat) diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 65486c670b78..133715d04bdc 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -402,6 +402,7 @@ def __init__(self, context_phase): self.context_phase = context_phase self.context_done = asyncio.Event() self.context_cancelled = asyncio.Event() + self.generation_cancelled = asyncio.Event() async def send_context_request(self, request, server=None, hooks=None, req_id=None): try: @@ -431,9 +432,13 @@ async def send_context_request(self, request, server=None, hooks=None, req_id=No async def send_generation_request(self, request, server=None, hooks=None, req_id=None): async def stream(): - await self.context_done.wait() - for frame in _generation_frames(): - yield frame + try: + await self.context_done.wait() + for frame in _generation_frames(): + yield frame + except asyncio.CancelledError: + self.generation_cancelled.set() + raise return stream() @@ -742,3 +747,55 @@ async def context_phase() -> None: handler.cancel() with contextlib.suppress(asyncio.CancelledError): await handler + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("asgi_version", ["2.3", "2.4"]) +async def test_disconnect_before_the_messages_body_starts_cancels_context( + monkeypatch, schedule_style, asgi_version +) -> None: + """Response cleanup must run even if the streaming iterator never starts.""" + from starlette.requests import ClientDisconnect + + _short_waits(monkeypatch) + release_context = asyncio.Event() + workers = _Workers(_waits_for(release_context)) + server = _messages_server(schedule_style, workers) + request = AnthropicMessagesRequest( + model=_MODEL, + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "hi"}], + ) + raw_request = Request({"type": "http", "headers": []}) + raw_request.state.server_arrival_time = 0 + response = await server.anthropic_messages(request, raw_request) + response_started = asyncio.Event() + sent = [] + + async def send(message): + sent.append(message) + response_started.set() + if asgi_version == "2.4": + raise OSError("client disconnected") + await asyncio.Event().wait() + + async def receive(): + await response_started.wait() + return {"type": "http.disconnect"} + + try: + scope = {"type": "http", "asgi": {"version": "3.0", "spec_version": asgi_version}} + if asgi_version == "2.4": + with pytest.raises(ClientDisconnect): + await response(scope, receive, send) + else: + await response(scope, receive, send) + assert [message["type"] for message in sent] == ["http.response.start"] + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + if schedule_style == "generation_first": + await asyncio.wait_for(workers.generation_cancelled.wait(), 2) + assert not workers.context_done.is_set() + finally: + release_context.set() From e8999ced24ea7045b4a39e5e69b655eb845b7001 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:42:06 +0000 Subject: [PATCH 4/4] fix: release conditional generation reservation on context cancellation Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_service.py | 2 +- .../test_openai_disagg_service.py | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index c62d62b31f8e..2de1717f2bfa 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/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 diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 46b063e71584..ed47990871b7 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -966,6 +966,43 @@ async def _gen_response(*_args, **_kwargs): ) +@pytest.mark.asyncio +async def test_context_cancellation_releases_conditional_generation_reservation(monkeypatch): + """Cancelling CTX must release the generation server reserved before it.""" + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + service = _make_service("context_first") + service._config.conditional_disagg_config = ConditionalDisaggConfig(max_local_prefill_length=32) + service._coordinator.get_disagg_request_id = AsyncMock(return_value=101) + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + router = KvCacheAwareRouter(server_role=ServerRole.GENERATION, servers=[]) + router.get_next_server = AsyncMock( + return_value=("gen:9001", {"match_length": 0, "num_tokens": 96}) + ) + router.finish_request = AsyncMock() + service._gen_router = router + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {})) + context_started = asyncio.Event() + + async def context_response(request, **kwargs): + context_started.set() + await asyncio.Event().wait() + + service._ctx_client.send_request = AsyncMock(side_effect=context_response) + request = CompletionRequest(model="model", prompt=[1] * 96, stream=True) + pending = asyncio.create_task(service._send_disagg_request(request)) + try: + await asyncio.wait_for(context_started.wait(), 2) + finally: + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + + router.get_next_server.assert_awaited_once_with(request, req_id=101) + router.finish_request.assert_awaited_once_with(request, success=False, req_id=101) + service._gen_client.send_request.assert_not_awaited() + + @pytest.mark.asyncio async def test_context_retry_preserves_generation_reservation_id(): service = _make_service("context_first")