From daa2863e89a482aecb66fc30643f59cb66d25a7d Mon Sep 17 00:00:00 2001 From: mayuriphad Date: Tue, 18 Aug 2026 17:30:56 +0530 Subject: [PATCH] fix(caching): guard RedisCache.disconnect against None connection pool in cluster mode get_redis_connection_pool() returns None by design when startup_nodes is present, since RedisCluster builds its own per-node pools rather than sharing a BlockingConnectionPool. RedisCache.disconnect() dereferenced that unconditionally, so every proxy shutdown in Redis cluster mode raised AttributeError and aborted the rest of proxy_shutdown_event(), skipping the jwt handler close, db writer close, billing metrics flush, and langfuse flush that run after the cache disconnect call. Guards the pool dereference, gives RedisClusterCache its own disconnect() override that tears down the cluster client via aclose(), and wraps the proxy_shutdown_event() cache disconnect call in try/except so one cache teardown failure can no longer abort the rest of shutdown. Fixes #37137 --- litellm/caching/redis_cache.py | 3 +- litellm/caching/redis_cluster_cache.py | 5 +++ litellm/proxy/proxy_server.py | 5 ++- .../test_litellm/caching/test_redis_cache.py | 15 +++++++ .../caching/test_redis_cluster_cache.py | 42 ++++++++++++++++++- 5 files changed, 67 insertions(+), 3 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f1c80eaacbe..90adda6b1cc 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1370,7 +1370,8 @@ class RedisCache(BaseCache): self.redis_client.flushall() async def disconnect(self): - await self.async_redis_conn_pool.disconnect(inuse_connections=True) + if self.async_redis_conn_pool is not None: + await self.async_redis_conn_pool.disconnect(inuse_connections=True) try: self.redis_client.close() except Exception as e: diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 12d285ca5a8..52a3a5696de 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -56,6 +56,11 @@ class RedisClusterCache(RedisCache): async_redis_cluster_client: Final = self.init_async_client() return await async_redis_cluster_client.mget_nonatomic(keys=keys) + async def disconnect(self): + if self.redis_async_redis_cluster_client is not None: + await self.redis_async_redis_cluster_client.aclose() + await super().disconnect() + async def test_connection(self) -> dict: """ Test the Redis Cluster connection. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 990682f10a5..cac8a9d6893 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -908,7 +908,10 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N await prisma_client.disconnect() if litellm.cache is not None: - await litellm.cache.disconnect() + try: + await litellm.cache.disconnect() + except Exception as e: + verbose_proxy_logger.debug("Error disconnecting litellm.cache: %s", e) await jwt_handler.close() diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 487a64797d1..16f527a5350 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -465,6 +465,21 @@ def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ mock_client.delete.assert_called_once_with(expected) +@pytest.mark.asyncio +async def test_disconnect_with_no_connection_pool_does_not_raise(): + """RedisCache.disconnect() must not blow up when async_redis_conn_pool is None, + which is the case in cluster mode (get_redis_connection_pool returns None for + startup_nodes configs). Regression test for + https://github.com/BerriAI/litellm/issues/37137.""" + redis_cache = RedisCache.__new__(RedisCache) + redis_cache.async_redis_conn_pool = None + redis_cache.redis_client = MagicMock() + + await redis_cache.disconnect() + + redis_cache.redis_client.close.assert_called_once() + + def _closed_port() -> int: """A port with nothing listening, so Redis calls fail fast and deterministically.""" import socket diff --git a/tests/test_litellm/caching/test_redis_cluster_cache.py b/tests/test_litellm/caching/test_redis_cluster_cache.py index 372425aa9fa..01ee4a9604b 100644 --- a/tests/test_litellm/caching/test_redis_cluster_cache.py +++ b/tests/test_litellm/caching/test_redis_cluster_cache.py @@ -1,5 +1,7 @@ import json -from unittest.mock import MagicMock, patch +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -116,6 +118,44 @@ def test_cache_init_creates_redis_cache_without_cluster_config( assert not isinstance(cache.cache, RedisClusterCache) +@pytest.mark.asyncio +@patch("litellm._redis.init_redis_cluster") +async def test_disconnect_closes_cluster_client_without_raising(mock_init_redis_cluster): + """Shutting down a proxy in cluster mode used to always raise AttributeError, + since RedisCache.disconnect() unconditionally dereferences async_redis_conn_pool, + which is None by design in cluster mode. Regression test for + https://github.com/BerriAI/litellm/issues/37137.""" + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 6379}], + password="hello", + ) + assert cache.async_redis_conn_pool is None + + mock_cluster_client = AsyncMock() + cache.redis_async_redis_cluster_client = mock_cluster_client + cache.redis_client = MagicMock() + + await cache.disconnect() + + mock_cluster_client.aclose.assert_awaited_once() + + +@pytest.mark.asyncio +@patch("litellm._redis.init_redis_cluster") +async def test_disconnect_without_cluster_client_does_not_raise(mock_init_redis_cluster): + """disconnect() must be a no-op for the cluster client when it was never + initialized (e.g. shutdown before any request touched Redis).""" + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 6379}], + password="hello", + ) + cache.redis_client = MagicMock() + + await cache.disconnect() + + cache.redis_client.close.assert_called_once() + + @pytest.mark.parametrize( "startup_nodes, env_var, expected_cache_type", [