refactor(redis): encapsulate router cooldown and usage stores

.github/workflows/test-unit.yml

@@ -167,6 +167,7 @@ jobs:
               tests/unit/google_genai
               tests/unit/router_strategy
               tests/unit/router_utils
+              tests/unit/redis
               tests/unit/proxy/common_utils/test_cache_aware_routing.py
               tests/unit/enterprise/enterprise_callbacks/send_emails
               tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py

litellm/caching/dual_cache.py

@@ -12,7 +12,7 @@ import asyncio
 import itertools
 import logging
 import time
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
 from threading import Lock
 from typing import TYPE_CHECKING, Any, Final
@@ -82,8 +82,10 @@ class DualCache(BaseCache):
         default_redis_ttl: float | None = None,
         default_redis_batch_cache_expiry: float | None = None,
         default_max_redis_batch_cache_size: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
+        clock: Callable[[], float] | None = None,
     ) -> None:
         super().__init__()
+        self._clock: Final = clock if clock is not None else lambda: time.time()
         # If in_memory_cache is not provided, use the default InMemoryCache
         self.in_memory_cache = in_memory_cache or InMemoryCache()
         # If redis_cache is not provided, use the default RedisCache
@@ -221,7 +223,7 @@ class DualCache(BaseCache):
             if None not in result or self.redis_cache is None or local_only:
                 return result
 
-            sublist_keys, previous_access_times = self._reserve_redis_batch_keys(time.time(), keys, result)
+            sublist_keys, previous_access_times = self._reserve_redis_batch_keys(self._clock(), keys, result)
             if len(sublist_keys) == 0:
                 return result
 
@@ -326,7 +328,7 @@ class DualCache(BaseCache):
             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)
+        return self._reserve_redis_batch_keys(self._clock(), 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:
@@ -350,7 +352,7 @@ class DualCache(BaseCache):
         previous_access_times: dict[str, float | None] = {}
         if None in result and self.redis_cache is not None and local_only is False:
             if throttle_redis:
-                redis_keys, previous_access_times = self._reserve_redis_batch_keys(time.time(), keys, result)
+                redis_keys, previous_access_times = self._reserve_redis_batch_keys(self._clock(), keys, result)
             else:
                 redis_keys = [key for key, value in zip(keys, result) if value is None]
         return PendingBatchRead(

litellm/redis/stores/_keys.py

@@ -0,0 +1,4 @@
+from typing import Final
+
+ROUTER_TPM_KEY: Final = "global_router:{id}:{model}:tpm:{current_minute}"
+ROUTER_RPM_KEY: Final = "global_router:{id}:{model}:rpm:{current_minute}"

litellm/redis/stores/_routing_prefetch.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import asyncio
+import itertools
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+from litellm._internal_context import service_target
+from litellm._logging import verbose_router_logger
+from litellm.caching.dual_cache import DualCache
+from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
+
+_PREFETCH_SLOT: Final = "routing_read"
+
+
+@dataclass(frozen=True, slots=True)
+class TierRead:
+    cache: DualCache
+    keys: tuple[str, ...]
+
+
+def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead:
+    return TierRead(cache, tuple(keys))
+
+
+def routing_target(reads: Sequence[TierRead]) -> str:
+    if len(reads) < 2 or not reads[1].keys:
+        return "router_cooldowns"
+    return "router_cooldowns_usage" if reads[0].keys else "router_usage"
+
+
+async def _backfill(read: TierRead, values: Mapping[str, object]) -> None:
+    pending: Final = await read.cache._prepare_batch_get(  # pyright: ignore[reportPrivateUsage]  # store owns tier
+        list(read.keys), local_only=True
+    )
+    missed: Final = {
+        key: values[key]
+        for key, local in zip(read.keys, pending.result)
+        if local is None and values.get(key) is not None
+    }
+    await read.cache._apply_batch_get(pending, missed)  # pyright: ignore[reportPrivateUsage]  # store owns tier
+
+
+@dataclass(frozen=True, slots=True)
+class ArmedRead:
+    keys: frozenset[str]
+    fetched: frozenset[str]
+    result: BatchResult[Mapping[str, object]]
+    reservations: tuple[tuple[TierRead, Mapping[str, float | None]], ...]
+
+    def release(self) -> None:
+        for read, previous in self.reservations:
+            read.cache._rollback_redis_batch_key_reservations(  # pyright: ignore[reportPrivateUsage]  # store owns tier
+                dict(previous)
+            )
+
+    async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
+        if future.cancelled() or future.exception() is not None:
+            self.release()
+            return
+        try:
+            for read, _ in self.reservations:
+                await _backfill(read, future.result())
+        except Exception:
+            self.release()
+            raise
+
+
+def arm(reads: Sequence[TierRead]) -> None:
+    request: Final = active_request_redis_batches()
+    redis: Final = reads[0].cache.redis_cache if reads else None
+    if request is None or redis is None or _PREFETCH_SLOT in request.prefetched:
+        return
+    reserved: Final = tuple((read, read.cache.reserve_redis_batch_reads(read.keys)) for read in reads)
+    due_reads: Final = tuple(TierRead(read.cache, tuple(due)) for read, (due, _) in reserved)
+    due: Final = tuple(itertools.chain.from_iterable(read.keys for read in due_reads))
+    if not due:
+        return
+    with service_target(routing_target(due_reads)):
+        result: Final = request.batch(redis).mget(due)
+    prefetch: Final = ArmedRead(
+        keys=frozenset(itertools.chain.from_iterable(read.keys for read in reads)),
+        fetched=frozenset(due),
+        result=result,
+        reservations=tuple((read, reservation[1]) for read, (_, reservation) in zip(due_reads, reserved)),
+    )
+    result.on_settled(prefetch.settle)
+    request.prefetched[_PREFETCH_SLOT] = prefetch
+
+
+def armed() -> bool:
+    request: Final = active_request_redis_batches()
+    return request is not None and _PREFETCH_SLOT in request.prefetched
+
+
+def _take(needed: Sequence[str]) -> ArmedRead | None:
+    request: Final = active_request_redis_batches()
+    if request is None:
+        return None
+    taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
+    if isinstance(taken, ArmedRead) and taken.keys.issuperset(needed):
+        return taken
+    if isinstance(taken, ArmedRead):
+        taken.release()
+    return None
+
+
+async def _read_prefetched(reads: Sequence[TierRead]) -> Sequence[Sequence[object | None] | None] | None:
+    prefetch: Final = _take(tuple(itertools.chain.from_iterable(read.keys for read in reads)))
+    if prefetch is None:
+        return None
+    try:
+        values: Final = await prefetch.result
+    except Exception as e:  # noqa: BLE001  # shared read applies the tiers' existing Redis fallback
+        verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
+        return None
+    pending_reads: Final = tuple(
+        [
+            await read.cache._prepare_batch_get(  # pyright: ignore[reportPrivateUsage]  # store owns tier
+                list(read.keys), local_only=True
+            )
+            for read in reads
+        ]
+    )
+    if any(
+        any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None)
+        for read, pending in zip(reads, pending_reads)
+    ):
+        return None
+    return [
+        await read.cache._apply_batch_get(  # pyright: ignore[reportPrivateUsage]  # store owns tier
+            pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None}
+        )
+        for read, pending in zip(reads, pending_reads)
+    ]
+
+
+async def read(reads: Sequence[TierRead], parent_otel_span: Span | None) -> Sequence[Sequence[object | None] | None]:
+    with service_target(routing_target(reads)):
+        prefetched: Final = await _read_prefetched(reads)
+        if prefetched is not None:
+            return prefetched
+        return await DualCache.async_batch_get_cache_shared(
+            tuple((read.cache, list(read.keys)) for read in reads), parent_otel_span=parent_otel_span
+        )

litellm/redis/stores/router_cooldowns.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import functools
+import time
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+from typing_extensions import ReadOnly, TypedDict
+
+from litellm._internal_context import service_target
+from litellm._logging import verbose_router_logger
+from litellm.caching.dual_cache import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.redis_cache import RedisCache
+from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
+
+from . import _routing_prefetch
+from ._routing_prefetch import TierRead
+
+
+class CooldownCacheValue(TypedDict):
+    exception_received: ReadOnly[str]
+    status_code: ReadOnly[str]
+    timestamp: ReadOnly[float]
+    cooldown_time: ReadOnly[float]
+
+
+@dataclass(frozen=True, slots=True)
+class CooldownPolicy:
+    redis_read_interval: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
+    backfill_ttl: int = 600
+    max_memory_entries: int = 200
+    max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
+    max_corrected_ttl: float = 60.0
+
+
+POLICY: Final = CooldownPolicy()
+ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
+ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]]
+
+
+@dataclass(frozen=True, slots=True)
+class RoutingRead:
+    cooldowns: ActiveCooldowns
+    selector_values: Sequence[object | None] | None
+
+
+@functools.lru_cache(maxsize=1024)
+def _cooldown_key(model_id: str) -> str:
+    return "deployment:" + model_id + ":cooldown"
+
+
+class RouterCooldowns(Protocol):
+    def attach_redis(self, connection: RedisCache | None) -> None: ...
+    def flush_local_counters(self) -> None: ...
+    def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
+    def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ...
+    def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None: ...
+    async def routing_read(
+        self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None
+    ) -> RoutingRead: ...
+    async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
+    def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
+    def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ...
+    def evict_local(self, model_id: str) -> None: ...
+
+
+class RouterCooldownStore:
+    def __init__(self, connection: RedisCache | None, *, clock: Callable[[], float] | None = None) -> None:
+        self._clock: Final = clock if clock is not None else lambda: time.time()
+        self._memory: Final = InMemoryCache(
+            max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl, clock=self._clock
+        )
+        self._cache: Final = DualCache(
+            redis_cache=connection,
+            in_memory_cache=self._memory,
+            default_redis_batch_cache_expiry=POLICY.redis_read_interval,
+            default_max_redis_batch_cache_size=POLICY.max_batch_entries,
+            clock=self._clock,
+        )
+        self._fails: Final = DualCache(redis_cache=connection, in_memory_cache=InMemoryCache(clock=self._clock))
+
+    def attach_redis(self, connection: RedisCache | None) -> None:
+        self._cache.attach_redis_cache(connection)
+        self._fails.attach_redis_cache(connection)
+
+    def flush_local_counters(self) -> None:
+        self._fails.in_memory_cache.flush_cache()
+
+    def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
+        with service_target(ROUTER_COOLDOWNS_TARGET):
+            self._cache.set_cache(key=_cooldown_key(model_id), value=value, ttl=ttl)
+
+    def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
+        base: Final = f"deployment:{model_id}:allowed_fails"
+        key: Final = f"{base}:{suffix}" if suffix else base
+        with service_target(ROUTER_COOLDOWNS_TARGET):
+            try:
+                return self._fails.increment_cache(key=key, value=1, ttl=ttl)
+            except Exception as e:  # noqa: BLE001  # preserve worker-local counting during a Redis outage
+                verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
+                local: Final[object] = self._fails.get_cache(key=key, local_only=True)
+                return local if isinstance(local, int) else 0
+
+    def _tier_read(self, model_ids: Sequence[str]) -> TierRead:
+        return TierRead(self._cache, tuple(_cooldown_key(model_id) for model_id in model_ids))
+
+    def arm_routing_read(self, model_ids: Sequence[str], selector: TierRead | None) -> None:
+        _routing_prefetch.arm((self._tier_read(model_ids), *(() if selector is None else (selector,))))
+
+    async def routing_read(
+        self, model_ids: Sequence[str], selector: TierRead | None, parent_otel_span: Span | None
+    ) -> RoutingRead:
+        reads: Final = (self._tier_read(model_ids), *(() if selector is None else (selector,)))
+        results: Final = await _routing_prefetch.read(reads, parent_otel_span)
+        return RoutingRead(
+            cooldowns=self._active_from_results(model_ids, results[0]),
+            selector_values=None if selector is None else results[1],
+        )
+
+    def _corrected_active_cooldown(
+        self, model_id: str, result: Mapping[str, object], current_time: float
+    ) -> CooldownCacheValue | None:
+        value: Final = CooldownCacheValue(
+            **result  # pyright: ignore[reportArgumentType]  # preserve main's unvalidated dict codec
+        )
+        remaining: Final = value["timestamp"] + value["cooldown_time"] - current_time
+        key: Final = _cooldown_key(model_id)
+        if remaining <= 0:
+            self._memory.delete_cache(key)
+            return None
+        expiry: Final = self._memory.ttl_dict.get(key)
+        if expiry is not None and expiry > current_time + remaining + 5:
+            self._memory.delete_cache(key)
+            self._memory.set_cache(key, result, ttl=min(remaining, POLICY.max_corrected_ttl))
+        return value
+
+    def _active_from_results(self, model_ids: Sequence[str], results: Sequence[object] | None) -> ActiveCooldowns:
+        if results is None:
+            return ()
+        now: Final = self._clock()
+        return tuple(
+            (model_id, active)
+            for model_id, result in zip(model_ids, results)
+            if isinstance(result, dict)
+            and result
+            and (active := self._corrected_active_cooldown(model_id, result, now)) is not None
+        )
+
+    async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
+        with service_target(ROUTER_COOLDOWNS_TARGET):
+            results: Final = await self._cache.async_batch_get_cache(
+                keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
+            )
+        return self._active_from_results(model_ids, results)
+
+    def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
+        with service_target(ROUTER_COOLDOWNS_TARGET):
+            results: Final = (
+                self._cache.batch_get_cache(
+                    keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
+                )
+                or ()
+            )
+        durations: Final = tuple(
+            CooldownCacheValue(**result)["cooldown_time"]  # pyright: ignore[reportArgumentType]  # preserve main's dict codec
+            for result in results
+            if isinstance(result, dict) and result
+        )
+        return (min(durations) if durations else None) or default
+
+    def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
+        with service_target(ROUTER_COOLDOWNS_TARGET):
+            results: Final = self._cache.batch_get_cache(
+                keys=[_cooldown_key(model_id) for model_id in model_ids], parent_otel_span=parent_otel_span
+            )
+        return self._active_from_results(model_ids, results)
+
+    def evict_local(self, model_id: str) -> None:
+        self._memory.delete_cache(_cooldown_key(model_id))

litellm/redis/stores/router_usage.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+from litellm._internal_context import with_service_target
+from litellm.caching.dual_cache import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.redis_cache import RedisCache
+from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
+from litellm.types.caching import RedisPipelineIncrementOperation
+
+from ._keys import ROUTER_RPM_KEY, ROUTER_TPM_KEY
+
+
+@dataclass(frozen=True, slots=True)
+class UsagePolicy:
+    redis_read_interval: float | None = None
+    memory_ttl: float | None = None
+    redis_ttl: float | None = None
+    backfill_ttl: int = 600
+    max_memory_entries: int = 200
+    max_batch_entries: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
+    counter_ttl: int = 60
+
+
+POLICY: Final = UsagePolicy()
+ROUTER_USAGE_TARGET: Final = "router_usage"
+CounterKind: TypeAlias = Literal["tpm", "rpm"]
+
+
+@dataclass(frozen=True, slots=True)
+class UsageCounter:
+    deployment_id: str
+    model: str | None
+    minute: str
+    kind: CounterKind
+
+
+def _key(counter: UsageCounter) -> str:
+    template: Final = ROUTER_TPM_KEY if counter.kind == "tpm" else ROUTER_RPM_KEY
+    return template.format(id=counter.deployment_id, model=counter.model, current_minute=counter.minute)
+
+
+class RouterUsage(Protocol):
+    def attach_redis(self, connection: RedisCache | None) -> None: ...
+    def flush_local_counters(self) -> None: ...
+    async def async_increment(
+        self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
+    ) -> float | None: ...
+    async def increment_deployment(
+        self,
+        deployment_id: str,
+        model: str | None,
+        minute: str,
+        tokens: float,
+        requests: int,
+        parent_otel_span: Span | None,
+    ) -> None: ...
+    async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None: ...
+
+
+class RouterUsageStore:
+    def __init__(self, connection: RedisCache | None) -> None:
+        self._cache: Final = DualCache(
+            redis_cache=connection,
+            in_memory_cache=InMemoryCache(
+                max_size_in_memory=POLICY.max_memory_entries, default_ttl=POLICY.backfill_ttl
+            ),
+            default_in_memory_ttl=POLICY.memory_ttl,
+            default_redis_ttl=POLICY.redis_ttl,
+            default_redis_batch_cache_expiry=POLICY.redis_read_interval,
+            default_max_redis_batch_cache_size=POLICY.max_batch_entries,
+        )
+
+    def attach_redis(self, connection: RedisCache | None) -> None:
+        self._cache.attach_redis_cache(connection)
+
+    def flush_local_counters(self) -> None:
+        self._cache.in_memory_cache.flush_cache()
+
+    @with_service_target(ROUTER_USAGE_TARGET)
+    async def async_increment(
+        self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
+    ) -> float | None:
+        return await self._cache.async_increment_cache(
+            key=_key(counter), value=value, ttl=ttl, parent_otel_span=parent_otel_span
+        )
+
+    @with_service_target(ROUTER_USAGE_TARGET)
+    async def increment_deployment(
+        self,
+        deployment_id: str,
+        model: str | None,
+        minute: str,
+        tokens: float,
+        requests: int,
+        parent_otel_span: Span | None,
+    ) -> None:
+        tpm: Final = UsageCounter(deployment_id, model, minute, "tpm")
+        rpm: Final = UsageCounter(deployment_id, model, minute, "rpm")
+        operations: Final = [
+            RedisPipelineIncrementOperation(key=_key(counter), increment_value=value, ttl=POLICY.counter_ttl)
+            for counter, value in ((tpm, tokens), (rpm, requests))
+        ]
+        values: Final = await self._cache.async_increment_cache_pipeline(
+            increment_list=operations, parent_otel_span=parent_otel_span
+        )
+        if values is not None and self._cache.redis_cache is not None:
+            for operation, value in zip(operations, values):
+                await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=POLICY.counter_ttl)
+
+    async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None:
+        return await self._cache.async_batch_get_cache(keys=[_key(counter) for counter in counters])

litellm/router_python_core.py

@@ -238,6 +238,8 @@ from litellm.llms.openai_like.model_info import (
     get_openai_compatible_model_info,
 )
 from litellm.main import Chat as LitellmChat
+from litellm.redis.stores.router_cooldowns import RouterCooldownStore
+from litellm.redis.stores.router_usage import RouterUsage, RouterUsageStore, UsageCounter
 from litellm.responses.bridge_check import keep_registered_mode
 from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
 from litellm.router_api import (
@@ -1743,6 +1745,7 @@ class PythonCore(DeprecatedRouterMembers):
         self._core_cache = DualCache(
             redis_cache=redis_cache, in_memory_cache=InMemoryCache()
         )  # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
+        self._core_usage_store: Final[RouterUsage] = RouterUsageStore(redis_cache)
         self._claude_code_session_router_cache: DualCache = DualCache(
             redis_cache=redis_cache,
             in_memory_cache=InMemoryCache(),
@@ -1848,7 +1851,7 @@ class PythonCore(DeprecatedRouterMembers):
         self.explicit_allowed_fails = allowed_fails
         self._core_cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS
         self._core_cooldown_cache = CooldownCache(
-            cache=self._core_cache, default_cooldown_time=self._core_cooldown_time
+            store=RouterCooldownStore(redis_cache), default_cooldown_time=self._core_cooldown_time
         )
         self._core_disable_cooldowns = disable_cooldowns
         self.fallback_on_unsupported_params = fallback_on_unsupported_params
@@ -2204,6 +2207,8 @@ class PythonCore(DeprecatedRouterMembers):
         and caching to just work.
         """
         self._core_cache.attach_redis_cache(cache)
+        self._core_usage_store.attach_redis(cache)
+        self._core_cooldown_cache.cooldown_store.attach_redis(cache)
         self._claude_code_session_router_cache.attach_redis_cache(cache)
         self._core_prompt_caching_pin_cache.attach_redis_cache(cache)
 
@@ -10060,35 +10065,17 @@ class PythonCore(DeprecatedRouterMembers):
         total_tokens: float,
         rpm_increment: int,
         parent_otel_span: Span | None,
-    ) -> str | None:
-        from litellm.types.caching import RedisPipelineIncrementOperation
-
+    ) -> UsageCounter | None:
         if total_tokens <= 0 and rpm_increment <= 0:
             return None
         if not self._cached_deployment_tracks_usage(deployment_id, model_group):
             return None
 
         current_minute: Final = get_utc_datetime().strftime("%H-%M")  # use the same timezone regardless of system clock
-        tpm_key: Final = RouterCacheEnum.TPM.value.format(
-            id=deployment_id, current_minute=current_minute, model=deployment_name
+        await self._core_usage_store.increment_deployment(
+            deployment_id, deployment_name, current_minute, total_tokens, rpm_increment, parent_otel_span
         )
-        rpm_key: Final = RouterCacheEnum.RPM.value.format(
-            id=deployment_id, current_minute=current_minute, model=deployment_name
-        )
-        pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [
-            RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value)
-            for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment))
-        ]
-        post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline(
-            increment_list=pipeline_operations,
-            parent_otel_span=parent_otel_span,
-        )
-        if post_increment_values is not None and self._core_cache.redis_cache is not None:
-            for operation, value in zip(pipeline_operations, post_increment_values):
-                await self._core_cache.async_set_cache(
-                    operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value
-                )
-        return tpm_key
+        return UsageCounter(deployment_id, deployment_name, current_minute, "tpm")
 
     def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool:
         deployment_info: Final = self._core_get_deployment(model_id=deployment_id)
@@ -10305,10 +10292,10 @@ class PythonCore(DeprecatedRouterMembers):
         current_minute: Final = dt.strftime("%H-%M")  # use the same timezone regardless of system clock
 
         ## RPM
-        rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
+        counter: Final = UsageCounter(id, deployment_name, current_minute, "rpm")
         with service_target(ROUTER_USAGE_TARGET):
-            await self._core_cache.async_increment_cache(
-                key=rpm_key,
+            await self._core_usage_store.async_increment(
+                counter=counter,
                 value=1,
                 parent_otel_span=parent_otel_span,
                 ttl=RoutingArgs.ttl.value,
@@ -13416,8 +13403,8 @@ class PythonCore(DeprecatedRouterMembers):
         """
         dt: Final = get_utc_datetime()
         current_minute: Final = dt.strftime("%H-%M")  # use the same timezone regardless of system clock
-        tpm_keys: Final[list[str]] = []
-        rpm_keys: Final[list[str]] = []
+        tpm_keys: Final[list[UsageCounter]] = []
+        rpm_keys: Final[list[UsageCounter]] = []
 
         model_list: Final = self._core_get_model_list(model_name=model_group)
         if model_list is None:  # no matching deployments
@@ -13430,28 +13417,16 @@ class PythonCore(DeprecatedRouterMembers):
             )  # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written.
             if id is None or litellm_model is None:
                 continue
-            tpm_keys.append(
-                RouterCacheEnum.TPM.value.format(
-                    id=id,
-                    model=litellm_model,
-                    current_minute=current_minute,
-                )
-            )
-            rpm_keys.append(
-                RouterCacheEnum.RPM.value.format(
-                    id=id,
-                    model=litellm_model,
-                    current_minute=current_minute,
-                )
-            )
+            tpm_keys.append(UsageCounter(id, litellm_model, current_minute, "tpm"))
+            rpm_keys.append(UsageCounter(id, litellm_model, current_minute, "rpm"))
         combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys
 
-        combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys)
+        combined_tpm_rpm_values: Final = await self._core_usage_store.async_batch_get(combined_tpm_rpm_keys)
         if combined_tpm_rpm_values is None:
             return None, None
 
-        tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)]
-        rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :]
+        tpm_usage_list: Final = combined_tpm_rpm_values[: len(tpm_keys)]
+        rpm_usage_list: Final = combined_tpm_rpm_values[len(tpm_keys) :]
 
         ## TPM
         tpm_usage: int | None = None
@@ -17653,6 +17628,8 @@ class PythonCore(DeprecatedRouterMembers):
 
     def flush_cache(self):
         self._host.set_global_cache(None)
+        self._core_usage_store.flush_local_counters()
+        self._core_cooldown_cache.cooldown_store.flush_local_counters()
         self._core_cache.flush_cache()
         session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache
         if session_in_memory_cache is not None:

litellm/router_utils/cooldown_cache.py

@@ -1,237 +1,54 @@
-"""
-Wrapper around router cache. Meant to handle model cooldown logic
-"""
+from __future__ import annotations
 
-import functools
 import time
-from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final
-
-from typing_extensions import TypedDict
-
-from litellm import verbose_logger
-from litellm._internal_context import service_target
-from litellm.caching.caching import DualCache
-from litellm.caching.in_memory_cache import InMemoryCache
-from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
-from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Final
 
 if TYPE_CHECKING:
-    from opentelemetry.trace import Span as _Span
+    from opentelemetry.trace import Span
 
-    Span = _Span
-else:
-    Span = Any
+from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
+from litellm.redis.stores.router_cooldowns import (
+    ROUTER_COOLDOWNS_TARGET,
+    CooldownCacheValue,
+    RouterCooldowns,
+)
 
-
-class CooldownCacheValue(TypedDict):
-    exception_received: str
-    status_code: str
-    timestamp: float
-    cooldown_time: float
-
-
-# Cap on the corrected in-memory TTL set in `_corrected_active_cooldown`: re-checks the
-# real remaining cooldown against Redis at least this often, so an entry that later gets
-# deleted or extended in Redis before its original deadline is still noticed promptly.
-_MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0
-ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
+__all__: Final = ("ROUTER_COOLDOWNS_TARGET", "CooldownCache", "CooldownCacheValue")
 
 
 class CooldownCache:
     def __init__(
-        self,
-        cache: DualCache,
-        default_cooldown_time: float,
-        redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
-    ):
-        self.cache = cache
-        self.default_cooldown_time = default_cooldown_time
-        self.in_memory_cache = InMemoryCache()
-        self._cooldown_store = DualCache(
-            in_memory_cache=self.in_memory_cache,
-            default_redis_batch_cache_expiry=redis_read_interval_seconds,
+        self, store: RouterCooldowns, default_cooldown_time: float, *, clock: Callable[[], float] | None = None
+    ) -> None:
+        self.cooldown_store: Final = store
+        self._clock: Final = clock if clock is not None else lambda: time.time()
+        self.default_cooldown_time: Final = default_cooldown_time
+        self.exception_masker: Final = SensitiveDataMasker(
+            visible_prefix=50, visible_suffix=0, mask_char="*", mask_short_values=False
         )
-        # Initialize the masker with custom settings for exception strings
-        self.exception_masker = SensitiveDataMasker(
-            visible_prefix=50,  # Show first 50 characters
-            visible_suffix=0,  # Show last 0 characters
-            mask_char="*",  # Use * for masking
-            mask_short_values=False,  # Truncate long messages only; keep short ones readable
-        )
-
-    @property
-    def cooldown_store(self) -> DualCache:
-        """
-        The cache cooldown entries live in, with the router's Redis attached on first use.
-
-        It is kept separate from the router-wide cache so that a key missing from memory is
-        re-read from Redis every `redis_read_interval_seconds` rather than on the router
-        cache's much longer batch interval, which is what lets a sibling replica see a
-        cooldown another replica wrote, and so that unrelated router keys cannot evict a
-        cooldown from the in-memory tier before it expires. Redis is attached lazily because
-        the router builds its cooldown cache before it wires up the shared Redis client.
-        """
-        self._cooldown_store.attach_redis_cache(self.cache.redis_cache)
-        return self._cooldown_store
-
-    def _common_add_cooldown_logic(
-        self, model_id: str, original_exception, exception_status, cooldown_time: float
-    ) -> tuple[str, CooldownCacheValue]:
-        try:
-            current_time: Final = time.time()
-            cooldown_key: Final = CooldownCache.get_cooldown_cache_key(model_id)
-
-            # Store the cooldown information for the deployment separately
-            cooldown_data: Final = CooldownCacheValue(
-                exception_received=self.exception_masker.mask_value(str(original_exception)),
-                status_code=str(exception_status),
-                timestamp=current_time,
-                cooldown_time=cooldown_time,
-            )
-
-            return cooldown_key, cooldown_data
-        except Exception as e:
-            verbose_logger.error("CooldownCache::_common_add_cooldown_logic - Exception occurred - %s", e)
-            raise e
 
     def add_deployment_to_cooldown(
-        self,
-        model_id: str,
-        original_exception: Exception,
-        exception_status: int,
-        cooldown_time: float | None,
-    ):
-        try:
-            #########################################################
-            # get cooldown time
-            # 1. If dynamic cooldown time is set for the model/deployment, use that
-            # 2. If no dynamic cooldown time is set, use the default cooldown time set on CooldownCache
-            _cooldown_time = cooldown_time
-            if _cooldown_time is None:
-                _cooldown_time = self.default_cooldown_time
-            #########################################################
-
-            cooldown_key, cooldown_data = self._common_add_cooldown_logic(
-                model_id=model_id,
-                original_exception=original_exception,
-                exception_status=exception_status,
-                cooldown_time=_cooldown_time,
-            )
-
-            # Set the cache with a TTL equal to the cooldown time
-            with service_target(ROUTER_COOLDOWNS_TARGET):
-                self.cooldown_store.set_cache(
-                    value=cooldown_data,
-                    key=cooldown_key,
-                    ttl=_cooldown_time,
-                )
-        except Exception as e:
-            verbose_logger.error("CooldownCache::add_deployment_to_cooldown - Exception occurred - %s", e)
-            raise e
-
-    @staticmethod
-    @functools.lru_cache(maxsize=1024)
-    def get_cooldown_cache_key(model_id: str) -> str:
-        return "deployment:" + model_id + ":cooldown"
-
-    def _corrected_active_cooldown(
-        self,
-        key: str,
-        result: Mapping[str, Any],
-        current_time: float,
-    ) -> CooldownCacheValue | None:
-        """
-        Return a CooldownCacheValue if the cooldown is still active, or None if it has expired.
-
-        Also corrects the in-memory TTL when DualCache promotes a Redis entry using the
-        default 600s TTL instead of the true remaining cooldown time.
-        """
-        cooldown_cache_value: Final = CooldownCacheValue(**result)  # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code
-        remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time
-        if remaining <= 0:
-            self.in_memory_cache.delete_cache(key)
-            return None
-        current_expiry: Final = self.in_memory_cache.ttl_dict.get(key)
-        if current_expiry is not None and current_expiry > current_time + remaining + 5:
-            corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS)
-            self.in_memory_cache.delete_cache(key)
-            self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
-        return cooldown_cache_value
+        self, model_id: str, original_exception: Exception, exception_status: int, cooldown_time: float | None
+    ) -> None:
+        ttl: Final = self.default_cooldown_time if cooldown_time is None else cooldown_time
+        value: Final = CooldownCacheValue(
+            exception_received=self.exception_masker.mask_value(str(original_exception)),
+            status_code=str(exception_status),
+            timestamp=self._clock(),
+            cooldown_time=ttl,
+        )
+        self.cooldown_store.put(model_id, value, ttl)
 
     async def async_get_active_cooldowns(
-        self, model_ids: list[str], parent_otel_span: Span | None
+        self, model_ids: Sequence[str], parent_otel_span: Span | None
     ) -> list[tuple[str, CooldownCacheValue]]:
-        # Generate the keys for the deployments
-        keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
+        return list(await self.cooldown_store.active(model_ids, parent_otel_span))
 
