mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 08:23:17 +00:00
fix: Apply max_connections configuration to Redis async client (#15797)
* fix: Apply max_connections configuration to Redis async client - Add max_connections to available Redis cluster kwargs - Add connection_pool parameter to get_redis_async_client() - Pass connection_pool to Redis client if provided - Prevents Redis connection exhaustion under high load * test: Add tests for Redis max_connections feature - Test max_connections is included in cluster kwargs - Test connection_pool parameter is properly passed to async client - Test async client works without connection_pool parameter All 3 tests pass successfully
This commit is contained in:
+6
-1
@@ -78,6 +78,7 @@ def _get_redis_cluster_kwargs(client=None):
|
||||
available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection
|
||||
available_args.append("gcp_service_account")
|
||||
available_args.append("gcp_ssl_ca_certs")
|
||||
available_args.append("max_connections")
|
||||
|
||||
return available_args
|
||||
|
||||
@@ -376,7 +377,7 @@ def get_redis_client(**env_overrides):
|
||||
|
||||
|
||||
def get_redis_async_client(
|
||||
**env_overrides,
|
||||
connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **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:
|
||||
@@ -447,6 +448,10 @@ def get_redis_async_client(
|
||||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_async_redis_sentinel(redis_kwargs)
|
||||
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
|
||||
|
||||
if connection_pool is not None:
|
||||
redis_kwargs["connection_pool"] = connection_pool
|
||||
|
||||
return async_redis.Redis(
|
||||
**redis_kwargs,
|
||||
)
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from litellm._redis import get_redis_url_from_environment
|
||||
from litellm._redis import get_redis_url_from_environment, _get_redis_cluster_kwargs, get_redis_async_client
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
import redis.asyncio as async_redis
|
||||
|
||||
def test_get_redis_url_from_environment_single_url(monkeypatch):
|
||||
"""Test when REDIS_URL is directly provided"""
|
||||
@@ -117,3 +119,51 @@ def test_get_redis_url_from_environment_missing_port(monkeypatch):
|
||||
|
||||
# Check the error message
|
||||
assert "Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified" in str(excinfo.value)
|
||||
|
||||
def test_max_connections_in_cluster_kwargs():
|
||||
"""Test that max_connections is included in Redis cluster kwargs"""
|
||||
kwargs = _get_redis_cluster_kwargs()
|
||||
assert "max_connections" in kwargs, "max_connections should be in available Redis cluster kwargs"
|
||||
|
||||
def test_get_redis_async_client_with_connection_pool():
|
||||
"""Test that connection_pool parameter is properly passed to Redis client"""
|
||||
# Create a mock connection pool
|
||||
mock_pool = MagicMock(spec=async_redis.BlockingConnectionPool)
|
||||
|
||||
# Mock the Redis client creation
|
||||
with patch('litellm._redis.async_redis.Redis') as mock_redis, \
|
||||
patch('litellm._redis._get_redis_client_logic') as mock_logic:
|
||||
|
||||
# Configure mock to return basic redis kwargs
|
||||
mock_logic.return_value = {
|
||||
"host": "localhost",
|
||||
"port": 6379,
|
||||
"db": 0
|
||||
}
|
||||
|
||||
# Call get_redis_async_client with connection_pool
|
||||
get_redis_async_client(connection_pool=mock_pool)
|
||||
|
||||
# Verify Redis was called with connection_pool in kwargs
|
||||
call_kwargs = mock_redis.call_args[1]
|
||||
assert "connection_pool" in call_kwargs, "connection_pool should be passed to Redis client"
|
||||
assert call_kwargs["connection_pool"] == mock_pool, "connection_pool should match the provided pool"
|
||||
|
||||
def test_get_redis_async_client_without_connection_pool():
|
||||
"""Test that Redis client works without connection_pool parameter"""
|
||||
with patch('litellm._redis.async_redis.Redis') as mock_redis, \
|
||||
patch('litellm._redis._get_redis_client_logic') as mock_logic:
|
||||
|
||||
# Configure mock to return basic redis kwargs
|
||||
mock_logic.return_value = {
|
||||
"host": "localhost",
|
||||
"port": 6379,
|
||||
"db": 0
|
||||
}
|
||||
|
||||
# Call get_redis_async_client without connection_pool
|
||||
get_redis_async_client()
|
||||
|
||||
# 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"
|
||||
|
||||
Reference in New Issue
Block a user