From 7d21ac9300e57de9757d5f323a7fcbfa8001ca3c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 20:17:21 +0000 Subject: [PATCH] fix(streaming): drain sync (boto3) streams in one thread to stop bursty delivery --- .../litellm_core_utils/streaming_handler.py | 88 ++++++++++++++++--- .../test_streaming_handler.py | 85 ++++++++++++++++++ 2 files changed, 161 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 60dbf7c644a..27a252d0e03 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -9,6 +9,7 @@ import traceback from dataclasses import dataclass from typing import ( Any, + AsyncGenerator, AsyncIterator, Callable, Dict, @@ -63,18 +64,70 @@ _SYNC_ITER_EXHAUSTED = object() _GCHUNK_FIELDS: frozenset = frozenset(GChunk.__annotations__) -def _next_sync_or_exhausted(it: Any) -> Any: - """ - Call next(it) from a thread and return _SYNC_ITER_EXHAUSTED on StopIteration. +@dataclass(frozen=True, slots=True) +class _SyncStreamError: + """Carries an exception raised while draining a synchronous stream across the + thread boundary so it can be re-raised on the consuming event loop.""" - asyncio.to_thread re-raises thread exceptions inside a coroutine, where PEP 479 - converts StopIteration to RuntimeError before any except clause can catch it. - Returning a sentinel instead keeps StopIteration out of the coroutine boundary. + exc: BaseException + + +async def _aiter_sync_stream_in_thread(sync_stream: Any) -> AsyncGenerator[Any, None]: """ + Drain a synchronous iterator (e.g. a boto3 Bedrock event stream) in a single + background thread, delivering items to the event loop through an asyncio.Queue. + + Dispatching one asyncio.to_thread per chunk pays thread-pool scheduling latency + on every chunk; during that latency the upstream socket buffers several SSE + events that then arrive back-to-back, so clients see stalls followed by bursts. + Iterating the whole stream in one thread lets each chunk reach the consumer as + soon as the provider emits it, restoring the provider's native cadence. + + A dedicated daemon thread is used rather than a pooled executor so a long-lived + stream never occupies a slot in the shared/default thread pool that the rest of + the event loop relies on for unrelated blocking work. + """ + loop = asyncio.get_running_loop() + queue: asyncio.Queue[Any] = asyncio.Queue() + stop = threading.Event() + + def _deliver(item: Any) -> None: + """ + Hand one item to the consumer's loop. A consumer may abandon the stream + (client disconnect) and let its event loop close while this thread is + still blocked inside next(); call_soon_threadsafe then raises, so stop + draining rather than crash the thread. + """ + if stop.is_set(): + return + try: + loop.call_soon_threadsafe(queue.put_nowait, item) + except RuntimeError: + stop.set() + + def _drain() -> None: + try: + for item in sync_stream: + if stop.is_set(): + break + _deliver(item) + except Exception as exc: + _deliver(_SyncStreamError(exc)) + finally: + _deliver(_SYNC_ITER_EXHAUSTED) + + thread = threading.Thread(target=_drain, daemon=True) + thread.start() try: - return next(it) - except StopIteration: - return _SYNC_ITER_EXHAUSTED + while True: + item = await queue.get() + if item is _SYNC_ITER_EXHAUSTED: + return + if isinstance(item, _SyncStreamError): + raise item.exc + yield item + finally: + stop.set() def is_async_iterable(obj: Any) -> bool: @@ -127,6 +180,7 @@ class CustomStreamWrapper: self.custom_llm_provider = custom_llm_provider self.logging_obj: LiteLLMLoggingObject = logging_obj self.completion_stream = completion_stream + self._threaded_sync_stream: AsyncGenerator[Any, None] | None = None self.sent_first_chunk = False self.sent_last_chunk = False self._stream_created_time: float = time.time() @@ -223,6 +277,8 @@ class CustomStreamWrapper: return self async def aclose(self): + threaded_sync_stream = self._threaded_sync_stream + self._threaded_sync_stream = None if self.completion_stream is not None: stream_to_close = self.completion_stream self.completion_stream = None @@ -230,6 +286,14 @@ class CustomStreamWrapper: # Without this, CancelledError is thrown into every await during # task group cancellation, preventing HTTP connection release. with anyio.CancelScope(shield=True): + if threaded_sync_stream is not None: + try: + await threaded_sync_stream.aclose() + except BaseException as e: + verbose_logger.debug( + "CustomStreamWrapper.aclose: error closing threaded sync stream: %s", + e, + ) try: if hasattr(stream_to_close, "aclose"): await stream_to_close.aclose() @@ -1982,9 +2046,9 @@ class CustomStreamWrapper: if isinstance(self.completion_stream, str) or isinstance(self.completion_stream, bytes): chunk = self.completion_stream else: - chunk = await asyncio.to_thread(_next_sync_or_exhausted, self.completion_stream) # type: ignore[arg-type] - if chunk is _SYNC_ITER_EXHAUSTED: - raise StopAsyncIteration + if self._threaded_sync_stream is None: + self._threaded_sync_stream = _aiter_sync_stream_in_thread(self.completion_stream) + chunk = await self._threaded_sync_stream.__anext__() if chunk is not None and chunk != b"": processed_chunk = self.chunk_creator(chunk=chunk) if processed_chunk is None: diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 514714136fd..454069926a1 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2660,6 +2660,91 @@ async def test_custom_stream_wrapper_anext_exhaustion_raises_stop_async_iteratio pytest.fail(f"PEP 479 regression: StopIteration leaked as RuntimeError: {e}") +@pytest.mark.asyncio +async def test_custom_stream_wrapper_anext_drains_sync_iterator_concurrently( + logging_obj: Logging, +): + """ + Regression for #34502: a synchronous (boto3-style) stream must be drained in a + single background thread so the provider keeps emitting chunks while the async + consumer is still processing earlier ones. With the previous per-chunk + asyncio.to_thread dispatch the sync iterator only advances when the consumer + calls __anext__, so a slow consumer stalls the provider and chunks arrive in + bursts. This test asserts the provider races ahead of a slow consumer. + """ + + def _make_chunk(content: str) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-concurrent", + created=int(time.time()), + model="test-model", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content=content, + role="assistant", + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields={}, + usage=None, + ) + + class TimestampingIterator: + """Fast sync producer that records when each chunk leaves the iterator.""" + + def __init__(self, chunks, produced_at: list): + self._it = iter(chunks) + self._produced_at = produced_at + + def __iter__(self): + return self + + def __next__(self): + chunk = next(self._it) # raises StopIteration when exhausted + self._produced_at.append(time.monotonic()) + return chunk + + n_chunks = 5 + produced_at: list = [] + wrapper = CustomStreamWrapper( + completion_stream=TimestampingIterator( + [_make_chunk(str(i)) for i in range(n_chunks)], produced_at + ), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="cached_response", + ) + + consumer_delay = 0.05 + consumed_at: list = [] + async for _ in wrapper: + consumed_at.append(time.monotonic()) + await asyncio.sleep(consumer_delay) # simulate a slow downstream client + + assert len(produced_at) == n_chunks + assert len(consumed_at) >= n_chunks # finalize appends a trailing finish chunk + # The single-thread pump produces every chunk while the consumer is still + # draining the first couple, so production completes well before the consumer + # reaches the tail. Per-chunk dispatch would interleave the two, making the + # last production timestamp land right before each corresponding consumption. + reference_consumption = consumed_at[n_chunks - 2] + assert max(produced_at) < reference_consumption, ( + f"provider did not race ahead of slow consumer " + f"(last produced at {max(produced_at):.3f}, " + f"reference consumption at {reference_consumption:.3f})" + ) + + # Azure streaming chunks that reproduce issue #24221: # Azure sends an initial chunk with prompt_filter_results and choices=[], # then a chunk with role='assistant' and content='', then content chunks.