-        with service_target(ROUTER_COOLDOWNS_TARGET):
-            results: Final = await self.cooldown_store.async_batch_get_cache(
-                keys=keys, parent_otel_span=parent_otel_span
-            )
-        return self.active_cooldowns_from_results(model_ids, results)
-
-    def active_cooldowns_from_results(
-        self, model_ids: list[str], results: Sequence[object] | None
-    ) -> list[tuple[str, CooldownCacheValue]]:
-        """The cooldowns still active in a `cooldown_store` batch read of `get_cooldown_cache_key(model_id)` per id."""
-        active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = []
-
-        if results is None or all(v is None for v in results):
-            return active_cooldowns
-
-        current_time: Final = time.time()
-        for model_id, result in zip(model_ids, results):
-            if result and isinstance(result, dict):
-                key = CooldownCache.get_cooldown_cache_key(model_id)
-                cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time)
-                if cooldown_cache_value is not None:
-                    active_cooldowns.append((model_id, cooldown_cache_value))
-
-        return active_cooldowns
+    def get_min_cooldown(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> float:
+        return self.cooldown_store.minimum(model_ids, self.default_cooldown_time, parent_otel_span)
 
     def get_active_cooldowns(
-        self, model_ids: list[str], parent_otel_span: Span | None
+        self, model_ids: Sequence[str], parent_otel_span: Span | None
     ) -> list[tuple[str, CooldownCacheValue]]:
-        # Generate the keys for the deployments
-        keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
-        # Retrieve the values for the keys using mget
-        with service_target(ROUTER_COOLDOWNS_TARGET):
-            results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
-
-        active_cooldowns: Final = []
-        current_time: Final = time.time()
-        for model_id, result in zip(model_ids, results):
-            if result and isinstance(result, dict):
-                key = CooldownCache.get_cooldown_cache_key(model_id)
-                cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time)
-                if cooldown_cache_value is not None:
-                    active_cooldowns.append((model_id, cooldown_cache_value))
-
-        return active_cooldowns
-
-    def get_min_cooldown(self, model_ids: list[str], parent_otel_span: Span | None) -> float:
-        """Return min cooldown time required for a group of model id's."""
-
-        # Generate the keys for the deployments
-        keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
-
-        # Retrieve the values for the keys using mget
-        with service_target(ROUTER_COOLDOWNS_TARGET):
-            results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
-
-        min_cooldown_time: float | None = None
-        # Process the results
-        for model_id, result in zip(model_ids, results):
-            if result and isinstance(result, dict):
-                cooldown_cache_value = CooldownCacheValue(**result)
-                if min_cooldown_time is None or cooldown_cache_value["cooldown_time"] < min_cooldown_time:
-                    min_cooldown_time = cooldown_cache_value["cooldown_time"]
-
-        return min_cooldown_time or self.default_cooldown_time
-
-
-# Usage example:
-# cooldown_cache = CooldownCache(cache=your_cache_instance, cooldown_time=your_cooldown_time)
-# cooldown_cache.add_deployment_to_cooldown(deployment, original_exception, exception_status)
-# active_cooldowns = cooldown_cache.get_active_cooldowns()
+        return list(self.cooldown_store.active_sync(model_ids, parent_otel_span))

litellm/router_utils/cooldown_handlers.py

@@ -16,9 +16,7 @@ from typing import TYPE_CHECKING, Any, Final
 from pydantic import TypeAdapter, ValidationError
 
 import litellm
-from litellm._internal_context import service_target
 from litellm._logging import verbose_router_logger
-from litellm.caching.dual_cache import DualCache
 from litellm.constants import (
     DEFAULT_COOLDOWN_TIME_SECONDS,
     DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS,
@@ -30,7 +28,7 @@ from litellm.exceptions import MidStreamFallbackError
 from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
 from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
 from litellm.router_deprecated import renamed_router_param
-from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue
+from litellm.router_utils.cooldown_cache import CooldownCacheValue
 from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
 from litellm.types.router import Deployment
 from litellm.types.utils import (
@@ -748,7 +746,7 @@ def should_cooldown_based_on_allowed_fails_policy(
     When *allowed_fails_override* / *cooldown_time_override* are supplied they
     take precedence over the router-level values (used by deployment-level overrides).
 
-    The counter lives in the router's shared ``DualCache`` (Redis when configured), so
+    The counter lives in the cooldown store (Redis when configured), so
     every worker process increments the same key and the threshold applies fleet-wide.
 
     When *cache_key_suffix* is supplied the fail counter is keyed as
@@ -771,26 +769,12 @@ def should_cooldown_based_on_allowed_fails_policy(
         else (ctx.settings().cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS)
     )
 
-    base_key: Final = f"deployment:{deployment}:allowed_fails"
-    cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key
-    updated_fails: Final = _increment_allowed_fails(cache=ctx.health().store, cache_key=cache_key, ttl=cooldown_time)
+    updated_fails: Final = ctx.health().cooldowns.cooldown_store.increment_allowed_fails(
+        deployment, cache_key_suffix, cooldown_time
+    )
     return updated_fails > allowed_fails
 
 
-def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int:
-    """
-    Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier
-    before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count.
-    """
-    with service_target(ROUTER_COOLDOWNS_TARGET):
-        try:
-            return cache.increment_cache(key=cache_key, value=1, ttl=ttl)
-        except Exception as e:  # noqa: BLE001  # a Redis outage must not stop failing deployments from cooling down
-            verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
-            local_fails: Final = cache.get_cache(key=cache_key, local_only=True)
-            return local_fails if isinstance(local_fails, int) else 0
-
-
 def _is_allowed_fails_set_on_router(
     ctx: "RoutingContext",
 ) -> bool:

litellm/router_utils/routing_read_batch.py

@@ -1,29 +1,13 @@
-"""
-One Redis round trip for the reads a request needs before a deployment can be picked.
-
-The cooldown filter (`CooldownCache`, its own `DualCache`) and usage-based selection
-(`LowestTPMLoggingHandler_v2`, the router cache) each issue their own MGET because they live in
-different objects. `RoutingReadBatch` fetches both key sets in one
-`DualCache.async_batch_get_cache_shared` while the healthy deployments are being resolved and hands
-the usage slice to the strategy, so selection does not read again.
-"""
-
-import asyncio
-import itertools
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Iterator
 from contextlib import contextmanager
 from contextvars import ContextVar
-from dataclasses import dataclass
-from types import MappingProxyType
 from typing import TYPE_CHECKING, Final
 
-from litellm._internal_context import service_target
 from litellm._logging import verbose_router_logger
-from litellm.caching.dual_cache import DualCache
-from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
+from litellm.redis.stores import _routing_prefetch
+from litellm.redis.stores._routing_prefetch import TierRead, selector_read
 from litellm.router_deprecated import renamed_router_param
 from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage
-from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span
@@ -31,127 +15,29 @@ if TYPE_CHECKING:
     from litellm.router_api import RoutingContext
 
 
-ROUTER_COOLDOWNS_USAGE_TARGET: Final = "router_cooldowns_usage"
 ROUTER_USAGE_TARGET: Final = "router_usage"
-_PREFETCH_SLOT: Final = "routing_read"
 
 
-def _routing_read_target(cooldown_keys: Sequence[str], usage_keys: Sequence[str]) -> str:
-    if not usage_keys:
-        return ROUTER_COOLDOWNS_TARGET
-    if not cooldown_keys:
-        return ROUTER_USAGE_TARGET
-    return ROUTER_COOLDOWNS_USAGE_TARGET
+def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: list) -> list[str]:
+    tpm_keys, rpm_keys = selector.usage_counter_keys(deployments)
+    return tpm_keys + rpm_keys
 
 
-async def _backfill_prefetched_cache(
-    cache: DualCache,
-    due_keys: tuple[str, ...],
-    values: Mapping[str, object],
-) -> None:
-    cache_keys: Final = list(due_keys)
-    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 = {
-        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)
+def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> TierRead | None:
+    return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments))
 
 
-@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
     @renamed_router_param("litellm_router_instance")
-    def arm(
-        ctx: "RoutingContext",
-        usage_selector: LowestTPMLoggingHandler_v2 | None,
-        deployments: list,
-    ) -> None:
-        request: Final = active_request_redis_batches()
-        redis_cache: Final = ctx.health().store.redis_cache
-        if request is None or redis_cache is None or _PREFETCH_SLOT in request.prefetched:
-            return
-        cooldown_keys: Final = tuple(
-            CooldownCache.get_cooldown_cache_key(model_id) for model_id in ctx.registry().names(kind="ids")
+    def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None:
+        ctx.health().cooldowns.cooldown_store.arm_routing_read(
+            ctx.registry().names(kind="ids"), _selector_read(usage_selector, deployments)
         )
-        usage_keys: Final = (
-            () if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments)))
-        )
-        keys: Final = (*cooldown_keys, *usage_keys)
-        cooldown_store: Final = ctx.health().cooldowns.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
-        with service_target(_routing_read_target(cooldown_due, usage_due)):
-            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:
-        request: Final = active_request_redis_batches()
-        return request is not None and _PREFETCH_SLOT in request.prefetched
-
-    @staticmethod
-    def take(needed: Sequence[str]) -> "RoutingPrefetch | None":
-        """The armed prefetch when it covers every key this read needs; taken once, so a retry reads fresh."""
-        request: Final = active_request_redis_batches()
-        if request is None:
-            return 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
+        return _routing_prefetch.armed()
 
 
 _active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar(
@@ -179,71 +65,22 @@ class RoutingReadBatch:
 
     @staticmethod
     def for_strategy(strategy: str | None, selector: object) -> "RoutingReadBatch | None":
-        """Usage-based routing reads its counters with the cooldown state; every other strategy reads only the
-        cooldown state, and only through this batch when the request armed a prefetch for it. Otherwise the
-        router's plain cooldown read stays in charge."""
         if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2):
             return RoutingReadBatch(usage_selector=selector)
         return RoutingReadBatch(usage_selector=None) if RoutingPrefetch.armed() else None
 
     @renamed_router_param("litellm_router_instance")
     async def async_get_cooldown_deployments(
-        self,
-        ctx: "RoutingContext",
-        healthy_deployments: list,
-        parent_otel_span: "Span | None",
+        self, ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None"
     ) -> list[str]:
-        """
-        `_async_get_cooldown_deployments`, with the strategy's tpm/rpm counters for
-        `healthy_deployments` fetched in the same MGET and kept as `prefetched_usage`.
-        """
-        model_ids: Final = ctx.registry().names(kind="ids")
-        cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
-        selector: Final = self.usage_selector
-        usage_keys: Final = (
-            () if selector is None else tuple(itertools.chain(*selector.usage_counter_keys(healthy_deployments)))
+        usage_read: Final = _selector_read(self.usage_selector, healthy_deployments)
+        read: Final = await ctx.health().cooldowns.cooldown_store.routing_read(
+            ctx.registry().names(kind="ids"), usage_read, parent_otel_span
         )
-        reads: Final = (
-            (ctx.health().cooldowns.cooldown_store, cooldown_keys),
-            *(() if selector is None else ((selector.router_cache, list(usage_keys)),)),
-        )
-        with service_target(_routing_read_target(cooldown_keys, usage_keys)):
-            results: Final = await self._read_prefetched(reads) or await DualCache.async_batch_get_cache_shared(
-                reads, parent_otel_span=parent_otel_span
-            )
-        cooldown_results: Final = results[0]
-        if selector is not None:
-            usage_values: Final = results[1]
+        if usage_read is not None:
             self.prefetched_usage = PrefetchedUsage(
-                keys=frozenset(usage_keys),
-                values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))),
+                keys=frozenset(usage_read.keys),
+                values=None if read.selector_values is None else dict(zip(usage_read.keys, read.selector_values)),
             )
-
-        cooldown_models: Final = ctx.health().cooldowns.active_cooldowns_from_results(model_ids, cooldown_results)
-        verbose_router_logger.debug("retrieve cooldown models: %s", cooldown_models)
-        return [model_id for model_id, _ in cooldown_models]
-
-    @staticmethod
-    async def _read_prefetched(
-        reads: Sequence[tuple[DualCache, list[str]]],
-    ) -> list[list[object | None] | None] | None:
-        """Serve the reads from the request's armed `RoutingPrefetch`, backfilling each cache's memory tier as
-        its own batch read would. None when nothing usable was armed or the prefetch failed."""
-        prefetch: Final = RoutingPrefetch.take(tuple(itertools.chain.from_iterable(keys for _, keys in reads)))
-        if prefetch is None:
-            return None
-        try:
-            values: Final = await prefetch.result
-        except Exception as e:  # noqa: BLE001  # the shared read below applies the caches' own Redis fallback
-            verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
-            return None
-        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 = {key: values.get(key) for key, local in zip(keys, pending.result) if local is None}
-            results.append(await cache._apply_batch_get(pending, missed))  # pyright: ignore[reportPrivateUsage]  # same two-step read as async_batch_get_cache_shared
-        return results
+        verbose_router_logger.debug("retrieve cooldown models: %s", read.cooldowns)
+        return [model_id for model_id, _ in read.cooldowns]

litellm/types/router.py

@@ -46,6 +46,8 @@ from litellm._uuid import uuid
 from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
 from litellm.litellm_core_utils.core_helpers import normalize_drop_params
 from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name
+from litellm.redis.stores._keys import ROUTER_RPM_KEY as _ROUTER_RPM_KEY
+from litellm.redis.stores._keys import ROUTER_TPM_KEY as _ROUTER_TPM_KEY
 from litellm.types.llms.base import LiteLLMBaseModel
 from litellm.types.router_weights import RouterWeights
 
@@ -1290,8 +1292,8 @@ class RoutingStrategy(enum.Enum):
 
 
 class RouterCacheEnum(enum.Enum):
-    TPM = "global_router:{id}:{model}:tpm:{current_minute}"
-    RPM = "global_router:{id}:{model}:rpm:{current_minute}"
+    TPM = _ROUTER_TPM_KEY
+    RPM = _ROUTER_RPM_KEY
     ITPM = "global_router:{id}:{model}:itpm:{current_minute}"
     OTPM = "global_router:{id}:{model}:otpm:{current_minute}"
 

tests/unit/caching/test_request_redis_batch_pre_call.py

@@ -30,16 +30,21 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
 )
 from litellm.proxy.utils import InternalUsageCache
 from litellm.router_api import RoutingContext
-from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache
+from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET
+from litellm.redis.stores.router_cooldowns import POLICY
+from litellm.redis.stores._routing_prefetch import ArmedRead, TierRead, routing_target
 from litellm.router_utils.routing_read_batch import (
-    ROUTER_COOLDOWNS_USAGE_TARGET,
     ROUTER_USAGE_TARGET,
     RoutingPrefetch,
-    _routing_read_target,  # pyright: ignore[reportPrivateUsage]  # the family rule under test
 )
 
 from .test_redis_batch import FakeClient, FakeRedisCache, replies
 
+
+def _cooldown_key(model_id: str) -> str:
+    return f"deployment:{model_id}:cooldown"
+
+
 _MODEL_GROUP = "claude"
 _FAR_FUTURE = 4_102_444_800.0  # 2100-01-01, a cooldown stamped then is still active
 
@@ -315,7 +320,7 @@ async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issue
     commands = client.pipelines[0].commands
     assert [c[0] for c in commands] == ["MGET", "EVALSHA", "EVALSHA"]
     mget_keys = set(commands[0][1:])
-    assert {CooldownCache.get_cooldown_cache_key("dep-a"), CooldownCache.get_cooldown_cache_key("dep-b")} <= mget_keys
+    assert {_cooldown_key("dep-a"), _cooldown_key("dep-b")} <= mget_keys
     assert any(":tpm:" in key for key in mget_keys) and any(":rpm:" in key for key in mget_keys)
     assert redis_cache.alone == []
 
@@ -326,23 +331,20 @@ async def test_a_cooldown_recorded_locally_after_the_prefetch_left_still_exclude
 
     def replies(command: tuple[Any, ...]) -> Any:
         if command[0] == "MGET":  # Redis holds a stale cooldown for dep-b and nothing for dep-a
-            return [
-                json.dumps(expired) if key == CooldownCache.get_cooldown_cache_key("dep-b") else None
-                for key in command[1:]
-            ]
+            return [json.dumps(expired) if key == _cooldown_key("dep-b") else None for key in command[1:]]
         return _lua_ok_replies(command)
 
     client = FakeClient(replies)
     redis_cache = FakeRedisCache(client)
     router = _router(redis_cache)
     cooldown_store = router.cooldown_cache.cooldown_store
-    assert cooldown_store.in_memory_cache is not None
 
     with request_redis_batch_scope():
         router.arm_routing_read_prefetch(_MODEL_GROUP, {})
