scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix: cache cooldown key (#16954)

Browse source 
There's no need to generate a key multiple times for the same model, cache it with a max limit.
This commit is contained in:
Alexsander Hamir 2025-11-21 19:11:17 -08:00 • committed by GitHub
parent 6881594632
commit f542011076
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 5 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/router_utils/cooldown_cache.py
View file

@ -2,6 +2,7 @@
Wrapper around router cache. Meant to handle model cooldown logic
"""
import functools
import time
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
@ -44,7 +45,7 @@ class CooldownCache:
) -> Tuple[str, CooldownCacheValue]:
try:
current_time = time.time()
cooldown_key = f"deployment:{model_id}:cooldown"
cooldown_key = CooldownCache.get_cooldown_cache_key(model_id)
# Store the cooldown information for the deployment separately
cooldown_data = CooldownCacheValue(
@ -104,8 +105,9 @@ class CooldownCache:
raise e
@staticmethod
@functools.lru_cache(maxsize=1024)
def get_cooldown_cache_key(model_id: str) -> str:
return f"deployment:{model_id}:cooldown"
return "deployment:" + model_id + ":cooldown"
async def async_get_active_cooldowns(
self, model_ids: List[str], parent_otel_span: Optional[Span]
@ -140,7 +142,7 @@ class CooldownCache:
self, model_ids: List[str], parent_otel_span: Optional[Span]
) -> List[Tuple[str, CooldownCacheValue]]:
# Generate the keys for the deployments
keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
keys = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
# Retrieve the values for the keys using mget
results = (
self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span)