phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(cooldown_cache.py): fix linting errors

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-24 17:11:32 -07:00
parent 33972cc79c
commit 5572ad7241
1 changed files with 11 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

36
litellm/router_utils/cooldown_cache.py
View file

@ -24,7 +24,7 @@ class CooldownCache:
def _common_add_cooldown_logic( def _common_add_cooldown_logic(
self, model_id: str, original_exception, exception_status, cooldown_time: float self, model_id: str, original_exception, exception_status, cooldown_time: float
) -> Tuple[str, dict]: ) -> Tuple[str, CooldownCacheValue]:
try: try:
current_time = time.time() current_time = time.time()
cooldown_key = f"deployment:{model_id}:cooldown" cooldown_key = f"deployment:{model_id}:cooldown"
@ -76,27 +76,9 @@ class CooldownCache:
) )
raise e raise e
async def async_add_deployment_to_cooldown(
self,
model_id: str,
original_exception: Exception,
exception_status: int,
cooldown_time: Optional[float],
):
cooldown_key, cooldown_data = self._common_add_cooldown_logic(
model_id=model_id, original_exception=original_exception
)
# Set the cache with a TTL equal to the cooldown time
self.cache.set_cache(
value=cooldown_data,
key=cooldown_key,
ttl=cooldown_time or self.default_cooldown_time,
)
async def async_get_active_cooldowns( async def async_get_active_cooldowns(
self, model_ids: List[str] self, model_ids: List[str]
) -> List[Tuple[str, dict]]: ) -> List[Tuple[str, CooldownCacheValue]]:
# Generate the keys for the deployments # Generate the keys for the deployments
keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
@ -106,12 +88,15 @@ class CooldownCache:
active_cooldowns = [] active_cooldowns = []
# Process the results # Process the results
for model_id, result in zip(model_ids, results): for model_id, result in zip(model_ids, results):
if result: if result and isinstance(result, dict):
active_cooldowns.append((model_id, result)) cooldown_cache_value = CooldownCacheValue(**result)
active_cooldowns.append((model_id, cooldown_cache_value))
return active_cooldowns return active_cooldowns
def get_active_cooldowns(self, model_ids: List[str]) -> List[Tuple[str, dict]]: def get_active_cooldowns(
self, model_ids: List[str]
) -> List[Tuple[str, CooldownCacheValue]]:
# Generate the keys for the deployments # Generate the keys for the deployments
keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
@ -121,8 +106,9 @@ class CooldownCache:
active_cooldowns = [] active_cooldowns = []
# Process the results # Process the results
for model_id, result in zip(model_ids, results): for model_id, result in zip(model_ids, results):
if result: if result and isinstance(result, dict):
active_cooldowns.append((model_id, result)) cooldown_cache_value = CooldownCacheValue(**result)
active_cooldowns.append((model_id, cooldown_cache_value))
return active_cooldowns return active_cooldowns