diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b3715333010..6027f3774fd 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -412,12 +412,19 @@ def _redis_timeout_error_types() -> tuple[type, ...]: ``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither. + ``MaxConnectionsError`` is the async cluster client's pool-exhausted signal (its nodes have + no blocking pool, so it is raised instantly instead of after a wait): a busy pool, not a + dead Redis, like the standalone pool's chained wait timeout. """ try: - from redis.exceptions import TimeoutError as RedisTimeoutError + from redis import exceptions as redis_exceptions except ImportError: return (TimeoutError, asyncio.TimeoutError) - return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError) + max_connections_error: Final = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # types-redis 4.6 stubs predate it + ) + return (redis_exceptions.TimeoutError, TimeoutError, asyncio.TimeoutError, max_connections_error) _MAX_EXCEPTION_CAUSE_DEPTH: Final = 20 @@ -604,9 +611,14 @@ def _redis_circuit_breaker_guard(method): """ @functools.wraps(method) - async def wrapper(self, *args, **kwargs): + async def wrapper(self: "RedisCache", *args, **kwargs): + async def guarded() -> object: + return await method(self, *args, **kwargs) + return await _run_under_circuit_breaker( - self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs) + self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator + wrapper.__name__, + lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator ) return wrapper @@ -690,6 +702,9 @@ class RedisCache(BaseCache): else: super().__init__() # defaults to 60s + def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: + return call() + def _setup_health_pings(self): """Setup async and sync health pings for Redis.""" # ASYNC HEALTH PING @@ -1070,7 +1085,9 @@ class RedisCache(BaseCache): litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor) return await executor(keys=keys, args=args, client=client) - return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute) + return await _run_under_circuit_breaker( + self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute) + ) return run_script diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index d0f3c11bf10..5e926fd2ed8 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -5,9 +5,17 @@ Key differences: - RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created """ -from typing import TYPE_CHECKING, Any, Final +import asyncio +import contextvars +import os +import weakref +from collections.abc import Awaitable, Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar +from litellm._logging import verbose_logger from litellm.caching.redis_cache import RedisCache +from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -23,9 +31,100 @@ else: Span = Any +_T: Final = TypeVar("_T") + + +def _connection_limit(max_connections: object) -> int | None: + try: + limit: Final = int(max_connections) # pyright: ignore[reportArgumentType] # config value, str or int + except (TypeError, ValueError): + return None + return limit if limit > 0 else None + + +_ClientSemaphores: TypeAlias = weakref.WeakValueDictionary[str, asyncio.Semaphore] + + +@dataclass(slots=True) +class _SemaphoreRegistry: + """Weak on both layers: a semaphore binds to its loop once a waiter queues on it, so a strong value would keep a + closed loop alive. In-flight ``ClusterConnectionGate.run`` calls hold the semaphore they acquired.""" + + _per_loop: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _ClientSemaphores]" = field( + default_factory=weakref.WeakKeyDictionary, repr=False + ) + + def get(self, client_key: str, limit: int) -> asyncio.Semaphore: + per_loop: Final = self._per_loop.setdefault(asyncio.get_running_loop(), weakref.WeakValueDictionary()) + existing: Final = per_loop.get(client_key) + if existing is not None: + return existing + created: Final = asyncio.Semaphore(limit) + per_loop[client_key] = created + return created + + +_Admission: TypeAlias = tuple[str, "asyncio.Task[object] | None"] +_SEMAPHORES: Final = _SemaphoreRegistry() +_ADMITTED: Final[contextvars.ContextVar[frozenset[_Admission]]] = contextvars.ContextVar( + "redis_cluster_gate_admitted", default=frozenset() +) + + +@dataclass(frozen=True, slots=True) +class ClusterConnectionGate: + """The async ``RedisCluster`` has no blocking pool, so past ``max_connections`` a node raises + ``MaxConnectionsError`` instead of waiting and the circuit breaker opens on a healthy Redis. This gate makes + guarded commands wait for a slot the way the standalone ``BlockingConnectionPool`` does. A command that finds no + slot within ``wait_timeout`` runs anyway, so a genuinely saturated pool still reaches the client's own error + handling and the breaker.""" + + limit: int + wait_timeout: float = REDIS_CONNECTION_POOL_TIMEOUT + + @classmethod + def from_max_connections(cls, max_connections: object) -> "ClusterConnectionGate | None": + limit: Final = _connection_limit(max_connections) + return None if limit is None else cls(limit=limit) + + async def run(self, call: Callable[[], Awaitable[_T]], client_key: str) -> _T: + admission: Final[_Admission] = (client_key, asyncio.current_task()) + if admission in _ADMITTED.get(): + return await call() + semaphore: Final = _SEMAPHORES.get(client_key, self.limit) + try: + await asyncio.wait_for(semaphore.acquire(), timeout=self.wait_timeout) + except asyncio.TimeoutError: + verbose_logger.debug( + "Redis Cluster: no free connection slot within %ss (max_connections=%s), sending anyway", + self.wait_timeout, + self.limit, + ) + return await call() + token: Final = _ADMITTED.set(_ADMITTED.get() | {admission}) + try: + return await call() + finally: + _ADMITTED.reset(token) + semaphore.release() + + class RedisClusterCache(RedisCache): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def __init__(self, *args, max_connections: object = None, connection_gate: object = None, **kwargs): + pool_kwargs: Final = {} if max_connections is None else {"max_connections": max_connections} + super().__init__(*args, **pool_kwargs, **kwargs) + self._connection_gate: Final = ( + connection_gate + if isinstance(connection_gate, ClusterConnectionGate) + else ClusterConnectionGate.from_max_connections( + os.environ.get("REDIS_MAX_CONNECTIONS") if max_connections is None else max_connections + ) + ) + + def _admit_async_call(self, call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]: + if self._connection_gate is None: + return call() + return self._connection_gate.run(call, self._get_async_client_cache_key()) def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index 46600e0bf60..5ebc3bf7336 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -489,6 +489,9 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) self.call_count = 0 + def _admit_async_call(self, call): + return call() + @_redis_circuit_breaker_guard async def boom(self): self.call_count += 1 diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 7cffa62d423..d27645d19fa 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -3,7 +3,7 @@ import time import types from collections.abc import Iterator from datetime import timedelta -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1327,6 +1327,32 @@ async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity(): await client.aclose() +@pytest.mark.asyncio +async def test_cluster_pool_exhaustion_is_a_timeout_failure_not_hard_connectivity(): + """The async cluster client has no blocking pool: a node whose max_connections are all busy raises + MaxConnectionsError at once. That is the cluster shape of a busy pool, so it must sit behind the + timeout duration gate like the standalone pool's wait timeout instead of opening the breaker.""" + from redis import exceptions as redis_exceptions + + max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs + ) + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + async def pool_exhausted() -> None: + raise max_connections_error() + + for _ in range(breaker.failure_threshold * 2): + with pytest.raises(max_connections_error): + await _run_under_circuit_breaker(breaker, "op", pool_exhausted) + + assert breaker.is_open() is False, "an exhausted client-side pool is a timeout, not a dead Redis" + + def test_timeout_classification_follows_the_explicit_cause_chain_only(): from redis.exceptions import ConnectionError as RedisConnectionError diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 68260bc6870..85422819926 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -1,20 +1,22 @@ import asyncio -from importlib import import_module +import gc import json import ssl +import weakref +from collections.abc import Awaitable, Callable +from importlib import import_module +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi.testclient import TestClient from redis.asyncio import Redis, RedisCluster from redis.asyncio.cluster import ClusterNode from redis.asyncio.connection import SSLConnection - -from litellm.caching.redis_cache import RedisCache -from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache @patch("litellm._redis.init_redis_cluster") @@ -73,9 +75,7 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): @patch("litellm._redis.get_redis_connection_pool") @patch("litellm._redis.get_redis_client") @patch.object(import_module("litellm.caching.redis_cache").RedisCache, "_setup_health_pings") -def test_cache_init_creates_cluster_cache_from_env_var( - mock_health, mock_get_client, mock_get_pool, monkeypatch -): +def test_cache_init_creates_cluster_cache_from_env_var(mock_health, mock_get_client, mock_get_pool, monkeypatch): """ Test that Cache() creates RedisClusterCache when REDIS_CLUSTER_NODES env var is set. @@ -153,9 +153,7 @@ def test_cache_init_creates_redis_cache_without_cluster_config( ), ], ) -def test_router_create_redis_cache_cluster_detection( - startup_nodes, env_var, expected_cache_type, monkeypatch -): +def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, expected_cache_type, monkeypatch): """ Test that Router._create_redis_cache() creates RedisClusterCache when either startup_nodes is in config or REDIS_CLUSTER_NODES env var is set. @@ -416,3 +414,186 @@ def test_evicted_async_client_is_closed_only_for_cluster(cluster: bool) -> None: client.aclose.assert_not_awaited() asyncio.run(exercise()) + + +class _InFlightCountingClient: + """Async cluster client double that records how many commands were in flight at once.""" + + def __init__(self, command_seconds: float) -> None: + self.command_seconds = command_seconds + self.in_flight = 0 + self.peak_in_flight = 0 + self.completed = 0 + + async def set(self, *args: object, **kwargs: object) -> bool: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + await asyncio.sleep(self.command_seconds) + self.in_flight -= 1 + self.completed += 1 + return True + + +def _cluster_set_cache_with_client( + client: _InFlightCountingClient, + max_connections: int | None = None, + connection_gate: ClusterConnectionGate | None = None, +) -> Callable[[str, str], Awaitable[object]]: + with patch("litellm._redis.init_redis_cluster"): + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 7000}], + max_connections=max_connections, + connection_gate=connection_gate, + ) + setattr(cache, "init_async_client", MagicMock(return_value=client)) + return cast( + "Callable[[str, str], Awaitable[object]]", cache.async_set_cache + ) # cast-ok: wrapped **kwargs are untyped + + +@pytest.mark.asyncio +async def test_cluster_cache_queues_commands_beyond_max_connections_instead_of_failing() -> None: + """A burst larger than max_connections used to hand every command to the node at once, so redis-py raised + MaxConnectionsError on the overflow and the circuit breaker opened. The overflow now waits for a slot.""" + client = _InFlightCountingClient(command_seconds=0.01) + set_cache = _cluster_set_cache_with_client(client, max_connections=3) + + await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20))) + + assert client.completed == 20 + assert client.peak_in_flight <= 3 + + +@pytest.mark.asyncio +async def test_cluster_cache_without_max_connections_does_not_throttle() -> None: + client = _InFlightCountingClient(command_seconds=0.01) + set_cache = _cluster_set_cache_with_client(client) + + await asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(20))) + + assert client.completed == 20 + assert client.peak_in_flight == 20 + + +@pytest.mark.asyncio +async def test_cluster_cache_sends_after_the_connection_wait_timeout() -> None: + """A command that cannot get a slot within the pool timeout is still sent, so a saturated pool keeps + surfacing through the client's own errors and the breaker rather than stalling forever.""" + client = _InFlightCountingClient(command_seconds=0.2) + set_cache = _cluster_set_cache_with_client( + client, connection_gate=ClusterConnectionGate(limit=1, wait_timeout=0.02) + ) + + await asyncio.wait_for(asyncio.gather(*(set_cache(f"key-{i}", "value") for i in range(3))), timeout=1.0) + + assert client.completed == 3 + assert client.peak_in_flight == 3 + + +@pytest.mark.asyncio +async def test_cluster_caches_sharing_a_client_share_the_connection_cap() -> None: + """init_async_client hands caches with identical settings the same client, so two caches must not each admit + a full max_connections worth of commands against that one pool.""" + client = _InFlightCountingClient(command_seconds=0.01) + first = _cluster_set_cache_with_client(client, max_connections=3) + second = _cluster_set_cache_with_client(client, max_connections=3) + + await asyncio.gather( + *(first(f"first-{i}", "value") for i in range(10)), *(second(f"second-{i}", "value") for i in range(10)) + ) + + assert client.completed == 20 + assert client.peak_in_flight <= 3 + + +@pytest.mark.asyncio +async def test_guarded_command_nested_in_a_guarded_command_reuses_the_slot() -> None: + """batch_cache_write flushes through async_set_cache_pipeline while it is itself admitted. With one slot the + inner command must run on the outer slot instead of waiting for the pool timeout.""" + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + + async def inner() -> str: + return "flushed" + + async def outer() -> str: + return await gate.run(inner, "client-a") + + assert await asyncio.wait_for(gate.run(outer, "client-a"), timeout=1.0) == "flushed" + + +@pytest.mark.asyncio +async def test_nested_guarded_command_on_another_client_waits_for_that_clients_slot() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + client_b_busy = asyncio.Event() + release_b = asyncio.Event() + peak_b = 0 + in_flight_b = 0 + + async def hold_b() -> None: + nonlocal peak_b, in_flight_b + in_flight_b += 1 + peak_b = max(peak_b, in_flight_b) + client_b_busy.set() + await release_b.wait() + in_flight_b -= 1 + + async def touch_b() -> None: + nonlocal peak_b, in_flight_b + in_flight_b += 1 + peak_b = max(peak_b, in_flight_b) + await asyncio.sleep(0) + in_flight_b -= 1 + + async def outer_on_a() -> None: + await client_b_busy.wait() + await gate.run(touch_b, "client-b") + + holder = asyncio.create_task(gate.run(hold_b, "client-b")) + nested = asyncio.create_task(gate.run(outer_on_a, "client-a")) + await client_b_busy.wait() + await asyncio.sleep(0.05) + release_b.set() + await asyncio.gather(holder, nested) + + assert peak_b == 1 + + +@pytest.mark.asyncio +async def test_child_task_spawned_inside_a_guarded_command_does_not_inherit_the_slot() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + peak = 0 + in_flight = 0 + + async def command() -> None: + nonlocal peak, in_flight + in_flight += 1 + peak = max(peak, in_flight) + await asyncio.sleep(0.02) + in_flight -= 1 + + spawned: list[asyncio.Task[None]] = [] + + async def parent() -> None: + spawned.append(asyncio.create_task(gate.run(command, "client-a"))) + await command() + + await gate.run(parent, "client-a") + await asyncio.gather(*spawned) + + assert peak == 1 + + +def test_closed_event_loop_is_collected_after_contended_gate_use() -> None: + gate = ClusterConnectionGate(limit=1, wait_timeout=5.0) + + async def command() -> None: + await asyncio.sleep(0.01) + + async def contend() -> "weakref.ref[asyncio.AbstractEventLoop]": + await asyncio.gather(gate.run(command, "client-a"), gate.run(command, "client-a")) + return weakref.ref(asyncio.get_running_loop()) + + loop_ref = asyncio.run(contend()) + gc.collect() + + assert loop_ref() is None diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index ad86c3c5267..01400384842 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -23,7 +23,7 @@ from __future__ import annotations import asyncio from datetime import datetime -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -830,31 +830,25 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 -@pytest.mark.asyncio -async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( - monkeypatch, -): - """A failing pipeline must invalidate every pending counter so the next - request reseeds from the DB (which already holds this request's cost) - instead of trusting a value the write may have partially applied.""" - from redis.exceptions import MaxConnectionsError +_PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( + "spend:key:hashed", + "spend:team:team-1", + "spend:team_member:user-1:team-1", + "spend:user:user-1", + "spend:end_user:eu-1", + "spend:tag:tag-a", + "spend:tag:tag-b", + "spend:org:org-1", +) + +def _warm_counters_behind_a_failing_pipeline(monkeypatch, pipeline_error: Exception) -> AsyncMock: counter_cache = ps.DualCache() - pending_keys = ( - "spend:key:hashed", - "spend:team:team-1", - "spend:team_member:user-1:team-1", - "spend:user:user-1", - "spend:end_user:eu-1", - "spend:tag:tag-a", - "spend:tag:tag-b", - "spend:org:org-1", - ) - for key in pending_keys: + for key in _PENDING_SPEND_COUNTER_KEYS: counter_cache.in_memory_cache.set_cache(key=key, value=1.0) fake_redis = AsyncMock() fake_redis.async_get_cache = AsyncMock(return_value=1.0) # counters warm - fake_redis.async_increment_pipeline = AsyncMock(side_effect=MaxConnectionsError()) + fake_redis.async_increment_pipeline = AsyncMock(side_effect=pipeline_error) fake_redis.async_increment = AsyncMock() fake_redis.async_delete_cache = AsyncMock() fake_redis.get_ttl = MagicMock(return_value=None) @@ -862,23 +856,63 @@ async def test_increment_spend_counters_pipeline_failure_invalidates_all_counter monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) monkeypatch.setattr(ps, "prisma_client", None) + return fake_redis - with pytest.raises(MaxConnectionsError): - await ps.increment_spend_counters( - token="hashed", - team_id="team-1", - user_id="user-1", - response_cost=0.5, - org_id="org-1", - end_user_id="eu-1", - tags=["tag-a", "tag-b"], - ) +async def _increment_every_spend_counter() -> None: + await ps.increment_spend_counters( + token="hashed", + team_id="team-1", + user_id="user-1", + response_cost=0.5, + org_id="org-1", + end_user_id="eu-1", + tags=["tag-a", "tag-b"], + ) + + +def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None: assert fake_redis.async_increment.await_count == 0 deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} - assert deleted_keys == set(pending_keys) - for key in pending_keys: - assert counter_cache.in_memory_cache.get_cache(key=key) is None + assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS) + for key in _PENDING_SPEND_COUNTER_KEYS: + assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( + monkeypatch, +): + """A failing pipeline must invalidate every pending counter so the next + request reseeds from the DB (which already holds this request's cost) + instead of trusting a value the write may have partially applied.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, RedisConnectionError("Redis is down")) + + with pytest.raises(RedisConnectionError): + await _increment_every_spend_counter() + + _assert_every_pending_counter_invalidated(fake_redis) + + +@pytest.mark.asyncio +async def test_increment_spend_counters_pool_exhaustion_invalidates_all_counters_without_failing( + monkeypatch, +): + """A cluster node with every connection busy raises MaxConnectionsError at once. That is a busy pool, the + same as the standalone pool's wait timeout, so the counters are invalidated and the request still succeeds.""" + from redis import exceptions as redis_exceptions + + max_connections_error = cast( # cast-ok: types-redis 4.6 stubs predate MaxConnectionsError + "type[Exception]", + redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs + ) + fake_redis = _warm_counters_behind_a_failing_pipeline(monkeypatch, max_connections_error()) + + await _increment_every_spend_counter() + + _assert_every_pending_counter_invalidated(fake_redis) # ---------------------------------------------------------------------------