Merge abcd363acf into a392a39627

litellm/scheduler.py

@@ -1,15 +1,20 @@
 import asyncio
 import enum
 import heapq
+import logging as _logging
 import time
 from collections.abc import Awaitable, Callable, Sequence
 from typing import Final, TypeAlias
+from typing import Protocol as _Protocol
 
 from pydantic import TypeAdapter
 
 from litellm import print_verbose
 from litellm._internal_context import with_service_target
+from litellm._logging import verbose_logger as _verbose_logger
 from litellm.caching.caching import DualCache, RedisCache
+from litellm.caching.redis_cache import RedisCircuitBreakerOpenError as _RedisCircuitBreakerOpenError
+from litellm.caching.redis_cache import log_redis_failure as _log_redis_failure
 from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL
 from litellm.exceptions import Timeout
 from litellm.types.llms.base import LiteLLMBaseModel
@@ -30,6 +35,26 @@ class FlowItem(LiteLLMBaseModel):
     model_name: str
 
 
+class _QueueCache(_Protocol):
+    async def async_get_cache(self, key: str, *, local_only: bool = False, raise_on_error: bool = False) -> object: ...
+
+    async def async_set_cache(self, key: str, value: object, *, local_only: bool = False) -> object: ...
+
+
+async def _read_queue(cache: _QueueCache, redis_cache: _QueueCache | None, key: str) -> object:
+    if redis_cache is None:
+        return await cache.async_get_cache(key=key, local_only=True)
+    from redis.exceptions import RedisError as _RedisError
+
+    try:
+        response: Final = await redis_cache.async_get_cache(key=key, raise_on_error=True)
+    except (_RedisError, _RedisCircuitBreakerOpenError) as error:
+        _log_redis_failure(_verbose_logger, _logging.ERROR, "Scheduler queue using local snapshot", error)
+        return await cache.async_get_cache(key=key, local_only=True)
+    await cache.async_set_cache(key=key, value=[] if response is None else response, local_only=True)
+    return response
+
+
 class Scheduler:
     cache: DualCache
 
@@ -153,7 +178,7 @@ class Scheduler:
         """
         if self.cache is not None:
             _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
-            response: Final = await self.cache.async_get_cache(key=_cache_key)
+            response: Final = await _read_queue(self.cache, self.cache.redis_cache, _cache_key)
             if not isinstance(response, list):
                 return []
             return _QUEUE_ENTRIES.validate_python(response)

tests/unit/test_scheduler.py

@@ -5,12 +5,20 @@ import asyncio
 import importlib
 import json
 import os
+import time
 from collections.abc import Sequence
 from typing import Final
 
 import pytest
+import redis
+import redis.asyncio
+from fakeredis import FakeRedis, FakeServer
+from fakeredis.aioredis import FakeConnection
+from redis.asyncio import BlockingConnectionPool
 
 import litellm
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.redis_cache import RedisCache
 from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys
 from litellm.utils import _invalidate_model_cost_lowercase_map
@@ -207,6 +215,53 @@ class _ExpiringCache:
         self.store[key] = value
 
 
+class _SchedulerClock:
+    now: float = 1000.0
+
+    def __call__(self) -> float:
+        return self.now
+
+
+@pytest.mark.asyncio
+async def test_waiter_observes_shared_expiry_after_reading_just_before_it(monkeypatch: pytest.MonkeyPatch) -> None:
+    clock: Final = _SchedulerClock()
+    server: Final = FakeServer()
+
+    class Pool(BlockingConnectionPool):
+        def __init__(self, **kwargs: object) -> None:
+            super().__init__(connection_class=FakeConnection, server=server, **kwargs)
+
+    monkeypatch.setattr(time, "time", clock)
+    monkeypatch.setattr(redis, "Redis", FakeRedis)
+    monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool)
+    cache: Final = RedisCache(host="scheduler-expiry")
+    scheduler: Final = Scheduler(redis_cache=cache)
+    scheduler.cache.in_memory_cache = InMemoryCache(clock=clock)
+    waiter: Final = FlowItem(priority=2, request_id="waiter", model_name="expiry")
+    key: Final = f"{SchedulerCacheKeys.queue.value}:{waiter.model_name}"
+    try:
+        await cache.async_set_cache(key, [(1, "dead-replica")])
+        assert await scheduler.get_queue(waiter.model_name) == [(1, "dead-replica")]
+        clock.now = 1000.1
+        await scheduler.add_request(waiter)
+        clock.now = 1005.05
+        assert not await scheduler.poll(waiter, [])
+        clock.now = 1005.2
+        server.connected = False
+        assert not await scheduler.poll(waiter, [])
+        server.connected = True
+        assert await cache.async_get_cache(key) is None
+        assert await scheduler.get_queue(waiter.model_name) == []
+        server.connected = False
+        assert await scheduler.get_queue(waiter.model_name) == []
+        server.connected = True
+        assert await scheduler.poll(waiter, [])
+        assert await scheduler.get_queue(waiter.model_name) == []
+    finally:
+        server.connected = True
+        await cache.disconnect()
+
+
 @pytest.mark.asyncio
 async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
     cache: Final = _ExpiringCache()