mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
fix(streaming): stop flushing the abandoned-stream log from aclose()
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".
This commit is contained in:
parent
17210cd7b0
commit
e06e7321b0
1 changed files with 19 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue