From 39f37918f34f158f562748a3312b0ca28333706b Mon Sep 17 00:00:00 2001 From: xyi74976-del <290725798+xyi74976-del@users.noreply.github.com> Date: Sat, 17 Oct 2026 01:37:31 +0800 Subject: [PATCH 1/2] fix(scheduler): atomically update shared Redis queues --- litellm/scheduler.py | 157 ++++++++++++++++++ .../test_priority_scheduler_queue_cleanup.py | 156 ++++++++++++++++- 2 files changed, 311 insertions(+), 2 deletions(-) diff --git a/litellm/scheduler.py b/litellm/scheduler.py index 73196d24e4d..871cff4d241 100644 --- a/litellm/scheduler.py +++ b/litellm/scheduler.py @@ -1,15 +1,28 @@ import asyncio import enum import heapq +import logging as _logging import time from collections.abc import Awaitable, Callable, Sequence +from functools import partial as _partial from typing import Final, TypeAlias +from typing import Protocol as _Protocol from pydantic import TypeAdapter from litellm import print_verbose from litellm._internal_context import with_service_target +from litellm._logging import verbose_router_logger as _verbose_router_logger from litellm.caching.caching import DualCache, RedisCache +from litellm.caching.redis_cache import ( + RedisCircuitBreakerOpenError as _RedisCircuitBreakerOpenError, +) +from litellm.caching.redis_cache import ( + _is_redis_health_failure, # pyright: ignore[reportPrivateUsage] # RedisCache's connectivity classifier has no public counterpart +) +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 @@ -17,6 +30,25 @@ from litellm.types.llms.base import LiteLLMBaseModel SCHEDULER_QUEUE_TARGET: Final = "scheduler_queue" QueueEntry: TypeAlias = tuple[int, str] _QUEUE_ENTRIES: Final = TypeAdapter(list[QueueEntry]) +_Queue: TypeAlias = tuple[QueueEntry, ...] +_QueueResult: TypeAlias = tuple[bool, _Queue] +_QUEUE_SNAPSHOT: Final = TypeAdapter(_Queue) +_RAW_QUEUE: Final = TypeAdapter[str | bytes | None](str | bytes | None) +_CAS_RESULT: Final = TypeAdapter(bool) +_QUEUE_READ: Final = "return redis.call('GET', KEYS[1])" +_QUEUE_COMPARE_SET: Final = """ +local current = redis.call('GET', KEYS[1]) +if (ARGV[1] == '0' and current ~= false) or + (ARGV[1] == '1' and current ~= ARGV[2]) then + return 0 +end +if ARGV[4] == '' then + redis.call('SET', KEYS[1], ARGV[3]) +else + redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[4]) +end +return 1 +""" class SchedulerCacheKeys(enum.Enum): @@ -30,6 +62,96 @@ class FlowItem(LiteLLMBaseModel): model_name: str +class _LocalQueueCache(_Protocol): + @property + def default_in_memory_ttl(self) -> float | None: ... + + async def async_get_cache(self, key: str, *, local_only: bool) -> object: ... + + async def async_set_cache(self, key: str, value: Sequence[QueueEntry], *, local_only: bool) -> None: ... + + +async def _read_local_queue(cache: _LocalQueueCache, key: str) -> _Queue: + local: Final = await cache.async_get_cache(key=key, local_only=True) + return _QUEUE_SNAPSHOT.validate_python(local or ()) + + +async def _write_local_queue(cache: _LocalQueueCache, key: str, queue: _Queue) -> None: + await cache.async_set_cache(key=key, value=list(queue), local_only=True) + + +def _enqueue(entry: QueueEntry, queue: _Queue) -> _QueueResult: + return False, queue if entry in queue else tuple(sorted((*queue, entry))) + + +def _remove(request_id: str, queue: _Queue) -> _QueueResult: + return False, tuple(sorted(item for item in queue if item[1] != request_id)) + + +def _poll(entry: QueueEntry, queue: _Queue) -> _QueueResult: + pending: Final = tuple(sorted(queue if entry in queue else (*queue, entry))) + return (True, pending[1:]) if pending[0] == entry else (False, pending) + + +def _redis_unavailable(error: Exception) -> bool: + return isinstance(error, _RedisCircuitBreakerOpenError) or _is_redis_health_failure(error) + + +async def _read_redis_queue(cache: RedisCache, key: str) -> tuple[str | bytes | None, _Queue]: + raw: Final = _RAW_QUEUE.validate_python(await cache.async_register_script(_QUEUE_READ)(keys=(key,), args=())) + return raw, () if raw is None else _QUEUE_SNAPSHOT.validate_json(raw) + + +async def _mutate_redis_queue( + cache: RedisCache, key: str, model_name: str, ttl: int | None, change: Callable[[_Queue], _QueueResult] +) -> _QueueResult: + script: Final = cache.async_register_script(_QUEUE_COMPARE_SET) + for _ in range(100): + raw, queue = await _read_redis_queue(cache, key) + admitted, updated = change(queue) + if updated == queue: + return admitted, updated + if _CAS_RESULT.validate_python( + await script( + keys=(key,), + args=( + "0" if raw is None else "1", + raw or "", + _QUEUE_SNAPSHOT.dump_json(updated), + "" if ttl is None else str(ttl), + ), + ) + ): + return admitted, updated + raise Timeout( + message="Scheduler queue remained busy after concurrent updates", model=model_name, llm_provider="openai" + ) + + +async def _shared_queue_result( + cache: _LocalQueueCache, + redis_cache: RedisCache, + key: str, + model_name: str, + change: Callable[[_Queue], _QueueResult], +) -> _QueueResult: + try: + return await _mutate_redis_queue( + redis_cache, + key, + model_name, + redis_cache.default_ttl if cache.default_in_memory_ttl is None else int(cache.default_in_memory_ttl), + change, + ) + except Exception as error: + if not _redis_unavailable(error): + raise + _log_redis_failure( + _verbose_router_logger, _logging.WARNING, "Scheduler using local queue while Redis is unavailable", error + ) + return change(await _read_local_queue(cache, key)) + + class Scheduler: cache: DualCache @@ -49,7 +171,22 @@ class Scheduler: 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 + @with_service_target(SCHEDULER_QUEUE_TARGET) + async def _update_shared_queue( + self, redis_cache: RedisCache, model_name: str, change: Callable[[_Queue], _QueueResult] + ) -> bool: + key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}" + result: Final = await _shared_queue_result(self.cache, redis_cache, key, model_name, change) + await _write_local_queue(self.cache, key, result[1]) + return result[0] + async def add_request(self, request: FlowItem): + redis_cache: Final = self.cache.redis_cache + if isinstance(redis_cache, RedisCache): + await self._update_shared_queue( + redis_cache, request.model_name, _partial(_enqueue, (request.priority, request.request_id)) + ) + return # We use the priority directly, as lower values indicate higher priority # get the queue queue: Final = await self.get_queue(model_name=request.model_name) @@ -78,6 +215,11 @@ class Scheduler: if len(health_deployments) > 0: return True + redis_cache: Final = self.cache.redis_cache + if isinstance(redis_cache, RedisCache): + return await self._update_shared_queue( + redis_cache, request.model_name, _partial(_poll, (request.priority, request.request_id)) + ) queue: Final = await self.get_queue(model_name=request.model_name) entry: Final = (request.priority, request.request_id) print_verbose(f"queue: {queue}, seeking {entry}") @@ -117,6 +259,10 @@ class Scheduler: Remove a specific request from the priority queue for a model. Used when a request times out while waiting in the queue. """ + redis_cache: Final = self.cache.redis_cache + if isinstance(redis_cache, RedisCache): + await self._update_shared_queue(redis_cache, model_name, _partial(_remove, request_id)) + return queue: Final = await self.get_queue(model_name=model_name) filtered_queue: Final = [item for item in queue if item[1] != request_id] heapq.heapify(filtered_queue) # restore heap invariant after filtering @@ -151,6 +297,17 @@ class Scheduler: Redis hands the queue back as JSON lists, so every entry is validated into the (priority, request_id) tuple the heap operations compare against. """ + redis_cache: Final = self.cache.redis_cache + if isinstance(redis_cache, RedisCache): + key: Final = f"{SchedulerCacheKeys.queue.value}:{model_name}" + try: + _, snapshot = await _read_redis_queue(redis_cache, key) + except Exception as error: + if not _redis_unavailable(error): + raise + return list(await _read_local_queue(self.cache, key)) + await _write_local_queue(self.cache, key, snapshot) + return list(snapshot) 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) diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 993da52fc67..35c21c8c3a8 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -3,11 +3,12 @@ from __future__ import annotations import asyncio import http.client import json +import os import re import threading import time import uuid -from collections.abc import Callable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from dataclasses import dataclass @@ -20,17 +21,23 @@ import httpx import openai import psutil import pytest +import pytest_asyncio import yaml from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.async_server import LoopbackServer, close_writer from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, openai_error, responses_reply from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.ports import allocate_port from integration._support.redis_process import OwnedRedis, owned_redis from integration._support.timing import assert_cut_short from integration._support.wire import Reply, Request, Wire, wire_server -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError from redis import Redis +from litellm.caching.redis_cache import RedisCache +from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys + T = TypeVar("T") Endpoint = Literal["chat", "chat-stream", "completions", "completions-stream", "queue", "queue-stream"] Instance = Literal["first", "second"] @@ -851,3 +858,148 @@ def test_sibling_worker_keeps_serving_prioritized_requests_after_a_worker_is_kil seconds=graceful_stop_seconds(), ) assert_once(received_markers(wire), (*(marker for _, marker, _ in outcomes), follow_up)) + + +@pytest_asyncio.fixture(loop_scope="function") +async def shared_scheduler_cache() -> AsyncIterator[RedisCache]: + cache: Final = RedisCache( + host=os.environ["REDIS_HOST"], + port=int(os.environ["REDIS_PORT"]), + namespace="scheduler-atomic-" + uuid.uuid4().hex, + ) + try: + yield cache + finally: + client: Final = cache.init_async_client() + await client.delete(cache.check_and_fix_namespace(f"{SchedulerCacheKeys.queue.value}:shared")) + await client.aclose() + cache.redis_client.close() + + +async def test_shared_scheduler_preserves_concurrent_admissions(shared_scheduler_cache: RedisCache) -> None: + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + await asyncio.gather( + first.add_request(FlowItem(priority=1, request_id="first", model_name="shared")), + second.add_request(FlowItem(priority=2, request_id="second", model_name="shared")), + ) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[1, "first"], [2, "second"]] + assert await first.get_queue("shared") == [(1, "first"), (2, "second")] + assert await second.get_queue("shared") == [(1, "first"), (2, "second")] + + +async def test_shared_scheduler_does_not_resurrect_peer_removal(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "removed"], [2, "waiting"]]) + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + await first.remove_request("removed", "shared") + await second.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[2, "waiting"], [3, "new"]] + + +async def test_shared_scheduler_observes_peer_pop(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "first"], [2, "second"]]) + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + assert await first.poll(FlowItem(priority=1, request_id="first", model_name="shared"), []) + assert await second.poll(FlowItem(priority=2, request_id="second", model_name="shared"), []) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [] + + +async def test_shared_scheduler_keeps_other_requests_during_concurrent_cleanup( + shared_scheduler_cache: RedisCache, +) -> None: + schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(16)) + await asyncio.gather( + *( + scheduler.add_request(FlowItem(priority=i, request_id=str(i), model_name="shared")) + for i, scheduler in enumerate(schedulers) + ) + ) + await asyncio.gather( + *(scheduler.remove_request(str(i), "shared") for i, scheduler in enumerate(schedulers) if i % 2 == 0) + ) + assert await schedulers[0].get_queue("shared") == [(i, str(i)) for i in range(16) if i % 2] + + +@pytest.fixture +def interrupted_scheduler_cache(shared_scheduler_cache: RedisCache) -> Iterator[tuple[RedisCache, threading.Event]]: + interrupted: Final = threading.Event() + + async def relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + upstream_reader, upstream_writer = await asyncio.open_connection( + os.environ["REDIS_HOST"], int(os.environ["REDIS_PORT"]) + ) + + async def forward(source: asyncio.StreamReader, destination: asyncio.StreamWriter) -> None: + try: + while chunk := await source.read(65536): + if interrupted.is_set(): + return + destination.write(chunk) + await destination.drain() + except ConnectionError: + return + finally: + await close_writer(destination) + + await asyncio.gather(forward(reader, upstream_writer), forward(upstream_reader, writer)) + + port: Final = allocate_port() + server: Final = LoopbackServer(port, relay) + server.start() + cache: Final = RedisCache(host="127.0.0.1", port=port, namespace=shared_scheduler_cache.namespace) + try: + yield cache, interrupted + finally: + cache.redis_client.close() + server.stop() + + +async def test_shared_scheduler_keeps_local_fallback_out_of_recovered_shared_queue( + shared_scheduler_cache: RedisCache, + interrupted_scheduler_cache: tuple[RedisCache, threading.Event], +) -> None: + cache, interrupted = interrupted_scheduler_cache + scheduler: Final = Scheduler(redis_cache=cache) + try: + await scheduler.add_request(FlowItem(priority=1, request_id="removed", model_name="shared")) + interrupted.set() + await scheduler.add_request(FlowItem(priority=2, request_id="local", model_name="shared")) + assert await scheduler.get_queue("shared") == [(1, "removed"), (2, "local")] + await Scheduler(redis_cache=shared_scheduler_cache).remove_request("removed", "shared") + interrupted.clear() + await scheduler.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) + assert await scheduler.get_queue("shared") == [(3, "new")] + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[3, "new"]] + finally: + interrupted.clear() + await cache.init_async_client().aclose() + + +async def test_shared_scheduler_preserves_invalid_queue_for_diagnosis(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"}) + scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) + with pytest.raises(ValidationError): + await scheduler.add_request(FlowItem(priority=1, request_id="first", model_name="shared")) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"} + + +@pytest.mark.parametrize("priority", [0, -1, 2**80]) +async def test_shared_scheduler_retains_priority_and_request_id( + shared_scheduler_cache: RedisCache, priority: int +) -> None: + scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) + scheduler.cache.update_cache_ttl(default_in_memory_ttl=17, default_redis_ttl=None) + await scheduler.add_request(FlowItem(priority=priority, request_id="请求-é", model_name="shared")) + assert await scheduler.get_queue("shared") == [(priority, "请求-é")] + assert ( + 0 + < shared_scheduler_cache.redis_client.ttl( + shared_scheduler_cache.check_and_fix_namespace("scheduler:queue:shared") + ) + <= 17 + ) From e3760f87174c4fb7fa7fff24ba619e5b68934d41 Mon Sep 17 00:00:00 2001 From: xyi74976-del <290725798+xyi74976-del@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:49:07 +0800 Subject: [PATCH 2/2] test(scheduler): collect shared queue coverage in Redis CI --- .circleci/scripts/classify_changes.sh | 2 +- .github/workflows/test-unit.yml | 23 +++ .../test_priority_scheduler_queue_cleanup.py | 156 +----------------- tests/integration/run.py | 1 + .../integration/sdk/test_dual_cache_redis.py | 152 +++++++++++++++++ tests/unit/test_circleci_path_filter.py | 3 + 6 files changed, 182 insertions(+), 155 deletions(-) diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index d2ab3478397..b470d110e43 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -14,7 +14,7 @@ outside_cost_map_set=false while IFS= read -r file || [ -n "$file" ]; do [ -n "$file" ] || continue case "$file" in - litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .github/workflows/test-unit.yml | .github/actions/detect-changes/* | .github/scripts/detect_changes.sh | .circleci/scripts/classify_changes.sh | tests/integration/run.py | pyproject.toml | uv.lock) + litellm/scheduler.py | tests/integration/sdk/test_dual_cache_redis.py | litellm/caching/dual_cache.py | litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .github/workflows/test-unit.yml | .github/actions/detect-changes/* | .github/scripts/detect_changes.sh | .circleci/scripts/classify_changes.sh | tests/integration/run.py | pyproject.toml | uv.lock) has_redis_compat=true ;; esac case "$file" in diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 0558875bdad..c3025b9dade 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -1366,6 +1366,16 @@ jobs: pull-requests: read runs-on: ubuntu-latest timeout-minutes: 25 + services: + redis: + image: redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 strategy: fail-fast: false matrix: @@ -1454,6 +1464,19 @@ jobs: --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \ --cov=./litellm --cov-report=xml:coverage.xml + - name: Run shared Redis SDK regressions + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + env: + REDIS_HOST: "127.0.0.1" + REDIS_PORT: "6379" + PYTHONPATH: ".:tests" + COVERAGE_CORE: sysmon + run: | + uv run --no-sync pytest tests/integration/sdk/test_dual_cache_redis.py \ + --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 \ + --cov=./litellm --cov-append --cov-report=xml:coverage.xml + - name: Upload redis-py ${{ matrix.redis-py }} coverage to Codecov if: always() && steps.changes.outputs.decision != 'skip' && matrix.redis-py == '5.3.1' continue-on-error: true diff --git a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py index 35c21c8c3a8..993da52fc67 100644 --- a/tests/integration/routing/test_priority_scheduler_queue_cleanup.py +++ b/tests/integration/routing/test_priority_scheduler_queue_cleanup.py @@ -3,12 +3,11 @@ from __future__ import annotations import asyncio import http.client import json -import os import re import threading import time import uuid -from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack from dataclasses import dataclass @@ -21,23 +20,17 @@ import httpx import openai import psutil import pytest -import pytest_asyncio import yaml from integration._support.client import Gateway, eventually, gateway_from_environment -from integration._support.async_server import LoopbackServer, close_writer from integration._support.database import read_rows from integration._support.openai_wire import answering_model_discovery, chat_reply, openai_error, responses_reply from integration._support.process import graceful_stop_seconds, owned_proxy_process -from integration._support.ports import allocate_port from integration._support.redis_process import OwnedRedis, owned_redis from integration._support.timing import assert_cut_short from integration._support.wire import Reply, Request, Wire, wire_server -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter from redis import Redis -from litellm.caching.redis_cache import RedisCache -from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys - T = TypeVar("T") Endpoint = Literal["chat", "chat-stream", "completions", "completions-stream", "queue", "queue-stream"] Instance = Literal["first", "second"] @@ -858,148 +851,3 @@ def test_sibling_worker_keeps_serving_prioritized_requests_after_a_worker_is_kil seconds=graceful_stop_seconds(), ) assert_once(received_markers(wire), (*(marker for _, marker, _ in outcomes), follow_up)) - - -@pytest_asyncio.fixture(loop_scope="function") -async def shared_scheduler_cache() -> AsyncIterator[RedisCache]: - cache: Final = RedisCache( - host=os.environ["REDIS_HOST"], - port=int(os.environ["REDIS_PORT"]), - namespace="scheduler-atomic-" + uuid.uuid4().hex, - ) - try: - yield cache - finally: - client: Final = cache.init_async_client() - await client.delete(cache.check_and_fix_namespace(f"{SchedulerCacheKeys.queue.value}:shared")) - await client.aclose() - cache.redis_client.close() - - -async def test_shared_scheduler_preserves_concurrent_admissions(shared_scheduler_cache: RedisCache) -> None: - first: Final = Scheduler(redis_cache=shared_scheduler_cache) - second: Final = Scheduler(redis_cache=shared_scheduler_cache) - await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) - await asyncio.gather( - first.add_request(FlowItem(priority=1, request_id="first", model_name="shared")), - second.add_request(FlowItem(priority=2, request_id="second", model_name="shared")), - ) - assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[1, "first"], [2, "second"]] - assert await first.get_queue("shared") == [(1, "first"), (2, "second")] - assert await second.get_queue("shared") == [(1, "first"), (2, "second")] - - -async def test_shared_scheduler_does_not_resurrect_peer_removal(shared_scheduler_cache: RedisCache) -> None: - await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "removed"], [2, "waiting"]]) - first: Final = Scheduler(redis_cache=shared_scheduler_cache) - second: Final = Scheduler(redis_cache=shared_scheduler_cache) - await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) - await first.remove_request("removed", "shared") - await second.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) - assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[2, "waiting"], [3, "new"]] - - -async def test_shared_scheduler_observes_peer_pop(shared_scheduler_cache: RedisCache) -> None: - await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "first"], [2, "second"]]) - first: Final = Scheduler(redis_cache=shared_scheduler_cache) - second: Final = Scheduler(redis_cache=shared_scheduler_cache) - await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) - assert await first.poll(FlowItem(priority=1, request_id="first", model_name="shared"), []) - assert await second.poll(FlowItem(priority=2, request_id="second", model_name="shared"), []) - assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [] - - -async def test_shared_scheduler_keeps_other_requests_during_concurrent_cleanup( - shared_scheduler_cache: RedisCache, -) -> None: - schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(16)) - await asyncio.gather( - *( - scheduler.add_request(FlowItem(priority=i, request_id=str(i), model_name="shared")) - for i, scheduler in enumerate(schedulers) - ) - ) - await asyncio.gather( - *(scheduler.remove_request(str(i), "shared") for i, scheduler in enumerate(schedulers) if i % 2 == 0) - ) - assert await schedulers[0].get_queue("shared") == [(i, str(i)) for i in range(16) if i % 2] - - -@pytest.fixture -def interrupted_scheduler_cache(shared_scheduler_cache: RedisCache) -> Iterator[tuple[RedisCache, threading.Event]]: - interrupted: Final = threading.Event() - - async def relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - upstream_reader, upstream_writer = await asyncio.open_connection( - os.environ["REDIS_HOST"], int(os.environ["REDIS_PORT"]) - ) - - async def forward(source: asyncio.StreamReader, destination: asyncio.StreamWriter) -> None: - try: - while chunk := await source.read(65536): - if interrupted.is_set(): - return - destination.write(chunk) - await destination.drain() - except ConnectionError: - return - finally: - await close_writer(destination) - - await asyncio.gather(forward(reader, upstream_writer), forward(upstream_reader, writer)) - - port: Final = allocate_port() - server: Final = LoopbackServer(port, relay) - server.start() - cache: Final = RedisCache(host="127.0.0.1", port=port, namespace=shared_scheduler_cache.namespace) - try: - yield cache, interrupted - finally: - cache.redis_client.close() - server.stop() - - -async def test_shared_scheduler_keeps_local_fallback_out_of_recovered_shared_queue( - shared_scheduler_cache: RedisCache, - interrupted_scheduler_cache: tuple[RedisCache, threading.Event], -) -> None: - cache, interrupted = interrupted_scheduler_cache - scheduler: Final = Scheduler(redis_cache=cache) - try: - await scheduler.add_request(FlowItem(priority=1, request_id="removed", model_name="shared")) - interrupted.set() - await scheduler.add_request(FlowItem(priority=2, request_id="local", model_name="shared")) - assert await scheduler.get_queue("shared") == [(1, "removed"), (2, "local")] - await Scheduler(redis_cache=shared_scheduler_cache).remove_request("removed", "shared") - interrupted.clear() - await scheduler.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) - assert await scheduler.get_queue("shared") == [(3, "new")] - assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[3, "new"]] - finally: - interrupted.clear() - await cache.init_async_client().aclose() - - -async def test_shared_scheduler_preserves_invalid_queue_for_diagnosis(shared_scheduler_cache: RedisCache) -> None: - await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"}) - scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) - with pytest.raises(ValidationError): - await scheduler.add_request(FlowItem(priority=1, request_id="first", model_name="shared")) - assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"} - - -@pytest.mark.parametrize("priority", [0, -1, 2**80]) -async def test_shared_scheduler_retains_priority_and_request_id( - shared_scheduler_cache: RedisCache, priority: int -) -> None: - scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) - scheduler.cache.update_cache_ttl(default_in_memory_ttl=17, default_redis_ttl=None) - await scheduler.add_request(FlowItem(priority=priority, request_id="请求-é", model_name="shared")) - assert await scheduler.get_queue("shared") == [(priority, "请求-é")] - assert ( - 0 - < shared_scheduler_cache.redis_client.ttl( - shared_scheduler_cache.check_and_fix_namespace("scheduler:queue:shared") - ) - <= 17 - ) diff --git a/tests/integration/run.py b/tests/integration/run.py index 7ab8d59c02c..f7e5abe7b1f 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -28,6 +28,7 @@ GROUPS: Final = MappingProxyType( GITHUB_FILES: Final = frozenset( { "tests/integration/database/test_roi_observed.py", + "tests/integration/sdk/test_dual_cache_redis.py", "tests/integration/database/image/test_component_image_serves_offline.py", "tests/integration/database/image/test_image_admin_mcp.py", "tests/integration/database/image/test_image_bedrock_realtime_extra.py", diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index ddd01480d36..e189a90c687 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/tests/integration/sdk/test_dual_cache_redis.py @@ -1,14 +1,21 @@ import asyncio import os +import threading import uuid +from collections.abc import AsyncIterator, Iterator from typing import Final from unittest.mock import patch import pytest +import pytest_asyncio +from integration._support.async_server import LoopbackServer, close_writer +from integration._support.ports import allocate_port +from pydantic import ValidationError from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache +from litellm.scheduler import FlowItem, Scheduler, SchedulerCacheKeys def _redis_cache() -> RedisCache: @@ -92,3 +99,148 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi assert in_loop_results == [[expected[key]] for key in in_loop_keys] assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] + + +@pytest_asyncio.fixture(loop_scope="function") +async def shared_scheduler_cache() -> AsyncIterator[RedisCache]: + cache: Final = RedisCache( + host=os.environ["REDIS_HOST"], + port=int(os.environ["REDIS_PORT"]), + namespace="scheduler-atomic-" + uuid.uuid4().hex, + ) + try: + yield cache + finally: + client: Final = cache.init_async_client() + await client.delete(cache.check_and_fix_namespace(f"{SchedulerCacheKeys.queue.value}:shared")) + await client.aclose() + cache.redis_client.close() + + +async def test_shared_scheduler_preserves_concurrent_admissions(shared_scheduler_cache: RedisCache) -> None: + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + await asyncio.gather( + first.add_request(FlowItem(priority=1, request_id="first", model_name="shared")), + second.add_request(FlowItem(priority=2, request_id="second", model_name="shared")), + ) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[1, "first"], [2, "second"]] + assert await first.get_queue("shared") == [(1, "first"), (2, "second")] + assert await second.get_queue("shared") == [(1, "first"), (2, "second")] + + +async def test_shared_scheduler_does_not_resurrect_peer_removal(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "removed"], [2, "waiting"]]) + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + await first.remove_request("removed", "shared") + await second.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[2, "waiting"], [3, "new"]] + + +async def test_shared_scheduler_observes_peer_pop(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", [[1, "first"], [2, "second"]]) + first: Final = Scheduler(redis_cache=shared_scheduler_cache) + second: Final = Scheduler(redis_cache=shared_scheduler_cache) + await asyncio.gather(first.get_queue("shared"), second.get_queue("shared")) + assert await first.poll(FlowItem(priority=1, request_id="first", model_name="shared"), []) + assert await second.poll(FlowItem(priority=2, request_id="second", model_name="shared"), []) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [] + + +async def test_shared_scheduler_keeps_other_requests_during_concurrent_cleanup( + shared_scheduler_cache: RedisCache, +) -> None: + schedulers: Final = tuple(Scheduler(redis_cache=shared_scheduler_cache) for _ in range(16)) + await asyncio.gather( + *( + scheduler.add_request(FlowItem(priority=i, request_id=str(i), model_name="shared")) + for i, scheduler in enumerate(schedulers) + ) + ) + await asyncio.gather( + *(scheduler.remove_request(str(i), "shared") for i, scheduler in enumerate(schedulers) if i % 2 == 0) + ) + assert await schedulers[0].get_queue("shared") == [(i, str(i)) for i in range(16) if i % 2] + + +@pytest.fixture +def interrupted_scheduler_cache(shared_scheduler_cache: RedisCache) -> Iterator[tuple[RedisCache, threading.Event]]: + interrupted: Final = threading.Event() + + async def relay(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + upstream_reader, upstream_writer = await asyncio.open_connection( + os.environ["REDIS_HOST"], int(os.environ["REDIS_PORT"]) + ) + + async def forward(source: asyncio.StreamReader, destination: asyncio.StreamWriter) -> None: + try: + while chunk := await source.read(65536): + if interrupted.is_set(): + return + destination.write(chunk) + await destination.drain() + except ConnectionError: + return + finally: + await close_writer(destination) + + await asyncio.gather(forward(reader, upstream_writer), forward(upstream_reader, writer)) + + port: Final = allocate_port() + server: Final = LoopbackServer(port, relay) + server.start() + cache: Final = RedisCache(host="127.0.0.1", port=port, namespace=shared_scheduler_cache.namespace) + try: + yield cache, interrupted + finally: + cache.redis_client.close() + server.stop() + + +async def test_shared_scheduler_keeps_local_fallback_out_of_recovered_shared_queue( + shared_scheduler_cache: RedisCache, + interrupted_scheduler_cache: tuple[RedisCache, threading.Event], +) -> None: + cache, interrupted = interrupted_scheduler_cache + scheduler: Final = Scheduler(redis_cache=cache) + try: + await scheduler.add_request(FlowItem(priority=1, request_id="removed", model_name="shared")) + interrupted.set() + await scheduler.add_request(FlowItem(priority=2, request_id="local", model_name="shared")) + assert await scheduler.get_queue("shared") == [(1, "removed"), (2, "local")] + await Scheduler(redis_cache=shared_scheduler_cache).remove_request("removed", "shared") + interrupted.clear() + await scheduler.add_request(FlowItem(priority=3, request_id="new", model_name="shared")) + assert await scheduler.get_queue("shared") == [(3, "new")] + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == [[3, "new"]] + finally: + interrupted.clear() + await cache.init_async_client().aclose() + + +async def test_shared_scheduler_preserves_invalid_queue_for_diagnosis(shared_scheduler_cache: RedisCache) -> None: + await shared_scheduler_cache.async_set_cache("scheduler:queue:shared", {"invalid": "queue"}) + scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) + with pytest.raises(ValidationError): + await scheduler.add_request(FlowItem(priority=1, request_id="first", model_name="shared")) + assert await shared_scheduler_cache.async_get_cache("scheduler:queue:shared") == {"invalid": "queue"} + + +@pytest.mark.parametrize("priority", [0, -1, 2**80]) +async def test_shared_scheduler_retains_priority_and_request_id( + shared_scheduler_cache: RedisCache, priority: int +) -> None: + scheduler: Final = Scheduler(redis_cache=shared_scheduler_cache) + scheduler.cache.update_cache_ttl(default_in_memory_ttl=17, default_redis_ttl=None) + await scheduler.add_request(FlowItem(priority=priority, request_id="请求-é", model_name="shared")) + assert await scheduler.get_queue("shared") == [(priority, "请求-é")] + assert ( + 0 + < shared_scheduler_cache.redis_client.ttl( + shared_scheduler_cache.check_and_fix_namespace("scheduler:queue:shared") + ) + <= 17 + ) diff --git a/tests/unit/test_circleci_path_filter.py b/tests/unit/test_circleci_path_filter.py index 0856115e39b..395cd4fa991 100644 --- a/tests/unit/test_circleci_path_filter.py +++ b/tests/unit/test_circleci_path_filter.py @@ -122,6 +122,9 @@ def test_classify_decisions(category: str, changed: list[str], expected: str) -> ("changed", "expected"), ( ("litellm/caching/redis_cache.py", "run"), + ("litellm/scheduler.py", "run"), + ("litellm/caching/dual_cache.py", "run"), + ("tests/integration/sdk/test_dual_cache_redis.py", "run"), ("tests/unit/caching/test_redis_cluster_cache.py", "run"), ("tests/integration/sdk/test_redis_cluster_iam_auth.py", "run"), (".github/workflows/test-unit.yml", "run"),