mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
f769aa4675
Cooldown entries rode the router-wide DualCache, which re-reads a key that is missing from memory at most once every 10s. A deployment benched on one replica therefore kept taking traffic on its siblings for up to 10 seconds, and the same shared in-memory tier could evict a live cooldown once 200 unrelated router keys crowded it out, which sent even the benching replica back to the dead deployment. CooldownCache now owns a DualCache over the router's Redis with a 1s read interval and an in-memory tier that only holds cooldown keys. Redis is attached lazily because the router builds the cooldown cache before it wires Redis up.
223 lines
9.6 KiB
Python
223 lines
9.6 KiB
Python
"""
|
|
Wrapper around router cache. Meant to handle model cooldown logic
|
|
"""
|
|
|
|
import functools
|
|
import time
|
|
from collections.abc import Mapping
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
|
|
from typing_extensions import TypedDict
|
|
|
|
from litellm import verbose_logger
|
|
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
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
Span = _Span | Any
|
|
else:
|
|
Span = Any
|
|
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
)
|
|
# 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
|
|
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
|
|
|
|
async def async_get_active_cooldowns(
|
|
self, model_ids: list[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]
|
|
|
|
results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)
|
|
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_active_cooldowns(
|
|
self, model_ids: list[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
|
|
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
|
|
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()
|