diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..0f202a3c13c 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -56,6 +56,7 @@ if TYPE_CHECKING: from redis.asyncio import Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline + from redis.client import Pipeline as SyncPipeline from litellm._service_logger import ServiceLogging @@ -95,6 +96,16 @@ class _AsyncRedisCommands(Protocol): def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + def zrem(self, name: str, *values: str) -> Awaitable[int]: ... + + +class _SyncRedisCommands(Protocol): + """Sync redis commands this cache issues; ``get_redis_client`` is untyped, so call sites need this protocol.""" + + def pipeline(self, transaction: bool = True) -> "SyncPipeline[bytes]": ... + + def zrem(self, name: str, *values: str) -> int: ... + _GENERIC_CALLER_MODULES: Final = frozenset( { @@ -127,26 +138,12 @@ _GENERIC_CALLER_FRAME_NAMES: Final = frozenset( ) _CALL_STACK_END_MODULES: Final = ("asyncio", "concurrent", "threading") -_INCREMENT_WITH_FLOOR_LUA: Final = ( - "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " - "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " - "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " - "return count" -) - _INCREMENT_IF_PRESENT_LUA: Final = ( "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])" ) -_LUA_COUNT: Final = TypeAdapter(int) _LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) -_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) - - -def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: - return _OPTIONAL_COUNTS.validate_python( - tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) - ) +_COUNTS: Final = TypeAdapter(tuple[int, ...]) def _is_generic_caller_frame(frame: FrameType) -> bool: @@ -829,6 +826,9 @@ class RedisCache(BaseCache): def _async_commands(self) -> _AsyncRedisCommands: return self.init_async_client() + def _sync_commands(self) -> _SyncRedisCommands: + return self.redis_client # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # get_redis_client is untyped + def check_and_fix_namespace(self, key: str) -> str: """ Make sure each key starts with the given namespace @@ -954,41 +954,44 @@ class RedisCache(BaseCache): raise e @_redis_circuit_breaker_guard_sync - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. + def track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None: + """Record ``member`` as in flight under ``key`` since ``started_at``, dropping members older than ``max_age``. - A counter whose key expired while a request was still in flight would otherwise be - recreated negative by that request's decrement. Clamping inside the same call is what - keeps it safe: a separate corrective write could land after another pod's increment and - erase it. - - The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was - created rather than ``ttl`` after it was last touched. Refreshing it on every touch - would keep a count a dead worker never decremented alive for as long as the group - takes traffic. Returns the resulting count. + Every request is its own member of a sorted set, so finishing one request can only ever + remove that request. A shared counter with a key-wide TTL forgot every request still in + flight when it expired, and their later decrements made a busy deployment read idle. A + member a dead worker never removed still stops counting ``max_age`` after it started. """ namespaced_key: Final = self.check_and_fix_namespace(key=key) - count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval - _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl - ) - return _LUA_COUNT.validate_python(count) + with self._sync_commands().pipeline(transaction=False) as pipe: + pipe.zremrangebyscore(namespaced_key, "-inf", f"({started_at - max_age}") + pipe.zadd(namespaced_key, {member: started_at}) + pipe.expire(namespaced_key, max_age) + pipe.execute() @_redis_circuit_breaker_guard_sync - def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + def untrack_in_flight(self, key: str, member: str) -> None: + self._sync_commands().zrem(self.check_and_fix_namespace(key=key), member) - ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller - cannot tell apart from "every counter is unset". A caller that has to fall back to its - own numbers when Redis is unreachable needs the failure, not a dict of zeros. + @_redis_circuit_breaker_guard_sync + def batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]: + """Count the members of each key that started at or after ``since``, in order. + + Raises when Redis cannot answer, because a caller that falls back to its own numbers + cannot tell a swallowed failure apart from "nothing is in flight". """ - namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] - return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + with self._sync_commands().pipeline(transaction=False) as pipe: + for key in key_list: + pipe.zcount(self.check_and_fix_namespace(key=key), since, "+inf") + return _COUNTS.validate_python(pipe.execute()) @_redis_circuit_breaker_guard - async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - """Async twin of ``batch_get_counts``, raising on failure the same way.""" - namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] - return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + async def async_batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]: + """Async twin of ``batch_count_in_flight``, raising on failure the same way.""" + async with self._async_commands().pipeline(transaction=False) as pipe: + for key in key_list: + pipe.zcount(self.check_and_fix_namespace(key=key), since, "+inf") + return _COUNTS.validate_python(await pipe.execute()) @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: @@ -1566,12 +1569,18 @@ class RedisCache(BaseCache): return float(result) @_redis_circuit_breaker_guard - async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: - """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" - _redis_client: Final = self._async_commands() + async def async_track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None: + """Async twin of ``track_in_flight``.""" namespaced_key: Final = self.check_and_fix_namespace(key=key) - count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) - return _LUA_COUNT.validate_python(count) + async with self._async_commands().pipeline(transaction=False) as pipe: + pipe.zremrangebyscore(namespaced_key, "-inf", f"({started_at - max_age}") + pipe.zadd(namespaced_key, {member: started_at}) + pipe.expire(namespaced_key, max_age) + await pipe.execute() + + @_redis_circuit_breaker_guard + async def async_untrack_in_flight(self, key: str, member: str) -> None: + await self._async_commands().zrem(self.check_and_fix_namespace(key=key), member) @_redis_circuit_breaker_guard async def async_increment_if_present( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 8e0c3d6a7e7..8efbecbb31d 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,5 +1,8 @@ import logging -from collections.abc import Mapping, Sequence +import time +import uuid +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from typing import Final from pydantic import TypeAdapter, ValidationError @@ -29,6 +32,7 @@ class _LitellmParams(TypedDict, total=False): class _CallKwargs(TypedDict, total=False): + litellm_call_id: ReadOnly[str | None] litellm_params: ReadOnly[_LitellmParams | None] @@ -45,11 +49,22 @@ _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) +@dataclass(frozen=True, slots=True) +class _CallRef: + model_group: str + deployment_id: str + call_id: str | None + + def _request_count_key(model_group: str, deployment_id: str) -> str: return f"{model_group}_request_count:{deployment_id}" -def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: +def _in_flight_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_in_flight:{deployment_id}" + + +def _call_ref(kwargs: Mapping[str, object]) -> _CallRef | None: if is_batch_retrieve_call_type(kwargs.get("call_type")): return None try: @@ -63,13 +78,12 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: deployment_id: Final = model_info.get("id") if model_info else None if model_group is None or deployment_id is None: return None - return model_group, str(deployment_id) + return _CallRef(model_group, str(deployment_id), call.get("litellm_call_id")) -def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: +def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: return tuple( - _request_count_key(model_group, str(deployment["model_info"]["id"])) - for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) ) @@ -116,9 +130,24 @@ class LeastBusyLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache): + def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.time): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) + self._clock = clock + self._worker_id = uuid.uuid4().hex + + def _attempt_member(self, shared_key: str, call_id: str, delta: int) -> str: + """Name this call's sorted-set member, unique per concurrent attempt sharing a ``litellm_call_id``. + + Callers can supply their own call id, so two concurrent calls may share one. Numbering the + attempts per worker keeps one member per call, and finishing either call removes only one. + """ + attempts_key: Final = f"{shared_key}:{call_id}:attempts" + memory: Final = self.router_cache.in_memory_cache + after: Final = int(memory.increment_cache(attempts_key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)) # pyright: ignore[reportUnknownMemberType] # InMemoryCache.increment_cache leaves key untyped + if after <= 0: + memory.delete_cache(attempts_key) + return f"{call_id}:{self._worker_id}:{after if delta > 0 else after + 1}" @with_service_target("router_usage") def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: @@ -160,15 +189,19 @@ class LeastBusyLoggingHandler(CustomLogger): def get_available_deployments( self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) + shared: Final = redis_cache.batch_count_in_flight( + [_in_flight_key(model_group, deployment_id) for deployment_id in deployment_ids], + self._clock() - IN_FLIGHT_COUNT_TTL_SECONDS, + ) except Exception as e: _warn_unreadable(model_group, e) else: return _least_busy(healthy_deployments, shared) + keys: Final = tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids) local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) return _least_busy(healthy_deployments, local) @@ -176,23 +209,27 @@ class LeastBusyLoggingHandler(CustomLogger): async def async_get_available_deployments( self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: - keys: Final = _request_count_keys(model_group, healthy_deployments) + deployment_ids: Final = _deployment_ids(healthy_deployments) redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) + shared: Final = await redis_cache.async_batch_count_in_flight( + [_in_flight_key(model_group, deployment_id) for deployment_id in deployment_ids], + self._clock() - IN_FLIGHT_COUNT_TTL_SECONDS, + ) except Exception as e: _warn_unreadable(model_group, e) else: return _least_busy(healthy_deployments, shared) + keys: Final = tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids) local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) return _least_busy(healthy_deployments, local) def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) + ref: Final = _call_ref(kwargs) if ref is None: return - key: Final = _request_count_key(*ref) + key: Final = _request_count_key(ref.model_group, ref.deployment_id) redis_cache: Final = self.router_cache.redis_cache try: local: Final = self.router_cache.increment_cache( @@ -200,17 +237,22 @@ class LeastBusyLoggingHandler(CustomLogger): ) if local < 0: self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: + if redis_cache is None or ref.call_id is None: return - redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id) + member: Final = self._attempt_member(shared_key, ref.call_id, delta) + if delta > 0: + redis_cache.track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS) + else: + redis_cache.untrack_in_flight(shared_key, member) except Exception as e: _warn_unwritable(key, e) async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: - ref: Final = _deployment_ref(kwargs) + ref: Final = _call_ref(kwargs) if ref is None: return - key: Final = _request_count_key(*ref) + key: Final = _request_count_key(ref.model_group, ref.deployment_id) redis_cache: Final = self.router_cache.redis_cache try: local: Final = await self.router_cache.async_increment_cache( @@ -218,8 +260,13 @@ class LeastBusyLoggingHandler(CustomLogger): ) if local is not None and local < 0: await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - if redis_cache is None: + if redis_cache is None or ref.call_id is None: return - await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + shared_key: Final = _in_flight_key(ref.model_group, ref.deployment_id) + member: Final = self._attempt_member(shared_key, ref.call_id, delta) + if delta > 0: + await redis_cache.async_track_in_flight(shared_key, member, self._clock(), IN_FLIGHT_COUNT_TTL_SECONDS) + else: + await redis_cache.async_untrack_in_flight(shared_key, member) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/integration/routing/test_redis_in_flight_tracking.py b/tests/integration/routing/test_redis_in_flight_tracking.py new file mode 100644 index 00000000000..d4b0c7a8ffc --- /dev/null +++ b/tests/integration/routing/test_redis_in_flight_tracking.py @@ -0,0 +1,61 @@ +import os +import time +import uuid +from collections.abc import Iterator +from typing import Final + +import pytest + +from litellm.caching.redis_cache import RedisCache + +MAX_AGE: Final = 600 + + +@pytest.fixture +def in_flight() -> Iterator[tuple[RedisCache, str, str]]: + cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + key: Final = f"in-flight-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_requests_finishing_after_newer_ones_started_leave_the_newer_ones_counted( + in_flight: tuple[RedisCache, str, str], +) -> None: + cache, key, _ = in_flight + now: Final = time.time() + + for call_id in ("old-1", "old-2", "old-3", "new-1", "new-2", "new-3"): + cache.track_in_flight(key, call_id, now, MAX_AGE) + for call_id in ("old-1", "old-2", "old-3"): + cache.untrack_in_flight(key, call_id) + cache.track_in_flight(key, "new-4", now, MAX_AGE) + + assert cache.batch_count_in_flight([key], since=now - MAX_AGE) == (4,) + + +def test_a_request_older_than_max_age_ages_out_without_touching_the_rest( + in_flight: tuple[RedisCache, str, str], +) -> None: + cache, key, namespaced_key = in_flight + now: Final = time.time() + + cache.track_in_flight(key, "orphan", now - MAX_AGE - 1, MAX_AGE) + cache.track_in_flight(key, "running", now - MAX_AGE + 60, MAX_AGE) + cache.track_in_flight(key, "new", now, MAX_AGE) + + assert cache.batch_count_in_flight([key], since=now - MAX_AGE) == (2,) + assert cache.redis_client.zcard(namespaced_key) == 2 + assert 0 < cache.redis_client.ttl(namespaced_key) <= MAX_AGE + + +@pytest.mark.asyncio +async def test_the_async_twins_share_the_same_state(in_flight: tuple[RedisCache, str, str]) -> None: + cache, key, _ = in_flight + now: Final = time.time() + + await cache.async_track_in_flight(key, "async-1", now, MAX_AGE) + cache.track_in_flight(key, "sync-1", now, MAX_AGE) + await cache.async_untrack_in_flight(key, "sync-1") + + assert await cache.async_batch_count_in_flight([key], since=now - MAX_AGE) == (1,) diff --git a/tests/integration/routing/test_redis_increment_with_floor.py b/tests/integration/routing/test_redis_increment_with_floor.py deleted file mode 100644 index d5535d30715..00000000000 --- a/tests/integration/routing/test_redis_increment_with_floor.py +++ /dev/null @@ -1,69 +0,0 @@ -import os -import uuid -from collections.abc import Iterator -from typing import Final - -import pytest - -from litellm.caching.redis_cache import RedisCache - -TTL: Final = 600 - - -@pytest.fixture -def counter() -> Iterator[tuple[RedisCache, str, str]]: - cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) - key: Final = f"increment-with-floor-{uuid.uuid4()}" - yield cache, key, cache.check_and_fix_namespace(key=key) - cache.delete_cache(key) - - -def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter: tuple[RedisCache, str, str]) -> None: - cache, key, _ = counter - - assert cache.increment_with_floor(key, 3, TTL) == 3 - assert cache.increment_with_floor(key, 2, TTL) == 5 - assert cache.batch_get_counts([key]) == (5,) - - -def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter: tuple[RedisCache, str, str]) -> None: - cache, key, _ = counter - - assert cache.increment_with_floor(key, 1, TTL) == 1 - assert cache.increment_with_floor(key, -5, TTL) == 0 - assert cache.batch_get_counts([key]) == (0,) - - -def test_traffic_never_pushes_a_counters_expiry_back_out(counter: tuple[RedisCache, str, str]) -> None: - cache, key, namespaced_key = counter - - cache.increment_with_floor(key, 1, TTL) - assert cache.redis_client.ttl(namespaced_key) > TTL - 60 - - cache.redis_client.expire(namespaced_key, 30) - cache.increment_with_floor(key, 1, TTL) - - assert cache.redis_client.ttl(namespaced_key) <= 30 - - -def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter: tuple[RedisCache, str, str]) -> None: - cache, key, namespaced_key = counter - - cache.increment_with_floor(key, 1, TTL) - cache.redis_client.expire(namespaced_key, 30) - - assert cache.increment_with_floor(key, -5, TTL) == 0 - assert cache.redis_client.ttl(namespaced_key) <= 30 - - -@pytest.mark.asyncio -async def test_the_async_counter_behaves_the_same_way(counter: tuple[RedisCache, str, str]) -> None: - cache, key, namespaced_key = counter - - assert await cache.async_increment_with_floor(key, 2, TTL) == 2 - assert await cache.async_batch_get_counts([key]) == (2,) - - cache.redis_client.expire(namespaced_key, 30) - - assert await cache.async_increment_with_floor(key, -9, TTL) == 0 - assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..4c52e93acc6 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -506,48 +506,92 @@ def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record( sync_batch_redis_cache.get_cache("lit7468") -def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): +def test_batch_count_in_flight_raises_where_batch_get_cache_reports_a_miss(redis_no_ping: None): """A caller that must fall back when Redis is unreachable needs the failure, not zeros. The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell - apart from "every counter is unset". Least-busy routing read that as an idle deployment and + apart from "nothing is in flight". Least-busy routing read that as an idle deployment and kept sending traffic to it instead of falling back to this worker's own in-flight counts. """ - assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + from fakeredis import FakeRedis, FakeServer + from redis.exceptions import ConnectionError as RedisConnectionError - with pytest.raises(OSError, match="redis unavailable"): - sync_batch_redis_cache.batch_get_counts(["lit7039"]) + unreachable = FakeServer() + unreachable.connected = False + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=FakeRedis(server=unreachable) + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + assert cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(RedisConnectionError): + cache.batch_count_in_flight(["lit7039"], since=0.0) + + +@pytest.fixture +def in_flight_redis() -> Iterator[RedisCache]: + from fakeredis import FakeRedis, FakeServer + from fakeredis.aioredis import FakeRedis as AsyncFakeRedis + + server = FakeServer() + with ( + patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")), + patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=FakeRedis(server=server) + ), + ): + cache = RedisCache(host="127.0.0.1", port=6379, namespace="ns") + with patch.object(cache, "init_async_client", return_value=AsyncFakeRedis(server=server)): + yield cache + + +IN_FLIGHT_START: Final = 1_700_000_000.0 +IN_FLIGHT_MAX_AGE: Final = 600 + + +def test_finishing_a_request_removes_only_that_request(in_flight_redis: RedisCache): + for call_id in ("old-1", "old-2", "old-3", "new-1"): + in_flight_redis.track_in_flight("dep-a", call_id, IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + for call_id in ("old-1", "old-2", "old-3", "old-3", "never-started"): + in_flight_redis.untrack_in_flight("dep-a", call_id) + + assert in_flight_redis.batch_count_in_flight(["dep-a"], since=IN_FLIGHT_START) == (1,) + + +def test_in_flight_counts_come_back_in_key_order_with_unknown_keys_as_zero(in_flight_redis: RedisCache): + in_flight_redis.track_in_flight("dep-a", "a-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-c", "c-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-c", "c-2", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + + assert in_flight_redis.batch_count_in_flight(["dep-c", "dep-b", "dep-a"], since=IN_FLIGHT_START) == (2, 0, 1) + + +def test_requests_that_started_before_the_cutoff_are_not_counted(in_flight_redis: RedisCache): + in_flight_redis.track_in_flight("dep-a", "older", IN_FLIGHT_START - 1, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-a", "at-cutoff", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + + assert in_flight_redis.batch_count_in_flight(["dep-a"], since=IN_FLIGHT_START) == (1,) + + +def test_tracking_a_request_drops_members_older_than_max_age(in_flight_redis: RedisCache): + in_flight_redis.track_in_flight("dep-a", "orphan", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-a", "edge", IN_FLIGHT_START + 1, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-a", "new", IN_FLIGHT_START + 1 + IN_FLIGHT_MAX_AGE, IN_FLIGHT_MAX_AGE) + + assert in_flight_redis.redis_client.zrange("ns:dep-a", 0, -1) == [b"edge", b"new"] + assert 0 < in_flight_redis.redis_client.ttl("ns:dep-a") <= IN_FLIGHT_MAX_AGE @pytest.mark.asyncio -async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): - """Async twin: the async batch read hides the same failure behind an empty dict.""" - failing_client = AsyncMock() - failing_client.mget.side_effect = OSError("redis unavailable") - with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point - "litellm._redis.get_redis_client", return_value=MagicMock() - ): - cache = RedisCache(host="127.0.0.1", port=6379) +async def test_async_in_flight_tracking_shares_the_sync_state(in_flight_redis: RedisCache): + await in_flight_redis.async_track_in_flight("dep-a", "async-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + in_flight_redis.track_in_flight("dep-a", "sync-1", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + await in_flight_redis.async_track_in_flight("dep-b", "async-2", IN_FLIGHT_START, IN_FLIGHT_MAX_AGE) + await in_flight_redis.async_untrack_in_flight("dep-a", "sync-1") - with patch.object(cache, "init_async_client", return_value=failing_client): - assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} - - with pytest.raises(OSError, match="redis unavailable"): - await cache.async_batch_get_counts(["lit7039"]) - - -@pytest.mark.parametrize("stored", [b"3", "3"]) -def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): - """Counters come back positionally, so an unset key has to stay a hole rather than shift the - rest of the row onto the wrong deployments, and a count has to survive whether the client - hands it back as bytes or as text.""" - with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point - "litellm._redis.get_redis_client", return_value=MagicMock() - ): - cache = RedisCache(host="127.0.0.1", port=6379) - cache.redis_client.mget.return_value = [stored, None, b"0"] - - assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + assert await in_flight_redis.async_batch_count_in_flight(["dep-a", "dep-b"], since=IN_FLIGHT_START) == (1, 1) + assert in_flight_redis.batch_count_in_flight(["dep-a", "dep-b"], since=IN_FLIGHT_START) == (1, 1) @pytest.fixture @@ -1033,8 +1077,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() ServiceLogging() - assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( - redis_breaker_state_counts()["closed"] + assert ( + REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) + == (redis_breaker_state_counts()["closed"]) ) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 5b3af388ae5..dc9a7752040 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,19 +1,24 @@ import asyncio import json import logging +import time +from collections.abc import Callable, Iterator from typing import Final +from unittest.mock import patch +import fakeredis +import fakeredis.aioredis import httpx import pytest import respx import litellm from litellm import Router - from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler from litellm.types.utils import ModelResponse @@ -21,102 +26,172 @@ GROUP: Final = "least-busy-group" DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] +START: Final = 1_700_000_000.0 -def _call_kwargs(deployment_id: str) -> dict[str, object]: - return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} +def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]: + return { + "litellm_call_id": call_id, + "litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}, + } -class SharedRedisCounters: - """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when - the key is created, and ordered reads that raise rather than invent a value.""" - +class Clock: def __init__(self) -> None: - self.counts: dict[str, int] = {} - self.ttls: dict[str, int] = {} + self.now = START - def count(self, key: str) -> int | None: - return self.counts.get(key) - - def expire(self, key: str) -> None: - self.counts.pop(key, None) - self.ttls.pop(key, None) - - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - incremented: Final = max(0, self.counts.get(key, 0) + value) - self.counts[key] = incremented - self.ttls.setdefault(key, ttl) - return incremented - - async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: - return self.increment_with_floor(key, value, ttl) - - def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - return tuple(self.counts.get(key) for key in key_list) - - async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - return self.batch_get_counts(key_list) + def __call__(self) -> float: + return self.now -def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: - cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double - return LeastBusyLoggingHandler(router_cache=cache) +@pytest.fixture +def redis_server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() + + +@pytest.fixture +def shared_redis(redis_server: fakeredis.FakeServer) -> Iterator[RedisCache]: + server: Final = redis_server + with ( + patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")), + patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=fakeredis.FakeRedis(server=server) + ), + ): + cache: Final = RedisCache(host="127.0.0.1", port=6379) + with patch.object(cache, "init_async_client", return_value=fakeredis.aioredis.FakeRedis(server=server)): + yield cache + + +class UnreachableRedis: + def __init__(self, error: Exception) -> None: + self.error = error + + def track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None: + raise self.error + + async def async_track_in_flight(self, key: str, member: str, started_at: float, max_age: int) -> None: + raise self.error + + def untrack_in_flight(self, key: str, member: str) -> None: + raise self.error + + async def async_untrack_in_flight(self, key: str, member: str) -> None: + raise self.error + + def batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]: + raise self.error + + async def async_batch_count_in_flight(self, key_list: list[str], since: float) -> tuple[int, ...]: + raise self.error + + +def _worker( + shared: RedisCache | UnreachableRedis | None, clock: Callable[[], float] = time.time +) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # UnreachableRedis is a duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache, clock=clock) @pytest.mark.asyncio -async def test_worker_routes_around_a_request_another_worker_started() -> None: - shared: Final = SharedRedisCounters() - streaming_worker: Final = _worker(shared) - picking_worker: Final = _worker(shared) +async def test_worker_routes_around_a_request_another_worker_started(shared_redis: RedisCache) -> None: + streaming_worker: Final = _worker(shared_redis) + picking_worker: Final = _worker(shared_redis) - picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "quick")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a", "quick"), None, None, None) - streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "stream")) assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B - await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + await streaming_worker.async_log_success_event(_call_kwargs("dep-a", "stream"), None, None, None) assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_sync_pick_reads_the_shared_counts() -> None: - shared: Final = SharedRedisCounters() - streaming_worker: Final = _worker(shared) - picking_worker: Final = _worker(shared) +def test_sync_pick_reads_the_shared_counts(shared_redis: RedisCache) -> None: + streaming_worker: Final = _worker(shared_redis) + picking_worker: Final = _worker(shared_redis) - picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "quick")) + picking_worker.log_success_event(_call_kwargs("dep-a", "quick"), None, None, None) - streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "stream")) assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B - streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + streaming_worker.log_failure_event(_call_kwargs("dep-a", "stream"), None, None, None) assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: - """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that - stuck count an hour after the key was created, which only works while nothing writes the TTL - again: a handler that refreshed it on every touch would keep the count alive for as long as - the group takes traffic, and the deployment would read busier than it is forever.""" - shared: Final = SharedRedisCounters() - worker: Final = _worker(shared) - key: Final = f"{GROUP}_request_count:dep-a" +@pytest.mark.asyncio +async def test_a_request_outliving_the_ttl_cannot_cancel_a_newer_request_when_it_finishes( + shared_redis: RedisCache, +) -> None: + """A shared counter forgot every in-flight request when its key expired, then each late + decrement cancelled a request that was really running, so a busy deployment read idle and + kept getting a full share of new traffic (#48994)""" + clock: Final = Clock() + worker: Final = _worker(shared_redis, clock) - worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "old")) + clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS + 1 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "new")) + await worker.async_log_success_event(_call_kwargs("dep-a", "old"), None, None, None) - assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B - shared.ttls[key] = 5 - worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(key) == 1 - assert shared.ttls == {key: 5} +@pytest.mark.asyncio +async def test_finishing_one_of_several_calls_sharing_a_call_id_keeps_the_others_counted( + shared_redis: RedisCache, +) -> None: + """Callers can pass their own litellm_call_id, so concurrent calls may share one. Tracking the + id itself merged them into one entry, and the first to finish hid the rest from other workers""" + worker: Final = _worker(shared_redis) + other_worker: Final = _worker(shared_redis) + picking_worker: Final = _worker(shared_redis) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + other_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "shared")) + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "b")) + await worker.async_log_success_event(_call_kwargs("dep-a", "shared"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + worker.log_success_event(_call_kwargs("dep-a", "shared"), None, None, None) + await other_worker.async_log_failure_event(_call_kwargs("dep-a", "shared"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +@pytest.mark.asyncio +async def test_a_request_nobody_finished_stops_counting_once_it_is_older_than_the_ttl( + shared_redis: RedisCache, +) -> None: + """A worker that dies mid-request never reports the end. Its request has to age out on its + own, without taking the deployment's other in-flight requests with it""" + clock: Final = Clock() + worker: Final = _worker(shared_redis, clock) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "orphan")) + clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS - 1 + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "running-b")) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "running-a")) + clock.now = START + IN_FLIGHT_COUNT_TTL_SECONDS + 1 + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + await worker.async_log_success_event(_call_kwargs("dep-b", "running-b"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B @pytest.mark.asyncio @@ -134,17 +209,9 @@ async def test_counts_stay_in_memory_without_redis() -> None: assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 -class UnavailableRedis(SharedRedisCounters): - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - raise ConnectionError("redis is down") - - def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - raise ConnectionError("redis is down") - - @pytest.mark.asyncio async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: - worker: Final = _worker(UnavailableRedis()) + worker: Final = _worker(UnreachableRedis(ConnectionError("redis is down"))) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) @@ -156,22 +223,6 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: - shared: Final = SharedRedisCounters() - worker: Final = _worker(shared) - - worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - shared.expire(f"{GROUP}_request_count:dep-a") - worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - - assert shared.count(f"{GROUP}_request_count:dep-a") == 0 - - worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - - assert shared.count(f"{GROUP}_request_count:dep-a") == 1 - assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B - - @pytest.mark.asyncio async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: worker: Final = _worker(None) @@ -188,27 +239,20 @@ async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B -def test_calls_without_a_deployment_are_ignored() -> None: - shared: Final = SharedRedisCounters() - worker: Final = _worker(shared) +def test_calls_without_a_deployment_are_ignored(shared_redis: RedisCache, redis_server: fakeredis.FakeServer) -> None: + worker: Final = _worker(shared_redis) worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) worker.log_pre_api_call(model="m", messages=[], kwargs={}) - assert shared.counts == {} - - -class OpenBreakerRedis(SharedRedisCounters): - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") - - def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: - raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") + assert fakeredis.FakeRedis(server=redis_server).dbsize() == 0 @pytest.mark.asyncio -async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None: - worker: Final = _worker(OpenBreakerRedis()) +async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( + caplog: pytest.LogCaptureFixture, +) -> None: + worker: Final = _worker(UnreachableRedis(RedisCircuitBreakerOpenError("Redis circuit breaker is open"))) with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) @@ -307,6 +351,12 @@ def _chat_stream(_: httpx.Request) -> httpx.Response: ) +async def _open_chat_stream(router: Router) -> CustomStreamWrapper: + return await router.acompletion( + model=ROUTER_GROUP, messages=[{"role": "user", "content": "write a poem"}], stream=True + ) + + @pytest.mark.asyncio async def test_router_spreads_open_chat_streams_across_least_busy_deployments( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -315,12 +365,7 @@ async def test_router_spreads_open_chat_streams_across_least_busy_deployments( route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_chat_stream) router: Final = _least_busy_router() - streams: Final = [ - await router.acompletion( - model=ROUTER_GROUP, messages=[{"role": "user", "content": "write a poem"}], stream=True - ) - for _ in ROUTER_DEPLOYMENT_IDS - ] + streams: Final = [await _open_chat_stream(router) for _ in ROUTER_DEPLOYMENT_IDS] assert _router_counts(router) == {"1": 1, "2": 1, "3": 1} assert sorted(stream._hidden_params["model_id"] for stream in streams) == list(ROUTER_DEPLOYMENT_IDS) @@ -379,9 +424,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts if use_async else router.get_available_deployment(model=ROUTER_GROUP, messages=None) ) - response: Final = await router.acompletion( - model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}] - ) + response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]) for _ in range(10): await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush() @@ -391,3 +434,25 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts assert response._hidden_params["model_id"] == "1" assert route.calls.last.request.headers["authorization"] == "Bearer key-1" assert _router_counts(router) == seeded + + +@pytest.mark.asyncio +async def test_router_spreads_open_streams_across_workers_through_the_shared_counts( + shared_redis: RedisCache, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + keys_used: Final[list[str]] = [] + + def record_key_and_stream(request: httpx.Request) -> httpx.Response: + keys_used.append(request.headers["authorization"]) + return _chat_stream(request) + + respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=record_key_and_stream) + routers: Final = tuple(_least_busy_router() for _ in ROUTER_DEPLOYMENT_IDS) + for router in routers: + router.cache.redis_cache = shared_redis + + for router in routers: + await _open_chat_stream(router) + + assert sorted(keys_used) == [f"Bearer key-{deployment_id}" for deployment_id in ROUTER_DEPLOYMENT_IDS]