-        cooldown_store.in_memory_cache.set_cache(
-            CooldownCache.get_cooldown_cache_key("dep-a"),
+        cooldown_store.put(
+            "dep-a",
             {"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60},
+            60,
         )
         picks = {
             (
@@ -366,8 +368,8 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro
     with request_redis_batch_scope() as request:
         router.arm_routing_read_prefetch(_MODEL_GROUP, {})
         armed = request.prefetched["routing_read"]
-        assert isinstance(armed, RoutingPrefetch)
-        request.prefetched["routing_read"] = RoutingPrefetch(
+        assert isinstance(armed, ArmedRead)
+        request.prefetched["routing_read"] = ArmedRead(
             keys=frozenset({"other"}),
             fetched=armed.fetched,
             result=armed.result,
@@ -389,8 +391,8 @@ async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservati
     router: Final = _router(redis_cache)
     cooldown_keys: Final = frozenset(
         {
-            CooldownCache.get_cooldown_cache_key("dep-a"),
-            CooldownCache.get_cooldown_cache_key("dep-b"),
+            _cooldown_key("dep-a"),
+            _cooldown_key("dep-b"),
         }
     )
 
@@ -424,8 +426,8 @@ async def test_a_failed_prefetch_falls_back_to_the_shared_read():
     assert len(redis_cache.alone) == 1
     cooldown_keys: Final = frozenset(
         {
-            CooldownCache.get_cooldown_cache_key("dep-a"),
-            CooldownCache.get_cooldown_cache_key("dep-b"),
+            _cooldown_key("dep-a"),
+            _cooldown_key("dep-b"),
         }
     )
     fallback_cooldown_mgets: Final = tuple(
@@ -462,7 +464,7 @@ async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_fam
 async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkeypatch):
     clock: Final = 1_000_000.0
     monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
-    cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
+    cooldown_key: Final = _cooldown_key("dep-a")
     active_cooldown: Final = {
         "exception_received": "429",
         "status_code": "429",
@@ -518,8 +520,8 @@ async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkey
 async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(monkeypatch):
     clock: Final = 1_000_000.0
     monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
-    dep_a_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
-    dep_b_key: Final = CooldownCache.get_cooldown_cache_key("dep-b")
+    dep_a_key: Final = _cooldown_key("dep-a")
+    dep_b_key: Final = _cooldown_key("dep-b")
     old_cooldown: Final = {
         "exception_received": "429",
         "status_code": "429",
@@ -549,27 +551,29 @@ async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misse
     redis_cache.store[dep_a_key] = old_cooldown
     redis_cache.store[dep_b_key] = redis_only_cooldown
     router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
-    memory_cache: Final = router.cooldown_cache.cooldown_store.in_memory_cache
-    assert memory_cache is not None
+    cooldown_store: Final = router.cooldown_cache.cooldown_store
 
     with request_redis_batch_scope() as request:
         router.arm_routing_read_prefetch(_MODEL_GROUP, {})
-        memory_cache.set_cache(dep_a_key, newer_memory_cooldown)
+        cooldown_store.put("dep-a", newer_memory_cooldown, 60)
+        redis_cache.store[dep_a_key] = old_cooldown
         await request.flush_all()
 
     prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET")
 
     assert len(prefetched_mgets) == 1
     assert frozenset(prefetched_mgets[0][1:]) == frozenset({dep_a_key, dep_b_key})
-    assert memory_cache.get_cache(dep_a_key) == newer_memory_cooldown
-    assert memory_cache.get_cache(dep_b_key) == redis_only_cooldown
+    assert dict(await cooldown_store.active(("dep-a", "dep-b"), None)) == {
+        "dep-a": newer_memory_cooldown,
+        "dep-b": redis_only_cooldown,
+    }
 
 
 @pytest.mark.asyncio
 async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(monkeypatch):
     clock: Final = 1_000_000.0
     monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
-    cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
+    cooldown_key: Final = _cooldown_key("dep-a")
     active_cooldown: Final = {
         "exception_received": "429",
         "status_code": "429",
@@ -624,7 +628,7 @@ async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(m
 async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkeypatch):
     clock: Final = 1_000_000.0
     monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
-    cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
+    cooldown_key: Final = _cooldown_key("dep-a")
     active_cooldown: Final = {
         "exception_received": "429",
         "status_code": "429",
@@ -636,14 +640,12 @@ async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkey
     redis_cache.store[cooldown_key] = active_cooldown
     router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
     cooldown_store: Final = router.cooldown_cache.cooldown_store
-    memory_cache: Final = cooldown_store.in_memory_cache
-    assert memory_cache is not None
-    memory_cache.set_cache(cooldown_key, active_cooldown)
+    cooldown_store.put("dep-a", active_cooldown, 60)
 
     with request_redis_batch_scope() as request:
         router.arm_routing_read_prefetch(_MODEL_GROUP, {})
         await request.flush_all()
-        memory_cache.delete_cache(cooldown_key)
+        cooldown_store.evict_local("dep-a")
         deployment: Final = await router.async_get_available_deployment(
             model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
         )
@@ -654,7 +656,7 @@ async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkey
     )
 
     assert len(prefetched_mgets) == 1
-    assert prefetched_mgets[0][1:] == (CooldownCache.get_cooldown_cache_key("dep-b"),)
+    assert prefetched_mgets[0][1:] == (_cooldown_key("dep-b"),)
     assert deployment["model_info"]["id"] == "dep-b"
     assert fallback_cooldown_mgets == ((cooldown_key,),)
 
@@ -689,8 +691,8 @@ async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admissi
     commands = client.pipelines[0].commands
     assert [c[0] for c in commands] == ["MGET", "EVALSHA"]
     assert set(commands[0][1:]) == {
-        CooldownCache.get_cooldown_cache_key("dep-a"),
-        CooldownCache.get_cooldown_cache_key("dep-b"),
+        _cooldown_key("dep-a"),
+        _cooldown_key("dep-b"),
     }
     assert redis_cache.alone == []
 
@@ -699,8 +701,8 @@ async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admissi
     with request_redis_batch_scope() as request:
         shuffle.arm_routing_read_prefetch(_MODEL_GROUP, {})
         armed = request.prefetched["routing_read"]
-        assert isinstance(armed, RoutingPrefetch)
-        assert armed.keys == {CooldownCache.get_cooldown_cache_key("dep-a")}  # no usage counters for shuffle
+        assert isinstance(armed, ArmedRead)
+        assert armed.keys == {_cooldown_key("dep-a")}  # no usage counters for shuffle
 
 
 @pytest.mark.asyncio
@@ -733,8 +735,8 @@ async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_re
     pipeline_mgets = [command for pipeline in client.pipelines for command in pipeline.commands if command[0] == "MGET"]
     alone_mgets = [keys for command, keys in redis_cache.alone if command == "MGET"]
     cooldown_keys = {
-        CooldownCache.get_cooldown_cache_key("dep-a"),
-        CooldownCache.get_cooldown_cache_key("dep-b"),
+        _cooldown_key("dep-a"),
+        _cooldown_key("dep-b"),
     }
     cooldown_mgets = [command[1:] for command in pipeline_mgets if cooldown_keys.intersection(command[1:])] + [
         keys for keys in alone_mgets if cooldown_keys.intersection(keys)
@@ -768,8 +770,8 @@ async def test_concurrent_requests_share_one_cooldown_read_per_interval():
     deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm())
     cooldown_keys: Final = frozenset(
         {
-            CooldownCache.get_cooldown_cache_key("dep-a"),
-            CooldownCache.get_cooldown_cache_key("dep-b"),
+            _cooldown_key("dep-a"),
+            _cooldown_key("dep-b"),
         }
     )
     cooldown_mgets: Final = tuple(
@@ -799,8 +801,7 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses
         if command[0] == "MGET":
             result = next(mget_results)
             return [
-                None if result is None or key != CooldownCache.get_cooldown_cache_key("dep-a") else json.dumps(result)
-                for key in command[1:]
+                None if result is None or key != _cooldown_key("dep-a") else json.dumps(result) for key in command[1:]
             ]
         return _lua_ok_replies(command)
 
@@ -818,7 +819,7 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses
     monkeypatch.setattr(
         dual_cache_module.time,
         "time",
-        lambda: first_time + cooldown_store.redis_batch_cache_expiry + 1,
+        lambda: first_time + POLICY.redis_read_interval + 1,
     )
     with request_redis_batch_scope():
         router.arm_routing_read_prefetch(_MODEL_GROUP, {})
@@ -827,8 +828,8 @@ async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses
         )
 
     cooldown_keys = {
-        CooldownCache.get_cooldown_cache_key("dep-a"),
-        CooldownCache.get_cooldown_cache_key("dep-b"),
+        _cooldown_key("dep-a"),
+        _cooldown_key("dep-b"),
     }
     cooldown_mgets = [
         command
@@ -854,7 +855,7 @@ async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeyp
     router = _router(redis_cache, routing_strategy="usage-based-routing-v2")
     cooldown_store = router.cooldown_cache.cooldown_store
     usage_cache = router.lowesttpm_logger_v2.router_cache
-    time_offset = cooldown_store.redis_batch_cache_expiry + 0.5
+    time_offset = POLICY.redis_read_interval + 0.5
 
     assert time_offset < usage_cache.redis_batch_cache_expiry
 
@@ -873,8 +874,8 @@ async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeyp
 
     cooldown_keys = frozenset(
         {
-            CooldownCache.get_cooldown_cache_key("dep-a"),
-            CooldownCache.get_cooldown_cache_key("dep-b"),
+            _cooldown_key("dep-a"),
+            _cooldown_key("dep-b"),
         }
     )
     second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET")
@@ -1085,7 +1086,7 @@ async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_c
     [
         (("cooldown:a",), (), ROUTER_COOLDOWNS_TARGET),
         ((), ("usage:a",), ROUTER_USAGE_TARGET),
-        (("cooldown:a",), ("usage:a",), ROUTER_COOLDOWNS_USAGE_TARGET),
+        (("cooldown:a",), ("usage:a",), "router_cooldowns_usage"),
     ],
 )
 def test_the_routing_read_family_follows_the_keys_that_are_actually_due(
@@ -1093,4 +1094,4 @@ def test_the_routing_read_family_follows_the_keys_that_are_actually_due(
 ):
     """A routing MGET is ``router_cooldowns`` when only cooldown keys go out, ``router_usage`` when the
     cooldowns were already in memory and only usage counters go out, and the combined family otherwise."""
-    assert _routing_read_target(cooldown_keys, usage_keys) == expected
+    assert routing_target((TierRead(DualCache(), cooldown_keys), TierRead(DualCache(), usage_keys))) == expected

tests/unit/conftest.py

@@ -4,17 +4,88 @@ import importlib
 import os
 import shutil
 import tempfile
-from collections.abc import Callable, Coroutine, Iterator
+from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator
 from dataclasses import dataclass, field
 from pathlib import Path
 from types import MappingProxyType
-from typing import Final, Literal, TypeAlias
+from typing import TYPE_CHECKING, Final, Literal, TypeAlias
 
 import boto3
 import httpx
 import pytest
 from pytest_socket import enable_socket, socket_allow_hosts
 
+if TYPE_CHECKING:
+    from litellm.caching.redis_cache import RedisCache
+
+RouterRedis: TypeAlias = tuple["RedisCache", asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]]
+RedisWriteGate: TypeAlias = tuple[asyncio.Event, asyncio.Event, asyncio.Event]
+
+
+@pytest.fixture
+def redis_write_gate() -> RedisWriteGate:
+    released: Final = asyncio.Event()
+    released.set()
+    return asyncio.Event(), released, asyncio.Event()
+
+
+@pytest.fixture
+async def router_redis(
+    redis_write_gate: RedisWriteGate,
+) -> AsyncIterator[RouterRedis]:
+    from unittest.mock import patch
+
+    from fakeredis import FakeRedis, FakeServer
+    from fakeredis.aioredis import FakeConnection
+    from opentelemetry import trace
+    from redis.asyncio import BlockingConnectionPool
+    from redis.exceptions import ConnectionError
+
+    from litellm._internal_context import current_service_target
+    from litellm.caching.redis_cache import RedisCache
+
+    commands: Final[asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]] = asyncio.Queue()
+    started, released, unavailable = redis_write_gate
+
+    class ObservedConnection(FakeConnection):
+        async def send_command(self, *args: object, **kwargs: object) -> None:  # kwargs-ok: Redis transport signature
+            commands.put_nowait((args, current_service_target(), getattr(trace.get_current_span(), "name", None)))
+            await super().send_command(*args, **kwargs)
+
+        async def send_packed_command(self, command: Iterable[bytes] | bytes | str, check_health: bool = True) -> None:
+            chunks: Final = (
+                (command.encode(),)
+                if isinstance(command, str)
+                else (command,)
+                if isinstance(command, bytes)
+                else command
+            )
+            wire: Final = b"".join(chunks)
+            if b"INCRBYFLOAT" in wire:
+                commands.put_nowait(
+                    ((wire,), current_service_target(), getattr(trace.get_current_span(), "name", None))
+                )
+                started.set()
+                await released.wait()
+                if unavailable.is_set():
+                    raise ConnectionError("scripted write outage")
+            await super().send_packed_command(command, check_health=check_health)
+
+    server: Final = FakeServer()
+    client: Final = FakeRedis(server=server)
+    pool: Final = BlockingConnectionPool(connection_class=ObservedConnection, server=server)
+    try:
+        with (
+            patch("redis.Redis.from_url", return_value=client),
+            patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool),
+        ):
+            yield RedisCache(url="redis://router-unit.example:6379/0"), commands
+    finally:
+        released.set()
+        await pool.aclose()
+        client.close()
+
+
 HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset(
     (
         "PATH",
@@ -289,7 +360,9 @@ def isolate_litellm_globals() -> Iterator[None]:
     original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks)
     original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
     original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
-    original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__}
+    original_clients: Final = {
+        name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__
+    }
     original_loggers: Final = {
         logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters))
         for logger in ALL_LOGGERS

tests/unit/redis/stores/test_router_cooldowns.py

@@ -0,0 +1,112 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm._internal_context import current_service_target
+from litellm.caching.redis_batch import request_redis_batch_scope
+from litellm.caching.redis_cache import RedisCache
+from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
+
+if TYPE_CHECKING:
+    from tests.unit.conftest import RouterRedis
+
+
+@dataclass
+class Clock:
+    now: float = 1000.0
+    failure: RuntimeError | None = None
+
+    def __call__(self) -> float:
+        if self.failure is not None:
+            raise self.failure
+        return self.now
+
+
+def _entry(timestamp: float = 1000.0, ttl: float = 120.0) -> CooldownCacheValue:
+    return {"exception_received": "limited", "status_code": "429", "timestamp": timestamp, "cooldown_time": ttl}
+
+
+@pytest.mark.parametrize("suffix", (None, "", "timeout"))
+def test_allowed_fails_preserves_local_count_during_redis_failure(suffix: str | None) -> None:
+    redis: Final = MagicMock(spec=RedisCache)
+    redis.increment_cache.side_effect = ConnectionError("offline")
+    clock: Final = Clock()
+    store: Final = RouterCooldownStore(redis, clock=clock)
+    assert store.increment_allowed_fails("dep", suffix, 30.0) == 1
+    assert store.increment_allowed_fails("dep", suffix, 30.0) == 2
+    clock.now += 31.0
+    assert store.increment_allowed_fails("dep", suffix, 30.0) == 1
+    expected: Final = "deployment:dep:allowed_fails" + (f":{suffix}" if suffix else "")
+    redis.increment_cache.assert_called_with(expected, 1, ttl=30.0)
+
+
+def test_sync_wire_write_and_reads_keep_service_target_and_float_ttl() -> None:
+    redis: Final = MagicMock(spec=RedisCache)
+    seen: Final[list[str | None]] = []
+    redis.set_cache.side_effect = lambda *args, **kw: seen.append(current_service_target())
+    redis.batch_get_cache.side_effect = lambda **kw: {key: _entry() for key in kw["key_list"]}
+    store: Final = RouterCooldownStore(redis, clock=Clock())
+    store.put("dep", _entry(ttl=0.5), 0.5)
+    redis.set_cache.assert_called_once_with("deployment:dep:cooldown", _entry(ttl=0.5), ttl=0.5)
+    reader: Final = RouterCooldownStore(redis, clock=Clock())
+    assert reader.active_sync(["dep"], None) == (("dep", _entry()),)
+    assert reader.minimum(["dep"], 60.0, None) == 120.0
+    assert seen == ["router_cooldowns"]
+    assert current_service_target() is None
+
+
+@pytest.mark.asyncio
+async def test_remote_backfill_expires_at_corrected_sixty_second_cap() -> None:
+    clock: Final = Clock()
+    redis: Final = MagicMock(spec=RedisCache)
+    redis.async_batch_get_cache = AsyncMock(return_value={"deployment:dep:cooldown": _entry()})
+    store: Final = RouterCooldownStore(redis, clock=clock)
+    assert await store.active(["dep"], None) == (("dep", _entry()),)
+    clock.now += 61.0
+    assert await store.active(["dep"], None) == (("dep", _entry()),)
+    assert redis.async_batch_get_cache.await_count == 2
+    clock.now += 60.0
+    assert await store.active(["dep"], None) == ()
+
+
+@pytest.mark.asyncio
+async def test_lazy_attach_and_failure_retry_do_not_throttle_next_read() -> None:
+    store: Final = RouterCooldownStore(None, clock=Clock())
+    assert await store.active(["dep"], None) == ()
+    redis: Final = MagicMock(spec=RedisCache)
+    redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("offline"))
+    store.attach_redis(redis)
+    assert await store.active(["dep"], None) == ()
+    redis.async_batch_get_cache.side_effect = None
+    redis.async_batch_get_cache.return_value = {"deployment:dep:cooldown": _entry()}
+    assert await store.active(["dep"], None) == (("dep", _entry()),)
+    assert redis.async_batch_get_cache.await_count == 2
+
+
+def test_minimum_preserves_zero_default_and_does_not_correct_expired_entries() -> None:
+    store: Final = RouterCooldownStore(None, clock=Clock())
+    store.put("zero", _entry(ttl=0), 600)
+    store.put("expired", _entry(timestamp=0, ttl=20), 600)
+    assert store.minimum(["zero"], 60, None) == 60
+    assert store.minimum(["expired"], 60, None) == 20
+    assert store.active_sync(["expired"], None) == ()
+
+
+@pytest.mark.asyncio
+async def test_prefetch_backfill_failure_allows_an_immediate_read_retry(router_redis: "RouterRedis") -> None:
+    redis, _ = router_redis
+    clock: Final = Clock()
+    value: Final = _entry()
+    writer: Final = RouterCooldownStore(redis, clock=clock)
+    reader: Final = RouterCooldownStore(redis, clock=clock)
+    writer.put("dep", value, 120.0)
+
+    with request_redis_batch_scope() as request:
+        reader.arm_routing_read(("dep",), None)
+        clock.failure = RuntimeError("backfill clock unavailable")
+        await request.flush_all()
+
+    clock.failure = None
+    assert await reader.active(("dep",), None) == (("dep", value),)

tests/unit/redis/stores/test_router_usage.py

@@ -0,0 +1,87 @@
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from litellm import Router
+from litellm._internal_context import current_service_target
+from litellm.caching.redis_cache import RedisCache
+from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("operation", ("flush_cache", "reset"))
+async def test_router_reset_clears_usage_and_failure_counters(operation: str) -> None:
+    router: Final = Router(model_list=[])
+    usage: Final = router._core_usage_store
+    cooldowns: Final = router._core_cooldown_cache.cooldown_store
+    counters: Final = (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
+    await usage.increment_deployment("dep", "model", "12-34", 7, 1, None)
+    assert await usage.async_batch_get(counters) == [7, 1]
+    assert cooldowns.increment_allowed_fails("dep", None, 30.0) == 1
+    getattr(router, operation)()
+    assert await usage.async_batch_get(counters) == [None, None]
+    assert cooldowns.increment_allowed_fails("dep", None, 30.0) == 1
+
+
+@pytest.mark.asyncio
+async def test_deployment_increment_preserves_pipeline_ttl_and_integer_mirror() -> None:
+    redis: Final = MagicMock(spec=RedisCache)
+    redis.async_increment_pipeline = AsyncMock(return_value=[7.75, 2.0])
+    store: Final = RouterUsageStore(redis)
+    await store.increment_deployment("dep", "anthropic/model", "12-34", 7.75, 2, None)
+    redis.async_increment_pipeline.assert_awaited_once_with(
+        increment_list=[
+            {"key": "global_router:dep:anthropic/model:tpm:12-34", "increment_value": 7.75, "ttl": 60},
+            {"key": "global_router:dep:anthropic/model:rpm:12-34", "increment_value": 2, "ttl": 60},
+        ],
+        parent_otel_span=None,
+    )
+    values: Final = await store.async_batch_get(
+        (UsageCounter("dep", "anthropic/model", "12-34", "tpm"), UsageCounter("dep", "anthropic/model", "12-34", "rpm"))
+    )
+    assert values == [7, 2]
+    assert all(type(value) is int for value in values)
+    redis.async_batch_get_cache.assert_not_called()
+    assert current_service_target() is None
+
+
+@pytest.mark.asyncio
+async def test_memory_only_usage_increments_accumulate_without_redis() -> None:
+    store: Final = RouterUsageStore(None)
+    await store.increment_deployment("dep", "model", "12-34", 2, 1, None)
+    await store.increment_deployment("dep", "model", "12-34", 5, 1, None)
+    assert await store.async_batch_get(
+        (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
+    ) == [7, 2]
+
+
+@pytest.mark.asyncio
+async def test_pre_call_rpm_increment_uses_same_namespace_and_service_target() -> None:
+    redis: Final = MagicMock(spec=RedisCache)
+    seen: Final[list[str | None]] = []
+
+    async def increment(*args: object, **kwargs: object) -> float:
+        seen.append(current_service_target())
+        return 4.0
+
+    redis.async_increment = AsyncMock(side_effect=increment)
+    store: Final = RouterUsageStore(None)
+    store.attach_redis(redis)
+    counter: Final = UsageCounter("dep", "model", "12-34", "rpm")
+    assert await store.async_increment(counter, value=1, ttl=60) == 4.0
+    assert seen == ["router_usage"]
+    redis.async_increment.assert_awaited_once_with(
+        "global_router:dep:model:rpm:12-34", 1, parent_otel_span=None, ttl=60, refresh_ttl=False
+    )
+
+
+@pytest.mark.asyncio
+async def test_failed_usage_pipeline_keeps_the_memory_increment_and_is_readable() -> None:
+    redis: Final = MagicMock(spec=RedisCache)
+    redis.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("offline"))
+    store: Final = RouterUsageStore(redis)
+    await store.increment_deployment("dep", "model", "12-34", 5, 1, None)
+    assert await store.async_batch_get(
+        (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
+    ) == [5, 1]

tests/unit/router_strategy/test_complexity_router.py

@@ -15992,9 +15992,7 @@ class TestHealthFallbackDispatch:
                     {"primary-id": {"is_healthy": True, "timestamp": time.time()}}
                 )
             else:
-                router.cooldown_cache.cooldown_store.delete_cache(
-                    router.cooldown_cache.get_cooldown_cache_key("primary-id")
-                )
+                router.cooldown_cache.cooldown_store.evict_local("primary-id")
             recovered: Final[dict[str, object]] = {"session_id": "outage"}
             assert await self._request(router, surface, stream, recovered) == "primary"
             assert recovered["routing_decision"]["routed_model"] == "primary"
@@ -16170,9 +16168,7 @@ class TestHealthFallbackDispatch:
                     {"primary-id": {"is_healthy": True, "timestamp": time.time()}}
                 )
             else:
-                router.cooldown_cache.cooldown_store.delete_cache(
-                    router.cooldown_cache.get_cooldown_cache_key("primary-id")
-                )
+                router.cooldown_cache.cooldown_store.evict_local("primary-id")
             recovered: Final[dict[str, object]] = {"session_id": "pinned"}
             assert await self._request(router, "chat", False, recovered) == "primary"
             assert recovered["routing_decision"]["cause"] == "session_affinity_pin"

tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
-from typing import Final
+from itertools import chain
+from typing import TYPE_CHECKING, Final
 from collections.abc import AsyncIterator, Mapping
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -21,6 +22,10 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
 )
 
 
+if TYPE_CHECKING:
+    from tests.unit.conftest import RouterRedis
+
+
 def _repinning_router(enabled: bool | None = True) -> litellm.Router:
     return litellm.Router(
         model_list=[
@@ -100,7 +105,7 @@ async def test_session_repin_requires_opt_in_unavailable_original_and_success(
         assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
         assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"}
         assert await router.cache.async_get_cache(user_key) == {"model_id": "original"}
-        router.cooldown_cache.cooldown_store.flush_cache()
+        router.cooldown_cache.cooldown_store.evict_local("original")
         await _record_replacement_success(router, "other")
         assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
     finally:
@@ -129,7 +134,7 @@ async def test_concurrent_successful_replacements_keep_first_winner() -> None:
         await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other"))
         winner: Final = await router.cache.async_get_cache(_repin_key())
         assert winner in ({"model_id": "replacement"}, {"model_id": "other"})
-        router.cooldown_cache.cooldown_store.flush_cache()
+        router.cooldown_cache.cooldown_store.evict_local("original")
         await _record_replacement_success(router, "original")
         assert await router.cache.async_get_cache(_repin_key()) == winner
     finally:
@@ -320,22 +325,12 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M
         router.discard()
 
 
-class _ReadCountingCache(InMemoryCache):
-    def __init__(self) -> None:
-        super().__init__()
-        self.read_keys: tuple[str, ...] = ()
-
-    def get_cache(self, key: str, **kwargs: object) -> object:
-        self.read_keys += (key,)
-        return super().get_cache(key, **kwargs)
-
-
 @pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"])
-async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None:
+async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str, router_redis: "RouterRedis") -> None:
     router: Final = _repinning_router(case != "disabled")
-    reads: Final = _ReadCountingCache()
+    redis, commands = router_redis
     try:
-        router.cooldown_cache.cooldown_store.in_memory_cache = reads
+        router.update_redis_cache(redis)
         callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
         if case == "group-disabled":
             callback.model_group_affinity_config = {"repin-group": []}
@@ -351,14 +346,13 @@ async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: st
             {"status": "completed"},
         )
         expected: Final = (
-            {
-                router.cooldown_cache.get_cooldown_cache_key(model_id)
-                for model_id in ("original", "replacement", "other")
-            }
+            {f"deployment:{model_id}:cooldown" for model_id in ("original", "replacement", "other")}
             if case == "active"
             else set()
         )
-        assert set(reads.read_keys) == expected
+        observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
+        reads: Final = frozenset(chain.from_iterable(args[1:] for args, _, _ in observed if args[0] == "MGET"))
+        assert reads == expected
     finally:
         router.discard()
 

tests/unit/router_utils/test_cooldown_cache.py

@@ -1,629 +1,43 @@
-"""
-Unit tests for CooldownCache exception masking functionality
-"""
-
-import time
-from unittest.mock import MagicMock
+from typing import Final
 
 import pytest
 
-# Add the parent directory to the system path
-from litellm._internal_context import current_service_target
-from litellm.caching.dual_cache import DualCache
-from litellm.caching.in_memory_cache import InMemoryCache
-from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
-from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
+from litellm.redis.stores.router_cooldowns import RouterCooldownStore
+from litellm.router_utils.cooldown_cache import CooldownCache
+
+_NOW: Final = 1000.0
 
 
-class TestCooldownCacheExceptionMasking:
-    """Test suite for CooldownCache exception masking functionality"""
-
-    @pytest.fixture
-    def cooldown_cache(self):
-        """Create a CooldownCache instance for testing"""
-        mock_dual_cache = MagicMock(spec=DualCache)
-        return CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)
-
-    def test_exception_masker_initialization(self, cooldown_cache):
-        """Test that the exception masker is properly initialized"""
-        assert isinstance(cooldown_cache.exception_masker, SensitiveDataMasker)
-        assert cooldown_cache.exception_masker.visible_prefix == 50
-        assert cooldown_cache.exception_masker.visible_suffix == 0
-        assert cooldown_cache.exception_masker.mask_char == "*"
-
-    def test_short_exception_string_not_masked(self, cooldown_cache):
-        """Test that short exception strings are not masked"""
-        short_exception = "Short error"
-        model_id = "test-model"
-        exception_status = 500
-        cooldown_time = 30.0
-
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=Exception(short_exception),
-            exception_status=exception_status,
-            cooldown_time=cooldown_time,
+@pytest.mark.parametrize("message", ("Short error", "Rate limit " + "secret" * 100))
+@pytest.mark.parametrize("ttl", (None, 0.5, 30.0))
+def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message: str, ttl: float | None) -> None:
+    store: Final = RouterCooldownStore(None, clock=lambda: _NOW)
+    cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW)
+    cache.add_deployment_to_cooldown("dep", Exception(message), 429, ttl)
+    expected_ttl: Final = 60.0 if ttl is None else ttl
+    expected_message: Final = message if len(message) <= 50 else message[:50] + "*" * (len(message) - 50)
+    assert cache.get_active_cooldowns(["dep"], None) == [
+        (
+            "dep",
+            {
+                "exception_received": expected_message,
+                "status_code": "429",
+                "timestamp": _NOW,
+                "cooldown_time": expected_ttl,
+            },
         )
-
-        # Short exception should not be masked
-        assert cooldown_data["exception_received"] == short_exception
-        assert cooldown_key == f"deployment:{model_id}:cooldown"
-
-    def test_long_exception_string_masked(self, cooldown_cache):
-        """Test that long exception strings are properly masked"""
-        # Create a long exception string that simulates prompt leakage
-        long_exception = (
-            "litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occurred - "
-            "No deployments available for selected model, Try again in 5 seconds. "
-            "Passed model=anthropic_claude_sonnet_4_v1_0. pre-call-checks=False, "
-            "cooldown_list=[('deepseek_r1-eastus', {'exception_received': "
-            "'litellm.RateLimitError: RateLimitError: Azure_aiException - "
-            '{"error":{"code":"Invalid input","status":422,"message":"invalid input error",'
-            '"details":[{"type":"model_attributes_type","loc":["body"],'
-            '"msg":"Tell me a story about a dragon and a princess in a magical kingdom '
-            "where the dragon is actually protecting the princess from an evil wizard "
-            'who wants to steal her magical powers and use them to conquer the world"}]}'
-        )
-
-        model_id = "test-model"
-        exception_status = 429
-        cooldown_time = 60.0
-
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=Exception(long_exception),
-            exception_status=exception_status,
-            cooldown_time=cooldown_time,
-        )
-
-        masked_exception = cooldown_data["exception_received"]
-
-        # Should start with first 50 characters
-        assert masked_exception.startswith(long_exception[:50])
-
-        # Should contain masking characters
-        assert "*" in masked_exception
-
-        # Should be same length (prefix + asterisks)
-        assert len(masked_exception) == len(long_exception)
-
-        # Should not contain the sensitive prompt content
-        assert "Tell me a story about a dragon" not in masked_exception
-        assert "magical kingdom" not in masked_exception
-
-        # Should preserve the error type information at the beginning (first 50 chars)
-        assert masked_exception.startswith("litellm.proxy.proxy_server._handle_llm_api_excepti")
-
-    def test_exception_with_api_keys_masked(self, cooldown_cache):
-        """Test that API keys in exceptions are properly masked"""
-        exception_with_key = (
-            "Authentication failed with api_key=sk-9876567890abcdefghijklmnopqrstuvwxyz "
-            "and token=bearer_token_123456789 for model gpt-4"
-        )
-
-        model_id = "test-model"
-        exception_status = 401
-        cooldown_time = 30.0
-
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=Exception(exception_with_key),
-            exception_status=exception_status,
-            cooldown_time=cooldown_time,
-        )
-
-        masked_exception = cooldown_data["exception_received"]
-
-        # Should mask the sensitive content while preserving structure
-        assert masked_exception.startswith("Authentication failed with api_key=sk-98765678")
-        assert "*" in masked_exception
-        assert len(masked_exception) == len(exception_with_key)
-
-    def test_cooldown_data_structure(self, cooldown_cache):
-        """Test that the cooldown data structure is correctly formed"""
-        exception_msg = "Test exception for structure validation"
-        model_id = "test-model"
-        exception_status = 500
-        cooldown_time = 45.0
-
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=Exception(exception_msg),
-            exception_status=exception_status,
-            cooldown_time=cooldown_time,
-        )
-
-        # Verify cooldown data structure
-        assert isinstance(cooldown_data, dict)
-        assert "exception_received" in cooldown_data
-        assert "status_code" in cooldown_data
-        assert "timestamp" in cooldown_data
-        assert "cooldown_time" in cooldown_data
-
-        # Verify data types
-        assert isinstance(cooldown_data["exception_received"], str)
-        assert isinstance(cooldown_data["status_code"], str)
-        assert isinstance(cooldown_data["timestamp"], float)
-        assert isinstance(cooldown_data["cooldown_time"], float)
-
-        # Verify values
-        assert cooldown_data["status_code"] == str(exception_status)
-        assert cooldown_data["cooldown_time"] == cooldown_time
-        assert cooldown_data["exception_received"] == exception_msg
-
-    def test_exception_object_conversion(self, cooldown_cache):
-        """Test that different exception types are properly converted to strings"""
-        # Test with different exception types
-        exceptions = [
-            ValueError("Invalid value provided"),
-            KeyError("Missing required key"),
-            RuntimeError("Runtime error occurred"),
-            Exception("Generic exception"),
-        ]
-
-        for exc in exceptions:
-            model_id = f"test-model-{exc.__class__.__name__}"
-
-            cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-                model_id=model_id,
-                original_exception=exc,
-                exception_status=500,
-                cooldown_time=30.0,
-            )
-
-            # Should successfully convert exception to string
-            assert isinstance(cooldown_data["exception_received"], str)
-            assert str(exc) == cooldown_data["exception_received"]  # Short exceptions not masked
-
-    def test_masking_preserves_error_debugging_info(self, cooldown_cache):
-        """Test that masking preserves essential debugging information"""
-        debugging_exception = (
-            "RateLimitError: Rate limit exceeded for model gpt-4. "
-            "Current usage: 1000 tokens/minute. Limit: 500 tokens/minute. "
-            "Request details: model=gpt-4, user_id=user123, "
-            "prompt='Write a comprehensive analysis of the economic implications "
-            "of artificial intelligence adoption in the healthcare sector, including "
-            "potential cost savings, job displacement, and regulatory challenges'"
-        )
-
-        model_id = "gpt-4-deployment"
-        exception_status = 429
-        cooldown_time = 120.0
-
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=Exception(debugging_exception),
-            exception_status=exception_status,
-            cooldown_time=cooldown_time,
-        )
-
-        masked_exception = cooldown_data["exception_received"]
-
-        # Should preserve error type and initial debugging info (first 50 chars)
-        assert masked_exception.startswith("RateLimitError: Rate limit exceeded for model gpt-")
-
-        # Should mask the prompt content
-        assert "Write a comprehensive analysis" not in masked_exception
-        assert "healthcare sector" not in masked_exception
-
-        # Should contain masking indicator
-        assert "*" in masked_exception
-
-    def test_error_handling_in_common_add_cooldown_logic(self, cooldown_cache):
-        """Test error handling in the _common_add_cooldown_logic method"""
-        # This test ensures that edge cases are properly handled
-        model_id = "test-model"
-
-        # Test with None exception (edge case) - should be handled gracefully
-        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
-            model_id=model_id,
-            original_exception=None,
-            exception_status=500,
-            cooldown_time=30.0,
-        )
-
-        # Should handle None by converting to string
-        assert cooldown_data["exception_received"] == "None"
-        assert cooldown_key == f"deployment:{model_id}:cooldown"
-
-    def test_custom_masker_settings(self):
-        """Test that custom masker settings work correctly"""
-        mock_dual_cache = MagicMock(spec=DualCache)
-
-        # Create cooldown cache and verify default settings
-        cache = CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)
-
-        # Test that we can access and verify the masker configuration
-        assert cache.exception_masker.visible_prefix == 50
-        assert cache.exception_masker.visible_suffix == 0
-        assert cache.exception_masker.mask_char == "*"
-
-        # Test masking behavior with these settings
-        long_string = "A" * 100  # 100 character string
-        masked = cache.exception_masker.mask_value(long_string)
-
-        # Should show first 50 characters, then all asterisks
-        expected = "A" * 50 + "*" * 50
-        assert masked == expected
-
-
-class TestCooldownCacheTTLCorrection:
-    def _make_cooldown_cache(self) -> CooldownCache:
-        in_memory = InMemoryCache()
-        dual_cache = DualCache(in_memory_cache=in_memory)
-        return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
-
-    def test_expired_entry_evicted_and_not_returned(self):
-        """
-        An entry with timestamp+cooldown_time in the past must be evicted from
-        in-memory cache and excluded from the active cooldown list.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "expired-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        expired_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - 120.0,
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
-
-        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert active == [], "Expired cooldown entry must not appear in active cooldowns"
-        assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
-
-    def test_active_entry_is_returned(self):
-        """
-        An entry whose cooldown window has not elapsed must appear in the active list.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "active-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        active_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time(),
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, active_value, ttl=60)
-
-        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert len(active) == 1
-        assert active[0][0] == model_id
-
-    def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self):
-        """
-        When DualCache backfills from Redis using the default 600s TTL, the in-memory
-        TTL must be corrected to min(remaining, 60) seconds.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "backfilled-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        remaining = 30.0
-        value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - (60.0 - remaining),
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, value, ttl=600)
-
-        before_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert before_expiry is not None
-
-        cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert after_expiry is not None
-        corrected_remaining = after_expiry - time.time()
-        assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
-        assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)"
-
-    @pytest.mark.asyncio
-    async def test_async_expired_entry_evicted(self):
-        """
-        Async path must also evict expired entries.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "async-expired"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        expired_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - 120.0,
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
-
-        active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert active == [], "Expired entry must not appear in async active cooldowns"
-        assert cc.in_memory_cache.get_cache(key) is None
-
-    @pytest.mark.asyncio
-    async def test_async_active_entry_is_returned(self):
-        """
-        Async counterpart of test_active_entry_is_returned: an entry whose cooldown
-        window has not elapsed must appear in the async active list too.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "async-active-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        active_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time(),
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, active_value, ttl=60)
-
-        active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert len(active) == 1
-        assert active[0][0] == model_id
-
-
-class TestCorrectedActiveCooldown:
-    def _make_cooldown_cache(self) -> CooldownCache:
-        in_memory = InMemoryCache()
-        dual_cache = DualCache(in_memory_cache=in_memory)
-        return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
-
-    def _entry(self, timestamp: float, cooldown_time: float) -> CooldownCacheValue:
-        return CooldownCacheValue(
-            exception_received="Rate limit",
-            status_code="429",
-            timestamp=timestamp,
-            cooldown_time=cooldown_time,
-        )
-
-    def test_expired_entry_returns_none_and_evicts(self):
-        cc = self._make_cooldown_cache()
-        key = "deployment:expired-dep:cooldown"
-        entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0)
-        cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
-
-        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
-
-        assert result is None
-        assert cc.in_memory_cache.get_cache(key) is None
-
-    def test_active_entry_within_window_returns_value(self):
-        cc = self._make_cooldown_cache()
-        key = "deployment:active-dep:cooldown"
-        entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
-        cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
-
-        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
-
-        assert result is not None
-        assert result["status_code"] == "429"
-
-    def test_inflated_ttl_is_corrected(self):
-        cc = self._make_cooldown_cache()
-        key = "deployment:backfilled-dep:cooldown"
-        remaining = 30.0
-        entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0)
-        cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
-
-        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
-
-        assert result is not None
-        corrected_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert corrected_expiry is not None
-        assert corrected_expiry - time.time() <= 60.0
-
-    def test_normal_ttl_not_modified(self):
-        cc = self._make_cooldown_cache()
-        key = "deployment:normal-dep:cooldown"
-        entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
-        cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
-        original_expiry = cc.in_memory_cache.ttl_dict.get(key)
-
-        cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
-
-        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert after_expiry == original_expiry
-
-
-class SharedRedisDouble:
-    """
-    In-process stand-in for RedisCache, shared by several DualCache instances so that
-    tests can model two proxy replicas talking to one Redis.
-    """
-
-    def __init__(self) -> None:
-        self.store: dict = {}  # mutable-ok: stands in for Redis' own mutable keyspace
-
-    def set_cache(self, key, value, **kwargs):
-        self.store[key] = value
-
-    async def async_set_cache(self, key, value, **kwargs):
-        self.store[key] = value
-
-    def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
-        return {key: self.store.get(key) for key in key_list}
-
-    async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
-        return {key: self.store.get(key) for key in key_list}
-
-
-class TestCooldownPropagationBetweenReplicas:
-    """
-    A cooldown written by one replica has to reach its siblings quickly. The router's own
-    DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads
-    get their own cache with a much shorter Redis read interval.
-    """
-
-    def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache:
-        router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis)
-        if read_interval is None:
-            return CooldownCache(cache=router_cache, default_cooldown_time=60.0)
-        return CooldownCache(
-            cache=router_cache,
-            default_cooldown_time=60.0,
-            redis_read_interval_seconds=read_interval,
-        )
-
-    @pytest.mark.asyncio
-    async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self):
-        redis = SharedRedisDouble()
-        replica_a = self._make_replica(redis, read_interval=0.25)
-        replica_b = self._make_replica(redis, read_interval=0.25)
-        model_id = "shared-deployment"
-
-        assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
-
-        replica_a.add_deployment_to_cooldown(
-            model_id=model_id,
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=60.0,
-        )
-
-        time.sleep(0.3)
-
-        active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
-        assert [model_id] == [entry[0] for entry in active], (
-            "sibling replica must pick up a cooldown written by another replica within the read interval"
-        )
-
-    @pytest.mark.asyncio
-    async def test_sibling_replica_sees_cooldown_within_default_read_interval(self):
-        redis = SharedRedisDouble()
-        replica_a = self._make_replica(redis)
-        replica_b = self._make_replica(redis)
-        model_id = "default-interval-deployment"
-
-        assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
-
-        replica_a.add_deployment_to_cooldown(
-            model_id=model_id,
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=60.0,
-        )
-
-        time.sleep(1.2)
-
-        active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
-        assert [model_id] == [entry[0] for entry in active], (
-            "the shipped default read interval must let a sibling replica see a cooldown about a second later"
-        )
-
-    def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self):
-        redis = SharedRedisDouble()
-        replica_a = self._make_replica(redis, read_interval=0.25)
-        replica_b = self._make_replica(redis, read_interval=0.25)
-        model_id = "sync-shared-deployment"
-
-        assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == []
-
-        replica_a.add_deployment_to_cooldown(
-            model_id=model_id,
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=60.0,
-        )
-
-        time.sleep(0.3)
-
-        active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None)
-        assert [model_id] == [entry[0] for entry in active]
-
-    @pytest.mark.asyncio
-    async def test_redis_attached_after_construction_is_still_used(self):
-        redis = SharedRedisDouble()
-        router_cache = DualCache(in_memory_cache=InMemoryCache())
-        writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25)
-        router_cache.attach_redis_cache(redis)
-        reader = self._make_replica(redis, read_interval=0.25)
-        model_id = "late-redis-deployment"
-
-        writer.add_deployment_to_cooldown(
-            model_id=model_id,
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=60.0,
-        )
-
-        active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None)
-        assert [model_id] == [entry[0] for entry in active], (
-            "a router that wires Redis after building its cooldown cache must still publish cooldowns to it"
-        )
-
-
-class TestCooldownSurvivesUnrelatedCacheTraffic:
-    @pytest.mark.asyncio
-    async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self):
-        router_cache = DualCache(in_memory_cache=InMemoryCache())
-        cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0)
-        model_id = "busy-router-deployment"
-
-        cc.add_deployment_to_cooldown(
-            model_id=model_id,
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=30.0,
-        )
-
-        for i in range(400):
-            router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i})
-
-        active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None)
-        assert [model_id] == [entry[0] for entry in active], (
-            "unrelated router cache traffic must not evict a cooldown that is still running"
-        )
-
-
-class TestCooldownStoreCallsDeclareTheirKeyFamily:
-    """Every cooldown store call runs inside ``service_target("router_cooldowns")`` so the
-    Redis service spans read ``redis.set router_cooldowns`` / ``redis.mget router_cooldowns``,
-    the sync paths included (the async MGET already did)."""
-
-    def _cooldown_cache_with_recording_store(self, seen: list[tuple[str, str | None]]) -> CooldownCache:
-        cc = CooldownCache(cache=DualCache(in_memory_cache=InMemoryCache()), default_cooldown_time=60.0)
-        store = MagicMock()
-
-        def _set_cache(**_kwargs):
-            seen.append(("set", current_service_target()))
-
-        def _batch_get_cache(**_kwargs):
-            seen.append(("mget", current_service_target()))
-            return []
-
-        store.set_cache.side_effect = _set_cache
-        store.batch_get_cache.side_effect = _batch_get_cache
-        cc._cooldown_store = store
-        return cc
-
-    def test_sync_cooldown_write_runs_under_router_cooldowns(self):
-        seen: list[tuple[str, str | None]] = []
-        cc = self._cooldown_cache_with_recording_store(seen)
-
-        cc.add_deployment_to_cooldown(
-            model_id="dep-1",
-            original_exception=Exception("Internal server error"),
-            exception_status=500,
-            cooldown_time=30.0,
-        )
-
-        assert seen == [("set", "router_cooldowns")]
-        assert current_service_target() is None
-
-    def test_sync_cooldown_reads_run_under_router_cooldowns(self):
-        seen: list[tuple[str, str | None]] = []
-        cc = self._cooldown_cache_with_recording_store(seen)
-
-        assert cc.get_active_cooldowns(["dep-1"], parent_otel_span=None) == []
-        assert cc.get_min_cooldown(["dep-1"], parent_otel_span=None) == 60.0
-
-        assert seen == [("mget", "router_cooldowns"), ("mget", "router_cooldowns")]
+    ]
+
+
+@pytest.mark.asyncio
+async def test_sync_and_async_router_cooldown_surfaces_share_active_state() -> None:
+    store: Final = RouterCooldownStore(None, clock=lambda: _NOW)
+    cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: _NOW)
+    cache.add_deployment_to_cooldown("dep", Exception("unavailable"), 503, 30.0)
+    expected: Final = [
+        ("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": _NOW, "cooldown_time": 30.0})
+    ]
+    assert cache.get_active_cooldowns(["absent", "dep"], None) == expected
+    assert await cache.async_get_active_cooldowns(["absent", "dep"], None) == expected
+    assert cache.get_min_cooldown(["dep"], None) == 30.0
+    assert cache.get_min_cooldown(["absent"], None) == 60.0

tests/unit/router_utils/test_cooldown_handlers.py

@@ -18,6 +18,8 @@ import litellm
 from litellm import Router
 from litellm._internal_context import current_service_target
 from litellm.caching.dual_cache import DualCache
+from litellm.caching.redis_cache import RedisCache
+from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
 from litellm.router_api import RoutingContext
@@ -27,7 +29,6 @@ from litellm.router_utils.cooldown_handlers import (
     _get_deployment_cooldown_policy,
     _has_explicit_allowed_fails_policy_for_exception,
     _has_fallback_group,
-    _increment_allowed_fails,
     _is_cooldown_required,
     _resolve_allowed_fails_from_policy,
     _should_cooldown_based_on_deployment_policy,
@@ -74,6 +75,7 @@ def _ctx_double(
     ctx.settings.return_value.allowed_fails_policy = None
     ctx.settings.return_value.allowed_fails_for.return_value = None
     ctx.health.return_value.store = store if store is not None else DualCache(in_memory_cache=InMemoryCache())
+    ctx.health.return_value.cooldowns = CooldownCache(RouterCooldownStore(None), cooldown_time)
     return ctx
 
 
@@ -360,85 +362,57 @@ class TestShouldCooldownBasedOnDeploymentPolicy:
 
 
 class TestShouldCooldownBasedOnAllowedFailsPolicy:
-    def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | MagicMock | None = None) -> MagicMock:
-        return _ctx_double(cooldown_time=cooldown_time, allowed_fails=0, store=cache)
+    def _make_router(self, cooldown_time: float = 60.0, store: RouterCooldowns | None = None) -> MagicMock:
+        ctx: Final = _ctx_double(cooldown_time=cooldown_time, allowed_fails=0)
+        ctx.health.return_value.cooldowns = CooldownCache(store or RouterCooldownStore(None), cooldown_time)
+        return ctx
 
     def test_cooldown_time_override_zero_is_not_falsy(self):
-        """cooldown_time_override=0 must be honored; it must not fall through to the router-level value."""
-        store = MagicMock()
-        store.increment_cache.return_value = 1
-        ctx = self._make_router(cooldown_time=60.0, cache=store)
-        exc = litellm.RateLimitError("429", "openai", "gpt-4")
-
-        should_cooldown_based_on_allowed_fails_policy(
-            ctx=ctx,
-            deployment="dep-1",
-            original_exception=exc,
-            allowed_fails_override=5,
-            cooldown_time_override=0.0,
-        )
-
-        increment_call = store.increment_cache.call_args
-        assert increment_call is not None
-        assert increment_call[1]["ttl"] == 0.0, (
-            "cooldown_time_override=0 should be used as TTL, not the router-level 60.0"
+        store: Final = MagicMock(spec=RouterCooldowns)
+        store.increment_allowed_fails.return_value = 1
+        ctx: Final = self._make_router(cooldown_time=60.0, store=store)
+        exc: Final = litellm.RateLimitError("429", "openai", "gpt-4")
+        assert not should_cooldown_based_on_allowed_fails_policy(
+            ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5, cooldown_time_override=0.0
         )
+        store.increment_allowed_fails.assert_called_once_with("dep-1", None, 0.0)
 
     def test_fail_counter_is_shared_across_router_instances(self):
-        """Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails."""
-        shared_cache = DualCache(in_memory_cache=InMemoryCache())
-        workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache))
-        exc = litellm.AuthenticationError("401", "openai", "gpt-4")
-
-        results = [
+        redis: Final = MagicMock(spec=RedisCache)
+        redis.increment_cache.side_effect = itertools.count(1)
+        workers: Final = tuple(self._make_router(store=RouterCooldownStore(redis)) for _ in range(2))
+        exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
+        results: Final = [
             should_cooldown_based_on_allowed_fails_policy(
-                ctx=workers[i % 2],
-                deployment="dep-1",
-                original_exception=exc,
-                allowed_fails_override=5,
+                ctx=workers[i % 2], deployment="dep-1", original_exception=exc, allowed_fails_override=5
             )
             for i in range(6)
         ]
-
         assert results == [False, False, False, False, False, True]
-        assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6
+        assert redis.increment_cache.call_count == 6
 
     def test_fleet_wide_count_from_redis_decides_cooldown(self):
-        """The Redis (fleet-wide) count decides, even when this process has only seen one failure."""
-        redis_cache = MagicMock()
-        redis_cache.increment_cache.return_value = 6
-        router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
-        exc = litellm.AuthenticationError("401", "openai", "gpt-4")
-
-        result = should_cooldown_based_on_allowed_fails_policy(
-            ctx=router,
-            deployment="dep-1",
-            original_exception=exc,
-            allowed_fails_override=5,
+        redis: Final = MagicMock(spec=RedisCache)
+        redis.increment_cache.return_value = 6
+        ctx: Final = self._make_router(store=RouterCooldownStore(redis))
+        exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
+        assert should_cooldown_based_on_allowed_fails_policy(
+            ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5
         )
-
-        assert result is True
-        redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
+        redis.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
 
     def test_redis_outage_falls_back_to_this_workers_count(self):
-        """When every Redis increment fails, the worker's own in-memory count must still cool the deployment down."""
-        redis_cache = MagicMock()
-        redis_cache.increment_cache.side_effect = ConnectionError("redis down")
-        router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
-        exc = litellm.AuthenticationError("401", "openai", "gpt-4")
-
-        results = [
+        redis: Final = MagicMock(spec=RedisCache)
+        redis.increment_cache.side_effect = ConnectionError("redis down")
+        ctx: Final = self._make_router(store=RouterCooldownStore(redis))
+        exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
+        assert [
             should_cooldown_based_on_allowed_fails_policy(
-                ctx=router,
-                deployment="dep-1",
-                original_exception=exc,
-                allowed_fails_override=5,
+                ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5
             )
             for _ in range(6)
-        ]
-
-        assert results == [False, False, False, False, False, True]
-        assert redis_cache.increment_cache.call_count == 6
+        ] == [False, False, False, False, False, True]
+        assert redis.increment_cache.call_count == 6
 
 
 class TestRoutingGroupCooldownAlternatives:
