perf(proxy): one post-call Redis pipeline per backend for spend, rate-limit, routing and response-cache writes

litellm/caching/caching.py

@@ -8,6 +8,7 @@
 #  Thank you users! We ❤️ you! - Krrish & Ishaan
 
 import ast
+import asyncio
 import hashlib
 import json
 import logging
@@ -30,10 +31,11 @@
 from .azure_blob_cache import AzureBlobCache
 from .base_cache import BaseCache
 from .disk_cache import DiskCache
-from .dual_cache import DualCache  # noqa: F401
+from .dual_cache import DualCache
 from .gcs_cache import GCSCache
 from .in_memory_cache import InMemoryCache
 from .qdrant_semantic_cache import QdrantSemanticCache
+from .redis_batch import active_post_call_redis_batch
 from .redis_cache import RedisCache, log_redis_failure
 from .redis_cluster_cache import RedisClusterCache
 from .redis_semantic_cache import RedisSemanticCache
@@ -68,6 +70,15 @@ def print_verbose(print_statement):
         pass
 
 
+def _ttl_seconds(raw: object) -> int | None:
+    if not isinstance(raw, (int, float, str)):
+        return None
+    try:
+        return int(raw)
+    except ValueError:
+        return None
+
+
 class CacheMode(str, Enum):
     default_on = "default_on"
     default_off = "default_off"
@@ -759,13 +770,48 @@ async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None =
                 await self.batch_cache_write(result, **kwargs)
             else:
                 cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
+                if await self._defer_set_to_post_call_batch(cache_key, cached_data, kwargs, dynamic_cache_object):
+                    return
                 if dynamic_cache_object is not None:
                     await dynamic_cache_object.async_set_cache(cache_key, cached_data, **kwargs)
                 else:
                     await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
         except Exception as e:
             self._log_add_cache_failure(e)
 
