Merge e3760f8717 into 4b1928387f

.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

.github/workflows/test-unit.yml

@@ -1369,6 +1369,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:
@@ -1457,6 +1467,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

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)

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",

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

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"),