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", [