@@ -591,31 +565,31 @@ class TestIncrementAllowedFailsServiceTarget:
         """The allowed_fails INCR is cooldown bookkeeping, so its service span must read
         ``redis.incr router_cooldowns`` rather than a bare ``redis.incr``."""
         seen: list[str | None] = []
-        cache = MagicMock(spec=DualCache)
+        cache = MagicMock(spec=RedisCache)
 
-        def _increment(**_kwargs):
+        def _increment(*_args, **_kwargs):
             seen.append(current_service_target())
             return 2
 
         cache.increment_cache.side_effect = _increment
 
-        assert _increment_allowed_fails(cache, "deployment:dep-1:fails", ttl=60.0) == 2
+        assert RouterCooldownStore(cache).increment_allowed_fails("dep-1", None, 60.0) == 2
         assert seen == ["router_cooldowns"]
         assert current_service_target() is None
 
     def test_in_memory_fallback_reads_under_the_same_target(self):
         seen: list[str | None] = []
-        cache = MagicMock(spec=DualCache)
-        cache.increment_cache.side_effect = ConnectionError("redis down")
+        cache = MagicMock(spec=RedisCache)
 
-        def _get(**_kwargs):
+        def _increment(*_args, **_kwargs):
             seen.append(current_service_target())
-            return 4
+            raise ConnectionError("redis down")
 
