fix(router): give the rust core's priority scheduler the router's redis cache (#50265)

The Rust core built its Scheduler without a Redis cache, so the priority queue lived only in each worker's memory and was never shared across workers or instances. Build it on the routing cache's Redis cache, as the Python core did

litellm/router_rust/invoke/endpoints.py

@@ -1907,10 +1907,9 @@ class RustEndpoints:
         if self._scheduler is None:
             from litellm.scheduler import Scheduler
 
-            self._scheduler = (
-                Scheduler(polling_interval=self._polling_interval)  # pyright: ignore[reportArgumentType]  # as configured
-                if self._polling_interval is not None
-                else Scheduler()
+            self._scheduler = Scheduler(
+                polling_interval=self._polling_interval,  # pyright: ignore[reportArgumentType]  # as configured
+                redis_cache=self._core.routing_cache.redis_cache,
             )
         return self._scheduler
 

tests/unit/test_router/core/test_construction_precall_characterization.py

@@ -13,6 +13,11 @@ from collections.abc import Callable, Generator, Mapping
 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._logging import verbose_router_logger
@@ -23,6 +28,7 @@ from litellm.router_api import ChangeSet, router_init_args
 from litellm.router_core import RouterCore
 from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
 from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
+from litellm.scheduler import FlowItem
 from litellm.types.router import (
     AlertingConfig,
     AllowedFailsPolicy,
@@ -620,6 +626,41 @@ def test_applying_a_redis_cache_attaches_it_to_the_routing_state_and_the_pins(ma
     assert core.registry().plugins("prompt_caching_pins").redis_cache is redis_cache
 
 
+@pytest.fixture
+def shared_redis(monkeypatch: pytest.MonkeyPatch) -> FakeServer:
+    server: Final = FakeServer()
+
+    class Pool(BlockingConnectionPool):
+        def __init__(self, **kwargs: object) -> None:  # kwargs-ok: forwarded to the redis pool
+            super().__init__(connection_class=FakeConnection, server=server, **kwargs)
+
+    monkeypatch.setattr(redis, "Redis", FakeRedis)
+    monkeypatch.setattr(redis.asyncio, "BlockingConnectionPool", Pool)
+    return server
+
+
+async def test_two_workers_on_one_redis_share_the_priority_queue(make_core: MakeCore, shared_redis: FakeServer) -> None:
+    first: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399)
+    second: Final = _one(make_core, redis_host="scheduler-shared", redis_port=6399)
+    await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g"))  # pyright: ignore[reportAttributeAccessIssue]
+    waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g")
+    await second.scheduler.add_request(waiter)  # pyright: ignore[reportAttributeAccessIssue]
+
+    assert not await second.scheduler.poll(waiter, [])  # pyright: ignore[reportAttributeAccessIssue]
+    await first.scheduler.remove_request("urgent", "g")  # pyright: ignore[reportAttributeAccessIssue]
+    assert await second.scheduler.poll(waiter, [])  # pyright: ignore[reportAttributeAccessIssue]
+
+
+async def test_without_redis_each_worker_keeps_its_own_queue(make_core: MakeCore) -> None:
+    first: Final = _one(make_core)
+    second: Final = _one(make_core)
+    await first.scheduler.add_request(FlowItem(priority=0, request_id="urgent", model_name="g"))  # pyright: ignore[reportAttributeAccessIssue]
+    waiter: Final = FlowItem(priority=5, request_id="waiter", model_name="g")
+    await second.scheduler.add_request(waiter)  # pyright: ignore[reportAttributeAccessIssue]
+
+    assert await second.scheduler.poll(waiter, [])  # pyright: ignore[reportAttributeAccessIssue]
+
+
 def test_cache_responses_installs_a_local_response_cache_on_the_host(make_core: MakeCore, host: ScriptedHost) -> None:
     core: Final = _one(make_core, cache_responses=True)
     assert isinstance(host.cache, litellm.Cache)