diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..53aedc773d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -445,6 +445,18 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80)) #### Logging callback constants #### REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) +# Backpressure + lifetime bounds for the /v1/messages streaming relay (see +# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is +# bounded so a slow client throttles the upstream pump instead of letting it +# buffer the whole response in memory; the detached-drain cap bounds how many +# post-disconnect drains may run concurrently so client behavior can't create +# unbounded worker state. +ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int( + os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024") +) +ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int( + os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100") +) LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 3458705c7f8..bff41b9acd2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -8,6 +8,10 @@ import httpx from pydantic import TypeAdapter from typing_extensions import TypedDict +from litellm.constants import ( + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE, +) from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -25,6 +29,13 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() # guidance and drop it again from the done callback. _UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks +# Rooted set of pumps still draining upstream AFTER their client disconnected. +# Bounds how many detached drains run at once so a burst of slow/abandoned +# streams can't pin unbounded worker memory; a pump over the cap bills what it +# already collected instead of continuing to drain. Only ever touched from the +# event loop, so a plain set + len() check needs no lock. +_DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains + INCOMPLETE_STREAM_ERROR_MESSAGE: Final = ( "Provider stream ended before emitting a message_stop event; " "the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated." @@ -51,6 +62,23 @@ def _is_terminal_stream_chunk(chunk: object) -> bool: return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk) +def _try_claim_detached_drain_slot() -> bool: + """Claim a detached-drain slot for the current task, bounding concurrency. + + Returns True if a slot was claimed (the caller may keep draining upstream + for billing) or False if the cap is already reached (the caller should stop + and bill what it has). Only touched from the event loop, so the check + + insert need no lock. + """ + if len(_DETACHED_STREAM_DRAINS) >= ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: + return False + current_task: Final = asyncio.current_task() + if current_task is not None: + _DETACHED_STREAM_DRAINS.add(current_task) + current_task.add_done_callback(_DETACHED_STREAM_DRAINS.discard) + return True + + def _incomplete_stream_error_sse_event() -> bytes: payload: Final = json.dumps( { @@ -205,9 +233,18 @@ class BaseAnthropicMessagesStreamingIterator: generating and billing the full response regardless of the client, so draining it to completion is what lets spend tracking see the real terminal ``message_delta`` / ``message_stop`` usage instead of a - truncated placeholder count. Chunks reach the client through a queue; - once the client goes away the pump only buffers for billing so the queue - can't grow unbounded. + truncated placeholder count. + + Chunks reach the client through a bounded queue. While the client is + connected the pump blocks on a full queue (racing the disconnect + signal), so a slow reader throttles the upstream read exactly as the old + direct ``yield`` did instead of letting the whole response buffer in + memory. Once the client goes away the pump stops enqueueing and only + keeps a single ``collected_chunks`` copy for billing, and the number of + such post-disconnect drains running at once is capped so client behavior + can't create unbounded worker state; over the cap the pump bills what it + has rather than draining further. Detached-drain lifetime is otherwise + bounded by the upstream stream/read timeout. An upstream failure (Bedrock read / decode / chunk-conversion error) that happens while the client is still connected is forwarded through @@ -217,63 +254,12 @@ class BaseAnthropicMessagesStreamingIterator: This method provides the common logic for both Anthropic and Bedrock implementations. """ - from litellm._logging import verbose_proxy_logger - - queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue() + queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue( + maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + ) client_detached: Final = asyncio.Event() - async def _pump_upstream() -> None: - collected_chunks: Final[list[bytes]] = [] - saw_terminal_event = False # rebind-ok: accumulates across the upstream loop - - async def _bill() -> None: - try: - await self._handle_streaming_logging(collected_chunks) - except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump - verbose_proxy_logger.warning( - "async_sse_wrapper billing failed after %d chunks: %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - - try: - async for chunk in completion_stream: - if self.completion_start_time is None: - self.completion_start_time = datetime.now() - saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) - encoded_chunk = self._convert_chunk_to_sse_format(chunk) - collected_chunks.append(encoded_chunk) - if not client_detached.is_set(): - queue.put_nowait(encoded_chunk) - except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend - if not client_detached.is_set(): - # Preserve the provider-specific failure: hand the original - # exception to the client-facing generator so it re-raises - # and the proxy's failure handling (status code, - # post_call_failure_hook) runs. The failure path owns - # logging here, so don't also success-bill. - queue.put_nowait(exc) - return - # Client already disconnected: nothing to propagate to and no - # failure hook will run, so salvage the partial spend instead - # of dropping the request entirely. - verbose_proxy_logger.warning( - "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", - len(collected_chunks), - type(exc).__name__, - exc, - ) - await _bill() - return - - if not client_detached.is_set(): - if not saw_terminal_event: - queue.put_nowait(_incomplete_stream_error_sse_event()) - queue.put_nowait(None) - await _bill() - - pump_task: Final = asyncio.create_task(_pump_upstream()) + pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached)) _UPSTREAM_PUMP_TASKS.add(pump_task) pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard) @@ -288,7 +274,130 @@ class BaseAnthropicMessagesStreamingIterator: finally: # Client-facing generator is being torn down (normal end, a # re-raised upstream error, or a disconnect GeneratorExit). Signal - # the pump to stop enqueueing so the queue can't grow unbounded, but - # let it keep draining upstream to its terminal usage event for - # accurate billing. + # the pump to stop enqueueing and unblock any backpressure-blocked + # put; the pump then either finishes billing or drains detached + # (subject to the cap) for accurate usage. client_detached.set() + + async def _bill_collected_chunks( + self, + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _handle_streaming_logging + ) -> None: + from litellm._logging import verbose_proxy_logger + + try: + await self._handle_streaming_logging(collected_chunks) + except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump + verbose_proxy_logger.warning( + "async_sse_wrapper billing failed after %d chunks: %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + + @staticmethod + async def _enqueue_for_client( + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + item: bytes | None | BaseException, + ) -> bool: + """Deliver one item to the client, applying backpressure. + + Returns True if the item was queued, False if the client disconnected + before there was room (the item is then dropped, since a gone client + can't receive it). Never blocks once the client has detached. + """ + if client_detached.is_set(): + return False + try: + queue.put_nowait(item) + return True + except asyncio.QueueFull: + pass + put_task: Final = asyncio.ensure_future(queue.put(item)) + detached_task: Final = asyncio.ensure_future(client_detached.wait()) + try: + await asyncio.wait(frozenset((put_task, detached_task)), return_when=asyncio.FIRST_COMPLETED) + finally: + if not detached_task.done(): + detached_task.cancel() + if put_task.done() and not put_task.cancelled(): + return True + put_task.cancel() + return False + + async def _pump_upstream_to_queue( + self, + completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict], + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + ) -> None: + """Drain the whole upstream into ``queue`` (backpressured) and bill once. + + Runs detached so a client disconnect can't interrupt the upstream read; + see ``async_sse_wrapper`` for the full rationale. Returns after billing. + """ + from litellm._logging import verbose_proxy_logger + + collected_chunks: Final[list[bytes]] = [] # mutable-ok: SSE billing buffer appended to across the drain + saw_terminal_event = False # rebind-ok: accumulates across the upstream loop + draining_detached = False # rebind-ok: set once this pump claims a detached-drain slot + try: + async for chunk in completion_stream: + if self.completion_start_time is None: + self.completion_start_time = datetime.now() + saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk) + encoded_chunk = self._convert_chunk_to_sse_format(chunk) + collected_chunks.append(encoded_chunk) + if not client_detached.is_set(): + await self._enqueue_for_client(queue, client_detached, encoded_chunk) + continue + # Client has gone: keep draining only to reach the terminal usage + # event for billing, but claim a detached-drain slot first; over + # the cap, bill what we have rather than pinning more memory. + if not draining_detached: + if not _try_claim_detached_drain_slot(): + verbose_proxy_logger.warning( + "async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial " + "chunks without draining the rest of the upstream stream", + ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS, + len(collected_chunks), + ) + await self._bill_collected_chunks(collected_chunks) + return + draining_detached = True + except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend + await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc) + return + + if not client_detached.is_set(): + if not saw_terminal_event: + await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event()) + await self._enqueue_for_client(queue, client_detached, None) + await self._bill_collected_chunks(collected_chunks) + + async def _handle_pump_upstream_error( + self, + queue: "asyncio.Queue[bytes | None | BaseException]", + client_detached: "asyncio.Event", + collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks + exc: BaseException, + ) -> None: + from litellm._logging import verbose_proxy_logger + + if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc): + # Preserve the provider-specific failure: the client-facing + # generator re-raises it and the proxy's failure handling (status + # code, post_call_failure_hook) runs. That path owns logging, so + # don't also success-bill. + return + # Client already gone (or disconnected before the error reached it): no + # failure hook will run, so salvage the partial spend instead of + # dropping the request entirely. + verbose_proxy_logger.warning( + "async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)", + len(collected_chunks), + type(exc).__name__, + exc, + ) + await self._bill_collected_chunks(collected_chunks) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 21eeb514094..4afbe8fe833 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -9,6 +9,7 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( INCOMPLETE_STREAM_ERROR_MESSAGE, BaseAnthropicMessagesStreamingIterator, @@ -440,3 +441,129 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_ assert len(received) == 2 # Partial spend was still recorded rather than the whole request being dropped. assert iterator.logged_chunks == received + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch): + """ + Regression: the relay queue is bounded, so a slow client throttles the + upstream read instead of letting the pump buffer the whole response in + memory. With a tiny queue and a client that reads a single chunk, the pump + must stall after producing only a queue's worth of chunks ahead, not race + to the end of a large stream. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2) + + total = 200 + produced = 0 + + async def _fast_stream(): + nonlocal produced + for i in range(total): + produced += 1 + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}} + + iterator = _make_iterator("test_backpressure_slow_client") + gen = iterator.async_sse_wrapper(_fast_stream()) + try: + await gen.__anext__() # read exactly one chunk, then stall + # Let the pump run as far as the bounded queue permits. + for _ in range(500): + await asyncio.sleep(0) + # Bounded by queue maxsize + the one in-flight put + the one delivered + # chunk; nowhere near the full 200-chunk stream. + assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}" + assert produced < total + finally: + await gen.aclose() + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(monkeypatch): + """ + Regression: when the concurrent detached-drain cap is already reached, a + pump whose client has disconnected must bill what it collected instead of + continuing to drain (and accumulating) the rest of a large upstream stream, + so slow/abandoned clients can't pin unbounded worker state. + + The cap slot set is pre-occupied so the single slot is unavailable when this + pump reaches its first post-disconnect chunk; that isolates the cap decision + from multi-pump scheduling races. + """ + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + # Occupy the only detached-drain slot with a placeholder task. + async def _hold_slot(): + await asyncio.sleep(3600) + + holder = asyncio.ensure_future(_hold_slot()) + streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder) + tail_reached = False + + async def _long_stream(): + nonlocal tail_reached + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + # These arrive only after the client has disconnected. + for i in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}} + tail_reached = True + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_full"), request_body={}) + try: + gen = iterator.async_sse_wrapper(_long_stream()) + await gen.__anext__() # message_start + await gen.__anext__() # first delta + await gen.aclose() # client disconnects; 100+ chunks remain upstream + + for _ in range(200): + if iterator.logged_chunks: + break + await asyncio.sleep(0) + + # Billed a small bounded partial (the prefix plus at most a queue's + # worth the pump ran ahead before disconnect) without draining the + # 100-chunk tail. The exact count depends on how far the bounded queue + # let the pump run ahead, so assert the bound, not an exact number. + assert iterator.logged_chunks, "capped pump never billed" + assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE + assert len(iterator.logged_chunks) < 100 + assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + assert tail_reached is False, "pump kept draining past the cap instead of stopping" + finally: + holder.cancel() + streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder) + + +@pytest.mark.asyncio +async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch): + """Complement to the cap test: with a slot free, a disconnected pump drains + the full upstream and bills the terminal usage, and releases its slot after.""" + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + + async def _stream(): + yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}} + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}} + for i in range(20): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"m{i}"}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}} + yield {"type": "message_stop"} + + iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_free"), request_body={}) + gen = iterator.async_sse_wrapper(_stream()) + await gen.__anext__() + await gen.__anext__() + await gen.aclose() + + for _ in range(300): + if iterator.logged_chunks: + break + await asyncio.sleep(0.01) + + assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks) + # Slot released once the drain finished. + assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0