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,28 @@ 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]]) -> 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)
+
+
 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 +1006,7 @@ class LLMCachingHandler:
 
         if litellm.cache is None:
             return
+        cache: Final = litellm.cache
 
         new_kwargs: Final = kwargs.copy()
         new_kwargs.update(
@@ -1004,24 +1028,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))

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