From 371444b989e63112aa63bbfd4cc3476e81cccf0d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:37 +0000 Subject: [PATCH] refactor(redis): move scheduler queue snapshots behind a typed store --- litellm/redis/stores/scheduler_queue.py | 71 +++++++++ litellm/scheduler.py | 55 +++++-- .../unit/redis/stores/test_scheduler_queue.py | 150 ++++++++++++++++++ tests/unit/test_scheduler.py | 51 ++---- 4 files changed, 277 insertions(+), 50 deletions(-) create mode 100644 litellm/redis/stores/scheduler_queue.py create mode 100644 tests/unit/redis/stores/test_scheduler_queue.py diff --git a/litellm/redis/stores/scheduler_queue.py b/litellm/redis/stores/scheduler_queue.py new file mode 100644 index 00000000000..3df2469eb6a --- /dev/null +++ b/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 diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 73196d24e4d..a0bc8021c44 100644 --- a/litellm/scheduler.py +++ b/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)) diff --git a/tests/unit/redis/stores/test_scheduler_queue.py b/tests/unit/redis/stores/test_scheduler_queue.py new file mode 100644 index 00000000000..ac5f0a4842c --- /dev/null +++ b/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")] diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index e0968e9b2cf..0583cbc74b8 100644 --- a/tests/unit/test_scheduler.py +++ b/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__()