From 9a5337c503e58f5f8f1abdde2840a3a37287a85e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 14 Sep 2024 16:20:27 -0700 Subject: [PATCH] fix cooldown message time --- litellm/router_utils/cooldown_cache.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 2d29a1a5d3e..376aced36df 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -83,7 +83,7 @@ class CooldownCache: keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results = await self.cache.async_batch_get_cache(keys=keys) + results = await self.cache.async_batch_get_cache(keys=keys) or [] active_cooldowns = [] # Process the results @@ -101,7 +101,7 @@ class CooldownCache: keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results = self.cache.batch_get_cache(keys=keys) + results = self.cache.batch_get_cache(keys=keys) or [] active_cooldowns = [] # Process the results @@ -119,17 +119,19 @@ class CooldownCache: keys = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results = self.cache.batch_get_cache(keys=keys) + results = self.cache.batch_get_cache(keys=keys) or [] - min_cooldown_time = self.default_cooldown_time + min_cooldown_time: Optional[float] = None # Process the results for model_id, result in zip(model_ids, results): if result and isinstance(result, dict): cooldown_cache_value = CooldownCacheValue(**result) # type: ignore - if cooldown_cache_value["cooldown_time"] < min_cooldown_time: + if min_cooldown_time is None: + min_cooldown_time = cooldown_cache_value["cooldown_time"] + elif cooldown_cache_value["cooldown_time"] < min_cooldown_time: min_cooldown_time = cooldown_cache_value["cooldown_time"] - return min_cooldown_time + return min_cooldown_time or self.default_cooldown_time # Usage example: