diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index c03928825e1b..ff80c53eec6a 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -1062,8 +1062,20 @@ async def reframe_openai_stream( ) -> AsyncIterator[str]: """Translate an OpenAI SSE stream into Anthropic SSE frames.""" reframer = AnthropicStreamReframer(model=model) + done = False + # The frames that close the message, held from [DONE] until the upstream + # has been read to its end. + closing_frames: List[str] = [] try: async for line in _iter_openai_sse_lines(openai_sse): + if done: + # The upstream is read to its end before the message is closed: + # the chat stream records the request's metrics after its last + # frame, and a consumer that stops reading at message_stop would + # leave it suspended at [DONE], to be closed there (GeneratorExit) + # without doing that work. A relayed byte stream likewise ends + # cleanly at the worker's EOF. + continue line = line.strip() if not line.startswith("data:"): continue @@ -1071,15 +1083,19 @@ async def reframe_openai_stream( if not data: continue if data == "[DONE]": - for frame in reframer.finish(): - yield frame - return + closing_frames = reframer.finish() + done = True + continue try: chunk = ChatCompletionStreamResponse(**json.loads(data)) except (json.JSONDecodeError, ValueError) as e: raise AnthropicResponseError("Malformed upstream stream chunk") from e for frame in reframer.process_chunk(chunk): yield frame + if done: + for frame in closing_frames: + yield frame + return # Upstream ended without [DONE]. Both streaming producers emit it only # after the response is complete, so its absence means the stream was # cut short. Calling finish() here would close the message with @@ -1097,5 +1113,11 @@ async def reframe_openai_stream( yield frame except Exception as e: # noqa: BLE001 - stream must end with an event logger.error(f"Anthropic stream reframing failed: {e}\n{traceback.format_exc()}") + if done: + # The response was complete at [DONE], so a failure while reading + # the rest of the upstream is logged and the message still closes. + for frame in closing_frames: + yield frame + return for frame in reframer.error("Internal server error"): yield frame diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 753fa4d2b86b..a0d96c5ceedb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2061,7 +2061,10 @@ async def chat_stream_generator( yield "data: [DONE]\n\n" await self._extract_metrics(res, raw_request) nvtx_mark("generation ends") - except: + except Exception: + # Only failures are logged: a consumer that stops reading closes + # the stream with GeneratorExit, and a client that disconnects + # cancels it, and neither is an error of this stream. logger.error(traceback.format_exc()) raise diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 2e9bb12fc656..bbbc80254eed 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -1132,6 +1132,122 @@ async def collect(): assert "".join(text_deltas) == "hello" +def test_stream_reframer_reads_the_upstream_to_its_end_after_done(): + """The upstream runs past [DONE], so the work it does there runs. + + The chat stream records the request's metrics after its last frame. A + reframer that stopped reading at [DONE] left it suspended there, to be + closed with GeneratorExit: every streamed /v1/messages request lost its + metrics, and the server logged a traceback for it. + """ + after_done = [] + + async def source(): + yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n" + yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + after_done.append("ran") + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + assert_event_invariants(events) + assert after_done == ["ran"] + + +def test_stream_reframer_reads_the_upstream_to_its_end_before_message_stop(): + """A consumer that stops reading at message_stop still lets the upstream end. + + A client may disconnect as soon as it has message_stop. Had message_stop + gone out before the upstream was drained, the chat stream would be left + suspended at [DONE] and its metrics never recorded. + """ + after_done = [] + + async def source(): + yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n" + yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + after_done.append("ran") + + async def read_until_message_stop(): + stream = reframe_openai_stream(source(), MODEL) + frames = [] + async for frame in stream: + frames.append(frame) + if frame.startswith("event: message_stop"): + break + await stream.aclose() + return frames + + events = parse_frames(asyncio.run(read_until_message_stop())) + assert_event_invariants(events) + assert after_done == ["ran"] + + +def test_stream_reframer_drops_anything_after_done(): + """Lines after [DONE] are read to reach the end, never reframed.""" + + async def source(): + yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n" + yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + yield f"data: {chunk({'content': 'late'}).model_dump_json()}\n\n" + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + assert_event_invariants(events) + text_deltas = [ + payload["delta"]["text"] + for name, payload in events + if name == "content_block_delta" and payload["delta"]["type"] == "text_delta" + ] + assert "".join(text_deltas) == "hello" + + +def test_stream_reframer_adds_no_event_when_the_upstream_fails_after_done(): + """message_stop stays the last event when reading past [DONE] fails.""" + + async def source(): + yield f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n" + yield f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + raise RuntimeError("failed after the last frame") + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + assert_event_invariants(events) + assert "error" not in [name for name, _ in events] + + +def test_stream_reframer_reads_a_byte_stream_to_its_end_when_done_is_split(): + """A relayed byte stream is read to its EOF however [DONE] is chunked.""" + body = ( + f"data: {chunk({'content': 'hello'}).model_dump_json()}\n\n" + f"data: {chunk({}, finish_reason='stop').model_dump_json()}\n\n" + "data: [DONE]\n\n" + ).encode() + split = body.index(b"[DONE]") + len(b"[DO") + after_last_chunk = [] + + async def source(): + yield body[:split] + yield body[split:] + after_last_chunk.append("ran") + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + assert_event_invariants(events) + assert after_last_chunk == ["ran"] + + def test_stream_reframer_splits_multibyte_characters(): """A chunk boundary landing inside a UTF-8 sequence must not corrupt output. diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index 6c04971f764b..7c2df325f79f 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -27,8 +27,10 @@ from fastapi.testclient import TestClient from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError +from tensorrt_llm.serve import openai_server as openai_server_module from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionStreamResponse, @@ -38,6 +40,7 @@ UsageInfo, ) from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.serve.render import chat as render_chat_module # These tests are CPU-only (no GPU, engine or sockets) and run in the # CPU-Generic CI stage, which selects with `-m cpu_only`. @@ -271,6 +274,168 @@ def test_messages_route_reframes_streaming_response(server_kind): assert backend.await_args.args[0].stream +class _StreamedGeneration: + """A streamed generation whose steps are already post-processed frames. + + After its steps it ends, raises ``then``, or (``hang=True``) waits forever + for the next step, as a generation does while the model is still running. + """ + + prompt_token_ids = [1, 2, 3] + + def __init__(self, steps, *, then=None, hang=False): + self._steps = list(steps) + self._then = then + self._hang = hang + + def __aiter__(self): + return self + + async def __anext__(self): + if self._steps: + frames = self._steps.pop(0) + return SimpleNamespace( + finished=not self._steps, + outputs=[SimpleNamespace(_postprocess_result=frames)], + ) + if self._then is not None: + raise self._then + if self._hang: + await asyncio.Event().wait() + raise StopAsyncIteration + + +class _ModelConfig: + model_type = "llama" + vocab_size = 1024 + + +async def _chat_stream(monkeypatch, generation): + """The body of a streamed /v1/chat/completions response, and the server. + + The real ``openai_chat`` handler runs on a server built with + ``object.__new__``; prompt rendering and the engine are stubbed, and the + engine's output is already post-processed into SSE frames. + """ + + async def render(**kwargs): + return "rendered prompt" + + monkeypatch.setattr(render_chat_module, "async_apply_chat_template", render) + server = object.__new__(OpenAIServer) + server.model = MODEL + server.allow_request_chat_template = False + server.model_config = _ModelConfig() + server.processor = None + server.tokenizer = SimpleNamespace(tokenizer=SimpleNamespace(vocab_size=1024)) + server.chat_template = None + server.tool_parser = None + server.tool_call_id_type = "random" + server.multimodal_server_config = None + server._is_visual_gen = False + server._adjusted_steady_clock = SimpleNamespace(now=time.monotonic) + server.generator = SimpleNamespace( + args=SimpleNamespace( + gather_generation_logits=False, + reasoning_parser=None, + backend="pytorch", + guided_decoding_backend=None, + num_postprocess_workers=1, + ), + generate_async=Mock(return_value=generation), + ) + server.await_disconnected = AsyncMock() + server._extract_metrics = AsyncMock() + request = ChatCompletionRequest( + model=MODEL, messages=[{"role": "user", "content": "hi"}], stream=True + ) + raw_request = SimpleNamespace(state=SimpleNamespace(), headers={}) + response = await server.openai_chat(request, raw_request) + assert isinstance(response, StreamingResponse), response.body + return response.body_iterator, server + + +@pytest.fixture +def server_logger(monkeypatch): + log = Mock() + monkeypatch.setattr(openai_server_module, "logger", log) + return log + + +def test_chat_stream_read_to_its_end_records_the_metrics(monkeypatch, server_logger): + """Past [DONE] the chat stream records the request's metrics.""" + + async def run(): + generation = _StreamedGeneration([["data: one\n\n"], ["data: two\n\n"]]) + stream, server = await _chat_stream(monkeypatch, generation) + frames = [frame async for frame in stream] + return frames, server + + frames, server = asyncio.run(run()) + + assert frames == ["data: one\n\n", "data: two\n\n", "data: [DONE]\n\n"] + server._extract_metrics.assert_awaited_once() + server_logger.error.assert_not_called() + + +def test_chat_stream_closed_by_its_consumer_logs_no_error(monkeypatch, server_logger): + """A consumer that stops reading (GeneratorExit) is not a stream error.""" + + async def run(): + generation = _StreamedGeneration([["data: one\n\n"], ["data: two\n\n"]]) + stream, server = await _chat_stream(monkeypatch, generation) + assert await stream.__anext__() == "data: one\n\n" + await stream.aclose() + return server + + server = asyncio.run(run()) + + server._extract_metrics.assert_not_awaited() + server_logger.error.assert_not_called() + + +def test_chat_stream_cancelled_by_a_disconnect_logs_no_error(monkeypatch, server_logger): + """A client that disconnects cancels the stream; that is not a stream error.""" + + async def run(): + generation = _StreamedGeneration([["data: one\n\n"]], hang=True) + stream, _ = await _chat_stream(monkeypatch, generation) + first_frame = asyncio.Event() + + async def consume(): + async for _ in stream: + first_frame.set() + + consumer = asyncio.create_task(consume()) + # Bounded, so a stream that fails before its first frame fails the + # test instead of hanging it. + await asyncio.wait_for(first_frame.wait(), timeout=5.0) + # The stream is now waiting on the generation for its next step. + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + asyncio.run(run()) + + server_logger.error.assert_not_called() + + +def test_chat_stream_logs_a_failure_and_raises_it(monkeypatch, server_logger): + """A failure of the generation itself is still logged and raised.""" + + async def run(): + generation = _StreamedGeneration([["data: one\n\n"]], then=RuntimeError("engine failed")) + stream, _ = await _chat_stream(monkeypatch, generation) + with pytest.raises(RuntimeError, match="engine failed"): + async for _ in stream: + pass + + asyncio.run(run()) + + server_logger.error.assert_called_once() + assert "engine failed" in server_logger.error.call_args.args[0] + + def test_standard_and_disagg_register_messages_route(monkeypatch, tmp_path): # The disagg server's register_routes mounts a prometheus multiprocess # collector, which errors unless PROMETHEUS_MULTIPROC_DIR points at a real