+    async def _defer_set_to_post_call_batch(
+        self,
+        cache_key: str,
+        cached_data: object,
+        kwargs: Mapping[str, object],
+        dynamic_cache_object: BaseCache | None,
+    ) -> bool:
+        """A plain SET on the Redis response cache rides the request's post-call pipeline with the counters,
+        instead of its own round trip. Anything with SET options keeps the direct path."""
+        if kwargs.get("nx"):
+            return False
+        ttl: Final = _ttl_seconds(kwargs.get("ttl"))
+        if isinstance(dynamic_cache_object, DualCache):
+            deferred: Final = await dynamic_cache_object.async_set_cache_post_call(cache_key, cached_data, ttl)
+            if deferred is None:
+                return False
+            deferred.on_settled(self._log_deferred_add_cache_failure)
+            return True
+        if dynamic_cache_object is not None or not isinstance(self.cache, RedisCache):
+            return False
+        batch: Final = active_post_call_redis_batch(self.cache)
+        if batch is None:
+            return False
+        batch.set(cache_key, cached_data, ttl).on_settled(self._log_deferred_add_cache_failure)
+        return True
+
+    def _log_deferred_add_cache_failure(self, future: asyncio.Future[None]) -> None:
+        if future.cancelled():
+            return
+        failure: Final = future.exception()
+        if isinstance(failure, Exception):
+            self._log_add_cache_failure(failure)
+
     def _convert_to_cached_embedding(
         self,
         embedding_response: Any,

litellm/caching/dual_cache.py

@@ -8,23 +8,23 @@
     - async_get_cache
 """
 
+import asyncio
 import itertools
 import logging
 import time
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Final
 
-if TYPE_CHECKING:
-    from litellm.types.caching import RedisPipelineIncrementOperation
-
 import litellm
 from litellm._logging import print_verbose, verbose_logger
 from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
+from litellm.types.caching import RedisPipelineIncrementOperation
 
 from .base_cache import BaseCache
 from .in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache
+from .redis_batch import BatchResult, RedisBatch, active_post_call_redis_batch
 from .redis_cache import RedisCache, RedisCircuitBreakerOpenError, log_redis_failure
 
 if TYPE_CHECKING:
@@ -59,6 +59,24 @@ class PendingBatchRead:
     previous_access_times: dict[str, float | None]
 
 
+@dataclass(frozen=True, slots=True)
+class DeclaredBatchRead:
+    """A ``async_batch_get_cache`` split in two: the memory half done, the Redis half declared on a ``RedisBatch``
+    so it rides that batch's next round trip, resolved later with ``async_resolve_batch_get``."""
+
+    keys: tuple[str, ...]
+    pending: PendingBatchRead
+    result: BatchResult[Mapping[str, object]] | None
+
+
+def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
+    if future.cancelled():
+        return
+    failure: Final = future.exception()
+    if failure is not None:
+        log_redis_failure(verbose_logger, logging.WARNING, "post-call Redis increment failed", failure)
+
+
 class DualCache(BaseCache):
     """
     DualCache is a cache implementation that updates both Redis and an in-memory cache simultaneously.
@@ -335,7 +353,7 @@ async def _prepare_batch_get(
         )
 
     async def _apply_batch_get(
-        self, pending: PendingBatchRead, redis_result: dict[str, object] | None, **kwargs: object
+        self, pending: PendingBatchRead, redis_result: Mapping[str, object] | None, **kwargs: object
     ) -> list[object | None]:
         if redis_result is None or all(v is None for v in redis_result.values()):
             return pending.result
@@ -349,6 +367,22 @@ async def _apply_batch_get(
                     await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
         return merged
 
+    async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
+        pending: Final = await self._prepare_batch_get(
+            list(keys),  # mutable-ok: the shared batch read takes a list
+            local_only=False,
+            throttle_redis=False,
+        )
+        return DeclaredBatchRead(
+            keys=tuple(keys),
+            pending=pending,
+            result=batch.mget(pending.redis_keys) if pending.redis_keys else None,
+        )
+
+    async def async_resolve_batch_get(self, declared: DeclaredBatchRead) -> list[object | None]:
+        redis_result: Final = None if declared.result is None else await declared.result
+        return await self._apply_batch_get(declared.pending, redis_result)
+
     async def async_batch_get_cache(
         self,
         keys: list,
@@ -468,6 +502,18 @@ async def async_set_cache(self, key, value, local_only: bool = False, **kwargs):
                 verbose_logger, logging.ERROR, "LiteLLM Cache: exception in async add_cache", e, with_traceback=True
             )
 
+    async def async_set_cache_post_call(self, key: str, value: object, ttl: float | None) -> BatchResult[None] | None:
+        """Memory now, the Redis SET on the request's post-call pipeline; None when no pipeline is open, so the
+        caller takes its direct path."""
+        batch: Final = None if self.redis_cache is None else active_post_call_redis_batch(self.redis_cache)
+        if batch is None:
+            return None
+        if self.in_memory_cache is not None:
+            await self.in_memory_cache.async_set_cache(
+                key, value, ttl=self.default_in_memory_ttl if ttl is None else ttl
+            )
+        return batch.set(key, value, ttl)
+
     # async_batch_set_cache
     async def async_set_cache_pipeline(
         self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs
@@ -535,6 +581,41 @@ async def async_increment_cache(
             )
             return result
 
+    async def async_increment_cache_post_call(
+        self,
+        key: str,
+        value: float,
+        ttl: int | None,
+        parent_otel_span: Span | None = None,
+    ) -> None:
+        """Memory is incremented now; the Redis increment rides the request's post-call pipeline when one is
+        open, and runs on its own as ``async_increment_cache`` otherwise."""
+        await self.async_increment_cache_pipeline_post_call(
+            (RedisPipelineIncrementOperation(key=key, increment_value=value, ttl=ttl),), parent_otel_span
+        )
+
+    async def async_increment_cache_pipeline_post_call(
+        self,
+        increment_list: Sequence["RedisPipelineIncrementOperation"],
+        parent_otel_span: Span | None = None,
+    ) -> None:
+        batch: Final = None if self.redis_cache is None else active_post_call_redis_batch(self.redis_cache)
+        operations: Final = list(increment_list)  # mutable-ok: both increment pipelines take a list
+        if batch is None:
+            await self.async_increment_cache_pipeline(operations, parent_otel_span=parent_otel_span)
+            return
+        try:
+            if self.in_memory_cache is not None:
+                await self.in_memory_cache.async_increment_pipeline(
+                    increment_list=operations, parent_otel_span=parent_otel_span
+                )
+        except Exception as e:  # noqa: BLE001  # same tolerance as async_increment_cache_pipeline
+            log_redis_failure(verbose_logger, logging.WARNING, "in-memory increment failed", e)
+        for operation in increment_list:
+            batch.increment(operation["key"], operation["increment_value"], operation["ttl"]).on_settled(
+                _log_deferred_increment_failure
+            )
+
     async def async_increment_cache_pipeline(
         self,
         increment_list: list["RedisPipelineIncrementOperation"],