fix(caching): make Redis Cluster commands wait for a pool slot instead of tripping the circuit breaker (#46488)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-14 15:27:00 -07:00 • committed by GitHub
parent f95e466bb7
commit ca2f99edac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 414 additions and 54 deletions

View file

@ -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

View file

@ -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]:
"""

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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)
# ---------------------------------------------------------------------------