diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2dc71abee3e..c6bcb235cec 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2068,7 +2068,7 @@ class CustomStreamWrapper: ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" if self.logging_obj is not None: - self._record_partial_usage_for_failure() + await asyncio.to_thread(self._record_partial_usage_for_failure) ## LOGGING asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) @@ -2076,10 +2076,10 @@ class CustomStreamWrapper: self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: - self._log_stream_failure_and_raise(e) + await self._alog_stream_failure_and_raise(e) return await self._finalize_completed_stream(cache_hit=cache_hit) except Exception as e: - self._log_stream_failure_and_raise(e) + await self._alog_stream_failure_and_raise(e) async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: @@ -2183,10 +2183,10 @@ class CustomStreamWrapper: # relies on aclose() or the best-effort __del__ guard. return processed_chunk - def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: + async def _alog_stream_failure_and_raise(self, e: Exception) -> NoReturn: traceback_exception: Final = traceback.format_exc() if self.logging_obj is not None: - self._record_partial_usage_for_failure() + await asyncio.to_thread(self._record_partial_usage_for_failure) ## LOGGING asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..0f77db3cdcf 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -183,7 +183,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons Awaited directly by the shielded cleanup rather than scheduled with create_task: the client is already gone so the extra latency is harmless, - and an unrooted task could be garbage-collected before it bills. + and an unrooted task could be garbage-collected before it bills. Assembly + is offloaded to a thread because a partial stream never carries a usage + chunk, so it re-tokenizes the whole delivered response with tiktoken, which + on a multi-MB stream blocks the event loop for hundreds of milliseconds. Returns True when a disconnect-time success event owns the request's max_parallel_requests slot release (one was dispatched here, or one had @@ -211,7 +214,8 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons ) messages: Final[object] = getattr(response, "messages", None) try: - partial_response: Final = litellm.stream_chunk_builder( + partial_response: Final = await asyncio.to_thread( + litellm.stream_chunk_builder, chunks=chunks, messages=messages if isinstance(messages, list) else None, logging_obj=logging_obj, 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 101935cac0a..3c8e7aeb94e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -3190,6 +3190,63 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): assert isinstance(logging_obj.model_call_details["response_cost"], float) +@pytest.mark.asyncio +async def test_async_stream_failure_recovers_partial_usage_off_the_event_loop(): + """Recovering partial usage re-tokenizes everything streamed so far. On a + multi-MB stream that costs hundreds of ms of tiktoken, so the async failure + path must not run it on the event loop thread and stall the whole worker. + """ + import threading + + class _FailingStream: + def __aiter__(self): + return self + + async def __anext__(self): + raise ValueError("upstream died mid-stream") + + logging_obj = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="partial-usage-offload", + function_id="1245", + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + + wrapper = CustomStreamWrapper( + completion_stream=_FailingStream(), + model="gpt-4o-mini", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-offload", + created=1742056047, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Rome fell", role="assistant"))], + ) + ] + + builder_threads = [] + real_builder = litellm.stream_chunk_builder + + def _record_thread(*args, **kwargs): + builder_threads.append(threading.get_ident()) + return real_builder(*args, **kwargs) + + with patch.object(litellm, "stream_chunk_builder", _record_thread): + with pytest.raises(Exception): + await wrapper.__anext__() + + assert len(builder_threads) == 1 + assert builder_threads[0] != threading.get_ident() + + def test_record_partial_usage_for_failure_noop_without_chunks(): """With no chunks delivered there is nothing billed to recover, so the failure stash must stay absent and not force a zero-usage row. diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index aacc7498ccb..e8d1e048899 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,6 +1,7 @@ import asyncio import copy import datetime +import threading from types import SimpleNamespace from typing import AsyncGenerator, Callable, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -5097,6 +5098,31 @@ class TestStreamingClientDisconnectBilling: await stream_iter.__anext__() return response + @pytest.mark.asyncio + async def test_disconnect_billing_assembles_off_the_event_loop(self): + """ + A partial stream carries no usage chunk, so assembling it re-tokenizes + the whole delivered response. On a multi-MB stream that is hundreds of + ms of tiktoken, which stalls every other request on the worker if it + runs on the event loop thread. + """ + response = await self._start_partial_stream() + builder_threads: list[int] = [] + real_builder = litellm.stream_chunk_builder + + def _record_thread(*args, **kwargs): + builder_threads.append(threading.get_ident()) + return real_builder(*args, **kwargs) + + with patch.object(litellm, "stream_chunk_builder", _record_thread): + billed = await _bill_partial_streamed_spend_on_disconnect( + {"litellm_logging_obj": response.logging_obj}, response + ) + + assert billed is True + assert len(builder_threads) == 1 + assert builder_threads[0] != threading.get_ident() + @pytest.mark.asyncio async def test_disconnect_bills_partial_streamed_spend(self): recorder = _RecordingSuccessLogger()