mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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
This commit is contained in:
parent
02dcc4d347
commit
daa2863e89
5 changed files with 67 additions and 3 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue