From 5fd9b44a8584bf384a8990eaf678975cf2184bae Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 06:13:01 +0000 Subject: [PATCH] perf(scheduler): share one Redis queue read per polling interval across waiters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/scheduler.py | 90 +++++++++++++++++++++++--- tests/unit/test_scheduler.py | 119 ++++++++++++++++++++++++++++++++++- 2 files changed, 199 insertions(+), 10 deletions(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 48be61283df..b3e9330dc76 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -3,7 +3,9 @@ import enum import heapq import logging as _logging import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType from typing import Final, TypeAlias from typing import Protocol as _Protocol @@ -24,6 +26,18 @@ QueueEntry: TypeAlias = tuple[int, str] _QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry]) +@dataclass(frozen=True, slots=True) +class _QueueRead: + value: object + from_redis: bool + + +@dataclass(frozen=True, slots=True) +class _QueueSnapshot: + read_at: float + entries: tuple[QueueEntry, ...] + + class SchedulerCacheKeys(enum.Enum): queue = "scheduler:queue" default_in_memory_ttl = DEFAULT_IN_MEMORY_TTL # cache queue in-memory for 5s when redis cache available @@ -41,27 +55,37 @@ class _QueueCache(_Protocol): 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: +def _validated_queue(response: object) -> tuple[QueueEntry, ...]: + if not isinstance(response, list): + return () + return tuple(_QUEUE_ENTRIES.validate_python(response)) + + +async def _read_queue(cache: _QueueCache, redis_cache: _QueueCache | None, key: str) -> _QueueRead: if redis_cache is None: - return await cache.async_get_cache(key=key, local_only=True) + return _QueueRead(value=await cache.async_get_cache(key=key, local_only=True), from_redis=False) 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) + return _QueueRead(value=await cache.async_get_cache(key=key, local_only=True), from_redis=False) await cache.async_set_cache(key=key, value=[] if response is None else response, local_only=True) - return response + return _QueueRead(value=response, from_redis=True) class Scheduler: cache: DualCache + clock: Callable[[], float] + _queue_snapshots: Mapping[str, _QueueSnapshot] + _queue_read_tasks: Mapping[str, asyncio.Task[_QueueRead]] def __init__( self, polling_interval: float | None = None, redis_cache: RedisCache | None = None, + clock: Callable[[], float] = time.monotonic, ): """ polling_interval: float or null - frequency of polling queue. Default is 3ms. @@ -73,6 +97,49 @@ class Scheduler: default_in_memory_ttl = SchedulerCacheKeys.default_in_memory_ttl.value self.cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=default_in_memory_ttl) self.polling_interval = polling_interval or DEFAULT_POLLING_INTERVAL # default to 3ms + self.clock = clock + self._queue_snapshots = MappingProxyType({}) + self._queue_read_tasks = MappingProxyType({}) + + async def _read_queue_with_snapshot(self, key: str, read_at: float) -> _QueueRead: + result: Final = await _read_queue(self.cache, self.cache.redis_cache, key) + if not result.from_redis: + return result + entries: Final = _validated_queue(result.value) + snapshot: Final = self._queue_snapshots.get(key) + if snapshot is None or read_at > snapshot.read_at: + self._queue_snapshots = MappingProxyType( + {**self._queue_snapshots, key: _QueueSnapshot(read_at=read_at, entries=entries)} + ) + return _QueueRead(value=list(entries), from_redis=True) + + def _clear_in_flight_read(self, task: asyncio.Task[_QueueRead]) -> None: + self._queue_read_tasks = MappingProxyType( + {key: current for key, current in self._queue_read_tasks.items() if current is not task} + ) + + async def _await_queue_read(self, task: asyncio.Task[_QueueRead]) -> _QueueRead: + try: + return await asyncio.shield(task) + finally: + if task.done(): + self._clear_in_flight_read(task) + + async def _get_redis_queue(self, key: str) -> tuple[QueueEntry, ...]: + now: Final = self.clock() + snapshot: Final = self._queue_snapshots.get(key) + if snapshot is not None and now - snapshot.read_at < self.polling_interval: + return snapshot.entries + task: Final = self._queue_read_tasks.get(key) + if task is None: + read_at: Final = now + new_task: Final = asyncio.create_task(self._read_queue_with_snapshot(key, read_at)) + self._queue_read_tasks = MappingProxyType({**self._queue_read_tasks, key: new_task}) + new_task.add_done_callback(self._clear_in_flight_read) + new_result: Final = await self._await_queue_read(new_task) + return _validated_queue(new_result.value) + shared_result: Final = await self._await_queue_read(task) + return _validated_queue(shared_result.value) async def add_request(self, request: FlowItem): # We use the priority directly, as lower values indicate higher priority @@ -178,10 +245,10 @@ class Scheduler: """ if self.cache is not None: _cache_key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}" + if self.cache.redis_cache is not None: + return list(await self._get_redis_queue(_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) + return list(_validated_queue(response.value)) return self.queue @with_service_target(SCHEDULER_QUEUE_TARGET) @@ -192,3 +259,10 @@ class Scheduler: 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) + if self.cache.redis_cache is not None: + self._queue_snapshots = MappingProxyType( + { + **self._queue_snapshots, + _cache_key: _QueueSnapshot(read_at=self.clock(), entries=tuple(queue)), + } + ) diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index fbc3c0e47c3..bb58c870ce0 100644 --- a/tests/unit/test_scheduler.py +++ b/tests/unit/test_scheduler.py @@ -235,7 +235,7 @@ async def test_waiter_observes_shared_expiry_after_reading_just_before_it(monkey monkeypatch.setattr(redis, "Redis", FakeRedis) monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool) cache: Final = RedisCache(host="scheduler-expiry") - scheduler: Final = Scheduler(redis_cache=cache) + scheduler: Final = Scheduler(redis_cache=cache, clock=clock) 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}" @@ -284,17 +284,132 @@ async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring(): class _JsonRoundTripRedisCache: - def __init__(self) -> None: + def __init__( + self, + read_started: asyncio.Event | None = None, + read_release: asyncio.Event | None = None, + ) -> None: self.store: dict[str, str] = {} + self.read_count = 0 + self.fail_reads = False + self.read_started = read_started + self.read_release = read_release async def async_get_cache(self, key: str, **kwargs: object) -> object: + self.read_count += 1 raw: Final = self.store.get(key) + if self.read_started is not None: + self.read_started.set() + if self.read_release is not None: + await self.read_release.wait() + if self.fail_reads: + raise redis.exceptions.RedisError("Redis unavailable") 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_concurrent_get_queue_calls_share_redis_read_and_return_independent_lists() -> None: + redis_cache: Final = _JsonRoundTripRedisCache() + redis_cache.store[f"{SchedulerCacheKeys.queue.value}:shared"] = json.dumps([(1, "queued")]) + scheduler: Final = Scheduler(redis_cache=redis_cache) + + queues: Final = await asyncio.gather(*(scheduler.get_queue("shared") for _ in range(20))) + + assert redis_cache.read_count == 1 + assert all(queue == [(1, "queued")] for queue in queues) + queues[0].append((2, "local")) + assert all(queue == [(1, "queued")] for queue in queues[1:]) + + +@pytest.mark.asyncio +async def test_get_queue_snapshot_expires_at_polling_interval() -> None: + redis_cache: Final = _JsonRoundTripRedisCache() + key: Final = f"{SchedulerCacheKeys.queue.value}:clocked" + redis_cache.store[key] = json.dumps([(1, "before")]) + clock: Final = _SchedulerClock() + scheduler: Final = Scheduler(polling_interval=1.0, redis_cache=redis_cache, clock=clock) + + assert await scheduler.get_queue("clocked") == [(1, "before")] + redis_cache.store[key] = json.dumps([(1, "after")]) + clock.now = 1000.5 + assert await scheduler.get_queue("clocked") == [(1, "before")] + assert redis_cache.read_count == 1 + clock.now = 1001.0 + assert await scheduler.get_queue("clocked") == [(1, "after")] + assert redis_cache.read_count == 2 + + +@pytest.mark.asyncio +async def test_save_queue_refreshes_local_snapshot_and_other_scheduler_reads_redis() -> None: + redis_cache: Final = _JsonRoundTripRedisCache() + scheduler: Final = Scheduler(redis_cache=redis_cache) + other_scheduler: Final = Scheduler(redis_cache=redis_cache) + queue: Final = [(1, "saved")] + + await scheduler.save_queue(queue, "saved") + assert await scheduler.get_queue("saved") == queue + assert redis_cache.read_count == 0 + assert await other_scheduler.get_queue("saved") == queue + assert redis_cache.read_count == 1 + + +@pytest.mark.asyncio +async def test_failed_redis_read_is_not_shared_after_redis_recovers() -> None: + redis_cache: Final = _JsonRoundTripRedisCache() + scheduler: Final = Scheduler(redis_cache=redis_cache) + redis_cache.fail_reads = True + + assert await scheduler.get_queue("recovering") == [] + assert redis_cache.read_count == 1 + redis_cache.fail_reads = False + redis_cache.store[f"{SchedulerCacheKeys.queue.value}:recovering"] = json.dumps([(1, "recovered")]) + + assert await scheduler.get_queue("recovering") == [(1, "recovered")] + assert redis_cache.read_count == 2 + + +@pytest.mark.asyncio +async def test_cancelling_one_get_queue_waiter_does_not_cancel_shared_redis_read() -> None: + read_started: Final = asyncio.Event() + read_release: Final = asyncio.Event() + redis_cache: Final = _JsonRoundTripRedisCache(read_started=read_started, read_release=read_release) + redis_cache.store[f"{SchedulerCacheKeys.queue.value}:blocked"] = json.dumps([(1, "queued")]) + scheduler: Final = Scheduler(redis_cache=redis_cache) + cancelled_waiter: Final = asyncio.create_task(scheduler.get_queue("blocked")) + remaining_waiter: Final = asyncio.create_task(scheduler.get_queue("blocked")) + await read_started.wait() + + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + read_release.set() + + assert await remaining_waiter == [(1, "queued")] + assert redis_cache.read_count == 1 + + +@pytest.mark.asyncio +async def test_older_in_flight_read_does_not_replace_saved_queue_snapshot() -> None: + read_started: Final = asyncio.Event() + read_release: Final = asyncio.Event() + redis_cache: Final = _JsonRoundTripRedisCache(read_started=read_started, read_release=read_release) + redis_cache.store[f"{SchedulerCacheKeys.queue.value}:write-race"] = json.dumps([(1, "before-save")]) + clock: Final = _SchedulerClock() + scheduler: Final = Scheduler(redis_cache=redis_cache, clock=clock) + pending_read: Final = asyncio.create_task(scheduler.get_queue("write-race")) + await read_started.wait() + + await scheduler.save_queue([(2, "saved")], "write-race") + read_release.set() + + assert await pending_read == [(1, "before-save")] + assert await scheduler.get_queue("write-race") == [(2, "saved")] + assert redis_cache.read_count == 1 + + @pytest.mark.asyncio async def test_second_replica_enqueues_behind_a_queue_decoded_from_redis(): redis_cache: Final = _JsonRoundTripRedisCache()