scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(proxy): increase spend counter cache capacity

Browse source
This commit is contained in:
wzj1228516103 2026-09-08 18:21:56 +08:00
parent 1af7a403c6
commit 8cb42ca7b4
2 changed files with 24 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/proxy/proxy_server.py
View file

@ -240,6 +240,7 @@ import litellm._redis
from litellm import Router
from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import (
_REALTIME_BODY_CACHE_SIZE,
@ -2301,7 +2302,11 @@ shared_aiohttp_session: Optional["ClientSession"] = None # Global shared sessio
user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
SPEND_COUNTER_CACHE_MAX_SIZE: Final = 10_000
spend_counter_cache: Final = DualCache(
in_memory_cache=InMemoryCache(max_size_in_memory=SPEND_COUNTER_CACHE_MAX_SIZE),
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value,
)
cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
model_max_budget_limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)

18
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -7542,6 +7542,24 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch):
# =====================================================================
def test_spend_counter_cache_keeps_active_counters_warm(monkeypatch):
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.proxy_server import spend_counter_cache
cache = InMemoryCache(
max_size_in_memory=spend_counter_cache.in_memory_cache.max_size_in_memory,
default_ttl=60,
)
monkeypatch.setattr(spend_counter_cache, "in_memory_cache", cache)
sentinel = "spend:key:sentinel"
cache.set_cache(key=sentinel, value=1.0)
for index in range(300):
cache.set_cache(key=f"spend:end_user:scope-{index}", value=1.0)
assert cache.get_cache(key=sentinel) == 1.0
@pytest.mark.asyncio
async def test_get_current_spend_reads_redis_first():
"""get_current_spend should prefer Redis over in-memory."""