scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix(proxy): global max_budget ignores budget_duration; enforce against the resettable proxy budget row (#33732)

Browse source 
* fix(proxy): enforce global max_budget against the resettable proxy budget row

The global proxy budget check compared litellm.max_budget against
SUM(spend) from the MonthlyGlobalSpend view, whose window is hardcoded
to a trailing 30 days. litellm.budget_duration was stored and reset on
a user row that enforcement never read, and startup budgeted the admin
user's own row (default_user_id) instead of the litellm-proxy-budget
aggregate row the spend writer increments per request. Net effect: 1d,
7d and 30d all behaved as a trailing 30 day cap that never reset on the
configured duration.

Startup now upserts the budget onto the litellm-proxy-budget row (and
zeroes lifetime accrual when first putting a row on a reset schedule),
enforcement loads global spend from that row, and ResetBudgetJob drops
the cached global spend accumulator when it resets that row so the cap
unblocks immediately after each window.

Fixes https://github.com/BerriAI/litellm/issues/31292

* refactor(proxy): address review nits on global proxy budget fix

Drop the redundant litellm_proxy_budget_name parameter from
_upsert_proxy_budget_with_reset_at_backfill; its only caller always passed
LITELLM_PROXY_BUDGET_NAME, and any other value would write the budget to a
row enforcement never reads.

Introduce GLOBAL_PROXY_SPEND_CACHE_KEY in constants.py and use it at every
site that previously built the key from litellm_proxy_admin_name (auth
loads, spend-writer increments, startup warm, reset-job invalidation), so
the reader and invalidator can no longer drift apart. The literal key value
is unchanged. Also drop the now-pointless litellm_proxy_admin_name
parameter from _warm_global_spend_cache and the proxy_server import from
the reset-job helper.
This commit is contained in:
ryan-crabbe-berri 2026-07-25 11:29:39 -07:00 • committed by GitHub
parent 8ce365511f
commit fe5cc1eb0c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 178 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/constants.py
View file

@ -1418,6 +1418,8 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int(
os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000)
)
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"
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"

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

