From e06e7321b02868e15d39c6b4f270384a17ee4e52 Mon Sep 17 00:00:00 2001 From: JingHao-Leon <102573344+JingHao-Leon@users.noreply.github.com> Date: Sun, 11 Oct 2026 16:26:57 +0800 Subject: [PATCH] fix(streaming): stop flushing the abandoned-stream log from aclose() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `router` job timed out twice at 20 minutes, always at the same test, so this was measured rather than guessed. Timing `tests/unit/test_router` locally with four workers: upstream/main 150s this branch, before this commit >1500s (did not finish) this branch, without the aclose flush 156s Same 12 failed / 1339 passed / 1339 passed / 32 errors in the two runs that complete, so nothing else moved. `aclose()` is litellm's own cleanup path — `Router` reaches it through `hasattr(model_response, "aclose")` on fallback and shutdown paths — so hooking the flush there turned every internal teardown into "assemble a response and submit a success log", each one running `stream_chunk_builder` (token counting) synchronously. That is a change in semantics for every caller, not a performance tweak, and it is what blew the budget. The flush stays reachable only through `close_stream()` and the context manager, where the caller has said it is abandoning the stream. The docstring records why the other two hooks are excluded so this is not re-added by a later "just cover aclose too". --- .../litellm_core_utils/streaming_handler.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f05cfcd51a4..dc409938d1f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -321,11 +321,25 @@ class CustomStreamWrapper: the proxy bills a partial stream on client disconnect. Only reachable through ``close_stream()`` / the context manager. It is - deliberately not wired into ``__del__``: a finalizer runs at an - arbitrary point on an arbitrary thread, possibly during interpreter - shutdown or after the event loop it needs is gone, so submitting - logging from there is neither reliable nor testable. A caller that - abandons a stream has to say so. + deliberately not wired into either of the other two hooks this class + already had: + + - ``__del__`` runs at an arbitrary point on an arbitrary thread, + possibly during interpreter shutdown or after the event loop it + needs is gone, so submitting logging from there is neither reliable + nor testable. It also races the success callbacks of a stream that is + still being drained. + - ``aclose()`` is litellm's own cleanup path. `Router` calls it on + fallback and shutdown paths (`hasattr(model_response, "aclose")`), so + treating it as "the consumer abandoned this stream" would assemble a + response and submit a success log on behalf of every internal + teardown. Measured cost: with the flush in ``aclose()`` the + ``tests/unit/test_router`` suite went from 150s to over 25 minutes, + because each one runs `stream_chunk_builder` (token counting) + synchronously. + + So a caller that abandons a stream has to say so, through + ``close_stream()``. """ if self._abandonment_logged or self.sent_last_chunk: return @@ -469,9 +483,6 @@ class CustomStreamWrapper: self._restore_consumer_correlation_context(guarded=True) async def aclose(self): - # Emit the abandoned-stream log before closing the provider stream: - # the chunks collected so far are what this call actually served. - self._flush_abandoned_stream_logging() # Restore the consumer's outer context only after the underlying # provider stream's own close (and its diagnostic logging below, if # closing fails) completes - not before - so those log lines still