From dffcdb55d4ee54e082676324d34b6086e32702b6 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay <34904987+lets-order-some-fries@users.noreply.github.com> Date: Tue, 15 Sep 2026 21:18:43 +0530 Subject: [PATCH] fix(router): let a longer cooldown extend an active one InMemoryCache.allow_ttl_override returns False while a key's TTL is still in the future, so set_cache keeps the expiry that is already recorded and silently drops the new one. Cooldowns go through that path, so the deadline a deployment gets is whichever failure happened to land first, not the longest one asked for. A deployment that took a 2s cooldown and then a 60s one on the next failure went back into rotation after 2 seconds, which is exactly when the provider is least likely to be ready for it. add_deployment_to_cooldown now deletes the in-memory entry before writing when the incoming cooldown ends later than the one on record. An equal or shorter cooldown leaves the entry alone, so a brief 500 cannot cut a long rate-limit bench short. Redis already overwrote its own TTL, so only the in-memory tier needed this. Verified with tests/test_litellm/router_utils/test_cooldown_cache.py, 26 passed. Removing the delete call makes test_longer_cooldown_extends_the_deadline fail on "the longer cooldown was dropped", so the regression is really covered Rebased onto litellm_internal_staging to pick up the separate cooldown store from #40025: the write goes through cooldown_store and the entry is read from the CooldownCache's own in_memory_cache rather than the router cache's --- litellm/router_utils/cooldown_cache.py | 9 ++++ .../unit/router_utils/test_cooldown_cache.py | 52 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index ef29f7d8fd3..4fd433d9391 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -118,6 +118,7 @@ class CooldownCache: ) # Set the cache with a TTL equal to the cooldown time + self._drop_in_memory_entry_if_extending(cooldown_key, _cooldown_time) self.cooldown_store.set_cache( value=cooldown_data, key=cooldown_key, @@ -127,6 +128,14 @@ class CooldownCache: verbose_logger.error("CooldownCache::add_deployment_to_cooldown - Exception occurred - %s", e) raise e + def _drop_in_memory_entry_if_extending(self, cooldown_key: str, new_cooldown_time: float) -> None: + """InMemoryCache keeps a live key's expiry, so a longer cooldown lands only if the entry is deleted first.""" + current_expiry: Final = self.in_memory_cache.ttl_dict.get(cooldown_key) + if current_expiry is None: + return + if float(current_expiry) < time.time() + float(new_cooldown_time): + self.in_memory_cache.delete_cache(cooldown_key) + @staticmethod @functools.lru_cache(maxsize=1024) def get_cooldown_cache_key(model_id: str) -> str: diff --git a/tests/unit/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py index 6f90fa8465f..fc59a2054f5 100644 --- a/tests/unit/router_utils/test_cooldown_cache.py +++ b/tests/unit/router_utils/test_cooldown_cache.py @@ -582,3 +582,55 @@ class TestCooldownSurvivesUnrelatedCacheTraffic: assert [model_id] == [entry[0] for entry in active], ( "unrelated router cache traffic must not evict a cooldown that is still running" ) + + +class TestCooldownExtension: + """A second failure asking for a longer cooldown must move the deadline out.""" + + @staticmethod + def _cache() -> CooldownCache: + return CooldownCache(cache=DualCache(), default_cooldown_time=60.0) + + def test_longer_cooldown_extends_the_deadline(self): + cc = self._cache() + key = CooldownCache.get_cooldown_cache_key("dep-a") + + cc.add_deployment_to_cooldown( + model_id="dep-a", original_exception=Exception("429"), exception_status=429, cooldown_time=2.0 + ) + first_expiry = cc.in_memory_cache.ttl_dict[key] + + started = time.time() + cc.add_deployment_to_cooldown( + model_id="dep-a", original_exception=Exception("429"), exception_status=429, cooldown_time=60.0 + ) + second_expiry = cc.in_memory_cache.ttl_dict[key] + + assert second_expiry > first_expiry, "the longer cooldown was dropped" + assert second_expiry - started == pytest.approx(60.0, abs=2.0) + + def test_shorter_cooldown_does_not_cut_an_active_one_short(self): + cc = self._cache() + key = CooldownCache.get_cooldown_cache_key("dep-b") + + cc.add_deployment_to_cooldown( + model_id="dep-b", original_exception=Exception("429"), exception_status=429, cooldown_time=60.0 + ) + long_expiry = cc.in_memory_cache.ttl_dict[key] + + cc.add_deployment_to_cooldown( + model_id="dep-b", original_exception=Exception("500"), exception_status=500, cooldown_time=1.0 + ) + + assert cc.in_memory_cache.ttl_dict[key] == long_expiry, "a brief cooldown shortened a long one" + + def test_first_cooldown_is_unaffected(self): + cc = self._cache() + key = CooldownCache.get_cooldown_cache_key("dep-c") + + started = time.time() + cc.add_deployment_to_cooldown( + model_id="dep-c", original_exception=Exception("429"), exception_status=429, cooldown_time=30.0 + ) + + assert cc.in_memory_cache.ttl_dict[key] - started == pytest.approx(30.0, abs=2.0)