diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/anthropic_adapter.py index cc07ae2f1b51..85785555dea0 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/anthropic_adapter.py @@ -24,6 +24,7 @@ chat templates, or the engine. """ +import asyncio import base64 import hashlib import json @@ -31,7 +32,7 @@ import traceback import uuid from dataclasses import dataclass, field -from typing import Any, AsyncIterator, Dict, List, Optional, Union +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union from fastapi.responses import JSONResponse @@ -52,6 +53,7 @@ AnthropicMessagesResponse, AnthropicMessageStartEvent, AnthropicMessageStopEvent, + AnthropicPingEvent, AnthropicSignatureDelta, AnthropicStopReason, AnthropicTextBlock, @@ -1073,10 +1075,12 @@ def error(self, message: str) -> List[str]: async def _iter_openai_sse_lines( - openai_sse: AsyncIterator[Union[str, bytes]], -) -> AsyncIterator[str]: + openai_sse: AsyncIterator[Optional[Union[str, bytes]]], +) -> AsyncIterator[Optional[str]]: """Yield complete lines from string frames or arbitrarily chunked bytes. + A None item, a tick from ``_with_ticks``, is passed through as None. + The in-process server produces complete SSE strings, while the disaggregated frontend relays raw ``aiohttp`` byte chunks. Network chunks need not align with either UTF-8 characters or SSE line boundaries, so @@ -1088,24 +1092,33 @@ async def _iter_openai_sse_lines( buffer = "" input_type: Optional[type] = None - async for payload in openai_sse: - if not isinstance(payload, (str, bytes)): - raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}") + try: + async for payload in openai_sse: + if payload is None: + yield None + continue + if not isinstance(payload, (str, bytes)): + raise TypeError(f"OpenAI SSE payload must be str or bytes, got {type(payload)!r}") - current_type = type(payload) - if input_type is None: - input_type = current_type - elif current_type is not input_type: - raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads") + current_type = type(payload) + if input_type is None: + input_type = current_type + elif current_type is not input_type: + raise TypeError("OpenAI SSE stream cannot mix str and bytes payloads") - if isinstance(payload, bytes): - buffer += decoder.decode(payload) - else: - buffer += payload + if isinstance(payload, bytes): + buffer += decoder.decode(payload) + else: + buffer += payload - while "\n" in buffer: - line, buffer = buffer.split("\n", 1) - yield line.rstrip("\r") + while "\n" in buffer: + line, buffer = buffer.split("\n", 1) + yield line.rstrip("\r") + finally: + # Leaving these lines early leaves the upstream too. + aclose = getattr(openai_sse, "aclose", None) + if aclose is not None: + await aclose() if input_type is bytes: buffer += decoder.decode(b"", final=True) @@ -1113,17 +1126,105 @@ async def _iter_openai_sse_lines( yield buffer.rstrip("\r") +# Seconds a stream may stay silent before it sends a ping. Clients and proxies +# reset their timers on any event; without one the stream is silent until the +# first token (queueing and prefill), and for as long as the model writes a tool +# call, which is sent only once complete. +STREAM_PING_INTERVAL_S = 10.0 + + +async def _with_ticks( + source: AsyncIterator[Any], timeout: Callable[[], Optional[float]] +) -> AsyncIterator[Any]: + """Yield the items of ``source``, and None each time ``timeout()`` seconds pass first. + + ``timeout()`` is asked before every wait; None waits for the next item. One + task reads ``source`` to its end, so a tick never interrupts a read and the + source is never resumed from another task; that task is cancelled when the + consumer stops. + """ + # One item at a time: the reader waits for the consumer, so a slow client + # holds back the source as reading it directly would. + items: asyncio.Queue = asyncio.Queue(maxsize=1) + end = object() + consumer_stopped = False + + async def read_all() -> None: + try: + async for item in source: + await items.put(item) + except asyncio.CancelledError: + if not consumer_stopped: + # Raised inside the source, not by the consumer stopping: the + # consumer still waits, and the end marker ends its stream + # instead of leaving it to pings. + await items.put(end) + raise + except Exception: + await items.put(end) + raise + else: + await items.put(end) + finally: + # The reader closes the source in every case: stopped while handing + # over an item, it leaves the source suspended at its yield. + aclose = getattr(source, "aclose", None) + if aclose is not None: + await aclose() + + reader = asyncio.create_task(read_all()) + try: + while True: + try: + item = await asyncio.wait_for(items.get(), timeout()) + except asyncio.TimeoutError: + yield None + continue + if item is end: + # Raises what reading the source raised. + await reader + return + yield item + finally: + consumer_stopped = True + reader.cancel() + # Return only once the reader has stopped and closed the source. + await asyncio.gather(reader, return_exceptions=True) + + async def reframe_openai_stream( openai_sse: AsyncIterator[Union[str, bytes]], model: str ) -> AsyncIterator[str]: - """Translate an OpenAI SSE stream into Anthropic SSE frames.""" + """Translate an OpenAI SSE stream into Anthropic SSE frames. + + Until [DONE], a ping goes out whenever no other event has for + ``STREAM_PING_INTERVAL_S`` seconds, before ``message_start`` too: the message + starts at the first upstream chunk, which carries the prompt's usage, so a + slow first token would otherwise leave the stream silent. Clients skip + pings wherever they appear. + """ 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] = [] + loop = asyncio.get_running_loop() + last_event = loop.time() + + def until_ping() -> Optional[float]: + if done: + return None + return max(0.0, STREAM_PING_INTERVAL_S - (loop.time() - last_event)) + + # The ticks wrap the upstream's payloads, not the lines split from them: + # one hand-off between the reading task and this one per payload. + lines = _iter_openai_sse_lines(_with_ticks(openai_sse, until_ping)) try: - async for line in _iter_openai_sse_lines(openai_sse): + async for line in lines: + if line is None: + yield anthropic_sse(AnthropicPingEvent()) + last_event = loop.time() + continue 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 @@ -1148,6 +1249,7 @@ async def reframe_openai_stream( raise AnthropicResponseError("Malformed upstream stream chunk") from e for frame in reframer.process_chunk(chunk): yield frame + last_event = loop.time() if done: for frame in closing_frames: yield frame @@ -1177,3 +1279,6 @@ async def reframe_openai_stream( return for frame in reframer.error("Internal server error"): yield frame + finally: + # Stops the task reading the upstream when this stream ends early. + await lines.aclose() diff --git a/tensorrt_llm/serve/anthropic_protocol.py b/tensorrt_llm/serve/anthropic_protocol.py index 1c331f75d661..929f04d520e0 100644 --- a/tensorrt_llm/serve/anthropic_protocol.py +++ b/tensorrt_llm/serve/anthropic_protocol.py @@ -501,6 +501,12 @@ class AnthropicMessageStopEvent(AnthropicBaseModel): type: Literal["message_stop"] = "message_stop" +class AnthropicPingEvent(AnthropicBaseModel): + """A keepalive: the stream format allows any number of them after message_start.""" + + type: Literal["ping"] = "ping" + + class AnthropicErrorEvent(AnthropicBaseModel): type: Literal["error"] = "error" error: AnthropicError diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index af3eaeb22d5d..9449c721e1c1 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -26,6 +26,7 @@ import pytest from pydantic import ValidationError +import tensorrt_llm.serve.anthropic_adapter as adapter from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, @@ -881,8 +882,11 @@ def run_reframer(chunks): def assert_event_invariants(events): - """First-principles invariants of the Anthropic streaming protocol.""" - names = [name for name, _ in events] + """First-principles invariants of the Anthropic streaming protocol. + + Pings may appear anywhere; the invariants hold for the other events. + """ + names = [name for name, _ in events if name != "ping"] assert names[0] == "message_start" assert names.count("message_start") == 1 assert names[-1] == "message_stop" @@ -1840,3 +1844,287 @@ def test_cached_tokens_exceeding_prompt_tokens_is_logged(monkeypatch): assert converted.input_tokens == 0 assert converted.cache_read_input_tokens == 99 assert any("exceeds prompt_tokens" in str(w) for w in warnings), warnings + + +# --------------------------------------------------------------------------- +# Keepalive pings while a stream is silent +# --------------------------------------------------------------------------- + + +def _line(c: ChatCompletionStreamResponse) -> str: + return f"data: {c.model_dump_json()}\n\n" + + +def _reframe_paced(monkeypatch, steps, interval): + """Reframe a source that waits ``pause`` seconds before each of its lines.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", interval) + + async def source(): + for pause, line in steps: + await asyncio.sleep(pause) + yield line + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + return parse_frames(asyncio.run(collect())) + + +def test_stream_pings_while_silent_after_message_start(monkeypatch): + """A silence longer than the interval is broken by pings, inside the message.""" + events = _reframe_paced( + monkeypatch, + [ + (0, _line(chunk({"content": "Writing the file."}))), + (0.3, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + names = [name for name, _ in events] + pings = [i for i, name in enumerate(names) if name == "ping"] + assert len(pings) >= 2 + assert names.index("message_start") < pings[0] < pings[-1] < names.index("message_delta") + assert all(events[i][1] == {"type": "ping"} for i in pings) + + +def test_stream_pings_while_waiting_for_the_first_token(monkeypatch): + """A slow first token is covered too: pings come first, then the message starts.""" + events = _reframe_paced( + monkeypatch, + [ + (0.3, _line(chunk({"content": "hello"}))), + (0, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + assert events[0] == ("ping", {"type": "ping"}) + + +def test_stream_pings_then_fails_before_the_first_token(monkeypatch): + """An upstream that fails before its first chunk still ends the stream with an error.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + await asyncio.sleep(0.1) + raise RuntimeError("engine failed") + yield # an async generator + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + names = [name for name, _ in parse_frames(asyncio.run(collect()))] + assert names[0] == "ping" + assert names[-1] == "error" + assert "message_start" not in names + + +def test_stream_pings_when_chunks_bring_no_event(monkeypatch): + """Upstream chunks that produce no event (an empty delta) do not count as activity.""" + steps = [(0, _line(chunk({"content": "Thinking about it."})))] + steps += [(0.02, _line(chunk({}))) for _ in range(15)] + steps += [(0, _line(chunk({}, finish_reason="stop"))), (0, "data: [DONE]\n\n")] + + events = _reframe_paced(monkeypatch, steps, interval=0.05) + + assert_event_invariants(events) + assert "ping" in [name for name, _ in events] + + +def test_stream_sends_no_ping_while_events_flow(monkeypatch): + """Every event restarts the wait: a stream longer than the interval but never silent has no ping. + + The 1 s interval is 50 times the 20 ms gap, so a loaded worker does not fake a silence. + """ + steps = [(0.02, _line(chunk({"content": f"word{i} "}))) for i in range(60)] + steps += [(0, _line(chunk({}, finish_reason="stop"))), (0, "data: [DONE]\n\n")] + + events = _reframe_paced(monkeypatch, steps, interval=1.0) + + assert_event_invariants(events) + assert "ping" not in [name for name, _ in events] + + +def test_stream_sends_no_ping_after_done(monkeypatch): + """After [DONE] the message is complete: a slow end of the upstream brings no ping.""" + events = _reframe_paced( + monkeypatch, + [ + (0, _line(chunk({"content": "hello"}))), + (0, _line(chunk({}, finish_reason="stop"))), + (0, "data: [DONE]\n\n"), + (0.3, "\n"), + ], + interval=0.05, + ) + + assert_event_invariants(events) + assert "ping" not in [name for name, _ in events] + + +@pytest.mark.parametrize("fails", [False, True]) +def test_a_slow_client_holds_back_the_upstream(monkeypatch, fails): + """The upstream is read only as fast as the client takes frames; its end or error still arrives.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 10) + words = [f"w{i} " for i in range(20)] + produced = [] + + async def source(): + for word in words: + produced.append(word) + yield _line(chunk({"content": word})) + if fails: + raise RuntimeError("engine failed") + yield _line(chunk({}, finish_reason="stop")) + yield "data: [DONE]\n\n" + + async def read_slowly(): + frames, lead = [], 0 + async for frame in reframe_openai_stream(source(), MODEL): + frames.append(frame) + taken = sum( + 1 + for name, payload in parse_frames(frames) + if name == "content_block_delta" and payload["delta"].get("text") in words + ) + lead = max(lead, len(produced) - taken) + await asyncio.sleep(0.005) + return parse_frames(frames), lead + + events, lead = asyncio.run(read_slowly()) + + assert lead <= 3, f"the upstream ran {lead} chunks ahead of the client" + texts = [p["delta"]["text"] for n, p in events if n == "content_block_delta"] + assert texts == words + assert events[-1][0] == ("error" if fails else "message_stop") + + +def test_a_client_that_leaves_while_a_frame_waits_closes_the_upstream(monkeypatch): + """The upstream is closed by the time the stream is, also while its next frame waits for the client.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 10) + closed = [] + + async def source(): + try: + for i in range(20): + yield _line(chunk({"content": f"w{i} "})) + finally: + closed.append(True) + + async def leave_after_one_frame(): + stream = reframe_openai_stream(source(), MODEL) + async for _ in stream: + break + # The reader fills the queue and then waits to hand over the next frame. + for _ in range(5): + await asyncio.sleep(0) + await stream.aclose() + return list(closed) + + assert asyncio.run(leave_after_one_frame()) == [True] + + +def test_stopping_the_stream_cancels_its_pending_read(monkeypatch): + """A client that goes away mid-silence stops the upstream read at once.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.05) + cancelled = [] + + async def source(): + yield _line(chunk({"content": "hello"})) + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + cancelled.append(True) + raise + yield _line(chunk({}, finish_reason="stop")) + + async def stop_at_first_ping(): + stream = reframe_openai_stream(source(), MODEL) + async for frame in stream: + if frame.startswith("event: ping"): + break + await stream.aclose() + # Checked while the loop still runs: asyncio.run cancels every task left at its exit. + for _ in range(100): + if cancelled: + break + await asyncio.sleep(0.01) + return list(cancelled) + + assert asyncio.run(asyncio.wait_for(stop_at_first_ping(), timeout=5)) == [True] + + +def test_stream_reads_its_upstream_from_one_task(monkeypatch): + """Pings never move the upstream between tasks. + + The upstream is the chat stream, and an async generator that holds a + task-bound context (a lock, a timeout or a cancel scope) across a yield + breaks if it is resumed from another task. + """ + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + readers = set() + + async def source(): + for line in [ + _line(chunk({"content": "one"})), + _line(chunk({"content": "two"})), + _line(chunk({}, finish_reason="stop")), + "data: [DONE]\n\n", + ]: + readers.add(asyncio.current_task()) + await asyncio.sleep(0.1) + yield line + + 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 "ping" in [name for name, _ in events] + assert len(readers) == 1 + + +def test_stream_ends_with_an_error_when_the_upstream_fails_mid_stream(monkeypatch): + """A failure while the stream pings still ends the message with an error event.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + yield _line(chunk({"content": "hello"})) + await asyncio.sleep(0.1) + raise RuntimeError("engine failed") + + async def collect(): + return [frame async for frame in reframe_openai_stream(source(), MODEL)] + + events = parse_frames(asyncio.run(collect())) + names = [name for name, _ in events] + assert "ping" in names + assert names[-1] == "error" + # Reported as the upstream's own failure, not as a stream that ended early. + assert events[-1][1]["error"]["message"] == "Internal server error" + + +def test_stream_ends_when_the_upstream_is_cancelled_on_its_own(monkeypatch): + """A cancellation raised inside the upstream ends the stream instead of leaving it to pings.""" + monkeypatch.setattr(adapter, "STREAM_PING_INTERVAL_S", 0.02) + + async def source(): + yield _line(chunk({"content": "hello"})) + await asyncio.sleep(0.1) + raise asyncio.CancelledError() + + async def collect(): + frames = [] + with pytest.raises(asyncio.CancelledError): + async for frame in reframe_openai_stream(source(), MODEL): + frames.append(frame) + assert len(frames) < 100, "the stream kept pinging after its upstream ended" + return frames + + events = parse_frames(asyncio.run(asyncio.wait_for(collect(), 5))) + assert "ping" in [name for name, _ in events]