scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix: make global proxy spend an atomic shared counter

Browse source 
Concurrent completions read-add-set the global proxy spend scalar in the management cache, so simultaneous callbacks lost increments. Track proxy-wide spend in the shared spend counter (spend:user:litellm-proxy-budget) instead, which increments atomically and reseeds from the durable proxy budget row.

Fixes #35567

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-02 02:30:13 +00:00
parent 23de7a15d9
commit 1da0836d17
9 changed files with 224 additions and 377 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -1420,7 +1420,7 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int(
)
LITELLM_PROXY_ADMIN_NAME = "default_user_id"
LITELLM_PROXY_BUDGET_NAME = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
GLOBAL_PROXY_SPEND_COUNTER_KEY = f"spend:user:{LITELLM_PROXY_BUDGET_NAME}"
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"

53
litellm/proxy/auth/user_api_key_auth.py
View file

@ -24,8 +24,7 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import (
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_BUDGET_NAME,
GLOBAL_PROXY_SPEND_COUNTER_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
)
from litellm.integrations.otel.model.config import is_otel_v2_enabled
@ -76,7 +75,6 @@ from litellm.proxy.auth.resolvers import CredentialRef, Principal
from litellm.proxy.auth.resolvers.store import IdentityStore
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
@ -492,33 +490,21 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
return valid_token
# Reusable coordinator for global spend to prevent cache stampede
_global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]")
async def _fetch_global_spend_with_event_coordination(
cache_key: str,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> Optional[float]:
async def _fetch_global_proxy_spend() -> float:
"""
Fetch global spend with event-driven coordination to prevent cache stampede.
Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done.
Read proxy-wide spend from the shared spend counter.
Reads the proxy budget aggregate user row, which accrues proxy-wide spend
per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``.
The counter is incremented atomically per request (``increment_spend_counters``)
so concurrent requests cannot lose an increment, and it reseeds from the
"litellm-proxy-budget" user row, which the spend writer accrues into and
ResetBudgetJob zeroes every ``litellm.budget_duration``.
"""
from litellm.proxy.proxy_server import get_current_spend
async def _load_global_spend() -> Optional[float]:
proxy_budget_row = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": LITELLM_PROXY_BUDGET_NAME}
)
return float(proxy_budget_row.spend) if proxy_budget_row is not None else None
return await _global_spend_coordinator.get_or_load(
cache_key=cache_key,
cache=user_api_key_cache, # pyright: ignore[reportArgumentType]
load_fn=_load_global_spend,
return await get_current_spend(
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
fallback_spend=0.0,
max_budget=litellm.max_budget,
)
@ -531,13 +517,7 @@ async def get_global_proxy_spend(
) -> Optional[float]:
global_proxy_spend = None
if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget
# Use event-driven coordination to prevent cache stampede
cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
global_proxy_spend = await _fetch_global_proxy_spend()
if global_proxy_spend is not None:
user_info = CallInfo(
user_id=litellm_proxy_admin_name,
@ -1975,13 +1955,8 @@ async def _user_api_key_auth_builder(
global_proxy_spend = None
if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget
cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY
with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
global_proxy_spend = await _fetch_global_proxy_spend()
if global_proxy_spend is not None:
call_info = CallInfo(

195
litellm/proxy/common_utils/cache_coordinator.py
View file

@ -1,195 +0,0 @@
"""
Event-driven cache coordinator to prevent cache stampede.
Use this when many requests can miss the same cache key at once (e.g. after
expiry or restart). Without coordination, they would all run the expensive
load (DB query, API call) in parallel and overload the backend.
This module ensures only one request performs the load; the rest wait for a
signal and then read the freshly cached value. Reuse it for any cache-aside
pattern: global spend, feature flags, config, or other shared read-through data.
"""
import asyncio
import time
from typing import Any, Awaitable, Callable, Optional, Protocol, TypeVar
from litellm._logging import verbose_proxy_logger
T = TypeVar("T")
class AsyncCacheProtocol(Protocol):
"""Protocol for cache backends used by EventDrivenCacheCoordinator.
Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params
before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate.
"""
async def async_get_cache(
self,
key: str,
parent_otel_span: Any = None,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
async def async_set_cache(
self,
key: str,
value: Any,
local_only: bool = False,
**kwargs: Any,
) -> Any: ...
class EventDrivenCacheCoordinator:
"""
Coordinates a single in-flight load per logical resource to prevent cache stampede.
Pattern:
- First request: loads data (e.g. DB query), caches it, then signals waiters.
- Other requests: wait for the signal, then read from cache.
Create one instance per resource (e.g. one for global spend, one for feature flags).
Args:
log_prefix: Prefix for debug log messages.
"""
def __init__(self, log_prefix: str = "[CACHE]"):
self._lock = asyncio.Lock()
self._event: Optional[asyncio.Event] = None
self._query_in_progress = False
self._log_prefix = log_prefix
async def _get_cached(self, cache_key: str, cache: AsyncCacheProtocol) -> Optional[Any]:
"""Return value from cache if present, else None."""
return await cache.async_get_cache(key=cache_key)
def _log_cache_hit(self, value: T) -> None:
if self._log_prefix:
verbose_proxy_logger.debug("%s Cache hit, value: %s", self._log_prefix, value)
def _log_cache_miss(self) -> None:
if self._log_prefix:
verbose_proxy_logger.debug("%s Cache miss", self._log_prefix)
async def _claim_role(self) -> Optional[asyncio.Event]:
"""
Under lock: return event to wait on if load is in progress, else set us as loader and return None.
"""
async with self._lock:
if self._query_in_progress and self._event is not None:
if self._log_prefix:
verbose_proxy_logger.debug("%s Load in flight, waiting for signal", self._log_prefix)
return self._event
self._query_in_progress = True
self._event = asyncio.Event()
if self._log_prefix:
verbose_proxy_logger.debug(
"%s Starting load (will signal others when done)",
self._log_prefix,
)
return None
async def _wait_for_signal_and_get(
self,
event: asyncio.Event,
cache_key: str,
cache: AsyncCacheProtocol,
) -> Optional[T]:
"""Wait for loader to finish, then read from cache."""
await event.wait()
if self._log_prefix:
verbose_proxy_logger.debug("%s Signal received, reading from cache", self._log_prefix)
value: Optional[T] = await cache.async_get_cache(key=cache_key)
if value is not None and self._log_prefix:
verbose_proxy_logger.debug(
"%s Cache filled by other request, value: %s",
self._log_prefix,
value,
)
elif value is None and self._log_prefix:
verbose_proxy_logger.debug("%s Signal received but cache still empty", self._log_prefix)
return value
async def _load_and_cache(
self,
cache_key: str,
cache: AsyncCacheProtocol,
load_fn: Callable[[], Awaitable[T]],
) -> Optional[T]:
"""Double-check cache, run load_fn, set cache, return value. Caller must call _signal_done in finally."""
value = await cache.async_get_cache(key=cache_key)
if value is not None:
if self._log_prefix:
verbose_proxy_logger.debug(
"%s Cache filled while acquiring lock, value: %s",
self._log_prefix,
value,
)
return value
if self._log_prefix:
verbose_proxy_logger.debug("%s Running load", self._log_prefix)
start = time.perf_counter()
value = await load_fn()
elapsed_ms = (time.perf_counter() - start) * 1000
if self._log_prefix:
verbose_proxy_logger.debug(
"%s Load completed in %.2fms, result: %s",
self._log_prefix,
elapsed_ms,
value,
)
await cache.async_set_cache(key=cache_key, value=value)
if self._log_prefix:
verbose_proxy_logger.debug("%s Result cached", self._log_prefix)
return value
async def _signal_done(self) -> None:
"""Reset loader state and signal all waiters."""
async with self._lock:
self._query_in_progress = False
if self._event is not None:
if self._log_prefix:
verbose_proxy_logger.debug("%s Signaling all waiting requests", self._log_prefix)
self._event.set()
self._event = None
async def get_or_load(
self,
cache_key: str,
cache: AsyncCacheProtocol,
load_fn: Callable[[], Awaitable[T]],
) -> Optional[T]:
"""
Return cached value or load it once and signal waiters.
- cache_key: Key to read/write in the cache.
- cache: Object with async_get_cache(key) and async_set_cache(key, value).
- load_fn: Async callable that performs the load (e.g. DB query). No args.
Return value is cached and returned. If it raises, waiters are
still signaled so they can retry or handle empty cache.
Returns the value from cache or from load_fn, or None if load failed or
cache was still empty after waiting.
"""
value = await self._get_cached(cache_key, cache)
if value is not None:
self._log_cache_hit(value)
return value
self._log_cache_miss()
event_to_wait = await self._claim_role()
if event_to_wait is not None:
return await self._wait_for_signal_and_get(event_to_wait, cache_key, cache)
try:
result = await self._load_and_cache(cache_key, cache, load_fn)
return result
finally:
await self._signal_done()

11
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -6,7 +6,6 @@ from typing import Any, Callable, List, Literal, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME
from litellm.proxy._types import (
LiteLLM_BudgetTableFull,
LiteLLM_EndUserTable,
@ -99,14 +98,6 @@ class ResetBudgetJob:
except Exception as e:
verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e)
@staticmethod
async def _invalidate_global_proxy_spend_cache() -> None:
"""Drop the cached global-proxy spend accumulator after the proxy
budget aggregate row is reset, so the next auth-time load reads the
zeroed row instead of a stale (potentially never-expiring) counter.
"""
await ResetBudgetJob._invalidate_user_api_key_cache_entry(GLOBAL_PROXY_SPEND_CACHE_KEY)
@staticmethod
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
"""Drop a stale management-cache entry so the next read fetches from DB.
@ -562,8 +553,6 @@ class ResetBudgetJob:
user_id = getattr(u, "user_id", None)
if user_id:
await self._invalidate_spend_counter(f"spend:user:{user_id}")
if user_id == LITELLM_PROXY_BUDGET_NAME:
await self._invalidate_global_proxy_spend_cache()
end_time = time.time()
if len(failed_users) > 0: # If any users failed to reset

51
litellm/proxy/proxy_server.py
View file

@ -232,7 +232,7 @@ from litellm.constants import (
DAYS_IN_A_MONTH,
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_MODEL_CREATED_AT_TIME,
GLOBAL_PROXY_SPEND_CACHE_KEY,
GLOBAL_PROXY_SPEND_COUNTER_KEY,
LITELLM_PROXY_ADMIN_NAME,
LITELLM_PROXY_BUDGET_NAME,
PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
@ -286,7 +286,6 @@ from litellm.proxy.auth.model_checks import (
get_team_models,
)
from litellm.proxy.auth.user_api_key_auth import (
_fetch_global_spend_with_event_coordination,
user_api_key_auth,
user_api_key_auth_websocket,
)
@ -1073,7 +1072,6 @@ async def proxy_startup_event(app: FastAPI):
ProxyStartupEvent._add_proxy_budget_to_db()
asyncio.create_task(
ProxyStartupEvent._warm_global_spend_cache(
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
)
@ -2492,6 +2490,18 @@ async def increment_spend_counters(
increment=cost,
)
async def _global_proxy_scope() -> None:
# The proxy-wide accumulator is the "litellm-proxy-budget" user row that
# db_spend_update_writer also accrues into, so it shares the user counter
# key space and reseeds from that row like any other user counter.
if GLOBAL_PROXY_SPEND_COUNTER_KEY in reserved_counter_keys:
return
await _init_and_increment_spend_counter(
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
source_cache_key=LITELLM_PROXY_BUDGET_NAME,
increment=cost,
)
scope_coros = tuple(
coro
for coro in (
@ -2499,6 +2509,7 @@ async def increment_spend_counters(
_team_scope(team_id) if team_id is not None else None,
_team_member_scope(user_id, team_id) if user_id is not None and team_id is not None else None,
_user_scope(user_id) if user_id is not None else None,
_global_proxy_scope() if litellm.max_budget > 0 and user_id != LITELLM_PROXY_BUDGET_NAME else None,
_increment_end_user_and_tag_spend_counters(
end_user_id=end_user_id,
tags=tags,
@ -2904,14 +2915,6 @@ async def update_cache(
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable),
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend = await user_api_key_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
if global_proxy_spend is None:
# do nothing if not in cache
return
elif response_cost is not None and global_proxy_spend is not None:
increment = global_proxy_spend + response_cost
values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment))
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update user spend in cache. "
@ -3072,27 +3075,15 @@ async def update_cache(
if tags is not None:
await _update_tag_cache()
global_proxy_spend_key = GLOBAL_PROXY_SPEND_CACHE_KEY
local_object_updates = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key)
shared_scalar_updates = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key)
if local_object_updates:
if values_to_update_in_cache:
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(local_object_updates),
cache_list=values_to_update_in_cache,
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
)
if shared_scalar_updates:
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(shared_scalar_updates),
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
)
)
def run_ollama_serve():
@ -7977,16 +7968,14 @@ class ProxyStartupEvent:
@classmethod
async def _warm_global_spend_cache(
cls,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
"""Warm global spend counter once at startup to reduce impact of first wave of requests."""
try:
cache_key = GLOBAL_PROXY_SPEND_CACHE_KEY
await _fetch_global_spend_with_event_coordination(
cache_key=cache_key,
user_api_key_cache=user_api_key_cache,
await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
)
except Exception as e:
verbose_proxy_logger.debug("Global spend cache warm-up at startup skipped or failed: %s", e)

64
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -5082,17 +5082,16 @@ class TestCheckKeyModelBudgetWithFallback:
@pytest.mark.asyncio
async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
async def test_global_proxy_spend_reads_resettable_proxy_budget_row(monkeypatch):
"""Regression for the global proxy budget ignoring budget_duration
(LIT-4309 / gh#31292): the enforced global spend must be loaded from the
"litellm-proxy-budget" user row, which the spend writer increments per
request and ResetBudgetJob zeroes every budget_duration. It must NOT be
loaded from the MonthlyGlobalSpend view, whose window is hardcoded to a
trailing 30 days and never resets on the configured duration."""
from litellm.proxy.auth.user_api_key_auth import (
_fetch_global_spend_with_event_coordination,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
import litellm.proxy.proxy_server as ps
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend
proxy_budget_row = MagicMock()
proxy_budget_row.spend = 42.5
@ -5101,12 +5100,11 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
prisma_client.db.query_raw = AsyncMock(
side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view")
)
monkeypatch.setattr(ps, "prisma_client", prisma_client)
monkeypatch.setattr(ps, "spend_counter_cache", DualCache())
monkeypatch.setattr(litellm, "max_budget", 100.0)
result = await _fetch_global_spend_with_event_coordination(
cache_key="default_user_id:spend",
user_api_key_cache=UserApiKeyCache(),
prisma_client=prisma_client,
)
result = await _fetch_global_proxy_spend()
assert result == 42.5
prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(
@ -5115,24 +5113,46 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
@pytest.mark.asyncio
async def test_global_proxy_spend_none_when_proxy_budget_row_missing():
async def test_global_proxy_spend_zero_when_proxy_budget_row_missing(monkeypatch):
"""Before the startup upsert creates the aggregate row, enforcement must
see None (no cap applied) rather than raising."""
from litellm.proxy.auth.user_api_key_auth import (
_fetch_global_spend_with_event_coordination,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
see no recorded spend (no cap applied) rather than raising."""
import litellm.proxy.proxy_server as ps
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
monkeypatch.setattr(ps, "prisma_client", prisma_client)
monkeypatch.setattr(ps, "spend_counter_cache", DualCache())
monkeypatch.setattr(litellm, "max_budget", 100.0)
result = await _fetch_global_spend_with_event_coordination(
cache_key="default_user_id:spend",
user_api_key_cache=UserApiKeyCache(),
prisma_client=prisma_client,
)
result = await _fetch_global_proxy_spend()
assert result is None
assert result == 0.0
@pytest.mark.asyncio
async def test_global_proxy_spend_reads_shared_counter_not_the_db_row(monkeypatch):
"""Once the counter is warm, the enforced value is the shared counter that
every pod increments atomically, so spend recorded by other pods since the
last DB flush is visible without another read of the lagging DB row."""
import litellm.proxy.proxy_server as ps
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.user_api_key_auth import _fetch_global_proxy_spend
proxy_budget_row = MagicMock()
proxy_budget_row.spend = 5.0
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row)
counter_cache = DualCache()
counter_cache.in_memory_cache.set_cache(key="spend:user:litellm-proxy-budget", value=9.0)
monkeypatch.setattr(ps, "prisma_client", prisma_client)
monkeypatch.setattr(ps, "spend_counter_cache", counter_cache)
monkeypatch.setattr(litellm, "max_budget", 100.0)
result = await _fetch_global_proxy_spend()
assert result == 9.0
@pytest.mark.asyncio

22
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -1367,14 +1367,14 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock
counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60)
def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache(
def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""Regression for LIT-4309: resetting the proxy-wide budget aggregate row
("litellm-proxy-budget") must also drop the cached global-spend
accumulator ("{admin}:spend") that _global_proxy_budget_check enforces
against. Without the invalidation, the cached value survives the DB reset
and the global cap keeps blocking requests for the whole next window."""
("litellm-proxy-budget") must also zero the shared spend counter that
_global_proxy_budget_check enforces against. Without the invalidation, the
cached value survives the DB reset and the global cap keeps blocking
requests for the whole next window."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
@ -1394,13 +1394,15 @@ def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache(
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.user_api_key_cache.async_delete_cache.assert_any_call(key="default_user_id:spend")
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:user:litellm-proxy-budget", value=0.0, ttl=60
)
def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_counter(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""The global-spend accumulator must only be dropped when the proxy
"""The global-spend accumulator must only be zeroed when the proxy
budget aggregate row itself resets, not on every user reset."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
@ -1422,8 +1424,8 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
assert not any(
call.kwargs.get("key") == "default_user_id:spend"
for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list
call.kwargs.get("key") == "spend:user:litellm-proxy-budget"
for call in counter_cache.in_memory_cache.set_cache.call_args_list
)

135
tests/test_litellm/proxy/proxy_server/test_spend_counters.py
View file

@ -506,6 +506,141 @@ async def test_increment_spend_counters_runs_scopes_concurrently(monkeypatch):
}
class _AtomicRedisCounter:
"""Redis double with INCRBYFLOAT semantics and a read barrier.
``async_get_cache`` releases only once ``expected`` callers have read the
counter, so every concurrent caller observes the same pre-increment value.
A read-add-write implementation collapses to a single increment under that
barrier; an atomic increment does not.
"""
def __init__(self, expected_readers: int):
self.values: dict[str, float] = {}
self.barrier = asyncio.Barrier(expected_readers)
async def async_get_cache(self, key, **kwargs):
if key.startswith("spend:"):
await self.barrier.wait()
return self.values.get(key)
async def async_increment(self, *, key, value, refresh_ttl=True, **kwargs):
self.values[key] = self.values.get(key, 0.0) + value
return self.values[key]
@pytest.mark.asyncio
async def test_increment_spend_counters_global_proxy_counter_keeps_concurrent_costs(
monkeypatch,
):
"""Regression for gh#35567: the proxy-wide accumulator enforced against
``litellm.max_budget`` must be an atomic shared increment. Two concurrent
completions that both read the counter before either writes must leave the
counter at the sum of their costs, not at one of them."""
import litellm
redis = _AtomicRedisCounter(expected_readers=2)
fake_cache = _make_spend_counter_cache(redis_get_value=None)
fake_cache.redis_cache.async_get_cache = redis.async_get_cache
fake_cache.redis_cache.async_increment = redis.async_increment
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None))
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
monkeypatch.setattr(litellm, "max_budget", 100.0)
await asyncio.gather(
ps.increment_spend_counters(token=None, team_id=None, user_id=None, response_cost=1.0),
ps.increment_spend_counters(token=None, team_id=None, user_id=None, response_cost=1.0),
)
assert redis.values["spend:user:litellm-proxy-budget"] == 2.0
@pytest.mark.asyncio
async def test_increment_spend_counters_skips_global_proxy_counter_without_max_budget(
monkeypatch,
):
"""No global cap configured means no proxy-wide row to accrue into, so the
counter must not be created (and must not cost an extra Redis round trip)."""
import litellm
fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=1.0)
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None))
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
monkeypatch.setattr(litellm, "max_budget", 0.0)
await ps.increment_spend_counters(token=None, team_id=None, user_id="u1", response_cost=1.0)
incremented_keys = {call.kwargs["key"] for call in fake_cache.redis_cache.async_increment.call_args_list}
assert incremented_keys == {"spend:user:u1"}
@pytest.mark.asyncio
async def test_increment_spend_counters_does_not_double_count_proxy_budget_user(
monkeypatch,
):
"""The proxy-wide accumulator lives in the user counter key space, so a
request attributed to that user id must increment it once, not twice."""
import litellm
fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=1.0)
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
monkeypatch.setattr(ps, "user_api_key_cache", _make_user_api_key_cache(get_value=None))
monkeypatch.setattr(ps, "prisma_client", None)
monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None))
monkeypatch.setattr(litellm, "max_budget", 100.0)
await ps.increment_spend_counters(
token=None, team_id=None, user_id="litellm-proxy-budget", response_cost=1.0
)
incremented_keys = [call.kwargs["key"] for call in fake_cache.redis_cache.async_increment.call_args_list]
assert incremented_keys == ["spend:user:litellm-proxy-budget"]
@pytest.mark.asyncio
async def test_update_cache_does_not_write_the_global_spend_scalar(monkeypatch):
"""gh#35567: the cached-object refresh must stay per-pod. Any shared write
of the proxy-wide spend from here is a read-add-write that loses concurrent
increments; the atomic counter owns that value now."""
import litellm
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
cached = {
"u1": CacheCodec.serialize(LiteLLM_UserTable(user_id="u1", spend=0.0), model_type=LiteLLM_UserTable),
"default_user_id:spend": 10.0,
}
fake_user_cache = _make_user_api_key_cache(get_side_effect=lambda key, **kwargs: cached.get(key))
monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
monkeypatch.setattr(litellm, "max_budget", 100.0)
await ps.update_cache(
token=None,
user_id="u1",
end_user_id=None,
team_id=None,
response_cost=1.0,
parent_otel_span=None,
tags=None,
)
await asyncio.sleep(0)
written_keys = [
key
for call in fake_user_cache.async_set_cache_pipeline.call_args_list
for key, _ in call.kwargs["cache_list"]
]
assert written_keys == ["u1"]
assert all(
call.kwargs.get("local_only") is True
for call in fake_user_cache.async_set_cache_pipeline.call_args_list
)
@pytest.mark.asyncio
async def test_increment_spend_counters_skips_reserved_counter_keys(monkeypatch):
"""Counters already reserved by a budget reservation are skipped, every

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

@ -4952,74 +4952,6 @@ async def test_spend_tracking_never_writes_the_auth_object_back():
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
@pytest.mark.asyncio
async def test_update_cache_global_proxy_spend_scalar_stays_shared():
"""
The proxy-wide spend estimate must keep flowing to Redis when the spend
writeback goes per-pod: the global max_budget check reads the
``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB
reloads, so keeping it pod-local would let traffic spread across replicas
exceed the proxy budget by roughly a factor of the replica count within a
cache TTL. Sharing this scalar is safe because it carries no limits or
permissions, so it cannot resurrect an invalidated auth blob.
"""
from litellm.caching.caching import DualCache
admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name
global_key = "{}:spend".format(admin_name)
async def fake_get(key, **kwargs):
if key == "user-lit":
return {"user_id": "user-lit", "spend": 1.0}
if key == global_key:
return 10.0
return None
original_cache = litellm.proxy.proxy_server.user_api_key_cache
cache = DualCache(default_in_memory_ttl=300)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
try:
with patch.object(
cache, "async_get_cache", new=AsyncMock(side_effect=fake_get)
):
with patch.object(
cache, "async_set_cache_pipeline", new=AsyncMock()
) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id="user-lit",
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
)
pending = [
t for t in asyncio.all_tasks() if t is not asyncio.current_task()
]
if pending:
await asyncio.wait(pending, timeout=5)
calls = mock_set_cache.await_args_list
local_keys = [
k
for c in calls
if c.kwargs.get("local_only") is True
for k, _ in c.kwargs["cache_list"]
]
shared_keys = [
k
for c in calls
if c.kwargs.get("local_only") is not True
for k, _ in c.kwargs["cache_list"]
]
assert "user-lit" in local_keys
assert global_key not in local_keys
assert shared_keys == [global_key]
finally:
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
@pytest.mark.asyncio
async def test_init_sso_settings_in_db():
"""