diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index f0bda3662f1c..ad4faf8de4f3 100644 --- a/docs/source/features/disagg-serving.md +++ b/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. diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index e67ec2e6b371..9aff10b4486e 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -109,11 +109,15 @@ def _encode_request_body(request: _SerializableRequest) -> bytes: """Encode a request forwarded to a worker. Fields go out under their wire names, e.g. ``schema``, which the model - holds as ``schema_``. + holds as ``schema_``. Only fields that were set go out, except a Responses + request's ``request_id``, which always does, even when it is the default + generated when the request was parsed: a worker names its response after + it, so every worker answers under the id this server holds. """ - return _msgpack_encoder.encode( - request.model_dump(mode="json", exclude_unset=True, by_alias=True) - ) + body = request.model_dump(mode="json", exclude_unset=True, by_alias=True) + if isinstance(request, ResponsesRequest): + body["request_id"] = request.request_id + return _msgpack_encoder.encode(body) _ResponseT = TypeVar("_ResponseT", bound=BaseModel) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 92a298d4260f..576bbaa3f717 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, Callable, Optional import aiohttp +import anyio import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError @@ -67,10 +68,12 @@ 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_CONTEXT_LENGTH_EXCEEDED, + 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 +84,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 +198,62 @@ def on_resp_done( )) +class _ResponsesStreamingResponse(ClosingStreamingResponse): + """Release the 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, 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: + pipeline = self._pipeline + pipeline.cancel() + if (pipeline.done() and not pipeline.cancelled() + and pipeline.exception() is None + and isinstance(pipeline.result(), StreamingResponse)): + with anyio.CancelScope(shield=True): + await pipeline.result().body_iterator.aclose() + + +def _guard_relayed_responses_stream(answer: Response) -> Response: + """``answer``, with a stream that fails early ending in an ``error`` event. + + The stream is the generation worker's own, response.created included, so + there is no response of this server's to report as failed. + """ + if not isinstance(answer, StreamingResponse): + return answer + return StreamingResponse( + content=guard_responses_stream(answer.body_iterator, + stream_error_event), + media_type="text/event-stream") + + +def _error_response_failure(response: Response) -> tuple[str, str]: + """The cause and message a stream reports an error response with. + + An over-length rejection keeps its code, context_length_exceeded, which + clients act on; any other error response is an upstream error. + """ + try: + payload = json.loads(response.body) + message = str(payload.get("message") or payload.get("detail") or payload) + except (json.JSONDecodeError, AttributeError, TypeError): + return STREAM_TERMINATION_UPSTREAM_ERROR, "Internal server error" + if payload.get("code") == CONTEXT_LENGTH_EXCEEDED_CODE: + return STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED, message + return STREAM_TERMINATION_UPSTREAM_ERROR, message + + def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: """Recover the worker's own error text without leaking the worker. @@ -348,7 +418,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"]) @@ -442,14 +512,10 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: response_or_generator = await entry_point(req, hooks) self._perf_metrics_collector.total_responses.inc() if req.stream: - stream = response_or_generator - if isinstance(req, ResponsesRequest): - # A relayed Responses stream that fails early ends with - # an `error` event; other protocols end on a sentinel. - stream = guard_responses_stream(stream, - stream_error_event) + # openai_responses ends a Responses stream that fails early + # with error events; other protocols end on a sentinel. return StreamingResponse( - content=stream, + content=response_or_generator, media_type="text/event-stream") # by_alias: e.g. text.format.schema goes out under its wire name. return JSONResponse( @@ -460,6 +526,100 @@ 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) + 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 _guard_relayed_responses_stream(pipeline.result()) + try: + # The opening names request.request_id, the id every worker is sent + # (_encode_request_body) and answers under. + 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 _guard_relayed_responses_stream(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, and for an over-length rejection + its code, context_length_exceeded (_error_response_failure); so does a + pipeline's stream that fails before its terminal event. 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 = (STREAM_TERMINATION_UPSTREAM_ERROR, str(e.detail)) + else: + failure = (None if isinstance(answer, StreamingResponse) else + _error_response_failure(answer)) + finally: + pipeline.cancel() + if failure is not None: + cause, detail = failure + for frame in processor.get_stream_failed_events( + cause, detail, len(opening)): + yield frame + return + # The guard reads the worker's stream as it comes, so it knows whether + # the terminal event went out; its events are numbered on with the + # worker's by relay_after_opening. + frames = relay_after_opening( + guard_responses_stream(answer.body_iterator, + processor.get_stream_failed_events), + len(opening)) + try: + async for frame in frames: + yield frame + except Exception: # noqa: BLE001 - the guard has reported it + # The guard has logged the failure and sent error and + # response.failed, unless the terminal event had gone out: either + # way the response is over, and so is the stream. + return + finally: + await frames.aclose() + async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" 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/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index f7d93a1471de..1c6cc7248d43 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -17,6 +17,7 @@ Union) import anyio +import pydantic_core from fastapi.responses import StreamingResponse from jinja2.exceptions import TemplateError from openai.types.responses import (ResponseCompletedEvent, @@ -59,6 +60,7 @@ from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.executor import (EngineDeadError, GenerationResult, RequestError) +from tensorrt_llm.executor.utils import CONTEXT_LENGTH_EXCEEDED_CODE from tensorrt_llm.inputs import PromptInputs, prompt_inputs from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.utils import resolve_hf_chat_template @@ -3009,6 +3011,35 @@ def _stream_terminal_event( ) +def _response_failed_frame(snapshot: dict, code: str, message: str, + sequence_number: int) -> str: + """A ``response.failed`` frame whose snapshot fails with ``code``. + + ``snapshot`` is a ``ResponsesResponse`` dump, which has no ``error``; the + SDK's Response, which the event carries, does. The SDK's ``ResponseError`` + does not list ``context_length_exceeded``, the code OpenAI fails an + over-length input with and the one clients act on. So the event is + validated with ``server_error`` in ``error.code``; any other code is then + written into its serialized form, with the serializer every frame uses. + """ + event = ResponseFailedEvent( + type="response.failed", + sequence_number=sequence_number, + response={ + **snapshot, "error": { + "code": "server_error", + "message": message + } + }, + ) + if code == "server_error": + return _format_sse_event(event) + payload = event.model_dump(mode="json", by_alias=True) + payload["response"]["error"]["code"] = code + return (f"event: {event.type}\n" + f"data: {pydantic_core.to_json(payload).decode()}\n\n") + + class ResponsesStreamingProcessor: def __init__( @@ -3133,7 +3164,9 @@ def get_stream_failed_events( ``events_sent`` is the number of frames that reached the wire, which the terminal events are numbered after; without it they continue this - processor's own numbering. + processor's own numbering. The failed snapshot's ``error.code`` is + ``context_length_exceeded`` for an over-length rejection and + ``server_error`` for every other cause. """ if events_sent is not None: self.sequence_number = events_sent @@ -3154,15 +3187,12 @@ def get_stream_failed_events( status="failed", usage=None, ).model_dump(by_alias=True) - # Set on the dump: the SDK Response carries `error`, ResponsesResponse - # does not. - snapshot["error"] = {"code": "server_error", "message": detail} - failed_event = self._send_event( - ResponseFailedEvent( - type="response.failed", - sequence_number=-1, - response=snapshot, - )) + code = (CONTEXT_LENGTH_EXCEEDED_CODE + if cause == STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED else + "server_error") + failed_event = _response_failed_frame(snapshot, code, detail, + self.sequence_number) + self.sequence_number += 1 return [error_event, failed_event] @@ -3174,6 +3204,11 @@ def get_stream_failed_events( STREAM_TERMINATION_ENGINE_ERROR = "engine_error" STREAM_TERMINATION_UPSTREAM_ERROR = "upstream_error" STREAM_TERMINATION_INTERNAL_ERROR = "internal_error" +# The engine rejected an input longer than the context window. This answers +# the request rather than breaking it: the stream fails with OpenAI's code for +# it, which clients act on (the Codex CLI compacts its history), and then ends +# normally. +STREAM_TERMINATION_CONTEXT_LENGTH_EXCEEDED = CONTEXT_LENGTH_EXCEEDED_CODE _SSE_EVENT_DELIMITER = "\n\n" # A terminal event frame always follows another frame, and JSON payloads @@ -3475,6 +3510,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. diff --git a/tests/unittest/disaggregated/test_disagg_responses_api.py b/tests/unittest/disaggregated/test_disagg_responses_api.py index 52a8ad9dcdce..a03afd26b8cb 100644 --- a/tests/unittest/disaggregated/test_disagg_responses_api.py +++ b/tests/unittest/disaggregated/test_disagg_responses_api.py @@ -20,17 +20,29 @@ import asyncio import base64 +import contextlib import json from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import AsyncMock, Mock, patch import aiohttp +import msgspec import numpy as np import pytest -from fastapi import HTTPException -from openai.types.responses import ResponseCreatedEvent +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import StreamingResponse +from openai.types.responses import ( + ResponseCompletedEvent, + ResponseCreatedEvent, + ResponseTextDeltaEvent, +) from pydantic import ValidationError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.llmapi.disagg_utils import DisaggServerConfig +from tensorrt_llm.serve import openai_disagg_server, responses_utils +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.serve.openai_client import _encode_request_body from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, @@ -49,10 +61,13 @@ UsageInfo, ) from tensorrt_llm.serve.responses_utils import ( + ResponsesStreamingProcessor, + ServerArrivalTimeMiddleware, create_response_non_store, finish_reason_mapping, responses_done_generator, ) +from tensorrt_llm.serve.router import Router pytestmark = pytest.mark.cpu_only @@ -526,6 +541,21 @@ async def send(request, hooks): assert (sent[0].prompt_token_ids, sent[0].prompt_token_ids_b64) == (None, None) +@pytest.mark.parametrize( + "fields", [{}, {"request_id": "resp_set_by_the_client"}], ids=["generated", "client_set"] +) +def test_every_worker_is_sent_the_request_id(fields): + """A worker names its response after request_id, so it goes out even when generated here. + + A stream this server opens itself names that id, and the workers must answer under it. + """ + request = ResponsesRequest(model="m", input="hi", **fields) + # The context request is a copy of the request, as the service makes it. + for forwarded in (request, request.model_copy(update={"stream": False})): + body = msgspec.msgpack.decode(_encode_request_body(forwarded)) + assert body["request_id"] == request.request_id + + def test_a_generation_first_worker_ignores_client_prompt_ids(): final = _responses_response(finish_reason="stop", status="completed") service = _gen_first_service(ctx_response=_responses_response(), gen_response=final) @@ -602,3 +632,871 @@ def test_a_models_failure_is_an_error_response(client, servers, status): response = asyncio.run(server.get_model()) assert response.status_code == status assert json.loads(response.body)["code"] == status + + +# --------------------------------------------------------------------------- +# Streamed requests: a slow context phase must not leave the client silent +# --------------------------------------------------------------------------- + +_MODEL = "test-model" +_SCHEDULE_STYLES = ["context_first", "generation_first"] +_OVER_LENGTH = context_length_exceeded_message(8, 9) + + +def _as_received(request: ResponsesRequest) -> ResponsesRequest: + """``request`` as a worker parses it from the body the service forwards.""" + return ResponsesRequest(**msgspec.msgpack.decode(_encode_request_body(request))) + + +def _response(request: ResponsesRequest, status: str) -> ResponsesResponse: + """A worker's response to ``request``, without output.""" + return ResponsesResponse.from_request( + request=request, + sampling_params=request.to_sampling_params(), + model_name=_MODEL, + created_time=0, + output=[], + status=status, + ) + + +def _sse(event) -> str: + return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n" + + +def _generation_frames(request: ResponsesRequest) -> list[str]: + """A generation worker's stream for ``request``: its opening, "hello", the completion. + + Its response objects are created at 0, so they differ from any built elsewhere. + """ + opening = ResponsesStreamingProcessor( + request, request.to_sampling_params(), model_name=_MODEL, create_time=0 + ).get_initial_responses() + delta = ResponseTextDeltaEvent( + type="response.output_text.delta", + sequence_number=2, + item_id="msg_1", + output_index=0, + content_index=0, + delta="hello", + logprobs=[], + ) + completed = ResponseCompletedEvent( + type="response.completed", + sequence_number=3, + response=_response(request, "completed").model_dump(by_alias=True), + ) + return opening + [_sse(delta), _sse(completed)] + + +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/responses"), + (), + status=status, + message=f"Error: {body}", + ) + + +class _Workers: + """A context and a generation worker behind the disaggregated service, without a network. + + Each parses its request from the body the service forwards. The context phase runs + ``context_phase`` and then answers with the KV-cache handoff. The generation worker + answers once the context phase has succeeded, as it does once it has the KV cache, and + streams its answer in pieces cut anywhere, as the transport delivers it. + """ + + 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() + self.generation_requests = [] + # What the generation worker answered: a response, or the frames of a stream. + self.response = None + self.frames = None + + async def send_context_request(self, request, server=None, hooks=None, req_id=None): + received = _as_received(request) + try: + await self.context_phase() + except asyncio.CancelledError: + self.context_cancelled.set() + raise + self.context_done.set() + response = _response(received, "incomplete") + response.finish_reason = "length" + response.prompt_token_ids = [1, 2, 3] + disagg_request_id = received.disaggregated_params.disagg_request_id + response.disaggregated_params = DisaggregatedParams( + request_type="context_only", + disagg_request_id=disagg_request_id, + ctx_request_id=disagg_request_id, + ) + return response + + async def send_generation_request(self, request, server=None, hooks=None, req_id=None): + received = _as_received(request) + self.generation_requests.append(received) + if not received.stream: + await self.context_done.wait() + self.response = _response(received, "completed") + return self.response + + async def stream(): + try: + yield None + await self.context_done.wait() + async for piece in self.generation_stream(received): + yield piece + 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, received: ResponsesRequest): + """The generation worker's stream, in pieces cut anywhere, as the transport delivers it.""" + self.frames = _generation_frames(received) + data = "".join(self.frames).encode() + for start in range(0, len(data), 61): + yield data[start : start + 61] + + +class _WorkersSilentAfterTheFirstOutput(_Workers): + """Workers whose generation worker is silent for a while after its first output. + + The silence outlasts the keepalive interval, so the worker's own keepalive sends its + response.in_progress again, with its own response object, before the rest of its stream. + """ + + async def generation_stream(self, received: ResponsesRequest): + frames = _generation_frames(received) + + async def produced(): + # The opening and the first output, then the rest after a silence. + for frame in frames[:3]: + yield frame + await asyncio.sleep(0.3) + for frame in frames[3:]: + yield frame + + self.frames = [] + async for frame in responses_utils.keep_responses_stream_alive(produced()): + self.frames.append(frame) + yield frame.encode() + + +class _WorkersFailingAfterTheFirstOutput(_Workers): + """Workers whose generation worker's connection fails once its first output is relayed.""" + + def __init__(self, context_phase): + super().__init__(context_phase) + self.output_relayed = asyncio.Event() + + async def generation_stream(self, received: ResponsesRequest): + self.frames = _generation_frames(received)[:3] + for frame in self.frames: + yield frame.encode() + await self.output_relayed.wait() + raise aiohttp.ClientPayloadError("Response payload is not completed") + + def on_send(self, message): + """An ``on_send`` that notes the first output reaching the client.""" + if message.get("body", b"").startswith(b"event: response.output_text.delta\n"): + self.output_relayed.set() + + +class _WorkersStillGenerating(_Workers): + """Workers whose generation worker sends its opening and then goes on generating.""" + + async def generation_stream(self, received: ResponsesRequest): + self.frames = _generation_frames(received)[:2] + for frame in self.frames: + yield frame.encode() + await asyncio.Event().wait() + + +def _responses_server(schedule_style: str, workers: _Workers) -> OpenAIDisaggServer: + """A disaggregated server in front of ``workers``, with its routes.""" + config = DisaggServerConfig(server_configs=[], schedule_style=schedule_style) + coordinator = DisaggCoordinatorService(config, client_factory=Mock()) + service = OpenAIDisaggregatedService(config, coordinator, 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._coordinator = coordinator + server._disagg_cluster_storage = None + server._service = service + server._perf_metrics_collector = Mock() + server._allow_request_chat_template = False + server._collect_perf_metrics = False + server.app = FastAPI() + server.app.add_middleware(ServerArrivalTimeMiddleware) + with patch("prometheus_client.multiprocess.MultiProcessCollector"): + server.register_routes() + return server + + +def _responses_route(server: OpenAIDisaggServer): + """The handler ``server`` routes /v1/responses to.""" + (route,) = [ + route for route in server.app.routes if getattr(route, "path", "") == "/v1/responses" + ] + return route.endpoint + + +async def _post_responses(app: FastAPI, *, stream=True, on_send=None, leave=None, **fields): + """Send one /v1/responses 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, "input": "hi", "stream": stream, **fields} + 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/responses", + "raw_path": b"/v1/responses", + "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 _without_number(data: dict) -> dict: + return {key: value for key, value in data.items() if key != "sequence_number"} + + +def _short_waits(monkeypatch): + monkeypatch.setattr(openai_disagg_server, "_RESPONSES_STREAM_START_WAIT_S", 0.05) + monkeypatch.setattr(responses_utils, "RESPONSES_STREAM_KEEPALIVE_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 + + +def _on_keepalive(event: asyncio.Event): + """An ``on_send`` that sets ``event`` once response.in_progress has gone out again.""" + in_progress = 0 + + def on_send(message): + nonlocal in_progress + if message.get("body", b"").startswith(b"event: response.in_progress\n"): + in_progress += 1 + if in_progress == 2: + event.set() + + return on_send + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_responses_stream_opens_and_keeps_alive_while_the_context_phase_runs( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + workers = _Workers(_waits_for(kept_alive)) + keepalive = _on_keepalive(kept_alive) + 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()) + keepalive(message) + + sent = await _post_responses(_responses_server(schedule_style, workers).app, on_send=on_send) + + assert _status(sent) == 200 + assert started_during_context_phase == [True] + events = _sse_events(sent) + types = [event for event, _ in events] + first_output = types.index("response.output_text.delta") + # The opening, kept alive while the context phase ran, then the generation + # worker's events without its own opening. + assert types[:2] == ["response.created", "response.in_progress"] + assert first_output > 2 + assert set(types[2:first_output]) == {"response.in_progress"} + relayed = [data for event, data in events[first_output:] if event != "response.in_progress"] + answered = [json.loads(frame.split("data: ", 1)[1]) for frame in workers.frames[2:]] + assert [_without_number(data) for data in relayed] == [ + _without_number(data) for data in answered + ] + # One sequence, for one response: the one the generation worker was asked for. + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + (generation_request,) = workers.generation_requests + assert {data["response"]["id"] for _, data in events if "response" in data} == { + generation_request.request_id + } + + +def _frame_events(frames: list[str]) -> list[tuple[str, dict]]: + """The (event, data) pairs of SSE frames.""" + return [ + (frame.split("\n", 1)[0].removeprefix("event: "), json.loads(frame.split("data: ", 1)[1])) + for frame in frames + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_every_keepalive_of_an_opened_stream_repeats_its_opening(monkeypatch, schedule_style): + """The worker's own keepalives, which carry its response object, never reach the client.""" + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + workers = _WorkersSilentAfterTheFirstOutput(_waits_for(kept_alive)) + + sent = await _post_responses( + _responses_server(schedule_style, workers).app, on_send=_on_keepalive(kept_alive) + ) + + assert _status(sent) == 200 + events = _sse_events(sent) + types = [event for event, _ in events] + assert types[:2] == ["response.created", "response.in_progress"] + first_output = types.index("response.output_text.delta") + # The worker kept its own stream alive while it was silent after its first output. + answered = _frame_events(workers.frames) + answered_types = [event for event, _ in answered] + assert ( + "response.in_progress" + in answered_types[answered_types.index("response.output_text.delta") :] + ) + # The client was kept alive meanwhile, and every response.in_progress it got is + # the opening's, only numbered on; none is one of the worker's. + assert "response.in_progress" in types[first_output:] + opening = _without_number(events[1][1]) + assert all( + _without_number(data) == opening + for event, data in events + if event == "response.in_progress" + ) + assert opening not in [ + _without_number(data) for event, data in answered if event == "response.in_progress" + ] + # The worker's other events follow, in one sequence, for one response. + relayed = [data for event, data in events[first_output:] if event != "response.in_progress"] + assert [_without_number(data) for data in relayed] == [ + _without_number(data) + for event, data in answered + if event not in ("response.created", "response.in_progress") + ] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + (generation_request,) = workers.generation_requests + assert {data["response"]["id"] for _, data in events if "response" in data} == { + generation_request.request_id + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_responses_stream_after_a_fast_context_phase_is_the_generation_stream( + schedule_style, +): + async def context_phase(): + pass + + workers = _Workers(context_phase) + sent = await _post_responses(_responses_server(schedule_style, workers).app) + + assert _status(sent) == 200 + assert [event for event, _ in _sse_events(sent)] == [ + "response.created", + "response.in_progress", + "response.output_text.delta", + "response.completed", + ] + assert _body(sent) == "".join(workers.frames) + + +_CONTEXT_FAILURES = [ + pytest.param(400, _OVER_LENGTH, id="over_length_400"), + pytest.param(503, "no capacity left", id="503"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize("status, worker_message", _CONTEXT_FAILURES) +async def test_responses_stream_with_a_fast_context_failure_is_an_http_error( + schedule_style, status, worker_message +): + async def context_phase(): + raise _worker_rejection(status, worker_message) + + streamed = await _post_responses(_responses_server(schedule_style, _Workers(context_phase)).app) + not_streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, stream=False + ) + + assert _status(streamed) == status + assert worker_message in _body(streamed) + # The error response of the same failure without a stream. + assert _status(not_streamed) == status + assert json.loads(_body(streamed)) == json.loads(_body(not_streamed)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +@pytest.mark.parametrize( + "status, worker_message, error_code, failed_code", + [ + # An over-length rejection keeps its code in both events: clients act on it, + # as an agent compacts its history and retries. + pytest.param( + 400, + _OVER_LENGTH, + "context_length_exceeded", + "context_length_exceeded", + id="over_length_400", + ), + pytest.param(503, "no capacity left", "upstream_error", "server_error", id="503"), + ], +) +async def test_responses_stream_with_a_late_context_failure_ends_with_response_failed( + monkeypatch, schedule_style, status, worker_message, error_code, failed_code +): + _short_waits(monkeypatch) + kept_alive = asyncio.Event() + wait = _waits_for(kept_alive) + + async def context_phase(): + await wait() + raise _worker_rejection(status, worker_message) + + streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, + on_send=_on_keepalive(kept_alive), + ) + kept_alive.set() + not_streamed = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, stream=False + ) + + assert _status(streamed) == 200 + # The stream ends after its events, as any stream does. + assert streamed[-1] == {"type": "http.response.body", "body": b"", "more_body": False} + events = _sse_events(streamed) + types = [event for event, _ in events] + assert types[:2] == ["response.created", "response.in_progress"] + assert set(types[2:-2]) == {"response.in_progress"} + assert types[-2:] == ["error", "response.failed"] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + # error and response.failed, as a stream reports any failure, with the + # message of the error response of the same failure without a stream. + error_response = json.loads(_body(not_streamed)) + message = error_response.get("detail") or error_response["message"] + assert worker_message in message + assert (events[-2][1]["code"], events[-2][1]["message"]) == (error_code, message) + failed = events[-1][1]["response"] + assert failed["id"] == events[0][1]["response"]["id"] + assert failed["status"] == "failed" + assert (failed["error"]["code"], failed["error"]["message"]) == (failed_code, message) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_an_opened_stream_whose_generation_stream_fails_ends_with_response_failed( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + + workers = _WorkersFailingAfterTheFirstOutput(context_phase) + sent = await _post_responses( + _responses_server(schedule_style, workers).app, on_send=workers.on_send + ) + + assert _status(sent) == 200 + # The stream ends after its events, as any stream does. + assert sent[-1] == {"type": "http.response.body", "body": b"", "more_body": False} + events = _sse_events(sent) + assert events[1][0] == "response.in_progress" + assert [event for event, _ in events if event != "response.in_progress"] == [ + "response.created", + "response.output_text.delta", + "error", + "response.failed", + ] + assert [data["sequence_number"] for _, data in events] == list(range(len(events))) + assert (events[-2][1]["code"], events[-2][1]["message"]) == ( + "upstream_error", + "ClientPayloadError", + ) + # The failed response is the one the stream opened with. + opened, failed = events[0][1]["response"], events[-1][1]["response"] + assert (failed["id"], failed["created_at"]) == (opened["id"], opened["created_at"]) + assert (failed["status"], failed["error"]["code"]) == ("failed", "server_error") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_relayed_generation_stream_that_fails_ends_with_an_error_event(schedule_style): + """After a fast context phase the worker's own opening is relayed, so only error is added.""" + + async def context_phase(): + pass + + workers = _WorkersFailingAfterTheFirstOutput(context_phase) + sent = [] + + def on_send(message): + sent.append(message) + workers.on_send(message) + + with contextlib.suppress(aiohttp.ClientPayloadError): + await _post_responses(_responses_server(schedule_style, workers).app, on_send=on_send) + + events = _sse_events(sent) + assert [event for event, _ in events] == [ + "response.created", + "response.in_progress", + "response.output_text.delta", + "error", + ] + assert [data["sequence_number"] for _, data in events] == [0, 1, 2, 3] + assert (events[-1][1]["code"], events[-1][1]["message"]) == ( + "upstream_error", + "ClientPayloadError", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_non_streamed_response_still_waits_for_the_whole_pipeline( + monkeypatch, schedule_style +): + _short_waits(monkeypatch) + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + + workers = _Workers(context_phase) + sent = await _post_responses(_responses_server(schedule_style, workers).app, stream=False) + + assert _status(sent) == 200 + assert json.loads(_body(sent)) == workers.response.model_dump(mode="json", by_alias=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_a_request_no_opening_can_be_built_for_is_answered_by_its_pipeline( + monkeypatch, schedule_style +): + """The opening is built from the request; one it cannot be built from waits as before.""" + _short_waits(monkeypatch) + rejection = "require 0 <= top_p <= 1, got top_p=2.0" + + async def context_phase(): + # Outlasts the start wait. + await asyncio.sleep(0.2) + raise _worker_rejection(400, rejection) + + sent = await _post_responses( + _responses_server(schedule_style, _Workers(context_phase)).app, top_p=2.0 + ) + + assert _status(sent) == 400 + assert json.loads(_body(sent)) == {"detail": rejection} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_client_leaving_after_the_responses_stream_starts_cancels_the_context_phase( + monkeypatch, schedule_style +): + _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_responses( + _responses_server(schedule_style, workers).app, on_send=on_send, leave=left + ) + + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + + +def _raw_request() -> Request: + raw_request = Request({"type": "http", "headers": []}) + raw_request.state.server_arrival_time = 0 + return raw_request + + +@pytest.mark.asyncio +@pytest.mark.parametrize("schedule_style", _SCHEDULE_STYLES) +async def test_cancelling_the_responses_handler_during_its_wait_cancels_the_context_phase( + monkeypatch, schedule_style +): + monkeypatch.setattr(openai_disagg_server, "_RESPONSES_STREAM_START_WAIT_S", 60) + context_started = asyncio.Event() + release_context = asyncio.Event() + + async def context_phase(): + context_started.set() + await release_context.wait() + + workers = _Workers(context_phase) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + pending = asyncio.create_task(handler(request, _raw_request())) + try: + await asyncio.wait_for(context_started.wait(), 2) + assert not pending.done() + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + await asyncio.wait_for(workers.context_cancelled.wait(), 2) + assert not workers.context_done.is_set() + finally: + release_context.set() + pending.cancel() + with contextlib.suppress(asyncio.CancelledError): + await pending + + +@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_responses_body_starts_cancels_the_context_phase( + monkeypatch, schedule_style, asgi_version +): + """The pipeline is cancelled even if the body is never iterated.""" + from starlette.requests import ClientDisconnect + + _short_waits(monkeypatch) + release_context = asyncio.Event() + workers = _Workers(_waits_for(release_context)) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + response = await asyncio.wait_for(handler(request, _raw_request()), 2) + 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_responses_body_starts_closes_the_stream_it_answered_with( + monkeypatch, schedule_style, asgi_version +): + """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)) + handler = _responses_route(_responses_server(schedule_style, workers)) + request = ResponsesRequest(model=_MODEL, input="hi", stream=True) + response = await asyncio.wait_for(handler(request, _raw_request()), 2) + pipeline = response._pipeline + 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((pipeline,), 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(pipeline.result(), StreamingResponse) + # Closed by the time the response is done, while the generation worker goes on. + assert workers.generation_closed.is_set() + + +def _frame(event_type: str, sequence_number: int, **fields) -> str: + data = {"type": event_type, "sequence_number": sequence_number, **fields} + return f"event: {event_type}\ndata: {json.dumps(data, separators=(',', ':'))}\n\n" + + +async def _transport(data: str, piece: int): + """``data`` as a relayed stream delivers it, in pieces of ``piece`` bytes.""" + encoded = data.encode() + for start in range(0, len(encoded), piece): + yield encoded[start : start + piece] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("piece", [1, 7, 4096]) +async def test_the_relay_drops_the_worker_opening_and_keepalives_and_numbers_the_rest(piece): + worker = "".join( + [ + _frame("response.created", 0), + _frame("response.in_progress", 1), + # Kept alive before its first output. + _frame("response.in_progress", 2), + _frame("response.output_text.delta", 3, delta="hi"), + # Kept alive after it. + _frame("response.in_progress", 4), + _frame("response.completed", 5), + ] + ) + + relayed = responses_utils.relay_after_opening(_transport(worker, piece), 2) + + assert [frame async for frame in relayed] == [ + _frame("response.output_text.delta", 2, delta="hi"), + _frame("response.completed", 3), + ] + + +@pytest.mark.asyncio +async def test_the_relay_numbers_a_failure_the_worker_reported_before_its_opening(): + worker = _frame("error", 0, code="upstream_error", message="ClientResponseError") + + relayed = responses_utils.relay_after_opening(_transport(worker, 4096), 3) + + assert [frame async for frame in relayed] == [ + _frame("error", 3, code="upstream_error", message="ClientResponseError") + ] + + +@pytest.mark.asyncio +async def test_the_relay_passes_on_an_event_its_stream_stopped_inside(): + cut = 'event: response.output_text.delta\ndata: {"type":"response.out' + worker = _frame("response.output_text.delta", 2, delta="hi") + cut + + relayed = responses_utils.relay_after_opening(_transport(worker, 7), 2) + + assert [frame async for frame in relayed] == [ + _frame("response.output_text.delta", 2, delta="hi"), + cut, + ] + + +@pytest.mark.asyncio +async def test_the_relay_closes_its_stream_when_it_stops_early(): + closed = asyncio.Event() + + async def worker(): + try: + yield _frame("response.output_text.delta", 2, delta="hi").encode() + yield _frame("response.completed", 3).encode() + finally: + closed.set() + + relayed = responses_utils.relay_after_opening(worker(), 2) + await relayed.__anext__() + await relayed.aclose() + + assert 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")