diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8d6c9dd925b..49bd10c37ce 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10460,6 +10460,10 @@ async def async_data_generator( stream_started = True 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 14196d4a1ac..9783e92a88b 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -4351,6 +4351,69 @@ 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[-1] >= num_chunks // 2, ( + f"Event loop starved while {num_chunks} buffered chunks were drained, " + 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" + ) + + def _has_nested_none_values(obj, path="root"): """ Recursively check if an object contains nested None values.