-        cache.get_cache.side_effect = _get
+        cache.increment_cache.side_effect = _increment
+        store = RouterCooldownStore(cache)
 
-        assert _increment_allowed_fails(cache, "deployment:dep-1:fails", ttl=60.0) == 4
-        assert seen == ["router_cooldowns"]
+        assert [store.increment_allowed_fails("dep-1", None, 60.0) for _ in range(4)] == [1, 2, 3, 4]
+        assert seen == ["router_cooldowns"] * 4
 
 
 @pytest.fixture()
@@ -860,136 +834,57 @@ class TestExceptionTypeCountersTrackedIndependently:
         rate_limit_exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4")
         ise = litellm.InternalServerError("Internal error", "openai", "gpt-4")
 
-        for _ in range(3):
+        assert [
             should_cooldown_based_on_allowed_fails_policy(
                 ctx=RoutingContext(router),
                 deployment="primary",
                 original_exception=rate_limit_exc,
-                allowed_fails_override=5,
+                allowed_fails_override=3,
                 cache_key_suffix="RateLimitError",
             )
-
-        rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
-        generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
-
-        assert rl_counter == 3, "RateLimitError counter should be 3"
-        assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments"
-
-        should_cooldown_based_on_allowed_fails_policy(
+            for _ in range(3)
+        ] == [False, False, False]
+        assert not should_cooldown_based_on_allowed_fails_policy(
             ctx=RoutingContext(router),
             deployment="primary",
             original_exception=ise,
-            allowed_fails_override=5,
+            allowed_fails_override=1,
+            cache_key_suffix="generic",
+        )
+        assert not should_cooldown_based_on_allowed_fails_policy(
+            ctx=RoutingContext(router),
+            deployment="primary",
+            original_exception=rate_limit_exc,
+            allowed_fails_override=4,
+            cache_key_suffix="RateLimitError",
+        )
+        assert should_cooldown_based_on_allowed_fails_policy(
+            ctx=RoutingContext(router),
+            deployment="primary",
+            original_exception=ise,
+            allowed_fails_override=1,
             cache_key_suffix="generic",
         )
-
-        generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
-        rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
-
-        assert generic_counter_after == 1, "generic counter should now be 1"
-        assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError"
 
 
 @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
 class TestCooldownCacheTTLCorrection:
-    def _make_cooldown_cache(self) -> CooldownCache:
-        in_memory = InMemoryCache()
-        dual_cache = DualCache(in_memory_cache=in_memory)
-        return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
-
-    def test_expired_entry_evicted_and_not_returned(self):
-        """
-        An entry with timestamp+cooldown_time in the past must be evicted from
-        in-memory cache and excluded from the active cooldown list.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "expired-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        expired_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - 120.0,
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
-
-        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert active == [], "Expired cooldown entry must not appear in active cooldowns"
-        assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
-
-    def test_active_entry_is_returned(self):
-        """
-        An entry whose cooldown window has not elapsed must appear in the active list.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "active-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        active_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time(),
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, active_value, ttl=60)
-
-        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert len(active) == 1
-        assert active[0][0] == model_id
-
-    def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self):
-        """
-        When DualCache backfills from Redis using the default 600s TTL, the in-memory
-        TTL must be corrected to min(remaining, 60) seconds.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "backfilled-deployment"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        remaining = 30.0
-        value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - (60.0 - remaining),
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, value, ttl=600)
-
-        before_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert before_expiry is not None
-
-        cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
-        assert after_expiry is not None
-        corrected_remaining = after_expiry - time.time()
-        assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
-        assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)"
-
+    @pytest.mark.parametrize(("timestamp", "active"), ((880.0, False), (1000.0, True)))
+    @pytest.mark.parametrize("asynchronous", (False, True))
     @pytest.mark.asyncio
-    async def test_async_expired_entry_evicted(self):
-        """
-        Async path must also evict expired entries.
-        """
-        cc = self._make_cooldown_cache()
-        model_id = "async-expired"
-        key = CooldownCache.get_cooldown_cache_key(model_id)
-
-        expired_value: CooldownCacheValue = {
-            "exception_received": "Rate limit",
-            "status_code": "429",
-            "timestamp": time.time() - 120.0,
-            "cooldown_time": 60.0,
-        }
-        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
-
-        active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
-
-        assert active == [], "Expired entry must not appear in async active cooldowns"
-        assert cc.in_memory_cache.get_cache(key) is None
+    async def test_only_unexpired_cooldowns_are_returned(self, timestamp: float, active: bool, asynchronous: bool):
+        store: Final = RouterCooldownStore(None, clock=lambda: 1000.0)
+        cc: Final = CooldownCache(store, default_cooldown_time=60.0)
+        value: Final = CooldownCacheValue(
+            exception_received="Rate limit", status_code="429", timestamp=timestamp, cooldown_time=60.0
+        )
+        store.put("deployment", value, 600)
+        result: Final = (
+            await cc.async_get_active_cooldowns(("deployment",), None)
+            if asynchronous
+            else cc.get_active_cooldowns(("deployment",), None)
+        )
+        assert result == ([("deployment", value)] if active else [])
 
 
 @pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")

tests/unit/router_utils/test_health_check_allowed_fails_integration.py

@@ -6,6 +6,7 @@ should increment the same counters used by allowed_fails_policy, using the
 actual exception type from the health check error.
 """
 
