Merge 6541003141 into a392a39627

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:

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)