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

litellm/proxy/proxy_server.py

@@ -10012,7 +10012,6 @@ 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({})
@@ -10225,7 +10224,6 @@ async def async_data_generator(
     request: Request | None = None,
     *,
     responses_stream_errors: bool = False,
-    now: Callable[[], float] = time.monotonic,
 ):
     verbose_proxy_logger.debug("inside generator")
     stream_completed = False
@@ -10289,11 +10287,8 @@ async def async_data_generator(
             else stream_iterator
         )
 
-        loop_released_at = now()  # rebind-ok: last time this stream let the event loop run other tasks
         async for item in stream_source:
-            if now() - loop_released_at >= _STREAM_LOOP_HOLD_SECONDS:
-                await asyncio.sleep(0)
-                loop_released_at = now()
+            await asyncio.sleep(0)
             if item is _STREAM_KEEPALIVE:
                 yield ": ping\n\n"
                 continue

tests/unit/proxy/proxy_server/test_streaming_helpers.py

@@ -2095,16 +2095,13 @@ def test_fast_serialize_simple_model_response_stream_keeps_served_service_tier()
 @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 on its own: another task must get
-    to run before the stream is over, and once the hold budget has been used
-    up, not on every chunk."""
+    the generator has to hand the loop back itself: another task must get to
+    run before the stream is over."""
     _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 per chunk
 
     async def chunks():
         for index in range(6):
-            clock[0] += 0.02
             order.append(f"chunk{index}")
             yield _simple_chunk(content=str(index))
 
@@ -2116,10 +2113,8 @@ async def test_async_data_generator_lets_other_tasks_run_during_a_stream_that_ne
         response=chunks(),
         user_api_key_dict=_user_auth(),
         request_data={"model": "gpt-4"},
-        now=lambda: clock[0],
     ):
         pass
     await other
 
     assert "other" in order[:-1], order
-    assert order.index("other") >= 2, order