From ea4ee8ff218676881c13c25193781f8390171aae Mon Sep 17 00:00:00 2001 From: Rabari Mihir <169349389+Mihir-Rabari@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:09:00 +0000 Subject: [PATCH 1/2] fix(proxy): yield to the event loop between buffered stream chunks async_data_generator drains a buffered chunk backlog without handing the loop back, because every await between chunks completes without suspending while upstream data is queued. One long streamed answer then monopolizes the worker for as long as the backlog lasts, stalling concurrent requests and the health endpoints. Hand the loop back once per chunk so a backlog interleaves with the rest of the worker. Fixes #46516 --- litellm/proxy/proxy_server.py | 4 ++ ...test_proxy_server_endpoints_and_startup.py | 58 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 60842648d16..f99c3804171 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10238,6 +10238,10 @@ async def async_data_generator( if item is _STREAM_KEEPALIVE: yield ": ping\n\n" continue + # Hand the event loop back once per chunk. A buffered backlog is + # otherwise drained back-to-back and starves every other task on + # the worker for as long as the backlog lasts + await asyncio.sleep(0) chunk = item if needs_per_chunk_hook: ### CALL HOOKS ### - modify outgoing data diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3ebe445d89b..9f6500304f5 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -4350,6 +4350,64 @@ async def test_async_data_generator_midstream_error(): mock_proxy_logging_obj.post_call_failure_hook.assert_not_called() +@pytest.mark.asyncio +async def test_async_data_generator_yields_to_event_loop_while_chunks_are_buffered() -> None: + """ + Draining a buffered chunk backlog must not monopolize the event loop. Other + scheduled tasks (health probes, concurrent requests) need to run between + chunks, so the probe task's progress must advance while the stream is + still being consumed + """ + from collections.abc import AsyncIterator + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.proxy_server import async_data_generator + from litellm.proxy.utils import ProxyLogging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + num_chunks = 200 + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_request_data = {"messages": [{"role": "user", "content": "test"}]} + + async def buffered_source() -> AsyncIterator[ModelResponseStream]: + for i in range(num_chunks): + yield ModelResponseStream( + id="chatcmpl-test-backlog", + choices=[StreamingChoices(index=0, delta=Delta(content=f"tok{i} "))], + model="gpt-test", + ) + + probe_progress = {"steps": 0} + + async def probe() -> None: + while probe_progress["steps"] < num_chunks: + probe_progress["steps"] += 1 + await asyncio.sleep(0) + + mock_proxy_logging_obj = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.needs_iterator_wrap.return_value = False + mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj), patch( + "litellm.proxy.proxy_server.llm_router", None + ): + + async def consume() -> list[int]: + return [ + probe_progress["steps"] + async for _ in async_data_generator( + buffered_source(), mock_user_api_key_dict, mock_request_data + ) + ] + + _, observed_progress = await asyncio.gather(probe(), consume()) + + assert observed_progress[0] < observed_progress[-1], ( + f"Event loop starved while {num_chunks} buffered chunks were drained, " + f"probe progress stayed at {observed_progress[0]} the whole stream" + ) + + def _has_nested_none_values(obj, path="root"): """ Recursively check if an object contains nested None values. From 64b27eea92ec79f0e467d7a3f433b354d220ab98 Mon Sep 17 00:00:00 2001 From: Rabari Mihir <169349389+Mihir-Rabari@users.noreply.github.com> Date: Fri, 16 Oct 2026 04:29:05 +0000 Subject: [PATCH 2/2] test(proxy): require sustained event-loop interleaving across the buffered drain The old assertion only compared the probe's first and last observed step, which passes when the drain yields once early and then monopolizes the loop again. The probe must now advance through at least half the backlog and make progress across at least a quarter of the chunk transitions, so a regression that starves other tasks mid-drain turns the test red Addresses the review on PR #46810 --- .../proxy/test_proxy_server_endpoints_and_startup.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 9f6500304f5..12c7783187a 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -4402,9 +4402,14 @@ async def test_async_data_generator_yields_to_event_loop_while_chunks_are_buffer _, observed_progress = await asyncio.gather(probe(), consume()) - assert observed_progress[0] < observed_progress[-1], ( + assert observed_progress[-1] >= num_chunks // 2, ( f"Event loop starved while {num_chunks} buffered chunks were drained, " - f"probe progress stayed at {observed_progress[0]} the whole stream" + f"probe only reached {observed_progress[-1]} steps" + ) + sustained_advances = sum(1 for a, b in zip(observed_progress, observed_progress[1:]) if b > a) + assert sustained_advances >= num_chunks // 4, ( + f"Probe advanced in only {sustained_advances} of {num_chunks} chunks, " + f"the drain yields once at the start and then monopolizes the loop again" )