diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9eb2bc93391..01048f35929 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10012,6 +10012,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({}) @@ -10224,6 +10225,7 @@ 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 @@ -10287,7 +10289,11 @@ 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() if item is _STREAM_KEEPALIVE: yield ": ping\n\n" continue diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 09d71cb4ead..9908d698a53 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/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 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.""" + _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)) + + 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"}, + now=lambda: clock[0], + ): + pass + await other + + assert "other" in order[:-1], order + assert order.index("other") >= 2, order