perf(router): honour the cooldown read interval in the routing prefetch

litellm/caching/dual_cache.py

@@ -326,6 +326,20 @@ def _reserve_redis_batch_keys(
 
         return sublist_keys, previous_access_times
 
+    def reserve_redis_batch_reads(self, keys: Sequence[str]) -> tuple[list[str], dict[str, float | None]]:
+        """Reserve the memory-missed keys whose throttled Redis reads are due, as a batch read would."""
+        if self.redis_cache is None:
+            return [], {}  # mutable-ok: API contract returns an empty list and dictionary
+        key_list: Final = list(keys)  # mutable-ok: batch_get_cache takes a list
+        memory: Final = self.in_memory_cache
+        in_memory_result: Final = (
+            None
+            if memory is None  # pyright: ignore[reportUnnecessaryComparison]  # handle an absent in-memory tier
+            else memory.batch_get_cache(key_list)
+        )
+        result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in key_list)
+        return self._reserve_redis_batch_keys(time.time(), key_list, result)
+
     def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None:
         with self._last_redis_batch_access_time_lock:
             for key, previous_time in previous_access_times.items():

litellm/router_utils/routing_read_batch.py

@@ -8,6 +8,7 @@
 the usage slice to the strategy, so selection does not read again.
 """
 
+import asyncio
 import itertools
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
@@ -35,13 +36,54 @@
 _PREFETCH_SLOT: Final = "routing_read"
 
 
+async def _backfill_prefetched_cache(
+    cache: DualCache,
+    due_keys: tuple[str, ...],
+    values: Mapping[str, object],
+) -> None:
+    cache_keys: Final = list(due_keys)  # mutable-ok: _prepare_batch_get takes a list
+    prepare_batch_get: Final = cache._prepare_batch_get  # pyright: ignore[reportPrivateUsage]  # memory backfill
+    pending: Final = await prepare_batch_get(cache_keys, local_only=True)
+    redis_values: Final = {  # mutable-ok: _apply_batch_get accepts a dictionary
+        key: values[key]
+        for key, local in zip(due_keys, pending.result)
+        if local is None and values.get(key) is not None
+    }
+    apply_batch_get: Final = cache._apply_batch_get  # pyright: ignore[reportPrivateUsage]  # cache backfill
+    await apply_batch_get(pending, redis_values)
+
+
 @dataclass(frozen=True, slots=True)
 class RoutingPrefetch:
     """The cooldown and usage keys of a model group, declared on the request's Redis batch before admission
     flushes it, so the routing read rides the same round trip as the rate limiter's Lua calls."""
 
     keys: frozenset[str]
+    fetched: frozenset[str]
     result: BatchResult[Mapping[str, object]]
+    reservations: tuple[tuple[DualCache, tuple[str, ...], dict[str, float | None]], ...]
+
+    def release(self) -> None:
+        for cache, _, previous_access_times in self.reservations:
+            cache._rollback_redis_batch_key_reservations(  # pyright: ignore[reportPrivateUsage]  # rollback
+                previous_access_times
+            )
+
+    async def _settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
+        if future.cancelled():
+            self.release()
+            return
+        if future.exception() is not None:
+            self.release()
+            return
+
+        values: Final = future.result()
+        try:
+            for cache, due_keys, _ in self.reservations:
+                await _backfill_prefetched_cache(cache, due_keys, values)
+        except Exception:
+            self.release()
+            raise
 
     @staticmethod
     def arm(
@@ -60,9 +102,28 @@ def arm(
             () if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments)))
         )
         keys: Final = (*cooldown_keys, *usage_keys)
-        request.prefetched[_PREFETCH_SLOT] = RoutingPrefetch(
-            keys=frozenset(keys), result=request.batch(redis_cache).mget(keys)
+        cooldown_store: Final = litellm_router_instance.cooldown_cache.cooldown_store
+        cooldown_due, cooldown_previous = cooldown_store.reserve_redis_batch_reads(cooldown_keys)
+        usage_cache: Final = None if usage_selector is None else usage_selector.router_cache
+        usage_reservation: Final = None if usage_cache is None else usage_cache.reserve_redis_batch_reads(usage_keys)
+        usage_due: Final = () if usage_reservation is None else tuple(usage_reservation[0])
+        due: Final = (*cooldown_due, *usage_due)
+        reservations: Final = (
+            (cooldown_store, tuple(cooldown_due), cooldown_previous),
+            *(
+                ()
+                if usage_cache is None or usage_reservation is None
+                else ((usage_cache, usage_due, usage_reservation[1]),)
+            ),
+        )
+        if not due:
+            return
+        result: Final = request.batch(redis_cache).mget(due)
+        prefetch: Final = RoutingPrefetch(
+            keys=frozenset(keys), fetched=frozenset(due), result=result, reservations=reservations
         )
+        result.on_settled(prefetch._settle)
+        request.prefetched[_PREFETCH_SLOT] = prefetch
 
     @staticmethod
     def armed() -> bool:
@@ -78,6 +139,8 @@ def take(needed: Sequence[str]) -> "RoutingPrefetch | None":
         armed: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
         if isinstance(armed, RoutingPrefetch) and armed.keys.issuperset(needed):
             return armed
+        if isinstance(armed, RoutingPrefetch):
+            armed.release()
         return None
 
 
@@ -149,6 +212,10 @@ async def _read_prefetched(
         results: Final[list[list[object | None] | None]] = []  # mutable-ok: filled per read below
         for cache, keys in reads:
             pending = await cache._prepare_batch_get(keys, local_only=True)  # pyright: ignore[reportPrivateUsage]  # same two-step read as async_batch_get_cache_shared
+            if any(
+                key not in prefetch.fetched for key, local_value in zip(keys, pending.result) if local_value is None
+            ):
+                return None
             missed = {  # mutable-ok: _apply_batch_get takes a dict
                 key: values.get(key) for key, local in zip(keys, pending.result) if local is None
             }

tests/unit/caching/test_dual_cache.py

@@ -2,6 +2,7 @@
 import logging
 import time
 import uuid
+from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -136,6 +137,35 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error():
     assert "shared_a" not in dual_cache.last_redis_batch_access_time
 
 
+def test_reserve_redis_batch_reads_reserves_memory_misses_and_can_be_rolled_back():
+    mock_redis: Final = MagicMock(spec=RedisCache)
+    dual_cache: Final = DualCache(
+        in_memory_cache=InMemoryCache(),
+        redis_cache=mock_redis,
+        default_redis_batch_cache_expiry=10,
+    )
+    dual_cache.in_memory_cache.set_cache("memory_key", "memory_value")
+
+    reserved, previous_access_times = dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"])
+
+    assert reserved == ["missing_key"]
+    assert previous_access_times == {"missing_key": None}
+    assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == ([], {})
+
+    dual_cache._rollback_redis_batch_key_reservations(previous_access_times)
+
+    assert dual_cache.reserve_redis_batch_reads(["memory_key", "missing_key"]) == (
+        ["missing_key"],
+        {"missing_key": None},
+    )
+
+
+def test_reserve_redis_batch_reads_returns_empty_without_redis():
+    dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=None)
+
+    assert dual_cache.reserve_redis_batch_reads(["missing_key"]) == ([], {})
+
+
 def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled():
     mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"})
     dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10)