diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index dc21db6a601..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 @@ -38,6 +39,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,6 +128,18 @@ 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 with service_target(ROUTER_COOLDOWNS_TARGET): self.cooldown_store.set_cache( @@ -130,6 +151,16 @@ class CooldownCache: verbose_logger.error("CooldownCache::add_deployment_to_cooldown - Exception occurred - %s", e) raise e + 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. + """ + 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) 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 98f2b254a3f..189156346b6 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 @@ -627,3 +629,154 @@ class TestCooldownStoreCallsDeclareTheirKeyFamily: assert cc.get_min_cooldown(["dep-1"], parent_otel_span=None) == 60.0 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 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, clock: _Clock): + cc: Final = self._cache() + _cool(cc, 2.0) + clock.advance(1.0) + _cool(cc, 60.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 + ) + _cool(writer, 60.0, status=429) + clock.advance(0.3) + _cool(writer, 1.0, status=500) + + 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 + ) + 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 + + 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 + ) + 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) + + 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 + + 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)