fix(redis): reset only the failed node on a cluster client timeout, not the whole client

litellm/_redis.py

@@ -665,8 +665,16 @@ def get_redis_async_client(
         cluster_kwargs.setdefault("health_check_interval", REDIS_CLUSTER_HEALTH_CHECK_INTERVAL)
         cluster_kwargs.setdefault("socket_keepalive", True)
 
+        # A single node's client-side timeout must reset only that node's connections,
+        # not tear down the whole cluster client for every concurrent caller.
+        from litellm.caching.redis_cluster_node_isolation import (
+            get_litellm_async_redis_cluster_class,
+        )
+
+        async_redis_cluster_class: Final = get_litellm_async_redis_cluster_class()
+
         # Create async RedisCluster with IAM token as password if available
-        cluster_client: Final = async_redis.RedisCluster(
+        cluster_client: Final = async_redis_cluster_class(
             startup_nodes=new_startup_nodes,
             **cluster_kwargs,
         )

litellm/caching/redis_cluster_node_isolation.py

@@ -0,0 +1,173 @@
+"""Bounds the blast radius of a single node's transient connection error on the async
+Redis Cluster client.
+
+redis-py's ``RedisCluster._execute_command`` responds to a ``ConnectionError`` or
+``TimeoutError`` on ANY one node by tearing down every node's connections and flipping
+the client into "needs reinitialization", which forces every other concurrent caller
+sharing this client through one reinit lock until the whole cluster topology is
+re-walked. Under real proxy load, a client-side socket timeout on a single node is a
+routine event (the event loop was too busy to read the response before ``socket_timeout``
+elapsed) and does not mean the cluster's topology moved, so treating it as a full-cluster
+event turns one slow node into a proxy-wide latency spike while Redis itself stays
+healthy -- confirmed live: pausing one of three local cluster nodes made every concurrent
+command against the other two, untouched nodes stall for the full pause duration too.
+
+``get_litellm_async_redis_cluster_class`` returns a ``RedisCluster`` subclass that resets
+only the node that actually failed (mirroring what a plain, non-cluster Redis client
+already does when one of its pooled connections errors), leaving every other node's
+connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered,
+retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
+topology changed.
+"""
+
+import asyncio
+from typing import TYPE_CHECKING, Final, Protocol
+
+from litellm._logging import verbose_logger
+
+if TYPE_CHECKING:
+    from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType
+
+
+class _ClusterNodeAttrs(Protocol):
+    """The subset of ``redis.asyncio.cluster.ClusterNode`` this override reads. redis-py
+    ships no resolvable stub for these members under the repo's current types-redis pin,
+    so a plain attribute access resolves every downstream use to ``Unknown`` under strict
+    mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
+    own logic fully typed without a banned ``typing.cast``."""
+
+    async def execute_command(
+        self,
+        *args: object,
+        **kwargs: object,  # kwargs-ok: mirrors redis-py's own ClusterNode.execute_command signature, a raw command dispatch with no fixed keyword contract
+    ) -> object: ...
+    async def disconnect(self) -> None: ...
+
+
+class _NodesManagerAttrs(Protocol):
+    _moved_exception: object
+
+    def get_node_from_slot(
+        self, slot: int, read_from_replicas: bool, load_balancing_strategy: object
+    ) -> _ClusterNodeAttrs: ...
+
+
+class _ClusterAttrs(Protocol):
+    RedisClusterRequestTTL: int
+    reinitialize_counter: int
+    reinitialize_steps: int
+    read_from_replicas: bool
+    load_balancing_strategy: object
+    nodes_manager: _NodesManagerAttrs
+
+    def get_node(self, node_name: str) -> _ClusterNodeAttrs: ...
+    async def _determine_slot(self, *args: object) -> int: ...
+    async def aclose(self) -> None: ...
+
+
+#: redis-py versions this override's copied ``_execute_command`` body has been verified
+#: against. A version outside this set may have changed the method's structure in a way
+#: 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"})
+
+
+def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
+    """Builds the ``RedisCluster`` subclass with the per-node isolation fix.
+
+    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
+    submodules are cached in ``sys.modules`` after the first import.
+    """
+    import redis
+    from redis.asyncio.cluster import (
+        RedisCluster as _BaseAsyncRedisCluster,  # pyright: ignore[reportUnknownVariableType]  # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
+    )
+    from redis.cluster import get_node_name
+    from redis.commands import READ_COMMANDS
+    from redis.exceptions import (
+        AskError,
+        BusyLoadingError,
+        ClusterDownError,
+        ClusterError,
+        MaxConnectionsError,
+        MovedError,
+        SlotNotCoveredError,
+        TryAgainError,
+    )
+    from redis.exceptions import ConnectionError as _RedisConnectionError
+    from redis.exceptions import TimeoutError as _RedisTimeoutError
+
+    if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
+        verbose_logger.warning(
+            "redis-py %s is not in the set this cluster-teardown-storm fix was verified "
+            "against (%s). The per-node-isolation override may not match the installed library's "
+            "real _execute_command behavior.",
+            redis.__version__,
+            sorted(_VERIFIED_REDIS_VERSIONS),
+        )
+
+    class LiteLLMAsyncRedisCluster(
+        _BaseAsyncRedisCluster  # pyright: ignore[reportUntypedBaseClass]  # same stale-stub gap as the import above; the base class itself is unresolvable, not this subclass's own code
+    ):
+        async def _execute_command(
+            self,
+            target_node: _ClusterNodeAttrs,
+            *args: object,
+            **kwargs: object,  # kwargs-ok: overrides redis-py's own **kwargs signature; the keyword contract is defined by the Redis command being dispatched, not by this method
+        ) -> object:
+            cluster: _ClusterAttrs = self
+            node = target_node
+
+            asking = moved = False
+            redirect_addr: str | None = None
+            ttl = cluster.RedisClusterRequestTTL
+
+            while ttl > 0:
+                ttl -= 1
+                try:
+                    if asking:
+                        assert redirect_addr is not None
+                        node = cluster.get_node(node_name=redirect_addr)
+                        await node.execute_command("ASKING")
+                        asking = False
+                    elif moved:
+                        slot = await cluster._determine_slot(*args)  # pyright: ignore[reportPrivateUsage]  # mirrors upstream's own un-overridden branch, which makes this identical private call from the same subclass
+                        node = cluster.nodes_manager.get_node_from_slot(
+                            slot,
+                            cluster.read_from_replicas and args[0] in READ_COMMANDS,
+                            (cluster.load_balancing_strategy if args[0] in READ_COMMANDS else None),
+                        )
+                        moved = False
+
+                    return await node.execute_command(*args, **kwargs)
+                except (BusyLoadingError, MaxConnectionsError):
+                    raise
+                except (_RedisConnectionError, _RedisTimeoutError):
+                    # Reset only the node that actually failed instead of the upstream
+                    # default (`await self.aclose()`, a full-cluster teardown that forces
+                    # every other concurrent caller through the shared reinit lock).
+                    await node.disconnect()
+                    raise
+                except (ClusterDownError, SlotNotCoveredError):
+                    await cluster.aclose()
+                    await asyncio.sleep(0.25)
+                    raise
+                except MovedError as e:
+                    cluster.reinitialize_counter += 1
+                    if cluster.reinitialize_steps and cluster.reinitialize_counter % cluster.reinitialize_steps == 0:
+                        await cluster.aclose()
+                        cluster.reinitialize_counter = 0
+                    else:
+                        cluster.nodes_manager._moved_exception = e  # pyright: ignore[reportPrivateUsage]  # mirrors upstream's own un-overridden branch; redis-py exposes no public setter for this
+                    moved = True
+                except AskError as e:
+                    redirect_addr = get_node_name(host=e.host, port=e.port)
+                    asking = True
+                except TryAgainError:
+                    if ttl < cluster.RedisClusterRequestTTL / 2:
+                        await asyncio.sleep(0.05)
+
+            raise ClusterError("TTL exhausted.")
+
+    return LiteLLMAsyncRedisCluster

tests/test_litellm/caching/test_redis_cluster_node_isolation.py

@@ -0,0 +1,133 @@
+"""Regression: a single cluster node's ConnectionError/TimeoutError must reset only that
+node's connections, not tear down the whole cluster client for every other concurrent
+caller. Live confirmation against a real 3-master local cluster (pausing one node with
+CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nodes
+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."""
+
+from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock
+
+import pytest
+from redis.exceptions import (
+    BusyLoadingError,
+    ClusterDownError,
+    MaxConnectionsError,
+    MovedError,
+)
+from redis.exceptions import (
+    ConnectionError as RedisConnectionError,
+)
+from redis.exceptions import TimeoutError as RedisTimeoutError
+
+from litellm.caching.redis_cluster_node_isolation import (
+    get_litellm_async_redis_cluster_class,
+)
+
+if TYPE_CHECKING:
+    from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType
+
+
+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)
+        self.disconnect = AsyncMock()
+
+
+class _FakeNodesManager:
+    def __init__(self, node_to_return: _FakeClusterNode) -> None:
+        self._moved_exception: object = None
+        self._node_to_return = node_to_return
+
+    def get_node_from_slot(
+        self, slot: int, read_from_replicas: bool, load_balancing_strategy: object
+    ) -> _FakeClusterNode:
+        return self._node_to_return
+
+
+def _build_cluster_instance() -> "_AsyncRedisClusterType":
+    cluster_cls = get_litellm_async_redis_cluster_class()
+    instance = cluster_cls.__new__(cluster_cls)
+    instance.RedisClusterRequestTTL = 1
+    instance.reinitialize_counter = 0
+    instance.reinitialize_steps = 5
+    instance.read_from_replicas = False
+    instance.load_balancing_strategy = None
+    instance.aclose = AsyncMock()
+    return instance
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("error_cls", [RedisConnectionError, RedisTimeoutError])
+async def test_node_level_error_resets_only_that_node_not_the_whole_client(error_cls: type[Exception]) -> None:
+    """The fix: a ConnectionError/TimeoutError must disconnect only the failing node
+    and must NOT call the client-wide aclose() that tears down every node."""
+    target_node = _FakeClusterNode("node-a", raises=error_cls("boom"))
+    instance = _build_cluster_instance()
+
+    with pytest.raises(error_cls):
+        await instance._execute_command(target_node, "GET", "k")
+
+    target_node.disconnect.assert_awaited_once()
+    instance.aclose.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None:
+    target_node = _FakeClusterNode("node-a", response=b"v")
+    instance = _build_cluster_instance()
+
+    result = await instance._execute_command(target_node, "GET", "k")
+
+    assert result == b"v"
+    target_node.disconnect.assert_not_awaited()
+    instance.aclose.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("error_cls", [BusyLoadingError, MaxConnectionsError])
+async def test_busy_loading_and_max_connections_reraise_without_any_reset(error_cls: type[Exception]) -> None:
+    """Unchanged from upstream: these say nothing about node health, so neither the
+    node nor the client should be reset."""
+    target_node = _FakeClusterNode("node-a", raises=error_cls("boom"))
+    instance = _build_cluster_instance()
+
+    with pytest.raises(error_cls):
+        await instance._execute_command(target_node, "GET", "k")
+
+    target_node.disconnect.assert_not_awaited()
+    instance.aclose.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_cluster_down_error_still_triggers_a_full_reinit() -> None:
+    """Unchanged from upstream: ClusterDownError is real evidence the topology
+    changed, so a full-client reinit (unlike a plain timeout) is still correct here."""
+    target_node = _FakeClusterNode("node-a", raises=ClusterDownError("boom"))
+    instance = _build_cluster_instance()
+
+    with pytest.raises(ClusterDownError):
+        await instance._execute_command(target_node, "GET", "k")
+
+    instance.aclose.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_moved_error_still_triggers_reinit_after_reinitialize_steps() -> None:
+    """Unchanged from upstream: repeated MOVED responses are real evidence of a
+    slot migration, so they should still force a full reinit every `reinitialize_steps`."""
+    target_node = _FakeClusterNode("node-a", raises=MovedError("1 127.0.0.1:7001"))
+    instance = _build_cluster_instance()
+    instance.reinitialize_steps = 1
+    instance.RedisClusterRequestTTL = 2
+    instance.nodes_manager = _FakeNodesManager(node_to_return=target_node)
+    instance._determine_slot = AsyncMock(return_value=0)
+
+    target_node.execute_command = AsyncMock(side_effect=[MovedError("1 127.0.0.1:7001"), b"v"])
+
+    result = await instance._execute_command(target_node, "GET", "k")
+
+    assert result == b"v"
+    instance.aclose.assert_awaited_once()
+    assert instance.reinitialize_counter == 0