mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(proxy): build redis usage cache from REDIS_* env when cache backend is not Redis (#32635)
Selecting a semantic (or any non-Redis-KV) response cache left redis_usage_cache unset, silently downgrading cross-pod rate limits, parallel-request limits, spend coordination, and the pod lock manager to per-pod in-memory state. Fall back to a standalone RedisCache built from REDIS_* environment variables, mirroring the existing use_redis_transaction_buffer escape hatch, which now shares the same helper. Resolves LIT-3861
This commit is contained in:
parent
94e8e69397
commit
d1ae9571ae
2 changed files with 122 additions and 8 deletions
|
|
@ -212,6 +212,7 @@ from contextlib import asynccontextmanager
|
|||
from functools import lru_cache
|
||||
|
||||
import litellm
|
||||
import litellm._redis
|
||||
from litellm import Router
|
||||
from litellm._logging import verbose_proxy_logger, verbose_router_logger
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
|
|
@ -3549,6 +3550,22 @@ def _apply_ssrf_general_settings(settings: Mapping[str, object]) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _build_redis_usage_cache_from_environment() -> RedisCache | None:
|
||||
"""
|
||||
Builds a standalone RedisCache from REDIS_* environment variables.
|
||||
|
||||
Lets the proxy's coordination Redis (cross-pod tpm/rpm rate limits, spend
|
||||
tracking, pod lock manager) run when the response-cache backend is not a
|
||||
plain Redis KV cache (e.g. a semantic cache, disk, or s3).
|
||||
|
||||
Returns None when no Redis host or url is set in the environment.
|
||||
"""
|
||||
redis_env_kwargs = litellm._redis._redis_kwargs_from_environment()
|
||||
if "host" not in redis_env_kwargs and "url" not in redis_env_kwargs:
|
||||
return None
|
||||
return RedisCache(**redis_env_kwargs)
|
||||
|
||||
|
||||
class ProxyConfig:
|
||||
"""
|
||||
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
||||
|
|
@ -3763,9 +3780,21 @@ class ProxyConfig:
|
|||
|
||||
litellm.cache = Cache(**cache_params)
|
||||
|
||||
if litellm.cache is not None and isinstance(litellm.cache.cache, (RedisCache, RedisClusterCache)):
|
||||
cache_backend = litellm.cache.cache if litellm.cache is not None else None
|
||||
if isinstance(cache_backend, (RedisCache, RedisClusterCache)):
|
||||
## INIT PROXY REDIS USAGE CLIENT ##
|
||||
redis_usage_cache = litellm.cache.cache
|
||||
redis_usage_cache = cache_backend
|
||||
elif redis_usage_cache is None:
|
||||
redis_usage_cache = _build_redis_usage_cache_from_environment()
|
||||
if redis_usage_cache is not None:
|
||||
verbose_proxy_logger.info(
|
||||
"Cache backend %s is not a Redis KV cache; built a standalone "
|
||||
"Redis from REDIS_* environment variables for usage tracking, "
|
||||
"rate limiting, and cross-pod coordination.",
|
||||
type(cache_backend).__name__,
|
||||
)
|
||||
|
||||
if redis_usage_cache is not None:
|
||||
spend_counter_cache.attach_redis_cache(
|
||||
redis_usage_cache,
|
||||
default_redis_ttl=litellm.default_redis_ttl,
|
||||
|
|
@ -7277,7 +7306,6 @@ class ProxyStartupEvent:
|
|||
Returns None when the buffer is disabled, or when no Redis host or url
|
||||
is set in the environment.
|
||||
"""
|
||||
from litellm._redis import _redis_kwargs_from_environment
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
_use_redis_transaction_buffer: bool | str | None = general_settings.get("use_redis_transaction_buffer", False)
|
||||
|
|
@ -7287,11 +7315,7 @@ class ProxyStartupEvent:
|
|||
if not _use_redis_transaction_buffer:
|
||||
return None
|
||||
|
||||
redis_env_kwargs = _redis_kwargs_from_environment()
|
||||
if "host" not in redis_env_kwargs and "url" not in redis_env_kwargs:
|
||||
return None
|
||||
|
||||
return RedisCache(**redis_env_kwargs)
|
||||
return _build_redis_usage_cache_from_environment()
|
||||
|
||||
@classmethod
|
||||
async def _initialize_semantic_tool_filter(
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import os
|
|||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import types
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
|
@ -23,6 +24,9 @@ sys.path.insert(
|
|||
) # Adds the parent directory to the system-path
|
||||
|
||||
import litellm
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
from litellm.caching.caching import RedisCache
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.proxy_server import app, initialize
|
||||
|
|
@ -9358,3 +9362,89 @@ def test_update_config_redacts_all_environment_variable_values(
|
|||
assert "db.internal" not in data["updated_values"]
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
class _EnvBuiltRedisCache(RedisCache):
|
||||
"""RedisCache stand-in that records its constructor kwargs and never
|
||||
opens a network connection, so tests can assert which connection params
|
||||
the proxy used to build its coordination Redis."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.init_kwargs = kwargs
|
||||
|
||||
|
||||
def _run_init_cache_with_backend(cache_backend, redis_env_kwargs):
|
||||
"""Run ProxyConfig._init_cache with a stubbed response-cache backend and a
|
||||
controlled REDIS_* environment, returning (redis_usage_cache,
|
||||
spend_counter redis, config-cache redis) as observed after the call."""
|
||||
mock_litellm_cache = MagicMock()
|
||||
mock_litellm_cache.cache = cache_backend
|
||||
fresh_spend_cache = DualCache()
|
||||
fresh_config_cache = types.SimpleNamespace(redis_cache=None)
|
||||
|
||||
with (
|
||||
patch.object(proxy_server_module, "redis_usage_cache", None),
|
||||
patch.object(proxy_server_module, "spend_counter_cache", fresh_spend_cache),
|
||||
patch.object(proxy_server_module, "user_api_key_cache", DualCache()),
|
||||
patch.object(proxy_server_module, "llm_router", None),
|
||||
patch.object(proxy_server_module, "litellm_config_cache", fresh_config_cache),
|
||||
patch.object(proxy_server_module, "RedisCache", _EnvBuiltRedisCache),
|
||||
patch(
|
||||
"litellm._redis._redis_kwargs_from_environment",
|
||||
return_value=redis_env_kwargs,
|
||||
),
|
||||
patch("litellm.Cache", return_value=mock_litellm_cache),
|
||||
):
|
||||
litellm.cache = None
|
||||
proxy_server_module.ProxyConfig()._init_cache(
|
||||
cache_params={"type": "qdrant-semantic"}
|
||||
)
|
||||
return (
|
||||
proxy_server_module.redis_usage_cache,
|
||||
fresh_spend_cache.redis_cache,
|
||||
fresh_config_cache.redis_cache,
|
||||
)
|
||||
|
||||
|
||||
def test_init_cache_non_redis_backend_builds_usage_redis_from_environment():
|
||||
"""A semantic (non-Redis-KV) response cache must not disable the proxy's
|
||||
coordination Redis: when REDIS_* env vars provide a connection,
|
||||
_init_cache builds a standalone usage cache so cross-pod rate limits,
|
||||
spend tracking, and the pod lock manager stay Redis-backed."""
|
||||
usage_cache, spend_redis, config_redis = _run_init_cache_with_backend(
|
||||
cache_backend=object(),
|
||||
redis_env_kwargs={"host": "coordination-redis", "port": "6379"},
|
||||
)
|
||||
|
||||
assert isinstance(usage_cache, _EnvBuiltRedisCache)
|
||||
assert usage_cache.init_kwargs["host"] == "coordination-redis"
|
||||
assert spend_redis is usage_cache
|
||||
assert config_redis is usage_cache
|
||||
|
||||
|
||||
def test_init_cache_non_redis_backend_without_redis_env_stays_in_memory():
|
||||
"""Without any REDIS_* connection info, a non-Redis response cache must
|
||||
leave the coordination Redis unset instead of building a broken client."""
|
||||
usage_cache, spend_redis, config_redis = _run_init_cache_with_backend(
|
||||
cache_backend=object(),
|
||||
redis_env_kwargs={},
|
||||
)
|
||||
|
||||
assert usage_cache is None
|
||||
assert spend_redis is None
|
||||
assert config_redis is None
|
||||
|
||||
|
||||
def test_init_cache_redis_backend_reuses_cache_backend_over_environment():
|
||||
"""When the response cache itself is a plain Redis KV cache, it must be
|
||||
reused as the coordination Redis; the REDIS_* environment fallback must
|
||||
not construct a second client."""
|
||||
redis_backend = _EnvBuiltRedisCache(host="cache-params-host")
|
||||
usage_cache, spend_redis, _ = _run_init_cache_with_backend(
|
||||
cache_backend=redis_backend,
|
||||
redis_env_kwargs={"host": "env-host"},
|
||||
)
|
||||
|
||||
assert usage_cache is redis_backend
|
||||
assert usage_cache.init_kwargs["host"] == "cache-params-host"
|
||||
assert spend_redis is redis_backend
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue