diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 92a298d4260f..b608413c5f69 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -21,9 +21,10 @@ import traceback import uuid from contextlib import asynccontextmanager -from typing import Callable, Optional +from typing import AsyncIterator, Awaitable, Callable, Optional, Union import aiohttp +import anyio import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError @@ -67,7 +68,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) @@ -108,6 +110,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): @@ -211,6 +222,31 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: return f"context worker rejected the request with status {error.status}" +class _MessagesStreamingResponse(ClosingStreamingResponse): + """Release the chat pipeline however the response ends, body iterated or not. + + A pipeline still running is cancelled, and the stream of one that has + answered is closed: the body may stop, or never start, before it has read + that stream to its end. + """ + + 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: + chat = self._chat + chat.cancel() + if (chat.done() and not chat.cancelled() + and chat.exception() is None + and isinstance(chat.result(), StreamingResponse)): + with anyio.CancelScope(shield=True): + await chat.result().body_iterator.aclose() + + class OpenAIDisaggServer: def __init__(self, config: DisaggServerConfig, @@ -469,9 +505,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 _MessagesStreamingResponse( + chat=chat, + content=reframe_openai_stream( + self._openai_stream_when_done(chat), + model=request.model), + ) + + 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 +549,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 +576,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/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_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 4984a1ad749e..0a317c26cdb0 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -11,21 +11,38 @@ # 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 fastapi.responses import StreamingResponse 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.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 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 +341,537 @@ 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() + self.generation_cancelled = asyncio.Event() + # Set once the generation worker's stream has ended, however it ended. + self.generation_closed = 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(): + try: + yield None + await self.context_done.wait() + async for frame in self.generation_stream(): + yield frame + except asyncio.CancelledError: + self.generation_cancelled.set() + raise + finally: + self.generation_closed.set() + + # Started before it is returned, as the HTTP client starts it, so that + # closing it ends it even before it is read. + response = stream() + await response.__anext__() + return response + + async def generation_stream(self): + """The generation worker's stream once the context phase is done.""" + for frame in _generation_frames(): + yield frame + + +class _WorkersStillGenerating(_Workers): + """Workers whose generation worker sends its first frame and then goes on generating.""" + + async def generation_stream(self): + yield _generation_frames()[0] + await asyncio.Event().wait() + + +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() + ) + 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 + 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"]) + 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) +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() + + 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) + + +@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 + + +@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() + + +@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_closes_the_stream_it_answered_with( + monkeypatch, schedule_style, asgi_version +) -> None: + """The pipeline answers with a stream after the response is returned, which no body reads.""" + from starlette.requests import ClientDisconnect + + _short_waits(monkeypatch) + release_context = asyncio.Event() + workers = _WorkersStillGenerating(_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 asyncio.wait_for(server.anthropic_messages(request, raw_request), 2) + chat = response._chat + answered = asyncio.Event() + sent = [] + + async def send(message): + sent.append(message) + # The pipeline answers while the response starts, before the body does. + release_context.set() + await asyncio.wait((chat,), timeout=2) + answered.set() + if asgi_version == "2.4": + raise OSError("client disconnected") + await asyncio.Event().wait() + + async def receive(): + await answered.wait() + return {"type": "http.disconnect"} + + 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"] + assert isinstance(chat.result(), StreamingResponse) + # Closed by the time the response is done, while the generation worker goes on. + assert workers.generation_closed.is_set() 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")