@ -23,7 +23,11 @@ from fastapi.security.api_key import APIKeyHeader
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.constants import (
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_BUDGET_NAME,
LITELLM_PROXY_MASTER_KEY_ALIAS,
)
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.runtime import phase_span, seed_request_identity
from litellm.litellm_core_utils.dd_tracing import tracer
@ -500,13 +504,16 @@ async def _fetch_global_spend_with_event_coordination(
"""
Fetch global spend with event-driven coordination to prevent cache stampede.
Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done.
Reads the proxy budget aggregate user row, which accrues proxy-wide spend
per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``.
"""
async def _load_global_spend() -> Optional[float]:
sql_query = """SELECT SUM(spend) AS total_spend FROM "MonthlyGlobalSpend";"""
response = await prisma_client.db.query_raw(query=sql_query)
val = response[0]["total_spend"]
return float(val) if val is not None else None
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,
@ -525,7 +532,7 @@ async def get_global_proxy_spend(
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 = "{}:spend".format(litellm_proxy_admin_name)
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,
@ -1979,7 +1986,7 @@ 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 = "{}:spend".format(litellm_proxy_admin_name)
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,

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

@ -6,6 +6,7 @@ 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,
@ -98,6 +99,14 @@ 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.
@ -553,6 +562,8 @@ 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

42
litellm/proxy/proxy_server.py
View file

@ -230,7 +230,9 @@ from litellm.constants import (
DAYS_IN_A_MONTH,
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_MODEL_CREATED_AT_TIME,
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_ADMIN_NAME,
LITELLM_PROXY_BUDGET_NAME,
PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
PROXY_BATCH_POLLING_ENABLED,
PROXY_BATCH_POLLING_INTERVAL,
@ -1054,10 +1056,9 @@ async def proxy_startup_event(app: FastAPI):
verbose_proxy_logger.debug("prisma_client: %s", prisma_client)
if prisma_client is not None and litellm.max_budget > 0:
ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name=litellm_proxy_admin_name)
ProxyStartupEvent._add_proxy_budget_to_db()
asyncio.create_task(
ProxyStartupEvent._warm_global_spend_cache(
litellm_proxy_admin_name=litellm_proxy_admin_name,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
@ -2002,7 +2003,7 @@ health_check_results: Dict[str, Union[int, List[Dict[str, Any]]]] = {}
background_health_check_loop_active = False
background_health_check_cycle_seq = 0
queue: List = []
litellm_proxy_budget_name = "litellm-proxy-budget"
litellm_proxy_budget_name = LITELLM_PROXY_BUDGET_NAME
litellm_proxy_admin_name = LITELLM_PROXY_ADMIN_NAME
ui_access_mode: Union[Literal["admin", "all"], Dict] = "all"
proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME
@ -2871,15 +2872,13 @@ async def update_cache(
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend = await user_api_key_cache.async_get_cache(
key="{}:spend".format(litellm_proxy_admin_name)
)
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(("{}:spend".format(litellm_proxy_admin_name), increment))
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. "
@ -3040,7 +3039,7 @@ async def update_cache(
if tags is not None:
await _update_tag_cache()
global_proxy_spend_key = "{}:spend".format(litellm_proxy_admin_name)
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)
@ -7772,29 +7771,32 @@ class ProxyStartupEvent:
)
@classmethod
def _add_proxy_budget_to_db(cls, litellm_proxy_budget_name: str):
def _add_proxy_budget_to_db(cls):
"""Adds a global proxy budget to db"""
if litellm.budget_duration is None:
raise Exception("budget_duration not set on Proxy. budget_duration is required to use max_budget.")
asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill(litellm_proxy_budget_name))
asyncio.create_task(cls._upsert_proxy_budget_with_reset_at_backfill())
@classmethod
async def _upsert_proxy_budget_with_reset_at_backfill(cls, litellm_proxy_budget_name: str) -> None:
async def _upsert_proxy_budget_with_reset_at_backfill(cls) -> None:
"""
Upsert the proxy admin user row with the configured max_budget /
budget_duration, then backfill budget_reset_at if currently NULL.
Upsert the proxy budget aggregate user row with the configured
max_budget / budget_duration, then backfill budget_reset_at if
currently NULL.
The backfill uses `WHERE budget_reset_at IS NULL` so it only fires
when the row pre-existed without a reset schedule (e.g. row created
via a different path before the proxy budget was configured). On
subsequent restarts it no-ops, so an active reset window is never
slid forward.
slid forward. It also zeroes spend at that moment: a row that was
never on a reset schedule holds lifetime accrual, which must not
gate the first duration window.
"""
await generate_key_helper_fn( # type: ignore
request_type="user",
table_name="user",
user_id=litellm_proxy_budget_name,
user_id=LITELLM_PROXY_BUDGET_NAME,
duration=None,
models=[],
aliases={},
@ -7817,10 +7819,13 @@ class ProxyStartupEvent:
try:
await UserRepository(prisma_client).table.update_many(
where={
"user_id": litellm_proxy_budget_name,
"user_id": LITELLM_PROXY_BUDGET_NAME,
"budget_reset_at": None,
},
data={"budget_reset_at": get_budget_reset_time(budget_duration=litellm.budget_duration)},
data={
"budget_reset_at": get_budget_reset_time(budget_duration=litellm.budget_duration),
"spend": 0,
},
)
except Exception as e:
verbose_proxy_logger.warning("Failed to backfill budget_reset_at on proxy admin row: %s", e)
@ -7828,13 +7833,12 @@ class ProxyStartupEvent:
@classmethod
async def _warm_global_spend_cache(
cls,
litellm_proxy_admin_name: str,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient,
) -> None:
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
try:
cache_key = "{}:spend".format(litellm_proxy_admin_name)
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,

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

@ -5059,6 +5059,60 @@ class TestCheckKeyModelBudgetWithFallback:
assert "model" not in request_data
@pytest.mark.asyncio
async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
"""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
proxy_budget_row = MagicMock()
proxy_budget_row.spend = 42.5
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row)
prisma_client.db.query_raw = AsyncMock(
side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view")
)
result = await _fetch_global_spend_with_event_coordination(
cache_key="default_user_id:spend",
user_api_key_cache=UserApiKeyCache(),
prisma_client=prisma_client,
)
assert result == 42.5
prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(
where={"user_id": "litellm-proxy-budget"}
)
@pytest.mark.asyncio
async def test_global_proxy_spend_none_when_proxy_budget_row_missing():
"""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
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
result = await _fetch_global_spend_with_event_coordination(
cache_key="default_user_id:spend",
user_api_key_cache=UserApiKeyCache(),
prisma_client=prisma_client,
)
assert result is None
@pytest.mark.asyncio
async def test_temp_budget_increase_applied_for_cached_key():
"""

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

@ -1367,6 +1367,66 @@ 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(
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."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["user"] = [
type(
"User",
(),
{
"spend": 150.0,
"budget_duration": "30d",
"budget_reset_at": now,
"id": "row-1",
"user_id": "litellm-proxy-budget",
},
)
]
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")
def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""The global-spend accumulator must only be dropped when the proxy
budget aggregate row itself resets, not on every user reset."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
now = datetime.now(timezone.utc)
mock_prisma_client.data["user"] = [
type(
"User",
(),
{
"spend": 50.0,
"budget_duration": "7d",
"budget_reset_at": now,
"id": "user-1",
"user_id": "alice",
},
)
]
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
)
def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch):
"""Team budget reset must clear the Redis spend counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)

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

@ -2642,6 +2642,11 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys():
This validates that generate_key_helper_fn is called with table_name="user"
which should prevent key creation in LiteLLM_VerificationToken table.
Also guards the row identity: the budget must land on the proxy-wide
aggregate row "litellm-proxy-budget" (the one the spend writer increments
per request), not the admin user's own row ("default_user_id"). Budgeting
the admin row leaves the global budget without a resettable counter.
"""
from unittest.mock import AsyncMock, patch
@ -2670,7 +2675,7 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys():
"litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper
):
# Call the function under test
ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name)
ProxyStartupEvent._add_proxy_budget_to_db()
# Allow async task to complete
import asyncio
@ -2696,9 +2701,13 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at():
Test that _upsert_proxy_budget_with_reset_at_backfill issues a conditional
update_many with `WHERE budget_reset_at IS NULL` to backfill the column on
rows that pre-existed without a reset schedule. Without this, the proxy
admin row stays at NULL and reset_budget_for_litellm_users never matches
budget row stays at NULL and reset_budget_for_litellm_users never matches
it (NULL < now() is unknown in SQL), so the global proxy budget never
resets.
The same conditional update must zero spend: a row that was never on a
reset schedule holds lifetime accrual, which must not gate the first
duration window.
"""
from unittest.mock import AsyncMock, MagicMock, patch
@ -2729,9 +2738,7 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at():
),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
):
await ProxyStartupEvent._upsert_proxy_budget_with_reset_at_backfill(
litellm_proxy_budget_name
)
await ProxyStartupEvent._upsert_proxy_budget_with_reset_at_backfill()
# Upsert ran with the configured budget
mock_generate_key_helper.assert_called_once()
@ -2750,6 +2757,8 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at():
assert isinstance(backfilled_reset_at, datetime)
assert backfilled_reset_at > datetime.now(timezone.utc)
assert backfill_call.kwargs["data"]["spend"] == 0
@pytest.mark.asyncio
async def test_custom_ui_sso_sign_in_handler_config_loading():