diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index dd1c152a421..43a41180d77 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -40,7 +40,7 @@ from .base_cache import BaseCache if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from redis.asyncio import Redis, RedisCluster + from redis.asyncio import BlockingConnectionPool, Redis, RedisCluster from redis.asyncio.client import Pipeline from redis.asyncio.cluster import ClusterPipeline @@ -48,12 +48,14 @@ if TYPE_CHECKING: cluster_pipeline = ClusterPipeline async_redis_client = Redis async_redis_cluster_client = RedisCluster + async_redis_conn_pool = BlockingConnectionPool Span = Union[_Span, Any] else: pipeline = Any cluster_pipeline = Any async_redis_client = Any async_redis_cluster_client = Any + async_redis_conn_pool = Any Span = Any @@ -96,6 +98,14 @@ def _get_call_stack_info(num_frames: int = 2) -> str: return "unknown" +def _running_event_loop_id() -> int | None: + """Identity of the running event loop, or None when called outside one.""" + try: + return id(asyncio.get_running_loop()) + except RuntimeError: + return None + + class RedisCircuitBreaker: """ Tracks Redis health for a RedisCache instance. @@ -232,7 +242,8 @@ class RedisCache(BaseCache): self.redis_client = get_redis_client(**redis_kwargs) self.redis_async_client: Optional[Union[async_redis_client, async_redis_cluster_client]] = None self.redis_kwargs = redis_kwargs - self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) + self.async_redis_conn_pool: async_redis_conn_pool | None = get_redis_connection_pool(**redis_kwargs) + self._async_conn_pool_loop_id: int | None = _running_event_loop_id() # redis namespaces self.namespace = namespace @@ -331,21 +342,34 @@ class RedisCache(BaseCache): kwargs_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16] return f"async-redis-client-{kwargs_hash}" + def _get_async_conn_pool(self) -> async_redis_conn_pool | None: + """ + Return this instance's connection pool, rebuilding it only when there is no + pool yet or the existing one belongs to a different event loop. + """ + from .._redis import get_redis_connection_pool + + loop_id = _running_event_loop_id() + if self.async_redis_conn_pool is None or loop_id != self._async_conn_pool_loop_id: + self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) + self._async_conn_pool_loop_id = loop_id + return self.async_redis_conn_pool + def init_async_client( self, ) -> Union[async_redis_client, async_redis_cluster_client]: from litellm import in_memory_llm_clients_cache - from .._redis import get_redis_async_client, get_redis_connection_pool + from .._redis import get_redis_async_client cache_key = self._get_async_client_cache_key() cached_client = in_memory_llm_clients_cache.get_cache(key=cache_key) if cached_client is not None: redis_async_client = cast(Union[async_redis_client, async_redis_cluster_client], cached_client) else: - # Create new connection pool and client for current event loop - self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) - redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) + redis_async_client = get_redis_async_client( + connection_pool=self._get_async_conn_pool(), **self.redis_kwargs + ) in_memory_llm_clients_cache.set_cache(key=cache_key, value=redis_async_client) self.redis_async_client = redis_async_client # type: ignore diff --git a/tests/test_litellm/caching/test_redis_connection_pool.py b/tests/test_litellm/caching/test_redis_connection_pool.py index c824d3e7a0e..98d6f001adb 100644 --- a/tests/test_litellm/caching/test_redis_connection_pool.py +++ b/tests/test_litellm/caching/test_redis_connection_pool.py @@ -4,11 +4,13 @@ Regression tests for Redis connection pool leak fixes (RC1-RC5). Tests are pure unit tests — no Redis server required. """ +import asyncio from unittest.mock import AsyncMock, MagicMock, patch import pytest import redis.asyncio as async_redis +import litellm from litellm._redis import get_redis_async_client, get_redis_connection_pool @@ -128,3 +130,109 @@ async def test_disconnect_idempotent(): await cache.disconnect() await cache.disconnect() # should not raise + + +def _make_url_redis_cache(): + """ + RedisCache over a URL config with a real (lazy) connection pool. + + Unlike _make_redis_cache above, the pool is NOT mocked — these tests assert on + pool *identity*, so a single shared mock would hide the behaviour under test. + BlockingConnectionPool opens no sockets until a command runs, so this stays + server-free. + """ + from litellm.caching.redis_cache import RedisCache + + with patch("litellm._redis.get_redis_client", return_value=MagicMock()), patch( + "litellm.caching.redis_cache.RedisCache._setup_health_pings" + ): + return RedisCache(url="redis://localhost:6379/0") + + +@pytest.fixture +def clean_client_cache(): + """Isolate the process-wide async client cache from other tests.""" + litellm.in_memory_llm_clients_cache.cache_dict.clear() + litellm.in_memory_llm_clients_cache.ttl_dict.clear() + yield litellm.in_memory_llm_clients_cache + litellm.in_memory_llm_clients_cache.cache_dict.clear() + litellm.in_memory_llm_clients_cache.ttl_dict.clear() + + +@pytest.mark.asyncio +async def test_client_cache_expiry_reuses_connection_pool(clean_client_cache): + """ + The cached async client expires on a TTL (default 600s). Building a fresh + connection pool for the replacement client abandons a pool of established + connections while the new pool opens its own, so a single RedisCache + transiently holds ~2x max_connections against the server every rotation — + which is how a proxy walks into "max number of clients reached". + + The replacement client must reuse the existing pool. + """ + cache = _make_url_redis_cache() + + first_client = cache.init_async_client() + pool = cache.async_redis_conn_pool + assert pool is not None + + # what TTL expiry does to the entry, on the same event loop + clean_client_cache.cache_dict.clear() + clean_client_cache.ttl_dict.clear() + + second_client = cache.init_async_client() + + assert second_client is not first_client, "expected a new client once the cached entry expired" + assert cache.async_redis_conn_pool is pool, "connection pool was rebuilt instead of reused" + assert second_client.connection_pool is pool, "replacement client did not attach to the existing pool" + + +@pytest.mark.asyncio +async def test_shared_pool_survives_outgoing_client_close(clean_client_cache): + """ + Reusing one pool across clients is only safe if closing the outgoing client + leaves the pool alone. redis-py guarantees this by setting + auto_close_connection_pool=False whenever a pool is passed in; pin it, because + losing it would let an expiring client tear the pool out from under live traffic. + """ + cache = _make_url_redis_cache() + + first_client = cache.init_async_client() + pool = cache.async_redis_conn_pool + + assert first_client.auto_close_connection_pool is False + + clean_client_cache.cache_dict.clear() + clean_client_cache.ttl_dict.clear() + second_client = cache.init_async_client() + + with patch.object(pool, "disconnect", new=AsyncMock()) as pool_disconnect: + await first_client.aclose() + + pool_disconnect.assert_not_awaited() + assert second_client.connection_pool is pool + + +def test_new_event_loop_gets_its_own_connection_pool(clean_client_cache): + """ + A pool's connections are bound to the loop that created them, so a different + event loop must never be handed the previous loop's pool. Guards the reuse + above from over-reaching. + """ + cache = _make_url_redis_cache() + loops = [] # held so a closed loop's id() cannot be recycled by the next one + pools = [] + + async def init_in_current_loop(): + cache.init_async_client() + return cache.async_redis_conn_pool + + for _ in range(3): + loop = asyncio.new_event_loop() + loops.append(loop) + pools.append(loop.run_until_complete(init_in_current_loop())) + + for loop in loops: + loop.close() + + assert len({id(p) for p in pools}) == 3, "a connection pool was shared across event loops"