Merge pull request #38385 from BerriAI/litellm_lit6184_sdk_async_redis_cache_write

fix(caching): flush async cache writes cancelled at event loop shutdown

litellm/caching/caching_handler.py

@@ -18,7 +18,7 @@ import asyncio
 import datetime
 import inspect
 import time
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
 
 from pydantic import BaseModel
@@ -27,6 +27,7 @@ import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.caching import InMemoryCache
 from litellm.caching.caching import S3Cache
+from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS
 from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
     update_response_metadata,
 )
@@ -124,6 +125,29 @@ def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") ->
     return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
 
 
+_PENDING_CACHE_WRITES: Final[set["asyncio.Task[None]"]] = set()  # mutable-ok: strong refs to pending write tasks
+
+
+async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], Awaitable[None]]) -> None:
+    try:
+        await write_factory()
+    except asyncio.CancelledError:
+        try:
+            await asyncio.wait_for(write_factory(), timeout=CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
+        except Exception as flush_error:  # noqa: BLE001  # shutdown flush failures are logged, never raised
+            verbose_logger.warning(
+                "LiteLLM Cache: pending cache write failed during event loop shutdown: %s", flush_error
+            )
+        raise
+
+
+def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]":
+    task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory))
+    _PENDING_CACHE_WRITES.add(task)
+    task.add_done_callback(_PENDING_CACHE_WRITES.discard)
+    return task
+
+
 def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
     """Read the caller-supplied ``cache_key`` off the request kwargs."""
     return request_kwargs.get("cache_key", None)
@@ -983,6 +1007,7 @@ class LLMCachingHandler:
 
         if litellm.cache is None:
             return
+        cache: Final = litellm.cache
 
         new_kwargs: Final = kwargs.copy()
         new_kwargs.update(
@@ -1004,24 +1029,24 @@ class LLMCachingHandler:
             ):
                 if (
                     isinstance(result, EmbeddingResponse)
-                    and litellm.cache is not None
-                    and not isinstance(litellm.cache.cache, S3Cache)  # s3 doesn't support bulk writing. Exclude.
+                    and not isinstance(cache.cache, S3Cache)  # s3 doesn't support bulk writing. Exclude.
                 ):
-                    asyncio.create_task(
-                        litellm.cache.async_add_cache_pipeline(
+                    create_cache_write_task(
+                        lambda: cache.async_add_cache_pipeline(
                             result, dynamic_cache_object=self.dual_cache, **new_kwargs
                         )
                     )
                 else:
-                    asyncio.create_task(
-                        litellm.cache.async_add_cache(
-                            result.model_dump_json(),
+                    result_json: Final = result.model_dump_json()
+                    create_cache_write_task(
+                        lambda: cache.async_add_cache(
+                            result_json,
                             dynamic_cache_object=self.dual_cache,
                             **new_kwargs,
                         )
                     )
             else:
-                asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs))
+                create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
 
     def sync_set_cache(
         self,

litellm/constants.py

@@ -381,6 +381,7 @@ AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_
 AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
 AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96))
 REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
+CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
 REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
 REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
 REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))

litellm/responses/streaming_iterator.py

@@ -573,13 +573,16 @@ class BaseResponsesAPIStreamingIterator:
         ):
             return
 
-        if litellm.cache is None:
+        cache: Final = litellm.cache
+        if cache is None:
             return
 
         cached_response: Final = response_obj.model_dump_json()
         if is_async:
-            cache_write_task: Final = asyncio.create_task(
-                litellm.cache.async_add_cache(
+            from litellm.caching.caching_handler import create_cache_write_task
+
+            cache_write_task: Final = create_cache_write_task(
+                lambda: cache.async_add_cache(
                     cached_response,
                     dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
                     **request_kwargs,
@@ -592,7 +595,7 @@ class BaseResponsesAPIStreamingIterator:
                 )
             )
         else:
-            litellm.cache.add_cache(
+            cache.add_cache(
                 cached_response,
                 dynamic_cache_object=getattr(caching_handler, "dual_cache", None),
                 **request_kwargs,

tests/test_litellm/caching/test_caching_handler.py

@@ -617,3 +617,44 @@ def test_request_kwargs_does_not_retain_logging_obj():
     assert "litellm_logging_obj" not in handler.request_kwargs
     assert handler.request_kwargs["messages"] == kwargs["messages"]
     assert handler.request_kwargs["model"] == "gpt-4o"
+
+
+def test_async_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
+    """
+    Regression test for the SDK losing async cache writes in short-lived scripts:
+    async_set_cache dispatched the write as a bare fire-and-forget task, so
+    asyncio.run cancelled it at loop close before the write landed (LIT-6184,
+    deterministic with hiredis installed). The write must survive loop shutdown.
+    """
+    import litellm
+
+    writes = []
+
+    class _SlowWriteCache:
+        supported_call_types = ["acompletion"]
+        cache = None
+
+        async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
+            await asyncio.sleep(0.2)
+            writes.append(result)
+
+    async def acompletion(**kwargs):
+        return None
+
+    handler = LLMCachingHandler(
+        original_function=acompletion,
+        request_kwargs={},
+        start_time=datetime.now(),
+    )
+    monkeypatch.setattr(litellm, "cache", _SlowWriteCache())
+
+    async def _short_lived_script():
+        await handler.async_set_cache(
+            result=litellm.ModelResponse(),
+            original_function=acompletion,
+            kwargs={},
+        )
+
+    asyncio.run(_short_lived_script())
+
+    assert len(writes) == 1

tests/test_litellm/responses/test_streaming_iterator.py

@@ -235,3 +235,73 @@ def test_sync_transport_error_before_completed_event_raises():
     with pytest.raises(httpx.ReadError):
         for _ in iterator:
             pass
+
+
+def test_stream_cache_write_completes_when_asyncio_run_closes_the_loop(monkeypatch):
+    """
+    Regression test for LIT-6184 on the /v1/responses streaming surface: the
+    completed-stream cache write was dispatched as a bare fire-and-forget task,
+    so asyncio.run cancelled it at loop close before the write landed. The
+    write must survive loop shutdown just like the chat-completions one.
+    """
+    import asyncio
+    from types import SimpleNamespace
+
+    import litellm
+    from litellm.types.utils import CallTypes
+
+    writes = []
+
+    class _SlowWriteCache:
+        async def async_add_cache(self, result, dynamic_cache_object=None, **kwargs):
+            await asyncio.sleep(0.2)
+            writes.append(result)
+
+        def add_cache(self, *args, **kwargs):
+            raise AssertionError("sync write must not run on the async path")
+
+    caching_handler = SimpleNamespace(
+        request_kwargs={
+            "model": "test-model",
+            "input": "hello",
+            "stream": True,
+            "caching": True,
+            "metadata": None,
+            "custom_llm_provider": "openai",
+        },
+        preset_cache_key="responses-stream-cache-key",
+        original_function=litellm.aresponses,
+        dual_cache=None,
+        _should_store_result_in_cache=lambda original_function, kwargs: True,
+    )
+    logging_obj = SimpleNamespace(
+        model_call_details={"litellm_params": {}},
+        _llm_caching_handler=caching_handler,
+    )
+    iterator = ResponsesAPIStreamingIterator(
+        response=httpx.Response(200),
+        model="test-model",
+        responses_api_provider_config=Mock(spec=BaseResponsesAPIConfig),
+        logging_obj=logging_obj,
+        request_data=caching_handler.request_kwargs,
+        call_type=CallTypes.aresponses.value,
+    )
+    iterator.completed_response = ResponseCompletedEvent(
+        type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
+        response=ResponsesAPIResponse(
+            id="resp_lit6184",
+            created_at=int(datetime.now().timestamp()),
+            status="completed",
+            model="test-model",
+            object="response",
+            output=[],
+        ),
+    )
+    monkeypatch.setattr(litellm, "cache", _SlowWriteCache())
+
+    async def _short_lived_script():
+        iterator._persist_completed_response_to_cache(is_async=True)
+
+    asyncio.run(_short_lived_script())
+
+    assert len(writes) == 1