diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d0cefcb6086..d4720294f4a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1738,7 +1738,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 d9f8e04ebda..afe76b11b93 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -996,7 +996,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: # noqa: BLE001 # cache teardown must not abort remaining shutdown steps + 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 c1e3240adb7..e59d033bf32 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -479,6 +479,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 0763b5110d5..33757c7359c 100644 --- a/tests/test_litellm/caching/test_redis_cluster_cache.py +++ b/tests/test_litellm/caching/test_redis_cluster_cache.py @@ -1,6 +1,8 @@ from importlib import import_module 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 @@ -117,6 +119,52 @@ def test_cache_init_creates_redis_cache_without_cluster_config( assert not isinstance(cache.cache, RedisClusterCache) +@pytest.mark.asyncio +@patch("redis.RedisCluster") +async def test_disconnect_closes_cluster_client_without_raising(mock_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 + + closed = False + + class _ClusterClient: + async def aclose(self): + nonlocal closed + closed = True + + cache.redis_async_redis_cluster_client = _ClusterClient() + cache.redis_client = MagicMock() + + await cache.disconnect() + + assert closed is True + + +@pytest.mark.asyncio +@patch("redis.RedisCluster") +async def test_disconnect_without_cluster_client_does_not_raise(mock_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), rather than + raising on the absent client or the None connection pool.""" + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 6379}], + password="hello", + ) + cache.redis_client = MagicMock() + + assert cache.redis_async_redis_cluster_client is None + assert cache.async_redis_conn_pool is None + + await cache.disconnect() + + @pytest.mark.parametrize( "startup_nodes, env_var, expected_cache_type", [ diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index deb7289d2d1..65670f1935e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -209,6 +209,31 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): await proxy_shutdown_event() +@pytest.mark.asyncio +async def test_proxy_shutdown_event_cache_disconnect_error_does_not_abort_shutdown(monkeypatch): + """A cache backend failing to disconnect must not skip the jwt_handler close + that runs after it, unlike a prisma disconnect failure which does abort. + """ + monkeypatch.setattr(ps, "prisma_client", None, raising=False) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + fake_cache = MagicMock() + fake_cache.disconnect = AsyncMock(side_effect=RuntimeError("redis gone")) + monkeypatch.setattr(litellm, "cache", fake_cache, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + await proxy_shutdown_event() + + assert fake_cache.disconnect.await_count == 1 + assert fake_jwt.close.await_count == 1 + + # --------------------------------------------------------------------------- # _flush_spend_logs_queue_on_shutdown # ---------------------------------------------------------------------------