refactor(redis): move scheduler queue snapshots behind a typed store

litellm/redis/stores/scheduler_queue.py

@@ -0,0 +1,71 @@
+import ast
+import json
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Final, Protocol, TypeAlias
+
+from pydantic import JsonValue, TypeAdapter, ValidationError
+
+from litellm.redis.transport.contract import RedisFailure, RedisTransport
+
+QUEUE_PREFIX: Final = "scheduler:queue"
+QueueEntry: TypeAlias = tuple[int, str]
+_ENTRIES: Final = TypeAdapter(list[QueueEntry])
+_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_OBJECT: Final = TypeAdapter(object)
+
+
+def queue_key(model_name: str) -> str:
+    return f"{QUEUE_PREFIX}:{model_name}"
+
+
+def queue_entries(value: object) -> Sequence[QueueEntry]:
+    return _ENTRIES.validate_python(value) if isinstance(value, list) else []
+
+
+def _decode(raw: bytes) -> object:
+    decoded: Final = raw.decode("utf-8")
+    try:
+        return _JSON.validate_json(decoded)
+    except ValueError:
+        return _OBJECT.validate_python(ast.literal_eval(decoded))
+
+
+@dataclass(frozen=True, slots=True)
+class InvalidQueue:
+    error: ValidationError
+
+
+class SchedulerQueueStore(Protocol):
+    async def read_queue(self, model_name: str) -> tuple[QueueEntry, ...] | None | InvalidQueue | RedisFailure: ...
+
+    async def replace_queue(
+        self, model_name: str, queue: Sequence[QueueEntry], *, ttl: timedelta | None
+    ) -> None | RedisFailure: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TransportSchedulerQueueStore:
+    transport: RedisTransport
+
+    async def read_queue(self, model_name: str) -> tuple[QueueEntry, ...] | None | InvalidQueue | RedisFailure:
+        raw: Final = await self.transport.get(queue_key(model_name))
+        if raw is None or isinstance(raw, RedisFailure):
+            return raw
+        try:
+            value: Final = _decode(raw)
+        except (ValueError, SyntaxError, UnicodeError) as error:
+            return RedisFailure(kind="decode", message=str(error))
+        if value is None:
+            return None
+        try:
+            return tuple(queue_entries(value))
+        except ValidationError as error:
+            return InvalidQueue(error)
+
+    async def replace_queue(
+        self, model_name: str, queue: Sequence[QueueEntry], *, ttl: timedelta | None
+    ) -> None | RedisFailure:
+        reply: Final = await self.transport.set_seconds(queue_key(model_name), json.dumps(list(queue)), ttl=ttl)
+        return reply if isinstance(reply, RedisFailure) else None

litellm/scheduler.py

@@ -3,7 +3,8 @@ import enum
 import heapq
 import time
 from collections.abc import Awaitable, Callable, Sequence
-from typing import Final, TypeAlias
+from datetime import timedelta
+from typing import Final
 
 from pydantic import TypeAdapter
 
@@ -12,15 +13,24 @@ from litellm._internal_context import with_service_target
 from litellm.caching.caching import DualCache, RedisCache
 from litellm.constants import DEFAULT_IN_MEMORY_TTL, DEFAULT_POLLING_INTERVAL
 from litellm.exceptions import Timeout
+from litellm.redis.stores.scheduler_queue import (
+    QUEUE_PREFIX,
+    InvalidQueue,
+    QueueEntry,
+    SchedulerQueueStore,
+    TransportSchedulerQueueStore,
+    queue_entries,
+    queue_key,
+)
+from litellm.redis.transport.contract import RedisFailure
 from litellm.types.llms.base import LiteLLMBaseModel
 
 SCHEDULER_QUEUE_TARGET: Final = "scheduler_queue"
-QueueEntry: TypeAlias = tuple[int, str]
-_QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry])
+_OBJECT: Final = TypeAdapter(object)
 
 
 class SchedulerCacheKeys(enum.Enum):
-    queue = "scheduler:queue"
+    queue = QUEUE_PREFIX
     default_in_memory_ttl = DEFAULT_IN_MEMORY_TTL  # cache queue in-memory for 5s when redis cache available
 
 