+from typing import Final
 from unittest.mock import ANY, patch
 
 import pytest
@@ -146,30 +147,26 @@ class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions:
 class TestHealthCheckCooldownIntegration:
     """Test that health check failures trigger cooldown via set_cooldown_deployments."""
 
-    def test_health_check_failure_increments_failed_calls(self):
-        """Health check failure should increment the failed_calls counter."""
+    def test_health_check_failure_increments_failed_calls(self) -> None:
         from litellm.router_utils.cooldown_handlers import (
             should_cooldown_based_on_allowed_fails_policy,
         )
 
-        router = Router(
+        router: Final = Router(
             model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
             allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3),
         )
 
-        timeout_exc = litellm.Timeout(message="Health check timeout", model="gpt-4", llm_provider="openai")
-
-        # First call: should not cooldown (1 <= 3)
-        result = should_cooldown_based_on_allowed_fails_policy(
-            ctx=RoutingContext(router),
-            deployment="deploy-1",
-            original_exception=timeout_exc,
+        timeout_exc: Final = litellm.Timeout(message="Health check timeout", model="gpt-4", llm_provider="openai")
+        results: Final = tuple(
+            should_cooldown_based_on_allowed_fails_policy(
+                ctx=RoutingContext(router),
+                deployment="deploy-1",
+                original_exception=timeout_exc,
+            )
+            for _ in range(4)
         )
-        assert result is False
-
-        # Check counter was incremented
-        current_fails = router.health().store.get_cache(key="deployment:deploy-1:allowed_fails")
-        assert current_fails == 1
+        assert results == (False, False, False, True)
 
     def test_health_check_failure_triggers_cooldown_at_threshold(self):
         """After exceeding allowed_fails threshold, deployment should enter cooldown."""

tests/unit/router_utils/test_routing_read_batch.py

@@ -64,9 +64,7 @@ async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_rou
     redis = _redis_answering({})
     router = _router(redis, "usage-based-routing-v2")
 
-    deployment = await router.async_get_available_deployment(
-        model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
-    )
+    deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
 
     assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
     assert _redis_key_families(redis) == [
@@ -162,9 +160,7 @@ async def test_batched_read_still_excludes_a_cooled_down_deployment():
     )
     router = _router(redis, "usage-based-routing-v2")
 
-    deployment = await router.async_get_available_deployment(
-        model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
-    )
+    deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
 
     assert deployment["model_info"]["id"] == "dep-a", "dep-b has the lowest tpm but is cooling down"
     assert redis.async_batch_get_cache.await_count == 1
@@ -181,9 +177,7 @@ async def test_batched_read_ignores_an_expired_cooldown():
     )
     router = _router(redis, "usage-based-routing-v2")
 
-    deployment = await router.async_get_available_deployment(
-        model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
-    )
+    deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
 
     assert deployment["model_info"]["id"] == "dep-b"
 
@@ -204,7 +198,11 @@ async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did():
 
     assert str(routed_error.value) == str(alone_error.value)
     assert len(routed.cache.last_redis_batch_access_time) == 0, "a failed read must not throttle the next one"
-    assert len(routed.cooldown_cache.cooldown_store.last_redis_batch_access_time) == 0
+    before: Final = redis.async_batch_get_cache.await_count
+    redis.async_batch_get_cache.side_effect = None
+    redis.async_batch_get_cache.return_value = {}
+    assert await routed.cooldown_cache.async_get_active_cooldowns(["dep-a"], None) == []
+    assert redis.async_batch_get_cache.await_count == before + 1
 
 
 @pytest.mark.asyncio
@@ -213,8 +211,6 @@ async def test_a_failed_batched_read_leaves_simple_shuffle_routing():
     redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
     router = _router(redis, "simple-shuffle")
 
-    deployment = await router.async_get_available_deployment(
-        model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
-    )
+    deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
 
     assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}

tests/unit/test_internal_context.py

@@ -32,6 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset(
         "litellm/llms/anthropic/prompt_cache_prediction.py",
         "litellm/proxy/common_utils/cache_coordinator.py",
         "litellm/proxy/common_utils/user_api_key_cache.py",
+        "litellm/redis/stores/_routing_prefetch.py",
         "litellm/utils.py",
     }
 )

tests/unit/test_router/test_router.py

@@ -110,6 +110,7 @@ import inspect
 from typing import List, Optional
 
 if TYPE_CHECKING:
+    from tests.unit.conftest import RedisWriteGate, RouterRedis
     from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
     from litellm.integrations.otel.logger import OpenTelemetryV2
@@ -1410,17 +1411,6 @@ _BATCH_NOT_FOUND = {
 }
 
 
-async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]:
-    loop = asyncio.get_event_loop()
-    deadline = loop.time() + timeout
-    while loop.time() < deadline:
-        keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:"))
-        if keys:
-            return keys
-        await asyncio.sleep(0.05)
-    return []
-
-
 @pytest.mark.asyncio
 async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch):
     """
@@ -1467,11 +1457,13 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m
         )
         response = await router.aretrieve_batch(batch_id=_BATCH_ID)
         payload = await collector.retrieve_batch_payload()
-        usage_keys = await _router_usage_keys(router)
+        requested_usage: Final = await router.get_model_group_usage(_BATCH_GROUP)
+        unrelated_usage: Final = await router.get_model_group_usage(_UNRELATED_BATCH_GROUP)
 
     assert response.id == _BATCH_ID
     assert payload["model_group"] == _BATCH_GROUP
-    assert usage_keys == []
+    assert requested_usage == (None, None)
+    assert unrelated_usage == (None, None)
 
 
 @pytest.mark.parametrize(
@@ -24307,11 +24299,9 @@ def test_router_keeps_wif_secret_pointers_unresolved(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_failure_rpm_increment_declares_the_router_usage_key_family():
+async def test_failure_rpm_increment_declares_the_router_usage_key_family(router_redis: "RouterRedis"):
     """The RPM bump a failed call still earns is router usage bookkeeping, so its Redis span
     reads ``redis.incr router_usage`` rather than a bare ``redis.incr``."""
-    from unittest.mock import AsyncMock
-
     from litellm._internal_context import current_service_target
 
     router = Router(
@@ -24323,43 +24313,27 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family():
             }
         ]
     )
-    seen: list[str | None] = []
-
-    async def _increment(**_kwargs):
-        seen.append(current_service_target())
-
-    with patch.object(router.cache, "async_increment_cache", new=AsyncMock(side_effect=_increment)):
-        await router.async_deployment_callback_on_failure(
-            kwargs={
-                "call_type": "acompletion",
-                "litellm_params": {
-                    "metadata": {"deployment": "openai/gpt-4o", "model_group": "gpt-group"},
-                    "model_info": {"id": "dep-1"},
-                },
+    redis, commands = router_redis
+    router.update_redis_cache(redis)
+    await router.async_deployment_callback_on_failure(
+        kwargs={
+            "call_type": "acompletion",
+            "litellm_params": {
+                "metadata": {"deployment": "openai/gpt-4o", "model_group": "gpt-group"},
+                "model_info": {"id": "dep-1"},
             },
-            completion_response=None,
-            start_time=None,
-            end_time=None,
-        )
+        },
+        completion_response=None,
+        start_time=None,
+        end_time=None,
+    )
 
-    assert seen == ["router_usage"]
+    observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
+    assert tuple(
+        target for args, target, _ in observed if isinstance(args[0], bytes) and b"INCRBYFLOAT" in args[0]
+    ) == ("router_usage",)
     assert current_service_target() is None
 
-class _SpanRecordingInMemoryCache(InMemoryCache):
-    """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.active_span_names: list[str] = []
-
-    async def async_batch_get_cache(self, keys, **kwargs):
-        self.active_span_names.append(trace.get_current_span().name)
-        return await super().async_batch_get_cache(keys, **kwargs)
-
-    async def async_get_cache(self, key, **kwargs):
-        self.active_span_names.append(trace.get_current_span().name)
-        return await super().async_get_cache(key, **kwargs)
-
 
 @pytest.fixture
 def v2_span_exporter(monkeypatch):
@@ -24378,7 +24352,9 @@ def v2_span_exporter(monkeypatch):
 
 
 @pytest.mark.asyncio
-async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group(v2_span_exporter):
+async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_model_group(
+    v2_span_exporter, router_redis: "RouterRedis"
+):
     """Picking a deployment opens ``route {model_group}`` (the requested group, not the deployment
     it picks) under the server span, and the cooldown reads it issues run inside it, so their Redis
     spans nest there instead of lying flat under the request."""
@@ -24398,9 +24374,8 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
             },
         ]
     )
-    recording_cache = _SpanRecordingInMemoryCache()
-    router.cache.in_memory_cache = recording_cache
-    router.cooldown_cache.cooldown_store.in_memory_cache = recording_cache
+    redis, commands = router_redis
+    router.update_redis_cache(redis)
 
     with TracerProvider().get_tracer("test").start_as_current_span("POST /v1/chat/completions") as server_span:
         deployment = await router.async_get_available_deployment(model="gpt-group", request_kwargs={})
@@ -24410,7 +24385,8 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
     assert route_span.name == "route gpt-group"
     assert route_span.parent is not None and route_span.parent.span_id == server_span.get_span_context().span_id
     assert route_span.end_time is not None
-    assert recording_cache.active_span_names and set(recording_cache.active_span_names) == {"route gpt-group"}
+    observed: Final = tuple(commands.get_nowait() for _ in range(commands.qsize()))
+    assert tuple(span for args, _, span in observed if args[0] == "MGET") == ("route gpt-group",)
 
 
 def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]:
@@ -25165,7 +25141,7 @@ async def test_router_get_model_group_usage_wildcard_routes():
 
 @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
 @pytest.mark.asyncio
-async def test_call_router_callbacks_on_success():
+async def test_call_router_callbacks_on_success(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"):
     router = Router(
         model_list=[
             {
@@ -25176,30 +25152,26 @@ async def test_call_router_callbacks_on_success():
         ]
     )
 
-    with patch.object(router.cache, "async_increment_cache_pipeline", new=AsyncMock()) as mock_callback:
-        await router.acompletion(
-            model="gemini/gemini-2.5-flash",
-            messages=[{"role": "user", "content": "Hello, how are you?"}],
-            mock_response="Hello, I'm good.",
-        )
-        await asyncio.sleep(1)
-        assert mock_callback.call_count == 1
+    redis, _ = router_redis
+    started, _, _ = redis_write_gate
+    router.update_redis_cache(redis)
+    response: Final = await router.acompletion(
+        model="gemini/gemini-2.5-flash",
+        messages=[{"role": "user", "content": "Hello, how are you?"}],
+        mock_response="Hello, I'm good.",
+    )
+    await asyncio.wait_for(started.wait(), timeout=5)
+    client: Final = redis.init_async_client()
+    keys: Final = tuple(await client.keys("global_router:*"))
+    values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys}
+    assert len(keys) == 2
+    assert values == {b"tpm": response.usage.total_tokens, b"rpm": 1}
 
-        increment_list = mock_callback.call_args_list[0].kwargs["increment_list"]
-        assert len(increment_list) == 2
-
-        for increment in increment_list:
-            if "tpm" in increment["key"]:
-                assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:tpm")
-                assert increment["increment_value"] == 30
-            elif "rpm" in increment["key"]:
-                assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm")
-                assert increment["increment_value"] == 1
 
 @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
 @pytest.mark.serial
 @pytest.mark.asyncio
-async def test_call_router_callbacks_on_failure():
+async def test_call_router_callbacks_on_failure(router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"):
     router = Router(
         model_list=[
             {
@@ -25210,19 +25182,23 @@ async def test_call_router_callbacks_on_failure():
         ]
     )
 
-    with patch.object(router.cache, "async_increment_cache", new=AsyncMock()) as mock_callback:
-        with pytest.raises(litellm.RateLimitError):
-            await router.acompletion(
-                model="gemini/gemini-2.5-flash",
-                messages=[{"role": "user", "content": "Hello, how are you?"}],
-                mock_response="litellm.RateLimitError",
-                num_retries=0,
-            )
-        await asyncio.sleep(3)
-        print(mock_callback.call_args_list)
-        assert mock_callback.call_count == 1
+    redis, _ = router_redis
+    started, _, _ = redis_write_gate
+    router.update_redis_cache(redis)
+    with pytest.raises(litellm.RateLimitError):
+        await router.acompletion(
+            model="gemini/gemini-2.5-flash",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            mock_response="litellm.RateLimitError",
+            num_retries=0,
+        )
+    await asyncio.wait_for(started.wait(), timeout=5)
+    client: Final = redis.init_async_client()
+    keys: Final = tuple(await client.keys("global_router:*"))
+    values: Final = {key.rsplit(b":", 2)[-2]: float(await client.get(key)) for key in keys}
+    assert len(keys) == 1
+    assert values == {b"rpm": 1}
 
-        assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm")
 
 @pytest.mark.asyncio
 async def test_rate_limit_error_callback():

tests/unit/test_router/test_router_helper_utils.py

@@ -7,7 +7,7 @@ import weakref
 from collections.abc import Awaitable, Callable, Iterator, Mapping
 from datetime import datetime, timezone
 from types import MethodType
-from typing import Final
+from typing import TYPE_CHECKING, Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
@@ -35,6 +35,10 @@ from litellm.types.utils import (
 )
 
 
+if TYPE_CHECKING:
+    from tests.unit.conftest import RouterRedis
+
+
 def create_standard_logging_payload() -> StandardLoggingPayload:
     return StandardLoggingPayload(
         id="test_id",
@@ -497,7 +501,7 @@ async def test_deployment_callback_on_success(sync_mode):
 
 
 @pytest.mark.asyncio
-async def test_deployment_callback_on_success_counts_a_responses_websocket_turn():
+async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(router_redis: "RouterRedis"):
     import time
 
     model_list = [
@@ -508,6 +512,8 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
         }
     ]
     router = Router(model_list=model_list)
+    redis, _ = router_redis
+    router.update_redis_cache(redis)
     standard_logging_payload = create_standard_logging_payload()
     standard_logging_payload["total_tokens"] = 17
 
@@ -521,14 +527,15 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
             "standard_logging_object": standard_logging_payload,
         }
 
-    tpm_key = await router.deployment_callback_on_success(
+    counter = await router.deployment_callback_on_success(
         kwargs=turn_kwargs("_aresponses_websocket"),
         completion_response=None,
         start_time=time.time(),
         end_time=time.time(),
     )
-    assert tpm_key is not None
-    assert await router.cache.async_get_cache(key=tpm_key) == 17
+    assert counter is not None
+    tpm_key = f"global_router:ws-deployment:None:tpm:{counter.minute}"
+    assert await redis.async_get_cache(key=tpm_key) == 17
 
     realtime_key = await router.deployment_callback_on_success(
         kwargs=turn_kwargs("_arealtime"),
@@ -537,7 +544,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
         end_time=time.time(),
     )
     assert realtime_key is None
-    assert await router.cache.async_get_cache(key=tpm_key) == 17
+    assert await redis.async_get_cache(key=tpm_key) == 17
 
 
 async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]:
@@ -1120,7 +1127,7 @@ async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
 
     # The IO deployment is no longer skipped: its TPM routing counter is tracked.
     assert tpm_key is not None
-    assert await router.cache.async_get_cache(key=tpm_key) == 100
+    assert (await router.get_model_group_usage("opus"))[0] == 100
 
 
 @pytest.mark.asyncio
@@ -1778,8 +1785,8 @@ async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
     store: dict = {}
     worker_a = _rpm_tpm_router("lit-3058-workers")
     worker_b = _rpm_tpm_router("lit-3058-workers")
-    worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
-    worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
+    worker_a.update_redis_cache(_shared_redis_stub(store))
+    worker_b.update_redis_cache(_shared_redis_stub(store))
 
     messages = [{"role": "user", "content": "hi"}]
     tokens_on_a = 0

tests/unit/test_router/test_router_provider_endpoints.py

@@ -5,28 +5,28 @@ import json
 import uuid
 from collections.abc import Callable, Iterator
 from datetime import datetime, timezone
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import httpx
 import litellm
 import pytest
 import respx
 from litellm import CustomLogger, Router
-from litellm.caching.dual_cache import DualCache
-from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.router_utils.router_callbacks.track_deployment_metrics import (
     get_deployment_successes_for_current_minute,
 )
-from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.llms.openai import HttpxBinaryResponseContent
 from litellm.types.router import RoutingStrategy
 from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse
 from litellm.types.videos.utils import encode_video_id_with_provider
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 
+if TYPE_CHECKING:
+    from tests.unit.conftest import RedisWriteGate, RouterRedis
+
 EVENT_TIMEOUT_SECONDS: Final = 5
 ROUTING_SELECTIONS: Final = 40
 ROUTING_MESSAGES: Final = ({"role": "user", "content": "route this request"},)
@@ -534,86 +534,30 @@ def test_router_track_deployment_metrics_updates_observable_usage() -> None:
     assert router.cache.get_cache(key="metrics-deployment", local_only=True) == 2
 
 
-class _GatedIncrementCache(DualCache):
-    def __init__(self) -> None:
-        super().__init__(in_memory_cache=InMemoryCache())
-        self.first_increment_started = asyncio.Event()
-        self.release_first_increment = asyncio.Event()
-        self.deployment_success_incremented = asyncio.Event()
-        self.increment_calls = 0
-
-    def increment_cache(self, key: str, value: int, local_only: bool = False, **kwargs: object) -> int:
-        result: Final = super().increment_cache(key=key, value=value, local_only=local_only, **kwargs)
-        if key.endswith(":successes"):
-            self.deployment_success_incremented.set()
-        return result
-
-    async def async_increment_cache_pipeline(
-        self,
-        increment_list: list[RedisPipelineIncrementOperation],
-        local_only: bool = False,
-        parent_otel_span: object = None,
-        **kwargs: object,
-    ) -> list[float] | None:
-        self.increment_calls += 1
-        if self.increment_calls == 1:
-            self.first_increment_started.set()
-            await self.release_first_increment.wait()
-        return await super().async_increment_cache_pipeline(
-            increment_list=increment_list,
-            local_only=local_only,
-            parent_otel_span=parent_otel_span,
-            **kwargs,
-        )
-
-
-class _UnavailableIncrementCache(DualCache):
-    def __init__(self) -> None:
-        super().__init__(in_memory_cache=InMemoryCache())
-        self.first_increment_started = asyncio.Event()
-        self.release_first_increment = asyncio.Event()
-        self.deployment_success_incremented = asyncio.Event()
-        self.increment_calls = 0
-
-    def increment_cache(self, key: str, value: int, local_only: bool = False, **kwargs: object) -> int:
-        result: Final = super().increment_cache(key=key, value=value, local_only=local_only, **kwargs)
-        if key.endswith(":successes"):
-            self.deployment_success_incremented.set()
-        return result
-
-    async def async_increment_cache_pipeline(
-        self,
-        increment_list: list[RedisPipelineIncrementOperation],
-        local_only: bool = False,
-        parent_otel_span: object = None,
-        **kwargs: object,
-    ) -> list[float] | None:
-        self.increment_calls += 1
-        if self.increment_calls == 1:
-            self.first_increment_started.set()
-            await self.release_first_increment.wait()
-        raise RuntimeError("cache unavailable")
-
-
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("router_minute_pinned")
-async def test_router_success_callback_during_pre_header_increment_does_not_double_count() -> None:
+async def test_router_success_callback_during_pre_header_increment_does_not_double_count(
+    router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"
+) -> None:
+    capture: Final = _RouterLoggingCapture("lit-3058-race")
     router: Final = _rpm_tpm_router("lit-3058-race")
-    cache: Final = _GatedIncrementCache()
-    router.cache = cache
+    redis, _ = router_redis
+    started, released, _ = redis_write_gate
+    router.update_redis_cache(redis)
+    released.clear()
     request: Final = asyncio.create_task(
         router.acompletion(
             model="gpt-5-mini",
             messages=[{"role": "user", "content": "hi"}],
             mock_response="pong",
+            success_callback=[capture.async_log_success_event],
         )
     )
 
-    await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
-    await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS)
+    await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
+    await capture.next_event()
     assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1
-    assert cache.increment_calls == 1
-    cache.release_first_increment.set()
+    released.set()
     response: Final = await request
 
     assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
@@ -621,10 +565,16 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub
 
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("router_minute_pinned")
-async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp() -> None:
+async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting(
+    router_redis: "RouterRedis", redis_write_gate: "RedisWriteGate"
+) -> None:
+    capture: Final = _RouterLoggingCapture("lit-3058-fail")
     router: Final = _rpm_tpm_router("lit-3058-fail")
-    cache: Final = _UnavailableIncrementCache()
-    router.cache = cache
+    redis, _ = router_redis
+    started, released, unavailable = redis_write_gate
+    router.update_redis_cache(redis)
+    released.clear()
+    unavailable.set()
     metadata: Final[dict[str, object]] = {}
     request: Final = asyncio.create_task(
         router.acompletion(
@@ -632,21 +582,21 @@ async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp()
             messages=[{"role": "user", "content": "hi"}],
             mock_response="pong",
             metadata=metadata,
+            success_callback=[capture.async_log_success_event],
         )
     )
 
-    await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
-    await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS)
+    await asyncio.wait_for(started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
+    await capture.next_event()
     assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
     assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
-    assert cache.increment_calls == 1
-    cache.release_first_increment.set()
+    released.set()
     response: Final = await request
 
     assert response.usage.total_tokens == 30
-    assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata
-    assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100
-    assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
+    assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == response.usage.total_tokens
+    assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 99
+    assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
 
 
 ASSISTANT_RESPONSE: Final = {