From 756681907b266d1f257a44f85d6edd701e3072eb Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 11 Aug 2026 16:28:42 +0000 Subject: [PATCH] fix(router): ignore expired cooldown payloads in get_min_cooldown get_active_cooldowns and its async twin already drop a payload whose timestamp plus cooldown_time has elapsed, but get_min_cooldown still counted one, so a stale short cooldown shrank the retry-after time reported to callers. It now goes through the same _corrected_active_cooldown helper, which also evicts the dead in-memory entry, and it builds its keys with get_cooldown_cache_key instead of a duplicated f-string. Carries over the remaining piece of #34508 by InvisibleMan1306, whose stale-cooldown read fix otherwise landed upstream in #34416. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_utils/cooldown_cache.py | 19 +++++------ .../router_utils/test_cooldown_cache.py | 32 +++++++++++++++++++ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 9e7f457f631..c2fb1be8833 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -181,20 +181,21 @@ class CooldownCache: """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] + 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.cache.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"] + current_time: Final = time.time() + active_cooldown_times: Final = tuple( + cooldown_cache_value["cooldown_time"] + for key, result in zip(keys, results) + if result + and isinstance(result, dict) + and (cooldown_cache_value := self._corrected_active_cooldown(key, result, current_time)) is not None + ) - return min_cooldown_time or self.default_cooldown_time + return min(active_cooldown_times, default=0.0) or self.default_cooldown_time # Usage example: diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index a48402684b4..e01021ddd84 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -373,6 +373,38 @@ class TestCooldownCacheTTLCorrection: assert len(active) == 1 assert active[0][0] == model_id + def test_get_min_cooldown_ignores_expired_entry(self): + """ + get_min_cooldown must skip a payload whose cooldown window has elapsed, otherwise a + stale short cooldown keeps shrinking the retry-after time reported to callers. + """ + cc = self._make_cooldown_cache() + expired_model_id = "min-cooldown-expired" + active_model_id = "min-cooldown-active" + + cc.cache.in_memory_cache.set_cache( + CooldownCache.get_cooldown_cache_key(expired_model_id), + { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - 120.0, + "cooldown_time": 10.0, + }, + ttl=600, + ) + cc.cache.in_memory_cache.set_cache( + CooldownCache.get_cooldown_cache_key(active_model_id), + { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time(), + "cooldown_time": 120.0, + }, + ttl=120, + ) + + assert cc.get_min_cooldown(model_ids=[expired_model_id, active_model_id], parent_otel_span=None) == 120.0 + class TestCorrectedActiveCooldown: def _make_cooldown_cache(self) -> CooldownCache: