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 1/3] 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) From 203fe2d4443c9500781d0d4bff1162eeeffe6543 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay <34904987+lets-order-some-fries@users.noreply.github.com> Date: Sun, 25 Oct 2026 09:29:59 +0530 Subject: [PATCH 2/3] fix(router): keep the longer cooldown when a shorter one arrives Both reviews showed the never-shorten half did not hold. The old helper only kept ttl_dict, but the write after it still replaced the stored record, and _corrected_active_cooldown routes on that record's timestamp + cooldown_time. A 5 s cooldown re-cooled for 0.5 s went routable again after 0.5 s, and DualCache sent the short record to Redis as well, so a replica reading Redis saw it too. The test passed because it asserted on ttl_dict add_deployment_to_cooldown now reads this replica's in-memory record and skips the write when that record's timestamp + cooldown_time is at or after the new deadline, so the longer cooldown keeps its record and TTL in memory and in Redis. Otherwise it deletes the in-memory entry so the new TTL lands, then writes as before. Comparing the record rather than ttl_dict also covers an entry backfilled from Redis, whose in-memory TTL is capped at 60 s. A missing or malformed record counts as no live cooldown. Redis is not read, so a replica whose in-memory copy is missing or older than Redis can still overwrite a longer cooldown there TestCooldownExtension now drives time.time with a fake clock and asserts through get_active_cooldowns: a longer cooldown, an equal one (#40130), a shorter one read after its own deadline, the same shorter one seen by a fresh replica through Redis, a longer one backfilled from Redis behind its capped in-memory TTL, a first cooldown, and stale or malformed records --- litellm/router_utils/cooldown_cache.py | 34 +++- .../unit/router_utils/test_cooldown_cache.py | 169 ++++++++++++++---- 2 files changed, 161 insertions(+), 42 deletions(-) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index aa153f15d21..1ae61f91e8f 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -38,6 +38,14 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns" +def _record_deadline(stored: object) -> float | None: + match stored: + case {"timestamp": int() | float() as timestamp, "cooldown_time": int() | float() as cooldown_time}: + return timestamp + cooldown_time + case _: + return None + + class CooldownCache: def __init__( self, @@ -119,8 +127,19 @@ class CooldownCache: cooldown_time=_cooldown_time, ) + live_deadline: Final = self._in_memory_deadline(cooldown_key) + new_deadline: Final = cooldown_data["timestamp"] + cooldown_data["cooldown_time"] + if live_deadline is not None and live_deadline >= new_deadline: + verbose_logger.debug( + "CooldownCache::add_deployment_to_cooldown - keeping %s's cooldown until %s over a new one until %s", + cooldown_key, + live_deadline, + new_deadline, + ) + return + # InMemoryCache keeps a live key's TTL on overwrite, so the new TTL lands only on a deleted entry + self.in_memory_cache.delete_cache(cooldown_key) # Set the cache with a TTL equal to the cooldown time - self._drop_in_memory_entry_if_extending(cooldown_key, _cooldown_time) with service_target(ROUTER_COOLDOWNS_TARGET): self.cooldown_store.set_cache( value=cooldown_data, @@ -131,13 +150,12 @@ 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) + def _in_memory_deadline(self, cooldown_key: str) -> float | None: + """ + The in-memory record's `timestamp + cooldown_time`, the deadline routing reads. The in-memory TTL is not + used because it is capped for an entry backfilled from Redis. Redis itself is not read. + """ + return _record_deadline(self.in_memory_cache.get_cache(cooldown_key)) @staticmethod @functools.lru_cache(maxsize=1024) diff --git a/tests/unit/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py index 5d143206d38..601b2a0b5f4 100644 --- a/tests/unit/router_utils/test_cooldown_cache.py +++ b/tests/unit/router_utils/test_cooldown_cache.py @@ -3,6 +3,8 @@ Unit tests for CooldownCache exception masking functionality """ import time +from collections.abc import Mapping +from typing import Final from unittest.mock import MagicMock import pytest @@ -629,53 +631,152 @@ class TestCooldownStoreCallsDeclareTheirKeyFamily: assert seen == [("mget", "router_cooldowns"), ("mget", "router_cooldowns")] +_CLOCK_START: Final = 1_000_000.0 + + +class _Clock: + """Stands in for `time.time`, which CooldownCache and its InMemoryCache both read at call time.""" + + def __init__(self, now: float) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> _Clock: + fake: Final = _Clock(_CLOCK_START) + monkeypatch.setattr(time, "time", fake) + return fake + + +def _cool(cc: CooldownCache, cooldown_time: float, status: int = 429) -> None: + cc.add_deployment_to_cooldown( + model_id="dep", original_exception=Exception(str(status)), exception_status=status, cooldown_time=cooldown_time + ) + + +def _active_record(cc: CooldownCache) -> CooldownCacheValue | None: + active: Final = cc.get_active_cooldowns(["dep"], parent_otel_span=None) + return active[0][1] if active else None + + class TestCooldownExtension: - """A second failure asking for a longer cooldown must move the deadline out.""" + """A second cooldown on a deployment keeps whichever deadline is later, read the way routing reads it.""" @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") + def test_longer_cooldown_extends_the_deadline(self, clock: _Clock): + cc: Final = self._cache() + _cool(cc, 2.0) + clock.advance(1.0) + _cool(cc, 60.0) - cc.add_deployment_to_cooldown( - model_id="dep-a", original_exception=Exception("429"), exception_status=429, cooldown_time=2.0 + clock.advance(1.5) # t=2.5, past the first deadline (t=2) + assert _active_record(cc) is not None, "the longer cooldown was dropped" + clock.advance(58.4) # t=60.9 + assert _active_record(cc) is not None + clock.advance(0.2) # t=61.1, past the second deadline (t=61) + assert _active_record(cc) is None + + def test_equal_cooldown_moves_the_deadline_out(self, clock: _Clock): + cc: Final = self._cache() + _cool(cc, 1.0) + clock.advance(0.6) + _cool(cc, 1.0) + + clock.advance(0.6) # t=1.2, past the first deadline (t=1) + assert _active_record(cc) is not None, "the re-cooldown kept the first deadline" + clock.advance(0.5) # t=1.7, past the second deadline (t=1.6) + assert _active_record(cc) is None + + def test_shorter_cooldown_does_not_cut_an_active_one_short(self, clock: _Clock): + cc: Final = self._cache() + _cool(cc, 60.0, status=429) + clock.advance(0.3) + _cool(cc, 1.0, status=500) + + clock.advance(1.5) # t=1.8, past the short deadline (t=1.3), long before the first one (t=60) + record: Final = _active_record(cc) + assert record is not None, "a brief cooldown cut a long one short" + assert (record["timestamp"], record["cooldown_time"], record["status_code"]) == (_CLOCK_START, 60.0, "429") + + def test_shorter_cooldown_does_not_reach_a_sibling_replica(self, clock: _Clock): + redis: Final = SharedRedisDouble() + writer: Final = CooldownCache( + cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis), default_cooldown_time=60.0 ) - first_expiry = cc.in_memory_cache.ttl_dict[key] + _cool(writer, 60.0, status=429) + clock.advance(0.3) + _cool(writer, 1.0, status=500) - started = time.time() - cc.add_deployment_to_cooldown( - model_id="dep-a", original_exception=Exception("429"), exception_status=429, cooldown_time=60.0 + clock.advance(1.5) # t=1.8, past the short deadline + sibling: Final = CooldownCache( + cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis), default_cooldown_time=60.0 ) - second_expiry = cc.in_memory_cache.ttl_dict[key] + record: Final = _active_record(sibling) # nothing in its memory yet, so this reads Redis + assert record is not None, "the brief cooldown replaced the long one in Redis" + assert record["cooldown_time"] == 60.0 - 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 + def test_backfilled_longer_cooldown_is_not_shortened_through_its_capped_ttl(self, clock: _Clock): + redis: Final = SharedRedisDouble() + writer: Final = CooldownCache( + cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis), default_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 + reader: Final = CooldownCache( + cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis), default_cooldown_time=60.0 ) + _cool(writer, 300.0, status=429) + assert _active_record(reader) is not None # backfills the reader's memory, with its TTL capped at 60 s + _cool(reader, 120.0, status=500) - 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 + clock.advance(150.0) # past the 120 s deadline, inside the 300 s one + sibling: Final = CooldownCache( + cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis), default_cooldown_time=60.0 ) + record: Final = _active_record(sibling) + assert record is not None, "the capped in-memory TTL let a shorter cooldown replace a longer one in Redis" + assert record["cooldown_time"] == 300.0 - assert cc.in_memory_cache.ttl_dict[key] - started == pytest.approx(30.0, abs=2.0) + def test_first_cooldown_is_unaffected(self, clock: _Clock): + cc: Final = self._cache() + _cool(cc, 30.0) + + clock.advance(29.9) + assert _active_record(cc) is not None + clock.advance(0.2) + assert _active_record(cc) is None + + @pytest.mark.parametrize( + "stored", + [ + pytest.param( + { + "exception_received": "429", + "status_code": "429", + "timestamp": _CLOCK_START - 120.0, + "cooldown_time": 60.0, + }, + id="deadline-passed-under-a-live-ttl", + ), + pytest.param({"timestamp": _CLOCK_START}, id="no-cooldown-time"), + pytest.param({"timestamp": "now", "cooldown_time": 60.0}, id="non-numeric-timestamp"), + ], + ) + def test_stale_or_malformed_record_does_not_block_a_fresh_cooldown( + self, clock: _Clock, stored: Mapping[str, object] + ): + cc: Final = self._cache() + cc.in_memory_cache.set_cache(CooldownCache.get_cooldown_cache_key("dep"), dict(stored), ttl=600) + _cool(cc, 30.0) + + clock.advance(10.0) + record: Final = _active_record(cc) + assert record is not None, "the stored record blocked the new cooldown" + assert (record["timestamp"], record["cooldown_time"]) == (_CLOCK_START, 30.0) From 1f753e2468a47036e201b1406a1200c316a76e35 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay <34904987+lets-order-some-fries@users.noreply.github.com> Date: Sun, 25 Oct 2026 10:33:52 +0530 Subject: [PATCH 3/3] fix(router): keep the basedpyright gate level for the cooldown deadline read The in-memory read of the stored cooldown record goes through InMemoryCache.get_cache, which is unannotated, so the merged head added one reportUnknownMemberType and one reportUnknownArgumentType error and scripts/type_check_gate.py failed against its merge-base ceiling Validate the value into object with TypeAdapter, as the rest of the codebase does, and name the suppressed rule with its reason on the untyped call. No behaviour change, the 34 cooldown tests still pass --- litellm/router_utils/cooldown_cache.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index e393d27996b..57a68ecac5d 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -7,6 +7,7 @@ import time from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm import verbose_logger @@ -155,7 +156,10 @@ class CooldownCache: The in-memory record's `timestamp + cooldown_time`, the deadline routing reads. The in-memory TTL is not used because it is capped for an entry backfilled from Redis. Redis itself is not read. """ - return _record_deadline(self.in_memory_cache.get_cache(cooldown_key)) + stored: Final = TypeAdapter(object).validate_python( + self.in_memory_cache.get_cache(cooldown_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache.get_cache is unannotated + ) + return _record_deadline(stored) @staticmethod @functools.lru_cache(maxsize=1024)