mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(cache): Fix Redis cluster caching (#23480)
* fix redis cluster startup_nodes check order * add tests for redis cluster startup_nodes fix
This commit is contained in:
cohml
GitHub
parent
0be8085564
commit
3eeb14bf1a
2 changed files with 160 additions and 25 deletions
|
|
@ -222,8 +222,12 @@ def _get_redis_client_logic(**env_overrides):
|
|||
"REDIS_CLUSTER_NODES"
|
||||
)
|
||||
|
||||
# If startup_nodes resolved to None (not set by kwarg or env), remove the key
|
||||
# entirely so callers can rely on key presence as a reliable cluster-mode signal.
|
||||
if _startup_nodes is not None and isinstance(_startup_nodes, str):
|
||||
redis_kwargs["startup_nodes"] = json.loads(_startup_nodes)
|
||||
elif _startup_nodes is None:
|
||||
redis_kwargs.pop("startup_nodes", None)
|
||||
|
||||
_sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
|
||||
"REDIS_SENTINEL_NODES"
|
||||
|
|
@ -273,10 +277,14 @@ def _get_redis_client_logic(**env_overrides):
|
|||
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
redis_kwargs.pop("host", None)
|
||||
redis_kwargs.pop("port", None)
|
||||
redis_kwargs.pop("db", None)
|
||||
redis_kwargs.pop("password", None)
|
||||
# Only strip host/port/db/password when not routing to a cluster.
|
||||
# When startup_nodes is also present the cluster path takes priority and
|
||||
# needs the password for authentication.
|
||||
if not redis_kwargs.get("startup_nodes"):
|
||||
redis_kwargs.pop("host", None)
|
||||
redis_kwargs.pop("port", None)
|
||||
redis_kwargs.pop("db", None)
|
||||
redis_kwargs.pop("password", None)
|
||||
elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None:
|
||||
pass
|
||||
elif (
|
||||
|
|
@ -368,6 +376,10 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
|||
|
||||
def get_redis_client(**env_overrides):
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
return init_redis_cluster(redis_kwargs)
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
args = _get_redis_url_kwargs()
|
||||
url_kwargs = {}
|
||||
|
|
@ -377,9 +389,6 @@ def get_redis_client(**env_overrides):
|
|||
|
||||
return redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
if "startup_nodes" in redis_kwargs or get_secret("REDIS_CLUSTER_NODES") is not None: # type: ignore
|
||||
return init_redis_cluster(redis_kwargs)
|
||||
|
||||
# Check for Redis Sentinel
|
||||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_redis_sentinel(redis_kwargs)
|
||||
|
|
@ -392,21 +401,6 @@ def get_redis_async_client(
|
|||
**env_overrides,
|
||||
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
if connection_pool is not None:
|
||||
return async_redis.Redis(connection_pool=connection_pool)
|
||||
args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
|
||||
url_kwargs = {}
|
||||
for arg in redis_kwargs:
|
||||
if arg in args:
|
||||
url_kwargs[arg] = redis_kwargs[arg]
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(
|
||||
arg
|
||||
)
|
||||
)
|
||||
return async_redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
from redis.cluster import ClusterNode
|
||||
|
|
@ -469,6 +463,22 @@ def get_redis_async_client(
|
|||
|
||||
return cluster_client
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
if connection_pool is not None:
|
||||
return async_redis.Redis(connection_pool=connection_pool)
|
||||
args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
|
||||
url_kwargs = {}
|
||||
for arg in redis_kwargs:
|
||||
if arg in args:
|
||||
url_kwargs[arg] = redis_kwargs[arg]
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(
|
||||
arg
|
||||
)
|
||||
)
|
||||
return async_redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
# Check for Redis Sentinel
|
||||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_async_redis_sentinel(redis_kwargs)
|
||||
|
|
@ -482,9 +492,15 @@ def get_redis_async_client(
|
|||
)
|
||||
|
||||
|
||||
def get_redis_connection_pool(**env_overrides):
|
||||
def get_redis_connection_pool(
|
||||
**env_overrides,
|
||||
) -> Optional[async_redis.BlockingConnectionPool]:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
return None
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
pool_kwargs = {
|
||||
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
|
||||
|
|
@ -504,7 +520,6 @@ def get_redis_connection_pool(**env_overrides):
|
|||
connection_class = async_redis.SSLConnection
|
||||
redis_kwargs.pop("ssl", None)
|
||||
redis_kwargs["connection_class"] = connection_class
|
||||
redis_kwargs.pop("startup_nodes", None)
|
||||
return async_redis.BlockingConnectionPool(
|
||||
timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,15 @@
|
|||
from litellm._redis import get_redis_url_from_environment, _get_redis_cluster_kwargs, get_redis_async_client
|
||||
from litellm._redis import (
|
||||
get_redis_url_from_environment,
|
||||
_get_redis_cluster_kwargs,
|
||||
get_redis_async_client,
|
||||
get_redis_client,
|
||||
get_redis_connection_pool,
|
||||
)
|
||||
import json
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
|
||||
def test_get_redis_url_from_environment_single_url(monkeypatch):
|
||||
|
|
@ -167,3 +175,115 @@ def test_get_redis_async_client_without_connection_pool():
|
|||
# Verify Redis was called without connection_pool in kwargs
|
||||
call_kwargs = mock_redis.call_args[1]
|
||||
assert "connection_pool" not in call_kwargs, "connection_pool should not be in kwargs when not provided"
|
||||
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
def test_sync_client_prefers_cluster_over_url(mock_init_cluster, monkeypatch):
|
||||
"""
|
||||
Test get_redis_client returns RedisCluster when startup_nodes is present even if
|
||||
REDIS_URL is also set.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
mock_init_cluster.return_value = MagicMock(spec=redis.RedisCluster)
|
||||
|
||||
startup_nodes = [{"host": "cluster-node.example.com", "port": 6379}]
|
||||
get_redis_client(startup_nodes=startup_nodes)
|
||||
|
||||
mock_init_cluster.assert_called_once()
|
||||
call_kwargs = mock_init_cluster.call_args[0][0]
|
||||
assert (
|
||||
"startup_nodes" in call_kwargs
|
||||
), "startup_nodes must be forwarded to init_redis_cluster"
|
||||
|
||||
@patch("litellm._redis.async_redis.RedisCluster")
|
||||
def test_async_client_prefers_cluster_over_url(mock_cluster_cls, monkeypatch):
|
||||
"""
|
||||
Test (1) get_redis_async_client returns async RedisCluster when startup_nodes is present
|
||||
even if REDIS_URL is also set and (2) startup_nodes is forwarded to RedisCluster.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
|
||||
startup_nodes = [{"host": "cluster-node.example.com", "port": 6379}]
|
||||
get_redis_async_client(startup_nodes=startup_nodes)
|
||||
|
||||
mock_cluster_cls.assert_called_once()
|
||||
call_kwargs = mock_cluster_cls.call_args[1]
|
||||
assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster"
|
||||
assert len(call_kwargs["startup_nodes"]) == 1, "should forward exactly 1 cluster node"
|
||||
|
||||
|
||||
@patch("litellm._redis.async_redis.RedisCluster")
|
||||
def test_async_client_prefers_cluster_over_url_via_env_var(mock_cluster_cls, monkeypatch):
|
||||
"""
|
||||
Test get_redis_async_client returns async RedisCluster when REDIS_CLUSTER_NODES is set
|
||||
even if REDIS_URL is also set.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
monkeypatch.setenv(
|
||||
"REDIS_CLUSTER_NODES",
|
||||
json.dumps([{"host": "cluster-node.example.com", "port": 6379}]),
|
||||
)
|
||||
|
||||
get_redis_async_client()
|
||||
|
||||
mock_cluster_cls.assert_called_once()
|
||||
call_kwargs = mock_cluster_cls.call_args[1]
|
||||
assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to async RedisCluster"
|
||||
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
def test_sync_client_prefers_cluster_over_url_via_env_var(mock_init_cluster, monkeypatch):
|
||||
"""
|
||||
Test get_redis_client returns RedisCluster when REDIS_CLUSTER_NODES is set even if
|
||||
REDIS_URL is also set.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
monkeypatch.setenv(
|
||||
"REDIS_CLUSTER_NODES",
|
||||
json.dumps([{"host": "cluster-node.example.com", "port": 6379}]),
|
||||
)
|
||||
mock_init_cluster.return_value = MagicMock(spec=redis.RedisCluster)
|
||||
|
||||
get_redis_client()
|
||||
|
||||
mock_init_cluster.assert_called_once()
|
||||
call_kwargs = mock_init_cluster.call_args[0][0]
|
||||
assert "startup_nodes" in call_kwargs, "startup_nodes must be forwarded to init_redis_cluster"
|
||||
assert len(call_kwargs["startup_nodes"]) == 1
|
||||
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
def test_sync_client_preserves_password_for_cluster_when_url_also_set(mock_init_cluster, monkeypatch):
|
||||
"""
|
||||
Test _get_redis_client_logic does not strip password from redis_kwargs when
|
||||
startup_nodes is present even if REDIS_URL is also set.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
monkeypatch.setenv("REDIS_PASSWORD", "secret")
|
||||
mock_init_cluster.return_value = MagicMock(spec=redis.RedisCluster)
|
||||
|
||||
startup_nodes = [{"host": "cluster-node.example.com", "port": 6379}]
|
||||
get_redis_client(startup_nodes=startup_nodes)
|
||||
|
||||
mock_init_cluster.assert_called_once()
|
||||
call_kwargs = mock_init_cluster.call_args[0][0]
|
||||
assert "password" in call_kwargs, "password must not be stripped when routing to cluster"
|
||||
assert call_kwargs["password"] == "secret"
|
||||
|
||||
|
||||
def test_connection_pool_returns_none_for_cluster(monkeypatch):
|
||||
"""Test get_redis_connection_pool returns None when startup_nodes is present."""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://fallback-host:6379")
|
||||
startup_nodes = [{"host": "cluster-node.example.com", "port": 6379}]
|
||||
result = get_redis_connection_pool(startup_nodes=startup_nodes)
|
||||
assert result is None, "connection pool must be None for cluster mode"
|
||||
|
||||
|
||||
@patch("litellm._redis.redis.Redis.from_url")
|
||||
def test_sync_client_url_used_when_no_cluster(mock_from_url, monkeypatch):
|
||||
"""
|
||||
Test get_redis_client default to using URL path when no startup_nodes are provided.
|
||||
"""
|
||||
monkeypatch.setenv("REDIS_URL", "redis://plain-host:6379")
|
||||
monkeypatch.delenv("REDIS_CLUSTER_NODES", raising=False)
|
||||
|
||||
get_redis_client()
|
||||
|
||||
mock_from_url.assert_called_once()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue