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

@@ -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)),
+                    }
+                )

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()