@@ -152,18 +162,43 @@ class Scheduler:
         (priority, request_id) tuple the heap operations compare against.
         """
         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)
-            if not isinstance(response, list):
+            key: Final = queue_key(model_name)
+            local: Final = _OBJECT.validate_python(await self.cache.async_get_cache(key=key, local_only=True))
+            if isinstance(local, InvalidQueue):
+                raise local.error
+            if local is not None:
+                return list(queue_entries(local))
+            redis: Final = self.cache.redis_cache
+            if redis is None:
                 return []
-            return _QUEUE_ENTRIES.validate_python(response)
+            store: Final[SchedulerQueueStore] = TransportSchedulerQueueStore(redis.transport())
+            response: Final = await store.read_queue(model_name)
+            if isinstance(response, InvalidQueue):
+                await self._save_local_queue(response, model_name)
+                raise response.error
+            if response is None or isinstance(response, RedisFailure):
+                return []
+            await self._save_local_queue(response, model_name)
+            return list(response)
         return self.queue
 
+    async def _save_local_queue(self, queue: Sequence[QueueEntry] | InvalidQueue, model_name: str) -> None:
+        value: Final = queue if isinstance(queue, InvalidQueue) else list(queue)
+        await self.cache.async_set_cache(key=queue_key(model_name), value=value, local_only=True)
+
     @with_service_target(SCHEDULER_QUEUE_TARGET)
     async def save_queue(self, queue: list[QueueEntry], model_name: str) -> None:
         """
         Save the updated queue of the model group
         """
         if self.cache is not None:
-            _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}"
-            await self.cache.async_set_cache(key=_cache_key, value=queue)
+            await self._save_local_queue(queue, model_name)
+            redis: Final = self.cache.redis_cache
+            if redis is not None:
+                ttl: Final = (
+                    int(self.cache.default_in_memory_ttl)
+                    if self.cache.default_in_memory_ttl is not None
+                    else redis.default_ttl
+                )
+                store: Final[SchedulerQueueStore] = TransportSchedulerQueueStore(redis.transport())
+                await store.replace_queue(model_name, queue, ttl=timedelta(seconds=ttl))

tests/unit/redis/stores/test_scheduler_queue.py

@@ -0,0 +1,150 @@
+from datetime import timedelta
+from typing import Final, Protocol, cast
+from unittest.mock import patch
+
+import fakeredis
+import pytest
+from pydantic import ValidationError
+from redis.exceptions import ConnectionError as RedisConnectionError
+
+from litellm.redis.stores.scheduler_queue import TransportSchedulerQueueStore
+from litellm.redis.transport.contract import Arg, RedisFailure
+from litellm.scheduler import FlowItem, Scheduler
+from tests._support.fake_redis import fake_redis_cache
+
+
+class LegacyQueueCache(Protocol):
+    async def async_set_cache(self, key: str, value: object, *, ttl: int) -> bool | None: ...
+
+    async def async_get_cache(self, key: str) -> object: ...
+
+
+class RedisQueueWriter(Protocol):
+    async def set(
+        self, name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False
+    ) -> object: ...
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("seconds", [1.9, 0.9, 0, -0.9, -1.9])
+@patch("time.time", return_value=1800000000.0)
+async def test_queue_fractional_seconds_match_legacy_set_expiry(
+    _clock: object, server: fakeredis.FakeServer, seconds: float
+) -> None:
+    cache: Final = fake_redis_cache(server, namespace="tenant")
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    key: Final = "tenant:scheduler:queue:model"
+    await raw.set(key, b'[[2, "older"]]')
+    store: Final = TransportSchedulerQueueStore(cache.transport())
+    result: Final = await store.replace_queue("model", ((1, "new"),), ttl=timedelta(seconds=seconds))
+    if seconds >= 1:
+        assert result is None
+        assert await raw.get(key) == b'[[1, "new"]]'
+        assert await raw.pttl(key) == 1000
+    else:
+        assert isinstance(result, RedisFailure) and result.kind == "server"
+        assert await raw.get(key) == b'[[2, "older"]]'
+        assert await raw.pttl(key) == -1
+
+
+@pytest.mark.asyncio
+@patch("time.time", return_value=1800000000.0)
+async def test_queue_rolling_deploy_preserves_bytes_ttl_and_local_snapshot(
+    _clock: object, server: fakeredis.FakeServer
+) -> None:
+    cache: Final = fake_redis_cache(server, namespace="tenant")
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    scheduler: Final = Scheduler(redis_cache=cache)
+    legacy: Final = cast(
+        LegacyQueueCache, cache
+    )  # cast-ok: narrow the untyped legacy API for wire-compatibility probes
+    key: Final = "scheduler:queue:mødel"
+    with patch("time.time", return_value=1800000000.0):
+        await legacy.async_set_cache(key, [[2, "older-雪"]], ttl=5)
+        assert await scheduler.get_queue("mødel") == [(2, "older-雪")]
+        await scheduler.add_request(FlowItem(priority=1, request_id="new-雪", model_name="mødel"))
+        assert await raw.get("tenant:" + key) == b'[[1, "new-\\u96ea"], [2, "older-\\u96ea"]]'
+        assert await legacy.async_get_cache(key) == [[1, "new-雪"], [2, "older-雪"]]
+        assert await raw.pttl("tenant:" + key) == 5000
+        await raw.delete("tenant:" + key)
+        assert await scheduler.get_queue("mødel") == [(1, "new-雪"), (2, "older-雪")]
+        scheduler.cache.in_memory_cache.flush_cache()
+        assert await scheduler.get_queue("mødel") == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raw_value", [b"garbage", b"null", b"{}", b'"text"', b"\xff"])
+async def test_queue_unreadable_or_non_list_snapshot_is_empty(server: fakeredis.FakeServer, raw_value: bytes) -> None:
+    cache: Final = fake_redis_cache(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set("scheduler:queue:model", raw_value)
+    assert await Scheduler(redis_cache=cache).get_queue("model") == []
+
+
+@pytest.mark.asyncio
+async def test_queue_invalid_entries_still_raise_and_python_literals_still_read(server: fakeredis.FakeServer) -> None:
+    cache: Final = fake_redis_cache(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    await raw.set("scheduler:queue:valid", b"[(2, 'old')]")
+    await raw.set("scheduler:queue:invalid", b'[["bad", "id"]]')
+    scheduler: Final = Scheduler(redis_cache=cache)
+    assert await scheduler.get_queue("valid") == [(2, "old")]
+    with pytest.raises(ValidationError, match="int_parsing"):
+        await scheduler.get_queue("invalid")
+    await raw.set("scheduler:queue:invalid", b'[[1, "fixed"]]')
+    with pytest.raises(ValidationError, match="int_parsing"):
+        await scheduler.get_queue("invalid")
+    scheduler.cache.in_memory_cache.flush_cache()
+    assert await scheduler.get_queue("invalid") == [(1, "fixed")]
+
+
+@pytest.mark.asyncio
+async def test_queue_outage_keeps_local_writes_and_uses_shared_breaker() -> None:
+    server: Final = fakeredis.FakeServer()
+    cache: Final = fake_redis_cache(server)
+    scheduler: Final = Scheduler(redis_cache=cache)
+    server.connected = False
+    await scheduler.save_queue([(1, "local")], "model")
+    assert await scheduler.get_queue("model") == [(1, "local")]
+    scheduler.cache.in_memory_cache.flush_cache()
+    assert await scheduler.get_queue("model") == []
+    for _ in range(5):
+        await scheduler.get_queue("model")
+    result: Final = await TransportSchedulerQueueStore(cache.transport()).read_queue("model")
+    assert isinstance(result, RedisFailure) and result.kind == "breaker_open"
+
+
+@pytest.mark.asyncio
+@patch("time.time", return_value=1800000000.0)
+async def test_queue_late_redis_attachment_uses_backend_default_ttl(
+    _clock: object, server: fakeredis.FakeServer
+) -> None:
+    scheduler: Final = Scheduler()
+    await scheduler.save_queue([(1, "local")], "model")
+    cache: Final = fake_redis_cache(server, namespace="tenant")
+    scheduler.cache.attach_redis_cache(cache)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    assert await raw.get("tenant:scheduler:queue:model") is None
+    with patch("time.time", return_value=1800000000.0):
+        await scheduler.save_queue([(2, "shared")], "model")
+        assert await raw.pttl("tenant:scheduler:queue:model") == cache.default_ttl * 1000
+
+
+@pytest.mark.asyncio
+async def test_queue_unknown_applied_write_is_not_retried_or_rolled_back(server: fakeredis.FakeServer) -> None:
+    cache: Final = fake_redis_cache(server, namespace="tenant")
+    original_set: Final = cast(type[RedisQueueWriter], fakeredis.FakeAsyncRedis).set  # cast-ok: typed redis-py boundary
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    scheduler: Final = Scheduler(redis_cache=cache)
+
+    async def applied_then_disconnected(
+        name: str, value: Arg, *, ex: timedelta | None = None, px: int | None = None, nx: bool = False
+    ) -> None:
+        await original_set(raw, name, value, ex=ex, px=px, nx=nx)
+        raise RedisConnectionError("response lost after write")
+
+    with patch.object(fakeredis.FakeAsyncRedis, "set", side_effect=applied_then_disconnected) as write:
+        await scheduler.save_queue([(1, "written")], "model")
+        assert write.call_count == 1
+    assert await raw.get("tenant:scheduler:queue:model") == b'[[1, "written"]]'
+    assert await scheduler.get_queue("model") == [(1, "written")]

tests/unit/test_scheduler.py

@@ -3,17 +3,18 @@
 
 import asyncio
 import importlib
-import json
 import os
 from collections.abc import Sequence
 from typing import Final
 
+import fakeredis
 import pytest
 
 import litellm
 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
+from tests._support.fake_redis import fake_redis_cache
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
 
@@ -35,17 +36,7 @@ async def test_scheduler_diff_model_names():
 
 @pytest.mark.asyncio
 async def test_scheduler_poll_persists_queue_to_cache():
-    class StubRedisCache:
-        def __init__(self):
-            self.store = {}
-
-        async def async_get_cache(self, key, **kwargs):
-            return self.store.get(key)
-
-        async def async_set_cache(self, key, value, **kwargs):
-            self.store[key] = value
-
-    redis_cache = StubRedisCache()
+    redis_cache = fake_redis_cache(fakeredis.FakeServer())
     scheduler = Scheduler(redis_cache=redis_cache)
 
     item1 = FlowItem(priority=0, request_id="10", model_name="gpt-3.5-turbo")
@@ -56,7 +47,8 @@ async def test_scheduler_poll_persists_queue_to_cache():
     await scheduler.poll(request=item1, health_deployments=[])
 
     queue_key = f"{SchedulerCacheKeys.queue.value}:{item1.model_name}"
-    updated_queue = redis_cache.store[queue_key]
+    updated_queue = await redis_cache.async_get_cache(queue_key)
+    assert updated_queue is not None
     assert updated_queue[0][1] == "11"
 
 
@@ -196,25 +188,16 @@ async def test_poll_during_cooldown_admits_an_erased_request_that_outranks_the_q
     assert await scheduler.get_queue("sched-model") == [(2, "still-queued")]
 
 
-class _ExpiringCache:
-    def __init__(self) -> None:
-        self.store: dict[str, object] = {}
-
-    async def async_get_cache(self, key: str, **kwargs: object) -> object:
-        return self.store.get(key)
-
-    async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
-        self.store[key] = value
-
-
 @pytest.mark.asyncio
 async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
-    cache: Final = _ExpiringCache()
+    server: Final = fakeredis.FakeServer()
+    cache: Final = fake_redis_cache(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
     scheduler: Final = Scheduler(redis_cache=cache)
     scheduler.cache.in_memory_cache.cache_dict.clear()
 
     async def no_healthy_deployments_after_the_key_expired() -> Sequence[object]:
-        cache.store.clear()
+        await raw.flushall()
         scheduler.cache.in_memory_cache.cache_dict.clear()
         return ()
 
@@ -227,22 +210,9 @@ async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring():
     assert await scheduler.get_queue("sched-model") == []
 
 
-
-class _JsonRoundTripRedisCache:
-    def __init__(self) -> None:
-        self.store: dict[str, str] = {}
-
-    async def async_get_cache(self, key: str, **kwargs: object) -> object:
-        raw: Final = self.store.get(key)
-        return None if raw is None else json.loads(raw)
-
-    async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
-        self.store[key] = json.dumps(value)
-
-
 @pytest.mark.asyncio
 async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis():
-    redis_cache: Final = _JsonRoundTripRedisCache()
+    redis_cache: Final = fake_redis_cache(fakeredis.FakeServer())
     replica_a: Final = Scheduler(redis_cache=redis_cache)
     replica_b: Final = Scheduler(redis_cache=redis_cache)
     waiting_on_a: Final = FlowItem(priority=1, request_id="waiting-on-a", model_name="sched-model")
@@ -257,6 +227,7 @@ async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis():
     await replica_b.remove_request(request_id="waiting-on-a", model_name="sched-model")
     assert await replica_b.get_queue("sched-model") == []
 
+
 class _PausesAfterEnqueueScheduler(Scheduler):
     def __init__(self) -> None:
         super().__init__()