mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(caching): keep a node timeout from forcing a cluster-wide topology reinit on redis-py 8.x (#39349)
* fix(caching): keep node timeout from forcing cluster reinit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(caching): describe the 8.x timeout-tolerant wrapper in the module docstring Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(caching): format redis cluster isolation wrapper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): keep concurrent reinit requests when tolerating a node timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(caching): cover redis cluster redirect branches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(caching): let overlapping tolerated timeouts release their own reinit requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
5dd3fdbc3d
commit
fe770700f4
2 changed files with 336 additions and 37 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue