diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 6fb84df331e..beb11fad4ee 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -139,6 +139,13 @@ 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) -> None: + try: + if self.redis_async_client is not None: # pyright: ignore[reportUnknownMemberType] # redis generics + await self.redis_async_client.aclose() # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue] # redis stubs leave aclose unknown + finally: + 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 b48b5ab1c1e..9b9bea83197 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1297,7 +1297,10 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N await drain_post_call_redis_batches() 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.warning("Error disconnecting litellm.cache: %s", e) await jwt_handler.close() diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..4b7ee4b0aba 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -425,6 +425,23 @@ 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_kwargs = {} + mock_redis_client = MagicMock() + redis_cache.redis_client = mock_redis_client + + await redis_cache.disconnect() + + mock_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/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..10f91136f2e 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -124,6 +124,69 @@ 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: MagicMock) -> None: + """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_client = cast(RedisCluster, _ClusterClient()) + cache.redis_client = MagicMock() + + await cache.disconnect() + + assert closed is True + + +@pytest.mark.asyncio +@patch("redis.RedisCluster") +async def test_disconnect_still_closes_redis_client_when_cluster_aclose_fails(mock_redis_cluster: MagicMock) -> None: + cache = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 6379}], + password="hello", + ) + cache.redis_async_client = MagicMock(aclose=AsyncMock(side_effect=RuntimeError("aclose failed"))) + mock_redis_client = MagicMock() + cache.redis_client = mock_redis_client + + with pytest.raises(RuntimeError, match="aclose failed"): + await cache.disconnect() + + mock_redis_client.close.assert_called_once() + + +@pytest.mark.asyncio +@patch("redis.RedisCluster") +async def test_disconnect_without_cluster_client_does_not_raise(mock_redis_cluster: MagicMock) -> None: + """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_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/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..d14fb8ece29 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -336,6 +336,33 @@ 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: pytest.MonkeyPatch, +) -> None: + """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 # ---------------------------------------------------------------------------