fix(proxy): streaming responses hand the event loop back every 50 ms instead of holding it for the whole buffered stream (#46925)

* fix(proxy): streaming responses hand the event loop back every 50 ms instead of holding it for the whole buffered stream

* fix(proxy): yield to the event loop on every streamed chunk, no hold budget or clock

* docs(proxy): say why the streaming loop yields on every chunk

* fix(proxy): streaming loop yields once it has held the event loop 50 ms, clock read inline

* docs(proxy): streaming loop comment names the 50 ms hold

* Revert "docs(proxy): streaming loop comment names the 50 ms hold"

This reverts commit cbf1371873.

* refactor(proxy): streaming loop hold budget is _STREAM_LOOP_HOLD_MS = 50

* Revert "refactor(proxy): streaming loop hold budget is _STREAM_LOOP_HOLD_MS = 50"

This reverts commit 356454def6.

* docs(proxy): streaming loop comment notes the yield is a floor, not a count of releases

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/proxy_server.py

@@ -10013,6 +10013,7 @@ def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]:
 _STREAM_KEEPALIVE: Final = object()
 
 _KEEPALIVE_MIN_SECONDS: Final = 1.0
+_STREAM_LOOP_HOLD_SECONDS: Final = 0.05
 _KEEPALIVE_MAX_SECONDS: Final = 300.0
 _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({})
 _EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
@@ -10288,7 +10289,14 @@ async def async_data_generator(
             else stream_iterator
         )
 
+        loop_released_at = time.monotonic()  # rebind-ok: last time this stream let other tasks run
         async for item in stream_source:
+            # Unblock the event loop: a stream whose chunks are already buffered never awaits
+            # between chunks, so without this yield one long stream holds the loop for its whole body.
+            # A real await in between may already have released the loop; this only guarantees it
+            if time.monotonic() - loop_released_at >= _STREAM_LOOP_HOLD_SECONDS:
+                await asyncio.sleep(0)
+                loop_released_at = time.monotonic()
             if item is _STREAM_KEEPALIVE:
                 yield ": ping\n\n"
                 continue

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -2090,3 +2090,36 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier()
 
     assert result is not None
     assert json.loads(result)["service_tier"] == "priority"
+
+
+@pytest.mark.asyncio
+async def test_async_data_generator_lets_other_tasks_run_during_a_stream_that_never_waits(monkeypatch):
+    """A stream whose chunks are already buffered never suspends on its own, so
+    the generator has to hand the loop back itself once it has held the loop
+    long enough: another task must get to run before the stream is over, and
+    not before the hold budget is used up."""
+    _patch_logging_flags(monkeypatch)
+    order: list[str] = []  # mutable-ok: records the interleaving of the two tasks
+    clock: list[float] = [0.0]  # mutable-ok: fake monotonic clock, advanced 20 ms per chunk
+    monkeypatch.setattr(ps.time, "monotonic", lambda: clock[0])
+
+    async def chunks():
+        for index in range(6):
+            clock[0] += 0.02
+            order.append(f"chunk{index}")
+            yield _simple_chunk(content=str(index))
+
+    async def other_task():
+        order.append("other")
+
+    other = asyncio.ensure_future(other_task())
+    async for _ in async_data_generator(
+        response=chunks(),
+        user_api_key_dict=_user_auth(),
+        request_data={"model": "gpt-4"},
+    ):
+        pass
+    await other
+
+    assert "other" in order[:-1], order
+    assert order.index("other") >= 3, order