diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index ae8c78709d9..0035801018b 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov retry-exhaustion) is unchanged from upstream, since those already carry real evidence the topology changed. -redis-py 8.x fixed this upstream with gentler machinery than this override's -``node.disconnect()`` (which also kills connections other coroutines are mid-operation -on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per -killed connection): it marks in-use connections for reconnect only after their current -operation completes, disconnects only the idle pooled ones, and defers reinitialization -to the outer retry loop. When the installed ``ClusterNode`` has that per-connection -recovery API, the factory returns the base ``RedisCluster`` unmodified. +redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream +still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent +caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full +teardown. For those versions the factory returns a thin wrapper around upstream's +``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError, +a third consecutive timeout on the same node, or a concurrent request from any other command +or ``aclose()`` still reinits). """ import asyncio @@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol): mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's own logic fully typed without a banned ``typing.cast``.""" + name: str + async def execute_command( self, *args: object, @@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol): #: this override can't see (Python won't error -- it'll just run our now-stale copy), so #: construction logs a loud warning rather than silently trusting an unverified copy. _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"}) +_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3 -def get_litellm_async_redis_cluster_class( +def get_litellm_async_redis_cluster_class( # noqa: C901 # supports redis-py version-specific cluster implementations cluster_node_class: type | None = None, + base_cluster_class: type | None = None, ) -> type["_AsyncRedisClusterType"]: - """Returns the base ``RedisCluster`` when the installed redis-py already recovers a - node-level connection error per-connection (8.x+), else builds the ``RedisCluster`` - subclass with the per-node isolation fix for older versions whose upstream branch - tears down the whole cluster client. + """Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already + recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster`` + subclass with the per-node isolation fix for older versions whose upstream branch tears + down the whole cluster client. - ``cluster_node_class`` exists for dependency injection in tests; production callers - leave it unset and the installed ``ClusterNode`` is used. + ``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests; + production callers leave them unset and the installed redis-py classes are used. Imported lazily because this module is reachable from a base ``import litellm`` while redis is not a base dependency. Cheap to call repeatedly: the underlying redis @@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class( from redis.exceptions import TimeoutError as _RedisTimeoutError node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode + base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster if hasattr(node_class, "update_active_connections_for_reconnect"): verbose_logger.debug( - "redis-py %s recovers a node-level connection error per-connection upstream; " - "using the base RedisCluster without litellm's node-isolation override.", + "redis-py %s recovers node connections per-connection upstream; using " + "LiteLLM's timeout-tolerant RedisCluster wrapper.", redis.__version__, ) - return _BaseAsyncRedisCluster + + class LiteLLMAsyncRedisClusterTimeoutTolerant( + base_class # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass] # the injected base class is selected at runtime + ): + def __init__( + self, + *args: object, + **kwargs: object, # kwargs-ok: passes redis-py's constructor kwargs through untouched + ) -> None: + self._litellm_initialize = False + self._litellm_reinit_requests = 0 + self._litellm_tolerated_timeouts = 0 + super().__init__(*args, **kwargs) + self._litellm_consecutive_timeouts: dict[ # mutable-ok: per-node counter updated on the command hot path + str, int + ] = {} + + @property + def _initialize(self) -> bool: + return self._litellm_initialize + + @_initialize.setter + def _initialize(self, value: bool) -> None: + if value: + self._litellm_reinit_requests += 1 + self._litellm_initialize = value + + async def _execute_command( + self, + target_node: _ClusterNodeAttrs, + *args: object, + **kwargs: object, # kwargs-ok: matches redis-py's own command dispatch signature + ) -> object: + outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts + pending_before: Final = self._litellm_initialize + try: + result: Final = await super()._execute_command(target_node, *args, **kwargs) + except _RedisTimeoutError: + timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1 + if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: + self._litellm_consecutive_timeouts.pop(target_node.name, None) + raise + self._litellm_consecutive_timeouts[target_node.name] = timeouts + self._litellm_tolerated_timeouts += 1 + if ( + not pending_before + and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before + ): + self._initialize = False + raise + if self._litellm_consecutive_timeouts: + self._litellm_consecutive_timeouts.pop(target_node.name, None) + return result + + return LiteLLMAsyncRedisClusterTimeoutTolerant if redis.__version__ not in _VERIFIED_REDIS_VERSIONS: verbose_logger.warning( diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index c16ceec8c31..9ebc3673ab6 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -5,20 +5,26 @@ CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nod stalling for the full pause duration before this fix, and zero after -- these tests pin the same behavior at the unit level so it can run without a live Redis Cluster.""" +import asyncio from typing import TYPE_CHECKING -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock, patch import pytest from redis.exceptions import ( + AskError, BusyLoadingError, ClusterDownError, + ClusterError, MaxConnectionsError, MovedError, + TryAgainError, ) from redis.exceptions import ( ConnectionError as RedisConnectionError, ) -from redis.exceptions import TimeoutError as RedisTimeoutError +from redis.exceptions import ( + TimeoutError as RedisTimeoutError, +) from litellm.caching.redis_cluster_node_isolation import ( get_litellm_async_redis_cluster_class, @@ -39,10 +45,35 @@ class _NodeClassWithoutPerConnectionRecovery: class _FakeClusterNode: def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None: self.name = name - self.execute_command = AsyncMock(side_effect=raises, return_value=response) + + async def execute_command(*args: object, **kwargs: object) -> object: + await asyncio.sleep(0) + if raises is not None: + raise raises + return response + + self.execute_command = AsyncMock(side_effect=execute_command) self.disconnect = AsyncMock() +class _Fake8xRedisCluster: + def __init__(self) -> None: + self._initialize = False + + async def _execute_command( + self, target_node: _FakeClusterNode, *args: object, **kwargs: object + ) -> object: + try: + return await target_node.execute_command(*args, **kwargs) + except (RedisConnectionError, RedisTimeoutError): + self._initialize = True + await asyncio.sleep(0) + raise + + async def aclose(self) -> None: + self._initialize = True + + class _FakeNodesManager: def __init__(self, node_to_return: _FakeClusterNode) -> None: self._moved_exception: object = None @@ -68,31 +99,198 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType": return instance -def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None: - """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level - connection error per-connection, the factory must NOT install the copied override, - whose node.disconnect() also kills connections other coroutines are mid-operation on.""" - from redis.asyncio.cluster import RedisCluster - +def _build_8x_cluster_instance() -> _Fake8xRedisCluster: cluster_cls = get_litellm_async_redis_cluster_class( - cluster_node_class=_NodeClassWithPerConnectionRecovery + cluster_node_class=_NodeClassWithPerConnectionRecovery, + base_cluster_class=_Fake8xRedisCluster, + ) + return cluster_cls() + + +def test_unverified_redis_version_logs_warning(caplog: pytest.LogCaptureFixture) -> None: + import redis + + with patch.object(redis, "__version__", "8.0.1"): + get_litellm_async_redis_cluster_class(cluster_node_class=_NodeClassWithoutPerConnectionRecovery) + + assert "not in the set this cluster-teardown-storm fix was verified against" in caplog.text + + +@pytest.mark.asyncio +async def test_single_timeout_does_not_request_topology_reinit() -> None: + error = RedisTimeoutError("timeout") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = error + instance = _build_8x_cluster_instance() + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_connection_error_preserves_upstream_topology_reinit() -> None: + error = RedisConnectionError("connection error") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = error + instance = _build_8x_cluster_instance() + + with pytest.raises(RedisConnectionError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_three_consecutive_timeouts_request_topology_reinit_and_reset_counter() -> None: + errors = [ + RedisTimeoutError("timeout-1"), + RedisTimeoutError("timeout-2"), + RedisTimeoutError("timeout-3"), + ] + fourth_error = RedisTimeoutError("timeout-4") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = [*errors, fourth_error] + instance = _build_8x_cluster_instance() + + for error in errors: + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + + assert instance._initialize is True + instance._initialize = False + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is fourth_error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_success_resets_consecutive_timeout_counter() -> None: + errors = [RedisTimeoutError("timeout-1"), RedisTimeoutError("timeout-2")] + final_error = RedisTimeoutError("timeout-3") + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = [*errors, b"value", final_error] + instance = _build_8x_cluster_instance() + + for error in errors: + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + assert instance._initialize is False + + result = await instance._execute_command(target_node, "GET", "k") + assert result == b"value" + assert instance._initialize is False + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is final_error + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_timeout_counters_are_per_node() -> None: + node_a_errors = [RedisTimeoutError("node-a-1"), RedisTimeoutError("node-a-2")] + node_b_error = RedisTimeoutError("node-b-1") + node_a = _FakeClusterNode("node-a") + node_b = _FakeClusterNode("node-b") + node_a.execute_command.side_effect = node_a_errors + node_b.execute_command.side_effect = node_b_error + instance = _build_8x_cluster_instance() + + for target_node, error in ( + (node_a, node_a_errors[0]), + (node_b, node_b_error), + (node_a, node_a_errors[1]), + ): + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + assert exc_info.value is error + + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_timeout_does_not_clear_concurrent_topology_reinit_request() -> None: + error = RedisTimeoutError("timeout") + instance = _build_8x_cluster_instance() + + async def request_reinit(*args: object, **kwargs: object) -> object: + await instance.aclose() + raise error + + target_node = _FakeClusterNode("node-a") + target_node.execute_command.side_effect = request_reinit + + with pytest.raises(RedisTimeoutError) as exc_info: + await instance._execute_command(target_node, "GET", "k") + + assert exc_info.value is error + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_tolerated_timeout_does_not_erase_concurrent_connection_error_reinit() -> None: + instance = _build_8x_cluster_instance() + failing_node = _FakeClusterNode("node-a", raises=RedisConnectionError("gone")) + slow_node = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow")) + + results = await asyncio.gather( + instance._execute_command(failing_node, "GET", "a"), + instance._execute_command(slow_node, "GET", "b"), + return_exceptions=True, ) - assert cluster_cls is RedisCluster + assert isinstance(results[0], RedisConnectionError) + assert isinstance(results[1], RedisTimeoutError) + assert instance._initialize is True -def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None: - """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(), - so those versions must keep litellm's per-node isolation override.""" - from redis.asyncio.cluster import RedisCluster +@pytest.mark.asyncio +async def test_overlapping_tolerated_timeouts_do_not_request_topology_reinit() -> None: + instance = _build_8x_cluster_instance() + node_a = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow-a")) + node_b = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow-b")) - cluster_cls = get_litellm_async_redis_cluster_class( - cluster_node_class=_NodeClassWithoutPerConnectionRecovery + results = await asyncio.gather( + instance._execute_command(node_a, "GET", "a"), + instance._execute_command(node_b, "GET", "b"), + return_exceptions=True, ) - assert cluster_cls is not RedisCluster - assert issubclass(cluster_cls, RedisCluster) - assert "_execute_command" in cluster_cls.__dict__ + assert all(isinstance(result, RedisTimeoutError) for result in results) + assert instance._initialize is False + + +@pytest.mark.asyncio +async def test_tolerated_timeout_does_not_clear_pending_reinit() -> None: + instance = _build_8x_cluster_instance() + instance._initialize = True + target_node = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow")) + + with pytest.raises(RedisTimeoutError): + await instance._execute_command(target_node, "GET", "k") + + assert instance._initialize is True + + +@pytest.mark.asyncio +async def test_success_returns_value_without_topology_reinit() -> None: + target_node = _FakeClusterNode("node-a", response=b"value") + instance = _build_8x_cluster_instance() + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + assert instance._initialize is False @pytest.mark.asyncio @@ -110,6 +308,48 @@ async def test_node_level_error_resets_only_that_node_not_the_whole_client(error instance.aclose.assert_not_awaited() +@pytest.mark.asyncio +async def test_moved_error_retries_without_full_reinit_before_threshold() -> None: + moved_error = MovedError("1 127.0.0.1:7001") + target_node = _FakeClusterNode("node-a") + target_node.execute_command = AsyncMock(side_effect=[moved_error, b"value"]) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + instance.nodes_manager = _FakeNodesManager(node_to_return=target_node) + instance._determine_slot = AsyncMock(return_value=0) + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + assert instance.nodes_manager._moved_exception is moved_error + instance.aclose.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_ask_error_sends_asking_and_retries_on_redirected_node() -> None: + ask_error = AskError("0 127.0.0.1:7001") + target_node = _FakeClusterNode("node-a") + target_node.execute_command = AsyncMock(side_effect=[ask_error, None, b"value"]) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + instance.get_node = Mock(return_value=target_node) + + result = await instance._execute_command(target_node, "GET", "k") + + assert result == b"value" + instance.get_node.assert_called_once_with(node_name="127.0.0.1:7001") + + +@pytest.mark.asyncio +async def test_try_again_error_exhausts_ttl() -> None: + target_node = _FakeClusterNode("node-a", raises=TryAgainError("try again")) + instance = _build_cluster_instance() + instance.RedisClusterRequestTTL = 2 + + with pytest.raises(ClusterError): + await instance._execute_command(target_node, "GET", "k") + + @pytest.mark.asyncio async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None: target_node = _FakeClusterNode("node-a", response=b"v")