From b267f1704ad012d41cc1be177e693d0454d40195 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 03:21:51 +0000 Subject: [PATCH 1/4] fix(scheduler): observe shared queue expiry before polling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/scheduler.py | 26 ++++++++++++++++- tests/unit/test_scheduler.py | 55 ++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 73196d24e4d..77b5bb724c7 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -1,15 +1,21 @@ 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 redis.exceptions import RedisError as _RedisError 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 +36,24 @@ 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) + 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 +177,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) diff --git a/tests/unit/test_scheduler.py b/tests/unit/test_scheduler.py index e0968e9b2cf..fbc3c0e47c3 100644 --- a/tests/unit/test_scheduler.py +++ b/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() From 809e3cbd75e7a97f24669bd323ff027a94f549f5 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 03:40:38 +0000 Subject: [PATCH 2/4] fix(scheduler): defer optional Redis exception import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/scheduler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 77b5bb724c7..48be61283df 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -8,7 +8,6 @@ from typing import Final, TypeAlias from typing import Protocol as _Protocol from pydantic import TypeAdapter -from redis.exceptions import RedisError as _RedisError from litellm import print_verbose from litellm._internal_context import with_service_target @@ -45,6 +44,8 @@ class _QueueCache(_Protocol): 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: From 5fd9b44a8584bf384a8990eaf678975cf2184bae Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 06:13:01 +0000 Subject: [PATCH 3/4] 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() From abcd363acfefd7a264692ac4e9529e393df47f93 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 31 Oct 2026 06:34:05 +0000 Subject: [PATCH 4/4] revert(scheduler): drop shared Redis queue read This reverts commit 5fd9b44a8584bf384a8990eaf678975cf2184bae 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, 10 insertions(+), 199 deletions(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index b3e9330dc76..48be61283df 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -3,9 +3,7 @@ import enum import heapq import logging as _logging import time -from collections.abc import Awaitable, Callable, Mapping, Sequence -from dataclasses import dataclass -from types import MappingProxyType +from collections.abc import Awaitable, Callable, Sequence from typing import Final, TypeAlias from typing import Protocol as _Protocol @@ -26,18 +24,6 @@ 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 @@ -55,37 +41,27 @@ class _QueueCache(_Protocol): async def async_set_cache(self, key: str, value: object, *, local_only: bool = False) -> 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: +async def _read_queue(cache: _QueueCache, redis_cache: _QueueCache | None, key: str) -> object: if redis_cache is None: - return _QueueRead(value=await cache.async_get_cache(key=key, local_only=True), from_redis=False) + 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 _QueueRead(value=await cache.async_get_cache(key=key, local_only=True), from_redis=False) + 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 _QueueRead(value=response, from_redis=True) + return response 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. @@ -97,49 +73,6 @@ 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 @@ -245,10 +178,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) - return list(_validated_queue(response.value)) + if not isinstance(response, list): + return [] + return _QUEUE_ENTRIES.validate_python(response) return self.queue @with_service_target(SCHEDULER_QUEUE_TARGET) @@ -259,10 +192,3 @@ 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 bb58c870ce0..fbc3c0e47c3 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, clock=clock) + 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}" @@ -284,132 +284,17 @@ async def test_wait_for_turn_during_cooldown_survives_the_queue_key_expiring(): class _JsonRoundTripRedisCache: - def __init__( - self, - read_started: asyncio.Event | None = None, - read_release: asyncio.Event | None = None, - ) -> None: + def __init__(self) -> 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()