scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(proxy): close streams before rebuilding the served output for logging (#51487)

Browse source 
* fix(proxy): close streams before rebuilding the served output for logging

The streaming iterator hook awaited _record_served_stream_output before it
returned, so every streamed response stayed open after its last event while
the proxy re-parsed all served chunks. That held the connection about 11 ms
plus 25 ms per 1,000 output tokens, and SDK clients wait for the close.

The parked deferred logging is now claimed at stream end and run in a rooted
background task after the served output is recorded, so logging still sees
the served text and the client sees the stream end right away.

* fix(proxy): cap background served-output rebuilds and fall back to waiting past it

Each background rebuild holds its served chunks until it runs, and the stream's
concurrency slot is already released. Past 64 pending rebuilds a stream waits
for its own rebuild as it did before, so overload keeps its backpressure and
no spend logging is dropped.
This commit is contained in:
maciej-berri 2026-11-14 14:32:48 -08:00 • committed by GitHub
parent bc3113fff2
commit 8366e15be0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 193 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
litellm/proxy/utils.py
View file

@ -404,6 +404,8 @@ class _TeamMembershipRow(TypedDict):
team_members_with_roles: ReadOnly[Iterable[_TeamMemberRow] | None]
_ROOTED_STREAM_LOGGING_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep tasks alive
_MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS: Final = 64
_VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0
_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0
_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"'
@ -5335,8 +5337,17 @@ class ProxyLogging:
# completed. unified_guardrail writes guardrail_information during
# its end-of-stream block (inside current_response), so by the time
# we reach this point the metadata is fully populated.
await ProxyLogging._record_served_stream_output(request_data, served_chunks)
ProxyLogging.fire_deferred_stream_logging(request_data)
# Rebuilding the served output only feeds logging, so it runs after the stream closes,
# unless enough rebuilds are already pending that this stream should wait like it used to.
if len(_ROOTED_STREAM_LOGGING_TASKS) >= _MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS:
await ProxyLogging._record_served_stream_output(request_data, served_chunks)
ProxyLogging.fire_deferred_stream_logging(request_data)
return
ProxyLogging._spawn_rooted(
ProxyLogging._log_served_stream(
request_data, served_chunks, ProxyLogging._take_deferred_stream_logging(request_data)
)
)
async def async_post_call_event_stream_iterator_hook(
self,
@ -5598,6 +5609,39 @@ class ProxyLogging:
texts: Final = await offload_token_count(served_stream_output_texts)(served_chunks)
record_served_output_texts(logging_obj.model_call_details, texts)
@staticmethod
async def _log_served_stream(
request_data: Mapping[str, object],
served_chunks: Sequence[object],
deferred_logging: Coroutine[object, object, object] | None,
) -> None:
try:
await ProxyLogging._record_served_stream_output(request_data, served_chunks)
finally:
if deferred_logging is not None:
_ = await deferred_logging
@staticmethod
def _spawn_rooted(work: Coroutine[object, object, None]) -> None:
task: Final = asyncio.create_task(work)
_ROOTED_STREAM_LOGGING_TASKS.add(task)
task.add_done_callback(_ROOTED_STREAM_LOGGING_TASKS.discard)
@staticmethod
def _take_deferred_stream_logging(request_data: Mapping[str, object]) -> Coroutine[object, object, object] | None:
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return None
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
)
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
if _deferred_cb is None or _args is None:
return None
logging_obj._on_deferred_stream_complete = None # pyright: ignore[reportAttributeAccessIssue] # slot parked by the stream iterators
logging_obj._deferred_stream_complete_args = None # pyright: ignore[reportAttributeAccessIssue] # slot parked by the stream iterators
return _deferred_cb(*_args)
@staticmethod
def fire_deferred_stream_logging(request_data: dict) -> None:
"""
@ -5608,17 +5652,9 @@ class ProxyLogging:
scheduling via create_task (which would race with unified_guardrail's
end-of-stream block). This method retrieves and fires them.
"""
logging_obj: Final = request_data.get("litellm_logging_obj")
if logging_obj is None:
return
_deferred_cb: Final[Callable[..., Coroutine[object, object, object]] | None] = getattr(
logging_obj, "_on_deferred_stream_complete", None
)
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
if _deferred_cb is not None and _args is not None:
logging_obj._on_deferred_stream_complete = None
logging_obj._deferred_stream_complete_args = None
asyncio.create_task(_deferred_cb(*_args))
deferred_logging: Final = ProxyLogging._take_deferred_stream_logging(request_data)
if deferred_logging is not None:
asyncio.create_task(deferred_logging)
_fire_deferred_stream_logging = fire_deferred_stream_logging

11
tests/unit/proxy/test_proxy_logging_hook_detection.py
View file

@ -1,3 +1,5 @@
import asyncio
import pytest
from fastapi import HTTPException
@ -591,6 +593,13 @@ async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeyp
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
logging_obj = _streaming_logging_obj()
served_texts_seen_by_logging = asyncio.get_running_loop().create_future()
async def deferred_logging() -> None:
served_texts_seen_by_logging.set_result(logging_obj.model_call_details.get(SERVED_OUTPUT_TEXTS_KEY))
logging_obj._on_deferred_stream_complete = deferred_logging
logging_obj._deferred_stream_complete_args = ()
async def fake_stream():
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
@ -606,7 +615,7 @@ async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeyp
delivered_text += choice.delta.content or ""
assert "zebra" not in delivered_text
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
assert await asyncio.wait_for(served_texts_seen_by_logging, timeout=5) == (delivered_text,)
@pytest.mark.asyncio

128
tests/unit/proxy/test_proxy_utils.py
View file

@ -3768,3 +3768,131 @@ async def test_query_raw_within_statement_timeout_leaves_other_database_errors_a
await query_raw_within_statement_timeout(prisma, "SELECT 1")
assert exc_info.value is failure
@pytest.mark.asyncio
async def test_stream_ends_without_waiting_for_served_output_and_logging_still_sees_it(monkeypatch):
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop
from litellm.proxy.utils import ProxyLogging
class PassThroughIterator(CustomLogger):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
async for chunk in response:
yield chunk
monkeypatch.setattr(litellm, "callbacks", [PassThroughIterator()])
logging_obj: Final = Logging(
model="claude-sonnet-4-5",
messages=[],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id="served-output-close",
function_id="served-output-close",
)
served_texts_seen_by_logging: Final = asyncio.get_running_loop().create_future()
async def deferred_logging() -> None:
served_texts_seen_by_logging.set_result(logging_obj.model_call_details.get(SERVED_OUTPUT_TEXTS_KEY))
logging_obj._on_deferred_stream_complete = deferred_logging
logging_obj._deferred_stream_complete_args = ()
events: Final = _anthropic_thinking_sse(50)
async def upstream():
for event in events:
yield event
async def drain() -> list[object]:
return [
chunk
async for chunk in ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
response=upstream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1", request_route="/v1/messages"),
request_data={"model": "claude-sonnet-4-5", "litellm_logging_obj": logging_obj},
)
]
offload_limiter: Final = _count_offload_limiter_for_this_loop()
holders: Final = tuple(object() for _ in range(int(offload_limiter.total_tokens)))
for holder in holders:
await offload_limiter.acquire_on_behalf_of(holder)
served: Final = await asyncio.wait_for(drain(), timeout=5)
for holder in holders:
offload_limiter.release_on_behalf_of(holder)
assert served == events
assert await asyncio.wait_for(served_texts_seen_by_logging, timeout=5) == ("final answer",)
@pytest.mark.asyncio
async def test_streams_wait_for_their_served_output_once_background_rebuilds_hit_the_cap(monkeypatch):
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
from litellm.litellm_core_utils.token_counter import _count_offload_limiter_for_this_loop
from litellm.proxy.utils import _MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS, ProxyLogging
class PassThroughIterator(CustomLogger):
async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):
async for chunk in response:
yield chunk
monkeypatch.setattr(litellm, "callbacks", [PassThroughIterator()])
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
events: Final = _anthropic_thinking_sse(1)
logging_objs: Final = tuple(
Logging(
model="claude-sonnet-4-5",
messages=[],
stream=True,
call_type="anthropic_messages",
start_time=datetime.now(),
litellm_call_id=f"served-output-cap-{index}",
function_id=f"served-output-cap-{index}",
)
for index in range(_MAX_BACKGROUND_SERVED_OUTPUT_REBUILDS + 1)
)
async def upstream():
for event in events:
yield event
async def drain(logging_obj: Logging) -> list[object]:
return [
chunk
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
response=upstream(),
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1", request_route="/v1/messages"),
request_data={"model": "claude-sonnet-4-5", "litellm_logging_obj": logging_obj},
)
]
offload_limiter: Final = _count_offload_limiter_for_this_loop()
holders: Final = tuple(object() for _ in range(int(offload_limiter.total_tokens)))
for holder in holders:
await offload_limiter.acquire_on_behalf_of(holder)
for logging_obj in logging_objs[:-1]:
assert await asyncio.wait_for(drain(logging_obj), timeout=5) == events
over_cap: Final = asyncio.ensure_future(drain(logging_objs[-1]))
await asyncio.sleep(0.2)
over_cap_waited: Final = not over_cap.done()
for holder in holders:
offload_limiter.release_on_behalf_of(holder)
assert over_cap_waited
assert await asyncio.wait_for(over_cap, timeout=5) == events
async def every_stream_recorded() -> None:
while any(SERVED_OUTPUT_TEXTS_KEY not in obj.model_call_details for obj in logging_objs):
await asyncio.sleep(0.01)
await asyncio.wait_for(every_stream_recorded(), timeout=10)

8
tests/unit/proxy/utils/proxy_logging/test_streaming_hooks.py
View file

@ -563,8 +563,12 @@ async def test_native_messages_stream_logging_fires_after_guardrail_end_of_strea
request_data=request_data,
):
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
async def logging_dispatched() -> None:
while len(events) < 2:
await asyncio.sleep(0.01)
await asyncio.wait_for(logging_dispatched(), timeout=5)
assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)]