From aa6b7fc447dc6aa0fe257f2b112e62ff644b9a5f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 01:56:33 -0700 Subject: [PATCH 1/9] fix(proxy): report live key and team spend to budget gauges within the key cache TTL --- litellm/proxy/auth/user_api_key_auth.py | 12 +++- .../spend_tracking/carried_budget_state.py | 17 ++++- .../test_user_api_key_auth_request_flow.py | 62 +++++++++++++++++++ 3 files changed, 89 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 292a147d200..2c4c0ec4780 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -168,7 +168,10 @@ from litellm.proxy.db.db_lookup_gate import bounded_db_lookup from litellm.proxy.db.db_span import db_span from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state +from litellm.proxy.spend_tracking.carried_budget_state import ( + carry_live_key_and_team_spend, + carry_team_and_user_budget_state, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -3572,6 +3575,7 @@ async def run_centralized_common_checks( """ from litellm.proxy.proxy_server import ( general_settings, + get_current_spend, litellm_proxy_admin_name, llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries llm_router, @@ -3964,6 +3968,12 @@ async def run_centralized_common_checks( user_obj=user_object, ) + if not route_skips_budget_checks(route=route): + await carry_live_key_and_team_spend( + valid_token=user_api_key_auth_obj, + team_object=typed_team_object, + read_spend=get_current_spend, + ) await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 0dfb38272b0..2adf4300099 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,6 +1,6 @@ """Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from types import MappingProxyType from typing import Final @@ -34,6 +34,21 @@ def carry_team_and_user_budget_state( ) +async def carry_live_key_and_team_spend( + valid_token: UserAPIKeyAuth, + team_object: LiteLLM_TeamTable | None, + read_spend: Callable[[str, float], Awaitable[float]], +) -> None: + if valid_token.max_budget is not None and valid_token.token is not None: + valid_token.spend = await read_spend( # rebind-ok: the budget gauges and spend header read this token + f"spend:key:{valid_token.token}", valid_token.spend or 0.0 + ) + if team_object is not None and team_object.max_budget is not None: + valid_token.team_spend = await read_spend( # rebind-ok: the team budget gauge reads this token + f"spend:team:{team_object.team_id}", team_object.spend or 0.0 + ) + + def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: budget_table: Final = org_table.litellm_budget_table valid_token.organization_alias = org_table.organization_alias diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index b7f3ded9291..5f4acff1661 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5084,6 +5084,68 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +@pytest.mark.asyncio +@pytest.mark.parametrize("model_input_cost", [None, 0.0], ids=["priced_model_reserves_budget", "zero_cost_model"]) +async def test_centralized_common_checks_carry_live_key_and_team_spend_on_a_token_with_stale_cached_spend( + model_input_cost: float | None, +): + import litellm.proxy.proxy_server as _proxy_server_mod + from starlette.datastructures import URL + + from litellm.proxy.utils import ProxyLogging + + zero_cost_router: Final = ( + None + if model_input_cost is None + else litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "input_cost_per_token": model_input_cost, + "output_cost_per_token": model_input_cost, + }, + } + ] + ) + ) + token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0) + user_api_key_cache: Final = UserApiKeyCache() + user_api_key_cache.set_cache( + key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0) + ) + spend_counter_cache: Final = DualCache() + spend_counter_cache.in_memory_cache.set_cache(key=f"spend:key:{token.token}", value=0.37) + spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t-live", value=1.25) + attrs: Final = { + **_proxy_attrs_for_centralized_checks(), + "llm_router": zero_cost_router, + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "spend_counter_cache": spend_counter_cache, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=user_api_key_cache), + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + request: Final = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1000}, + route="/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + assert (token.spend, token.team_spend) == (0.37, 1.25) + + def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock: row = MagicMock() row.dict = lambda: {"budget_id": budget_id, "max_budget": max_budget} From 09261ae336353ecb00044d8d502c3ff478b0aba5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 03:34:56 -0700 Subject: [PATCH 2/9] fix(prometheus): carry live spend for the user, org and customer budget gauges too The remaining budget gauges for users, orgs and customers read the same cached rows as the key and team gauges, so they read high for about a minute after a charge. Auth now carries the live counter for each budgeted entity in one snapshot that only the gauges read, floored at the spend auth loaded. The x-litellm-key-spend header and logging spend fields keep their current values. --- litellm/integrations/prometheus.py | 25 +++-- litellm/proxy/auth/user_api_key_auth.py | 6 +- .../spend_tracking/carried_budget_state.py | 65 ++++++++++-- litellm/types/proxy/auth/user_api_key_auth.py | 2 + litellm/types/proxy/carried_budget_state.py | 29 +++++ .../test_prometheus_carried_budget_state.py | 23 ++++ .../test_prometheus_user_team_metrics.py | 31 ++++++ .../test_user_api_key_auth_request_flow.py | 11 +- .../test_carried_budget_state.py | 100 ++++++++++++++++++ 9 files changed, 270 insertions(+), 22 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index bc8344fa686..f69a5f53d45 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -77,6 +77,7 @@ from litellm.types.integrations.prometheus import ( from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.proxy.carried_budget_state import ( KeyBudgetSnapshot, + LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -2221,6 +2222,8 @@ class PrometheusLogger(CustomLogger): _user_max_budget: Final = _metadata.get("user_api_key_user_max_budget", None) _user_email: Final = _metadata.get("user_api_key_user_email", None) _org_alias: Final = _metadata.get("user_api_key_org_alias", None) + _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) + _live_spend: Final = LiveBudgetSpend.from_metadata(_carried_metadata) or LiveBudgetSpend() # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2236,34 +2239,36 @@ class PrometheusLogger(CustomLogger): user_api_key_alias=user_api_key_alias, response_cost=response_cost, key_max_budget=_api_key_max_budget, - key_spend=_api_key_spend, - carried=KeyBudgetSnapshot.from_metadata(_metadata), + key_spend=_api_key_spend if _live_spend.key is None else _live_spend.key, + carried=KeyBudgetSnapshot.from_metadata(_carried_metadata), ), self._set_team_budget_metrics_after_api_request( user_api_team=user_api_team, user_api_team_alias=user_api_team_alias, - team_spend=_team_spend, + team_spend=_team_spend if _live_spend.team is None else _live_spend.team, team_max_budget=_team_max_budget, response_cost=response_cost, - carried=TeamBudgetSnapshot.from_metadata(_metadata), + carried=TeamBudgetSnapshot.from_metadata(_carried_metadata), ), self._set_user_budget_metrics_after_api_request( user_id=user_id, - user_spend=_user_spend, + user_spend=_user_spend if _live_spend.user is None else _live_spend.user, user_max_budget=_user_max_budget, response_cost=response_cost, - carried=UserBudgetSnapshot.from_metadata(_metadata), + carried=UserBudgetSnapshot.from_metadata(_carried_metadata), user_email=_user_email if isinstance(_user_email, str) else None, ), self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=response_cost, - carried=OrgBudgetSnapshot.from_metadata(_metadata), + carried=OrgBudgetSnapshot.from_metadata(_carried_metadata), org_alias=_org_alias if isinstance(_org_alias, str) else None, + live_spend=_live_spend.org, ), self._set_customer_budget_metrics_after_api_request( end_user_id=end_user_id, response_cost=response_cost, + live_spend=_live_spend.end_user, ), return_exceptions=True, ) @@ -4483,6 +4488,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, carried: OrgBudgetSnapshot | None = None, org_alias: str | None = None, + live_spend: float | None = None, ): """ Set org budget metrics after an LLM API request @@ -4500,7 +4506,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, org_alias=org_alias or "", - spend=carried.spend + response_cost, + spend=(carried.spend if live_spend is None else live_spend) + response_cost, max_budget=carried.max_budget, budget_reset_at=None, ) @@ -4593,6 +4599,7 @@ class PrometheusLogger(CustomLogger): self, end_user_id: str | None, response_cost: float, + live_spend: float | None = None, ): if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled(): return @@ -4618,7 +4625,7 @@ class PrometheusLogger(CustomLogger): budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=cached_customer.spend + response_cost, + spend=(cached_customer.spend if live_spend is None else live_spend) + response_cost, max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=None, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2c4c0ec4780..6510a889b79 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -169,7 +169,7 @@ from litellm.proxy.db.db_span import db_span from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.carried_budget_state import ( - carry_live_key_and_team_spend, + carry_live_budget_spend, carry_team_and_user_budget_state, ) from litellm.proxy.spend_tracking.spend_counter_batch import ( @@ -3969,9 +3969,11 @@ async def run_centralized_common_checks( ) if not route_skips_budget_checks(route=route): - await carry_live_key_and_team_spend( + await carry_live_budget_spend( valid_token=user_api_key_auth_obj, team_object=typed_team_object, + user_object=user_object, + end_user_object=end_user_object, read_spend=get_current_spend, ) await _reserve_budget_after_common_checks( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 2adf4300099..06cb792e78e 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,14 +1,18 @@ """Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" +import asyncio from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass from types import MappingProxyType from typing import Final +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.carried_budget_state import ( + LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -34,19 +38,61 @@ def carry_team_and_user_budget_state( ) -async def carry_live_key_and_team_spend( +@dataclass(frozen=True, slots=True) +class _BudgetedSpend: + counter_key: str + stored_spend: float + + +async def _read_live_spend( + budgeted: _BudgetedSpend | None, read_spend: Callable[[str, float], Awaitable[float]] +) -> float | None: + if budgeted is None: + return None + return max(await read_spend(budgeted.counter_key, budgeted.stored_spend), budgeted.stored_spend) + + +def _org_budgeted_spend(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> _BudgetedSpend | None: + org_snapshot: Final = valid_token.org_budget_snapshot + org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) + if org_snapshot is None or org_snapshot.max_budget is None or org_id is None: + return None + return _BudgetedSpend(f"spend:org:{org_id}", org_snapshot.spend) + + +def _end_user_budgeted_spend(end_user_object: LiteLLM_EndUserTable | None) -> _BudgetedSpend | None: + budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None + if end_user_object is None or budget_table is None or budget_table.max_budget is None: + return None + return _BudgetedSpend(f"spend:end_user:{end_user_object.user_id}", end_user_object.spend) + + +async def carry_live_budget_spend( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, + user_object: LiteLLM_UserTable | None, + end_user_object: LiteLLM_EndUserTable | None, read_spend: Callable[[str, float], Awaitable[float]], ) -> None: - if valid_token.max_budget is not None and valid_token.token is not None: - valid_token.spend = await read_spend( # rebind-ok: the budget gauges and spend header read this token - f"spend:key:{valid_token.token}", valid_token.spend or 0.0 - ) - if team_object is not None and team_object.max_budget is not None: - valid_token.team_spend = await read_spend( # rebind-ok: the team budget gauge reads this token - f"spend:team:{team_object.team_id}", team_object.spend or 0.0 - ) + budgeted: Final = ( + _BudgetedSpend(f"spend:key:{valid_token.token}", valid_token.spend or 0.0) + if valid_token.max_budget is not None and valid_token.token is not None + else None, + _BudgetedSpend(f"spend:team:{team_object.team_id}", team_object.spend or 0.0) + if team_object is not None and team_object.max_budget is not None + else None, + _BudgetedSpend(f"spend:user:{user_object.user_id}", user_object.spend or 0.0) + if user_object is not None and user_object.max_budget is not None + else None, + _org_budgeted_spend(valid_token=valid_token, team_object=team_object), + _end_user_budgeted_spend(end_user_object=end_user_object), + ) + key, team, user, org, end_user = await asyncio.gather( + *(_read_live_spend(budgeted=entity, read_spend=read_spend) for entity in budgeted) + ) + valid_token.live_budget_spend = LiveBudgetSpend( # rebind-ok: the request credential is pinned in place + key=key, team=team, user=user, org=org, end_user=end_user + ) def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: @@ -63,6 +109,7 @@ def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object] valid_token.team_budget_snapshot, valid_token.user_budget_snapshot, valid_token.org_budget_snapshot, + valid_token.live_budget_spend, ) return MappingProxyType( { diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 2c1c75fbe04..ec3cd514b60 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -16,6 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.user_roles import LitellmUserRoles from litellm.types.proxy.carried_budget_state import ( + LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -202,6 +203,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) + live_budget_spend: LiveBudgetSpend | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index 1a0c55a38ff..ba39eeefc59 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -64,3 +64,32 @@ class OrgBudgetSnapshot(_BudgetSnapshot): max_budget: float | None = Field( validation_alias="user_api_key_org_max_budget", serialization_alias="user_api_key_org_max_budget" ) + + +class LiveBudgetSpend(_BudgetSnapshot): + """The cross-replica spend counters the budget check compares, in-flight reservations included, read + only by the remaining budget gauges. ``None`` for an entity without a max budget.""" + + key: float | None = Field( + default=None, validation_alias="user_api_key_live_spend", serialization_alias="user_api_key_live_spend" + ) + team: float | None = Field( + default=None, + validation_alias="user_api_key_team_live_spend", + serialization_alias="user_api_key_team_live_spend", + ) + user: float | None = Field( + default=None, + validation_alias="user_api_key_user_live_spend", + serialization_alias="user_api_key_user_live_spend", + ) + org: float | None = Field( + default=None, + validation_alias="user_api_key_org_live_spend", + serialization_alias="user_api_key_org_live_spend", + ) + end_user: float | None = Field( + default=None, + validation_alias="user_api_key_end_user_live_spend", + serialization_alias="user_api_key_end_user_live_spend", + ) diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index e49e6baa8b0..e7ed767fc80 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -26,6 +26,7 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( ) from litellm.types.proxy.carried_budget_state import ( KeyBudgetSnapshot, + LiveBudgetSpend, TeamBudgetSnapshot, UserBudgetSnapshot, ) @@ -198,6 +199,28 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() +@pytest.mark.asyncio +async def test_gauges_subtract_the_live_spend_auth_carried_instead_of_the_cached_spend(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + token = _authed_token() + token.live_budget_spend = LiveBudgetSpend(key=3.0, team=30.0, user=8.0, org=60.0) + + await _emit(prometheus_logger, _request_metadata(token)) + + remaining = { + "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], + "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], + "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], + "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], + } + assert remaining == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + @pytest.mark.asyncio async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): _stub_gauges(prometheus_logger) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index ab0fb67b52f..cb9d9db93f1 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1065,6 +1065,37 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default assert _customer_sample("litellm_customer_max_budget_metric", "cust-default") == pytest.approx(3.0) +@pytest.mark.asyncio +async def test_customer_gauge_subtracts_the_live_spend_auth_carried(prometheus_logger, customer_metrics_enabled): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-live", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._increment_remaining_budget_metrics( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + litellm_params={"metadata": {"user_api_key_end_user_live_spend": 1.5}}, + response_cost=0.5, + end_user_id="cust-live", + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(1.0) + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( prometheus_logger, customer_metrics_enabled diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 5f4acff1661..fcca7c48bf7 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5081,12 +5081,17 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t "user_api_key_user_budget_reset_at": None, "user_api_key_user_table_max_budget": None, "user_api_key_user_alias": "Alice", + "user_api_key_live_spend": None, + "user_api_key_team_live_spend": 0.0, + "user_api_key_user_live_spend": None, + "user_api_key_org_live_spend": None, + "user_api_key_end_user_live_spend": None, } @pytest.mark.asyncio @pytest.mark.parametrize("model_input_cost", [None, 0.0], ids=["priced_model_reserves_budget", "zero_cost_model"]) -async def test_centralized_common_checks_carry_live_key_and_team_spend_on_a_token_with_stale_cached_spend( +async def test_centralized_common_checks_carry_live_budget_spend_beside_the_stale_cached_spend( model_input_cost: float | None, ): import litellm.proxy.proxy_server as _proxy_server_mod @@ -5143,7 +5148,9 @@ async def test_centralized_common_checks_carry_live_key_and_team_spend_on_a_toke for attribute, value in originals.items(): setattr(_proxy_server_mod, attribute, value) - assert (token.spend, token.team_spend) == (0.37, 1.25) + assert token.live_budget_spend is not None + assert (token.live_budget_spend.key, token.live_budget_spend.team) == (0.37, 1.25) + assert (token.spend, token.team_spend) == (0.0, 0.0) def _end_user_budget_row(budget_id: str, max_budget: float) -> MagicMock: diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 0bdf43b396c..2f83ef8b74f 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -2,6 +2,9 @@ from datetime import datetime, timezone +import pytest + +from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, @@ -11,11 +14,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, + carry_live_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) from litellm.types.proxy.carried_budget_state import ( KeyBudgetSnapshot, + LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -123,10 +128,105 @@ def test_snapshots_never_reach_the_serialized_token(): user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), ) token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) + token.live_budget_spend = LiveBudgetSpend(key=1.0) dumped = token.model_dump() assert "team_budget_snapshot" not in dumped assert "user_budget_snapshot" not in dumped assert "org_budget_snapshot" not in dumped + assert "live_budget_spend" not in dumped assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None + + +def _counter_reader(counters: dict[str, float]): + reads: list[str] = [] + + async def read_spend(counter_key: str, fallback_spend: float) -> float: + reads.append(counter_key) + return counters.get(counter_key, fallback_spend) + + return read_spend, reads + + +def _budgeted_token() -> UserAPIKeyAuth: + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + return token + + +@pytest.mark.asyncio +async def test_live_budget_spend_reads_every_budgeted_counter_and_leaves_the_header_spend_alone(): + """Each counter holds committed spend plus another request's 0.33 reservation, which only the gauges see.""" + token = _budgeted_token() + read_spend, _ = _counter_reader( + { + "spend:key:hashed": 1.63, + "spend:team:t1": 20.63, + "spend:user:u1": 5.63, + "spend:org:o1": 40.63, + "spend:end_user:cust-1": 0.83, + } + ) + + await carry_live_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0), + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0) + ), + read_spend=read_spend, + ) + + expected = LiveBudgetSpend(key=1.63, team=20.63, user=5.63, org=40.63, end_user=0.83) + assert token.live_budget_spend == expected + assert LiveBudgetSpend.from_metadata(carried_budget_metadata(token)) == expected + assert (token.spend, token.team_spend) == (1.0, None) + + +@pytest.mark.asyncio +async def test_live_budget_spend_never_reads_below_the_spend_auth_loaded(): + token = _budgeted_token() + read_spend, _ = _counter_reader({"spend:key:hashed": 0.2, "spend:org:o1": 12.0}) + + await carry_live_budget_spend( + valid_token=token, team_object=None, user_object=None, end_user_object=None, read_spend=read_spend + ) + + assert token.live_budget_spend == LiveBudgetSpend(key=1.0, org=40.0) + + +@pytest.mark.asyncio +async def test_live_budget_spend_skips_entities_without_a_max_budget(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) + read_spend, reads = _counter_reader({}) + + await carry_live_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0), + user_object=LiteLLM_UserTable(user_id="u1", spend=5.0), + end_user_object=LiteLLM_EndUserTable(user_id="cust-1", blocked=False, spend=0.5), + read_spend=read_spend, + ) + + assert reads == [] + assert token.live_budget_spend == LiveBudgetSpend() + + +@pytest.mark.asyncio +async def test_live_budget_spend_reads_the_team_org_when_the_key_has_none(): + token = UserAPIKeyAuth(token="hashed", team_id="t1") + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + read_spend, _ = _counter_reader({"spend:org:team-org": 41.0}) + + await carry_live_budget_spend( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"), + user_object=None, + end_user_object=None, + read_spend=read_spend, + ) + + assert token.live_budget_spend == LiveBudgetSpend(org=41.0) From f0066404b7668272a8ed315b3057a2a28784d350 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 06:25:37 -0700 Subject: [PATCH 3/9] fix(prometheus): show committed spend in the remaining budget gauges Reads twin counters that take only charged cost, so a worst-case budget hold never shows as spent and a request that finishes last cannot set a gauge back. The twins exist only while the Prometheus logger is registered, and every path that resets a spend counter drops them. Failed requests now set the org gauge from the carried org row and committed spend, not the cached org row. --- litellm/integrations/prometheus.py | 174 ++++++++++++----- litellm/proxy/auth/user_api_key_auth.py | 12 +- .../proxy/common_utils/reset_budget_job.py | 13 +- .../key_management_endpoints.py | 6 +- litellm/proxy/proxy_server.py | 59 +++++- .../spend_tracking/carried_budget_state.py | 60 ++---- .../proxy/spend_tracking/committed_spend.py | 144 ++++++++++++++ litellm/types/proxy/auth/user_api_key_auth.py | 4 +- litellm/types/proxy/carried_budget_state.py | 38 ++-- .../test_request_redis_batch_post_call.py | 2 +- .../test_prometheus_carried_budget_state.py | 180 ++++++++++++++++- .../test_prometheus_user_team_metrics.py | 19 +- .../test_user_api_key_auth_request_flow.py | 47 ++--- .../common_utils/test_reset_budget_job.py | 62 +++++- .../test_key_management_endpoints.py | 6 + .../proxy/proxy_server/test_spend_counters.py | 82 ++++++-- .../test_carried_budget_state.py | 102 +++++----- .../spend_tracking/test_committed_spend.py | 183 ++++++++++++++++++ .../test_spend_counter_batch.py | 9 +- ...test_proxy_server_endpoints_and_startup.py | 5 +- 20 files changed, 949 insertions(+), 258 deletions(-) create mode 100644 litellm/proxy/spend_tracking/committed_spend.py create mode 100644 tests/unit/proxy/spend_tracking/test_committed_spend.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f69a5f53d45..9f6b26021c5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -76,8 +76,8 @@ from litellm.types.integrations.prometheus import ( ) from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, KeyBudgetSnapshot, - LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -359,6 +359,60 @@ def _get_budget_metrics_per_request_timeout() -> float: return parsed +def _stored_spend(value: object) -> float: + return float(value) if isinstance(value, (int, float)) else 0.0 + + +def _committed_floor(stored: float, committed: float | None) -> float: + return stored if committed is None else max(stored, committed) + + +def _max_known_spend(first: float | None, second: float | None) -> float | None: + if first is None: + return second + return first if second is None else max(first, second) + + +def _committed_spend_at_auth(metadata: Mapping[str, object]) -> CommittedBudgetSpend: + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: + return auth.committed_budget_spend + return CommittedBudgetSpend() + + +async def _committed_spend_now( + token: str | None, team_id: str | None, user_id: str | None, org_id: str | None, end_user_id: str | None +) -> CommittedBudgetSpend: + try: + from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys, read_committed_budget_spend + + return await read_committed_budget_spend( + spend_counter_cache, + EntitySpendCounterKeys.of( + token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id + ), + ) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e) + return CommittedBudgetSpend() + + +def _committed_spend_before_request( + at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float +) -> CommittedBudgetSpend: + def before(at_auth_spend: float | None, now_spend: float | None) -> float | None: + return _max_known_spend(at_auth_spend, None if now_spend is None else now_spend - response_cost) + + return CommittedBudgetSpend( + key=before(at_auth.key, now.key), + team=before(at_auth.team, now.team), + user=before(at_auth.user, now.user), + org=before(at_auth.org, now.org), + end_user=before(at_auth.end_user, now.end_user), + ) + + def _get_proxy_llm_router() -> Router | None: try: from litellm.proxy.proxy_server import llm_router @@ -2223,7 +2277,7 @@ class PrometheusLogger(CustomLogger): _user_email: Final = _metadata.get("user_api_key_user_email", None) _org_alias: Final = _metadata.get("user_api_key_org_alias", None) _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) - _live_spend: Final = LiveBudgetSpend.from_metadata(_carried_metadata) or LiveBudgetSpend() + _committed_at_auth: Final = _committed_spend_at_auth(_carried_metadata) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2233,47 +2287,61 @@ class PrometheusLogger(CustomLogger): # so dropping one slow per-request emission only loses sub-cron real-time # detail, not correctness. budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout() - gather_coro: Final = asyncio.gather( - self._set_api_key_budget_metrics_after_api_request( - user_api_key=user_api_key, - user_api_key_alias=user_api_key_alias, + + async def emit_budget_metrics() -> tuple[object, ...]: + committed: Final = _committed_spend_before_request( + at_auth=_committed_at_auth, + now=await _committed_spend_now( + token=user_api_key, + team_id=user_api_team, + user_id=user_id, + org_id=user_api_key_org_id, + end_user_id=end_user_id, + ), response_cost=response_cost, - key_max_budget=_api_key_max_budget, - key_spend=_api_key_spend if _live_spend.key is None else _live_spend.key, - carried=KeyBudgetSnapshot.from_metadata(_carried_metadata), - ), - self._set_team_budget_metrics_after_api_request( - user_api_team=user_api_team, - user_api_team_alias=user_api_team_alias, - team_spend=_team_spend if _live_spend.team is None else _live_spend.team, - team_max_budget=_team_max_budget, - response_cost=response_cost, - carried=TeamBudgetSnapshot.from_metadata(_carried_metadata), - ), - self._set_user_budget_metrics_after_api_request( - user_id=user_id, - user_spend=_user_spend if _live_spend.user is None else _live_spend.user, - user_max_budget=_user_max_budget, - response_cost=response_cost, - carried=UserBudgetSnapshot.from_metadata(_carried_metadata), - user_email=_user_email if isinstance(_user_email, str) else None, - ), - self._set_org_budget_metrics_after_api_request( - org_id=user_api_key_org_id, - response_cost=response_cost, - carried=OrgBudgetSnapshot.from_metadata(_carried_metadata), - org_alias=_org_alias if isinstance(_org_alias, str) else None, - live_spend=_live_spend.org, - ), - self._set_customer_budget_metrics_after_api_request( - end_user_id=end_user_id, - response_cost=response_cost, - live_spend=_live_spend.end_user, - ), - return_exceptions=True, - ) + ) + return await asyncio.gather( + self._set_api_key_budget_metrics_after_api_request( + user_api_key=user_api_key, + user_api_key_alias=user_api_key_alias, + response_cost=response_cost, + key_max_budget=_api_key_max_budget, + key_spend=_committed_floor(_stored_spend(_api_key_spend), committed.key), + carried=KeyBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_team_budget_metrics_after_api_request( + user_api_team=user_api_team, + user_api_team_alias=user_api_team_alias, + team_spend=_committed_floor(_stored_spend(_team_spend), committed.team), + team_max_budget=_team_max_budget, + response_cost=response_cost, + carried=TeamBudgetSnapshot.from_metadata(_carried_metadata), + ), + self._set_user_budget_metrics_after_api_request( + user_id=user_id, + user_spend=_committed_floor(_stored_spend(_user_spend), committed.user), + user_max_budget=_user_max_budget, + response_cost=response_cost, + carried=UserBudgetSnapshot.from_metadata(_carried_metadata), + user_email=_user_email if isinstance(_user_email, str) else None, + ), + self._set_org_budget_metrics_after_api_request( + org_id=user_api_key_org_id, + response_cost=response_cost, + carried=OrgBudgetSnapshot.from_metadata(_carried_metadata), + org_alias=_org_alias if isinstance(_org_alias, str) else None, + committed_spend=committed.org, + ), + self._set_customer_budget_metrics_after_api_request( + end_user_id=end_user_id, + response_cost=response_cost, + committed_spend=committed.end_user, + ), + return_exceptions=True, + ) + try: - results: Final = await asyncio.wait_for(gather_coro, timeout=budget_metrics_timeout) + results: Final = await asyncio.wait_for(emit_budget_metrics(), timeout=budget_metrics_timeout) except asyncio.TimeoutError: verbose_logger.debug( "[Non-Blocking] Prometheus: per-request budget metric emission " @@ -2796,9 +2864,23 @@ class PrometheusLogger(CustomLogger): label_context=PrometheusLabelFactoryContext(enum_values), ) self.set_llm_deployment_failure_metrics(kwargs) + failure_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python( + get_litellm_metadata_from_kwargs(kwargs) + ) + failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") + failure_committed: Final = _committed_spend_before_request( + at_auth=_committed_spend_at_auth(failure_metadata), + now=await _committed_spend_now( + token=None, team_id=None, user_id=None, org_id=user_api_key_org_id, end_user_id=None + ), + response_cost=0, + ) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=0, + carried=OrgBudgetSnapshot.from_metadata(failure_metadata), + org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None, + committed_spend=failure_committed.org, ) except Exception as e: verbose_logger.exception("prometheus Layer Error(): Exception occured - %s", e) @@ -4488,7 +4570,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, carried: OrgBudgetSnapshot | None = None, org_alias: str | None = None, - live_spend: float | None = None, + committed_spend: float | None = None, ): """ Set org budget metrics after an LLM API request @@ -4506,7 +4588,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, org_alias=org_alias or "", - spend=(carried.spend if live_spend is None else live_spend) + response_cost, + spend=_committed_floor(carried.spend, committed_spend) + response_cost, max_budget=carried.max_budget, budget_reset_at=None, ) @@ -4532,7 +4614,7 @@ class PrometheusLogger(CustomLogger): if org_info is None: return - _total_org_spend: Final = (org_info.spend or 0.0) + response_cost + _total_org_spend: Final = _committed_floor(org_info.spend or 0.0, committed_spend) + response_cost budget_table: Final = org_info.litellm_budget_table max_budget: Final = budget_table.max_budget if budget_table else None budget_reset_at: Final = getattr(budget_table, "budget_reset_at", None) if budget_table else None @@ -4599,7 +4681,7 @@ class PrometheusLogger(CustomLogger): self, end_user_id: str | None, response_cost: float, - live_spend: float | None = None, + committed_spend: float | None = None, ): if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled(): return @@ -4625,7 +4707,7 @@ class PrometheusLogger(CustomLogger): budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=(cached_customer.spend if live_spend is None else live_spend) + response_cost, + spend=_committed_floor(cached_customer.spend, committed_spend) + response_cost, max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=None, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6510a889b79..8c8b4e9f774 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -13,6 +13,7 @@ import re import secrets from collections.abc import Mapping, Sequence from datetime import datetime, timezone +from functools import partial from types import MappingProxyType from typing import Final, NamedTuple, Protocol, Union, cast @@ -169,9 +170,10 @@ from litellm.proxy.db.db_span import db_span from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.carried_budget_state import ( - carry_live_budget_spend, + carry_committed_budget_spend, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, read_committed_budget_spend from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -3575,7 +3577,6 @@ async def run_centralized_common_checks( """ from litellm.proxy.proxy_server import ( general_settings, - get_current_spend, litellm_proxy_admin_name, llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries llm_router, @@ -3583,6 +3584,7 @@ async def run_centralized_common_checks( model_max_budget_limiter, prisma_client, proxy_logging_obj, + spend_counter_cache, user_api_key_cache, user_custom_auth, ) @@ -3968,13 +3970,13 @@ async def run_centralized_common_checks( user_obj=user_object, ) - if not route_skips_budget_checks(route=route): - await carry_live_budget_spend( + if committed_spend_tracked() and not route_skips_budget_checks(route=route): + await carry_committed_budget_spend( valid_token=user_api_key_auth_obj, team_object=typed_team_object, user_object=user_object, end_user_object=end_user_object, - read_spend=get_current_spend, + read_committed=partial(read_committed_budget_spend, spend_counter_cache), ) await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..67194a719ef 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -51,6 +51,7 @@ from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_span import db_span, db_spanned from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.organization_repository import OrganizationRepository @@ -610,15 +611,17 @@ class ResetBudgetJob: try: from litellm.proxy.proxy_server import spend_counter_cache - spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) - if spend_counter_cache.redis_cache is not None: + for stale_key in (counter_key, *committed_spend_keys((counter_key,))): + spend_counter_cache.in_memory_cache.delete_cache(key=stale_key) + if spend_counter_cache.redis_cache is None: + continue try: - await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) + await spend_counter_cache.redis_cache.async_delete_cache(key=stale_key) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. " "Budget may be over-enforced until counter expires.", - counter_key, + stale_key, redis_err, ) except Exception as e: @@ -675,7 +678,7 @@ class ResetBudgetJob: match cache: case "spend counter": - await spend_counter_cache.async_delete_cache_keys(keys) + await spend_counter_cache.async_delete_cache_keys((*keys, *committed_spend_keys(keys))) case "user_api_key_cache": await user_api_key_cache.async_delete_cache_keys(keys) case _: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 55e59ca5776..d590e60e88d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -156,6 +156,7 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import ( ) from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -4055,7 +4056,7 @@ async def update_key_fn( ) if data.spend is not None: - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import invalidate_spend_counter, spend_counter_cache counter_key: Final = f"spend:key:{hash_token_if_needed(key)}" spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60) @@ -4070,6 +4071,7 @@ async def update_key_fn( counter_key, redis_err, ) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) asyncio.create_task( KeyManagementEventHooks.async_key_updated_hook( @@ -6756,6 +6758,7 @@ async def reset_key_spend_fn( try: from litellm.proxy.proxy_server import ( hash_token, + invalidate_spend_counter, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -6810,6 +6813,7 @@ async def reset_key_spend_fn( # pre-reset value. _counter_key: Final = f"spend:key:{hashed_api_key}" await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) + await invalidate_spend_counter(counter_key=committed_spend_key(_counter_key)) await _reset_key_budget_windows( prisma_client=prisma_client, hashed_api_key=hashed_api_key, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5e37ff1ba65..4f2f30716a4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -908,6 +908,13 @@ from litellm.proxy.spend_tracking.budget_reservation import ( release_unbound_budget_reservation, stamp_budget_reservation_actual_cost, ) +from litellm.proxy.spend_tracking.committed_spend import ( + EntitySpendCounterKeys, + committed_spend_key, + committed_spend_keys, + committed_spend_tracked, + seed_committed_spend, +) from litellm.proxy.spend_tracking.spend_capture_rate import ( run_scheduled_spend_capture_rate_check, ) @@ -915,6 +922,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import ( SPEND_COUNTERS_TARGET, PendingSpendIncrement, active_spend_counter_batch, + bind_spend_counter_keys, forget_spend_counter, post_call_counter_keys, read_batched_spend_counter, @@ -3527,6 +3535,18 @@ async def _increment_spend_counters_batched( ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" + committed_counter_keys: Final = ( + EntitySpendCounterKeys.of( + token=_spend_counter_token(token) if token is not None else None, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + ).tracked() + if response_cost is not None and response_cost > 0 and committed_spend_tracked() + else () + ) + bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys))) reservation_update: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -3549,9 +3569,7 @@ async def _increment_spend_counters_batched( # if a raw key somehow arrives, hash it; otherwise use as-is to # avoid double-hashing (budget checks read valid_token.token which # is single-hashed). - hashed_token: Final = ( - hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token - ) + hashed_token: Final = _spend_counter_token(key_token) key_counter_key: Final = f"spend:key:{hashed_token}" key_pending: Final[tuple[PendingSpendIncrement, ...]] = ( () @@ -3731,6 +3749,9 @@ async def _increment_spend_counters_batched( ) if project_id is not None else None, + _prepare_committed_spend_increments(counter_keys=committed_counter_keys, increment=cost) + if committed_counter_keys + else None, ) if coro is not None ) @@ -3806,6 +3827,27 @@ async def _reconcile_budget_reservation_for_counter_update( return _ReservationCounterUpdate(reserved_counter_keys=frozenset(reserved_counter_keys), pending=pending) +def _spend_counter_token(key_token: str) -> str: + return hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token + + +async def _committed_spend_seeded(counter_key: str) -> bool: + return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend( + spend_counter_cache=spend_counter_cache, prisma_client=prisma_client, counter_key=counter_key + ) + + +async def _prepare_committed_spend_increments( + counter_keys: tuple[str, ...], increment: float +) -> tuple[PendingSpendIncrement, ...]: + seeded: Final = await asyncio.gather(*(_committed_spend_seeded(counter_key) for counter_key in counter_keys)) + return tuple( + PendingSpendIncrement(counter_key=committed_spend_key(counter_key), increment=increment) + for counter_key, is_seeded in zip(counter_keys, seeded) + if is_seeded + ) + + async def _prepare_end_user_and_tag_spend_increments( end_user_id: str | None, tags: list[str] | None, @@ -4182,11 +4224,12 @@ _increment_spend_counter_cache: Final = increment_spend_counter_cache @with_service_target(SPEND_COUNTERS_TARGET) async def invalidate_spend_counter(counter_key: str) -> None: - forget_spend_counter(counter_key) - spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) - if spend_counter_cache.redis_cache is not None: - with contextlib.suppress(Exception): - await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) + for stale_key in (counter_key, *committed_spend_keys((counter_key,))): + forget_spend_counter(stale_key) + spend_counter_cache.in_memory_cache.delete_cache(key=stale_key) + if spend_counter_cache.redis_cache is not None: + with contextlib.suppress(Exception): + await spend_counter_cache.redis_cache.async_delete_cache(key=stale_key) _invalidate_spend_counter: Final = invalidate_spend_counter diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 06cb792e78e..e86ba782eef 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,8 +1,6 @@ """Pins the budget state auth resolved onto ``UserAPIKeyAuth`` and emits it as request metadata.""" -import asyncio from collections.abc import Awaitable, Callable, Mapping -from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -11,8 +9,9 @@ from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys from litellm.types.proxy.carried_budget_state import ( - LiveBudgetSpend, + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -38,61 +37,35 @@ def carry_team_and_user_budget_state( ) -@dataclass(frozen=True, slots=True) -class _BudgetedSpend: - counter_key: str - stored_spend: float - - -async def _read_live_spend( - budgeted: _BudgetedSpend | None, read_spend: Callable[[str, float], Awaitable[float]] -) -> float | None: - if budgeted is None: - return None - return max(await read_spend(budgeted.counter_key, budgeted.stored_spend), budgeted.stored_spend) - - -def _org_budgeted_spend(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> _BudgetedSpend | None: +def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None: org_snapshot: Final = valid_token.org_budget_snapshot - org_id: Final = valid_token.org_id or (team_object.organization_id if team_object is not None else None) - if org_snapshot is None or org_snapshot.max_budget is None or org_id is None: + if org_snapshot is None or org_snapshot.max_budget is None: return None - return _BudgetedSpend(f"spend:org:{org_id}", org_snapshot.spend) + return valid_token.org_id or (team_object.organization_id if team_object is not None else None) -def _end_user_budgeted_spend(end_user_object: LiteLLM_EndUserTable | None) -> _BudgetedSpend | None: +def _budgeted_end_user_id(end_user_object: LiteLLM_EndUserTable | None) -> str | None: budget_table: Final = end_user_object.litellm_budget_table if end_user_object is not None else None if end_user_object is None or budget_table is None or budget_table.max_budget is None: return None - return _BudgetedSpend(f"spend:end_user:{end_user_object.user_id}", end_user_object.spend) + return end_user_object.user_id -async def carry_live_budget_spend( +async def carry_committed_budget_spend( valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, end_user_object: LiteLLM_EndUserTable | None, - read_spend: Callable[[str, float], Awaitable[float]], + read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]], ) -> None: - budgeted: Final = ( - _BudgetedSpend(f"spend:key:{valid_token.token}", valid_token.spend or 0.0) - if valid_token.max_budget is not None and valid_token.token is not None - else None, - _BudgetedSpend(f"spend:team:{team_object.team_id}", team_object.spend or 0.0) - if team_object is not None and team_object.max_budget is not None - else None, - _BudgetedSpend(f"spend:user:{user_object.user_id}", user_object.spend or 0.0) - if user_object is not None and user_object.max_budget is not None - else None, - _org_budgeted_spend(valid_token=valid_token, team_object=team_object), - _end_user_budgeted_spend(end_user_object=end_user_object), - ) - key, team, user, org, end_user = await asyncio.gather( - *(_read_live_spend(budgeted=entity, read_spend=read_spend) for entity in budgeted) - ) - valid_token.live_budget_spend = LiveBudgetSpend( # rebind-ok: the request credential is pinned in place - key=key, team=team, user=user, org=org, end_user=end_user + counter_keys: Final = EntitySpendCounterKeys.of( + token=valid_token.token if valid_token.max_budget is not None else None, + team_id=team_object.team_id if team_object is not None and team_object.max_budget is not None else None, + user_id=user_object.user_id if user_object is not None and user_object.max_budget is not None else None, + org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object), + end_user_id=_budgeted_end_user_id(end_user_object=end_user_object), ) + valid_token.committed_budget_spend = await read_committed(counter_keys) # rebind-ok: request-scoped credential def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: @@ -109,7 +82,6 @@ def carried_budget_metadata(valid_token: UserAPIKeyAuth) -> Mapping[str, object] valid_token.team_budget_snapshot, valid_token.user_budget_snapshot, valid_token.org_budget_snapshot, - valid_token.live_budget_spend, ) return MappingProxyType( { diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py new file mode 100644 index 00000000000..45bad9e2eed --- /dev/null +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -0,0 +1,144 @@ +"""Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds. + +The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept +only while a ``PrometheusLogger`` is registered. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter + +import litellm +from litellm._internal_context import with_service_target +from litellm._logging import verbose_proxy_logger +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.proxy.utils import PrismaClient + +_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX) +_FETCHED_VALUES: Final = TypeAdapter(dict[str, float | None]) +_STORED_VALUE: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) + + +def committed_spend_key(counter_key: str) -> str: + return f"committed_{counter_key}" + + +def _has_committed_twin(counter_key: str) -> bool: + return counter_key.startswith(_BASE_COUNTER_PREFIXES) and ":window:" not in counter_key + + +def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: + return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if _has_committed_twin(counter_key)) + + +@dataclass(frozen=True, slots=True) +class EntitySpendCounterKeys: + key: str | None = None + team: str | None = None + user: str | None = None + org: str | None = None + end_user: str | None = None + + @classmethod + def of( + cls, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + ) -> EntitySpendCounterKeys: + return cls( + key=f"spend:key:{token}" if token else None, + team=f"spend:team:{team_id}" if team_id else None, + user=f"spend:user:{user_id}" if user_id else None, + org=f"spend:org:{org_id}" if org_id else None, + end_user=f"{END_USER_COUNTER_PREFIX}{end_user_id}" if end_user_id else None, + ) + + def tracked(self) -> tuple[str, ...]: + counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) + return tuple( + counter_key for counter_key in counter_keys if counter_key is not None and _has_committed_twin(counter_key) + ) + + +def committed_spend_tracked() -> bool: + return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) + + +def _in_memory_committed_value(spend_counter_cache: DualCache, twin_key: str) -> float | None: + return _STORED_VALUE.validate_python( + spend_counter_cache.in_memory_cache.get_cache(key=twin_key) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped cache API + ) + + +async def _read_committed_values(spend_counter_cache: DualCache, twin_keys: tuple[str, ...]) -> Mapping[str, float]: + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + stored: Final = {twin_key: _in_memory_committed_value(spend_counter_cache, twin_key) for twin_key in twin_keys} + return {twin_key: value for twin_key, value in stored.items() if value is not None} + fetched: Final = _FETCHED_VALUES.validate_python( + await redis_cache.async_batch_get_cache(key_list=list(twin_keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped cache API + ) + return {twin_key: value for twin_key, value in fetched.items() if value is not None} + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def read_committed_budget_spend( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + if not twin_keys: + return CommittedBudgetSpend() + try: + committed: Final = await _read_committed_values(spend_counter_cache=spend_counter_cache, twin_keys=twin_keys) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) + return CommittedBudgetSpend() + + def spend_of(counter_key: str | None) -> float | None: + return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None + + return CommittedBudgetSpend( + key=spend_of(counter_keys.key), + team=spend_of(counter_keys.team), + user=spend_of(counter_keys.user), + org=spend_of(counter_keys.org), + end_user=spend_of(counter_keys.end_user), + ) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def seed_committed_spend( + spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str +) -> bool: + db_spend: Final = await ( + SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) + if counter_key.startswith(END_USER_COUNTER_PREFIX) + else SpendCounterReseed.from_db(prisma_client=prisma_client, counter_key=counter_key) + ) + if db_spend is None: + return False + twin_key: Final = committed_spend_key(counter_key) + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + if _in_memory_committed_value(spend_counter_cache, twin_key) is None: + spend_counter_cache.in_memory_cache.set_cache(key=twin_key, value=db_spend) # pyright: ignore[reportUnknownMemberType] # untyped cache API + return True + try: + await redis_cache.async_set_max(key=twin_key, value=db_spend) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e) + return False + return True diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index ec3cd514b60..317a5ab7a19 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -16,7 +16,7 @@ from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.auth.user_roles import LitellmUserRoles from litellm.types.proxy.carried_budget_state import ( - LiveBudgetSpend, + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -203,7 +203,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True) - live_budget_spend: LiveBudgetSpend | None = Field(default=None, exclude=True) + committed_budget_spend: CommittedBudgetSpend | None = Field(default=None, exclude=True) matched_model_access_groups: list[str] | None = Field(default=None, exclude=True) budget_throttle_pct: float | None = Field(default=None, exclude=True) user: object | None = None # Expanded user object when expand=user is used diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index ba39eeefc59..255f7af779c 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -25,7 +25,7 @@ class _BudgetSnapshot(LiteLLMBaseModel): @classmethod def from_metadata(cls, metadata: Mapping[str, object]) -> Self | None: try: - return cls.model_validate(metadata) + return cls.model_validate(metadata, by_name=False) except ValidationError: return None @@ -66,30 +66,14 @@ class OrgBudgetSnapshot(_BudgetSnapshot): ) -class LiveBudgetSpend(_BudgetSnapshot): - """The cross-replica spend counters the budget check compares, in-flight reservations included, read - only by the remaining budget gauges. ``None`` for an entity without a max budget.""" +class CommittedBudgetSpend(LiteLLMBaseModel): + """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the + remaining budget gauges only. ``None`` when the entity has no committed counter.""" - key: float | None = Field( - default=None, validation_alias="user_api_key_live_spend", serialization_alias="user_api_key_live_spend" - ) - team: float | None = Field( - default=None, - validation_alias="user_api_key_team_live_spend", - serialization_alias="user_api_key_team_live_spend", - ) - user: float | None = Field( - default=None, - validation_alias="user_api_key_user_live_spend", - serialization_alias="user_api_key_user_live_spend", - ) - org: float | None = Field( - default=None, - validation_alias="user_api_key_org_live_spend", - serialization_alias="user_api_key_org_live_spend", - ) - end_user: float | None = Field( - default=None, - validation_alias="user_api_key_end_user_live_spend", - serialization_alias="user_api_key_end_user_live_spend", - ) + model_config = ConfigDict(frozen=True) + + key: float | None = None + team: float | None = None + user: float | None = None + org: float | None = None + end_user: float | None = None diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 4b73bb6dd34..2f6c8ebfe6d 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -488,7 +488,7 @@ async def test_a_spend_counter_whose_increment_failed_is_invalidated_not_trusted await flush_post_call_redis_batches() assert spend_cache.in_memory_cache.get_cache("spend:key:k1") is None - assert redis_cache.alone == [("DEL", "spend:key:k1")] + assert redis_cache.alone == [("DEL", "spend:key:k1"), ("DEL", "committed_spend:key:k1")] assert spend_cache.in_memory_cache.get_cache("spend:team:t1") == 7.5 diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index e7ed767fc80..8753062bd3c 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from litellm.caching.dual_cache import DualCache from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ( LiteLLM_BudgetTable, @@ -24,9 +25,10 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_organization_budget_state, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, KeyBudgetSnapshot, - LiveBudgetSpend, TeamBudgetSnapshot, UserBudgetSnapshot, ) @@ -56,6 +58,15 @@ def prometheus_logger(): return PrometheusLogger() +@pytest.fixture(autouse=True) +def spend_counter_cache(): + cache = DualCache() + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_cache", cache + ): + yield cache + + @pytest.fixture def getters(): """Every response-path object getter, patched where prometheus imports them from.""" @@ -138,6 +149,7 @@ def _request_metadata(token: UserAPIKeyAuth) -> dict: return { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(token), **carried_budget_metadata(token), + "user_api_key_auth": token, } @@ -199,21 +211,69 @@ async def test_authed_request_sets_every_gauge_without_any_object_getter(prometh prometheus_logger.litellm_user_budget_remaining_hours_metric.labels().set.assert_called_once() -@pytest.mark.asyncio -async def test_gauges_subtract_the_live_spend_auth_carried_instead_of_the_cached_spend(prometheus_logger, getters): - _stub_gauges(prometheus_logger) - token = _authed_token() - token.live_budget_spend = LiveBudgetSpend(key=3.0, team=30.0, user=8.0, org=60.0) +def _store_counters(spend_counter_cache: DualCache, key: float, team: float, user: float, org: float) -> None: + for counter_key, value in ( + ("spend:key:hashed", key), + ("spend:team:t1", team), + ("spend:user:u1", user), + ("spend:org:o1", org), + ): + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value) - await _emit(prometheus_logger, _request_metadata(token)) - remaining = { +def _store_committed(spend_counter_cache: DualCache, key: float, team: float, user: float, org: float) -> None: + for counter_key, value in ( + ("spend:key:hashed", key), + ("spend:team:t1", team), + ("spend:user:u1", user), + ("spend:org:o1", org), + ): + spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key(counter_key), value=value) + + +def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]: + return { "key": prometheus_logger.litellm_remaining_api_key_budget_metric.labels().set.call_args[0][0], "team": prometheus_logger.litellm_remaining_team_budget_metric.labels().set.call_args[0][0], "user": prometheus_logger.litellm_remaining_user_budget_metric.labels().set.call_args[0][0], "org": prometheus_logger.litellm_remaining_org_budget_metric.labels().set.call_args[0][0], } - assert remaining == { + + +def _token_committed_at_auth() -> UserAPIKeyAuth: + token = _authed_token() + token.committed_budget_spend = CommittedBudgetSpend(key=3.0, team=30.0, user=8.0, org=60.0) + return token + + +@pytest.mark.asyncio +async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finished_after_auth( + prometheus_logger, getters, spend_counter_cache +): + """This request's 2.0 plus 1.0 from another request that finished first, so the gauge is never set back.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=60.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { "key": pytest.approx(5.0), "team": pytest.approx(268.0), "user": pytest.approx(40.0), @@ -221,6 +281,108 @@ async def test_gauges_subtract_the_live_spend_auth_carried_instead_of_the_cached } +@pytest.mark.asyncio +async def test_gauges_keep_the_auth_time_spend_when_the_proxy_spend_counters_cannot_be_read( + prometheus_logger, getters, monkeypatch +): + _stub_gauges(prometheus_logger) + monkeypatch.delattr("litellm.proxy.proxy_server.spend_counter_cache") + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_a_budget_hold_on_the_enforcement_counters_never_shows_in_the_gauges( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=5.0, team=32.0, user=10.0, org=62.0) + _store_counters(spend_counter_cache, key=9.0, team=36.0, user=14.0, org=66.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +def _failed_request(metadata: dict) -> dict: + return { + "model": "openai/gpt-5.4-nano", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": { + "model": "openai/gpt-5.4-nano", + "model_id": "dep-1", + "model_group": "chat", + "api_base": "https://api.openai.com", + "custom_llm_provider": "openai", + "status": "failure", + "request_tags": [], + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + "metadata": { + "user_api_key_hash": "hashed", + "user_api_key_alias": "key-alias", + "user_api_key_team_id": "t1", + "user_api_key_team_alias": "team-alias", + "user_api_key_user_id": "u1", + "user_api_key_org_id": "o1", + "requester_metadata": None, + "user_api_key_end_user_id": None, + }, + }, + "exception": Exception("upstream failed"), + } + + +@pytest.mark.asyncio +async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_committed_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=65.0) + _store_counters(spend_counter_cache, key=9.0, team=36.0, user=14.0, org=90.0) + + await prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ) + + assert getters["get_org_object"].await_count == 0 + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + +@pytest.mark.asyncio +async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): + _stub_gauges(prometheus_logger) + metadata = { + "user_api_key_team_spend": 20.0, + "user_api_key_team_max_budget": 300.0, + "spend": 1000.0, + "max_budget": 1.0, + "budget_reset_at": None, + "user_alias": "not-the-user", + } + + await _emit(prometheus_logger, metadata) + + assert {name: getter.await_count for name, getter in getters.items()} == { + "get_key_object": 1, + "get_team_object": 1, + "get_user_object": 1, + "get_org_object": 1, + } + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(458.0) + + @pytest.mark.asyncio async def test_metadata_without_carried_state_still_fetches_each_object_once(prometheus_logger, getters): _stub_gauges(prometheus_logger) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index cb9d9db93f1..3c67fd5364d 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1066,11 +1066,18 @@ async def test_set_customer_budget_metrics_after_api_request_uses_cached_default @pytest.mark.asyncio -async def test_customer_gauge_subtracts_the_live_spend_auth_carried(prometheus_logger, customer_metrics_enabled): +async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cached_spend( + prometheus_logger, customer_metrics_enabled +): + """Auth saw 1.5 committed, and by emit the counter holds 2.5 with this request's 0.5 and another request's.""" import sys + from litellm.caching.dual_cache import DualCache from litellm.models.budget import LiteLLM_BudgetTable from litellm.models.end_user import LiteLLM_EndUserTable + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend mock_proxy_server = MagicMock() mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( @@ -1081,6 +1088,12 @@ async def test_customer_gauge_subtracts_the_live_spend_auth_carried(prometheus_l litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), ) ) + mock_proxy_server.spend_counter_cache = DualCache() + mock_proxy_server.spend_counter_cache.in_memory_cache.set_cache( + key=committed_spend_key("spend:end_user:cust-live"), value=2.5 + ) + token = UserAPIKeyAuth(token="hashed") + token.committed_budget_spend = CommittedBudgetSpend(end_user=1.5) with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): await prometheus_logger._increment_remaining_budget_metrics( @@ -1088,12 +1101,12 @@ async def test_customer_gauge_subtracts_the_live_spend_auth_carried(prometheus_l user_api_team_alias=None, user_api_key=None, user_api_key_alias=None, - litellm_params={"metadata": {"user_api_key_end_user_live_spend": 1.5}}, + litellm_params={"metadata": {"user_api_key_auth": token}}, response_cost=0.5, end_user_id="cust-live", ) - assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(1.0) + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5) @pytest.mark.asyncio diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index fcca7c48bf7..af7a181c83b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5081,52 +5081,38 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t "user_api_key_user_budget_reset_at": None, "user_api_key_user_table_max_budget": None, "user_api_key_user_alias": "Alice", - "user_api_key_live_spend": None, - "user_api_key_team_live_spend": 0.0, - "user_api_key_user_live_spend": None, - "user_api_key_org_live_spend": None, - "user_api_key_end_user_live_spend": None, } @pytest.mark.asyncio -@pytest.mark.parametrize("model_input_cost", [None, 0.0], ids=["priced_model_reserves_budget", "zero_cost_model"]) -async def test_centralized_common_checks_carry_live_budget_spend_beside_the_stale_cached_spend( - model_input_cost: float | None, +@pytest.mark.parametrize("prometheus_registered", [True, False], ids=["prometheus_logger", "no_prometheus_logger"]) +async def test_centralized_common_checks_carry_the_committed_spend_beside_the_stale_cached_spend( + monkeypatch: pytest.MonkeyPatch, prometheus_registered: bool ): + """The enforcement counters also hold budget holds, the committed twins only charged spend.""" import litellm.proxy.proxy_server as _proxy_server_mod from starlette.datastructures import URL + from litellm.integrations.prometheus import PrometheusLogger + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.proxy.utils import ProxyLogging + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend - zero_cost_router: Final = ( - None - if model_input_cost is None - else litellm.Router( - model_list=[ - { - "model_name": "gpt-4o", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "sk-test", - "input_cost_per_token": model_input_cost, - "output_cost_per_token": model_input_cost, - }, - } - ] - ) - ) + for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): + monkeypatch.setattr(litellm, callback_list, []) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)] if prometheus_registered else []) token: Final = UserAPIKeyAuth(api_key="sk-live-spend", max_budget=10.0, spend=0.0, team_id="t-live", team_spend=0.0) user_api_key_cache: Final = UserApiKeyCache() user_api_key_cache.set_cache( key="team_id:t-live", value=LiteLLM_TeamTableCachedObj(team_id="t-live", max_budget=20.0, spend=0.0) ) spend_counter_cache: Final = DualCache() - spend_counter_cache.in_memory_cache.set_cache(key=f"spend:key:{token.token}", value=0.37) - spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t-live", value=1.25) + for counter_key, held, committed in ((f"spend:key:{token.token}", 0.67, 0.37), ("spend:team:t-live", 1.55, 1.25)): + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=held) + spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key(counter_key), value=committed) attrs: Final = { **_proxy_attrs_for_centralized_checks(), - "llm_router": zero_cost_router, + "llm_router": None, "prisma_client": MagicMock(), "user_api_key_cache": user_api_key_cache, "spend_counter_cache": spend_counter_cache, @@ -5148,8 +5134,9 @@ async def test_centralized_common_checks_carry_live_budget_spend_beside_the_stal for attribute, value in originals.items(): setattr(_proxy_server_mod, attribute, value) - assert token.live_budget_spend is not None - assert (token.live_budget_spend.key, token.live_budget_spend.team) == (0.37, 1.25) + assert token.committed_budget_spend == ( + CommittedBudgetSpend(key=0.37, team=1.25) if prometheus_registered else None + ) assert (token.spend, token.team_spend) == (0.0, 0.0) diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 183dd7919c1..b83a8bd522c 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -1454,6 +1454,64 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_ counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:key:sk-abc") counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:key:sk-abc") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="committed_spend:key:sk-abc") + + +def test_reset_budget_for_keys_without_redis_clears_the_counter_and_its_committed_twin( + reset_budget_job, mock_prisma_client, monkeypatch +): + from litellm.caching.dual_cache import DualCache + + spend_counter_cache: Final = DualCache() + for counter_key in ("spend:key:sk-abc", "committed_spend:key:sk-abc"): + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=100.0) + fake_module: Final = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_cache = spend_counter_cache + fake_module.user_api_key_cache = DualCache() + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + mock_prisma_client.data["key"] = [ + type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": datetime.now(timezone.utc), + "id": "key-1", + "token": "sk-abc", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:sk-abc") is None + assert spend_counter_cache.in_memory_cache.get_cache(key="committed_spend:key:sk-abc") is None + + +def test_a_redis_failure_on_the_key_counter_still_deletes_its_committed_twin( + reset_budget_job, mock_prisma_client, monkeypatch +): + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + mock_prisma_client.data["key"] = [ + type( + "Key", + (), + { + "spend": 100.0, + "budget_duration": "30d", + "budget_reset_at": datetime.now(timezone.utc), + "id": "key-1", + "token": "sk-abc", + }, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) + + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:key:sk-abc") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="committed_spend:key:sk-abc") def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): @@ -1812,7 +1870,9 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma invalidated: Final = { key for call in counter_cache.async_delete_cache_keys.await_args_list for key in call.args[0] } - assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} + assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} | { + f"committed_spend:end_user:cust-{i:06d}" for i in range(population) + } evicted: Final = { key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0] } diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 859824e4a72..203967e15ce 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -8427,6 +8427,9 @@ async def test_reset_key_spend_success(monkeypatch): mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:spend:key:{hashed_key}", value=50.0, ttl=5 ) + mock_spend_counter_cache.in_memory_cache.delete_cache.assert_called_once_with( + key=f"committed_spend:key:{hashed_key}" + ) @pytest.mark.asyncio @@ -8743,6 +8746,9 @@ async def test_update_key_spend_updates_counter(monkeypatch): mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( key=f"spend:key:{hashed_key}", value=0.0, ttl=60 ) + mock_spend_counter_cache.in_memory_cache.delete_cache.assert_called_once_with( + key=f"committed_spend:key:{hashed_key}" + ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..1d3255a9be8 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -42,6 +42,7 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.integrations.prometheus import PrometheusLogger from litellm.models.end_user import LiteLLM_EndUserTable from litellm.models.tag import LiteLLM_TagTable from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj @@ -50,6 +51,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_loca from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key, committed_spend_keys from .conftest import normalize @@ -851,6 +853,54 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 +class _VerificationTokenTable: + def __init__(self, spend: float) -> None: + self._spend: Final = spend + + async def find_unique(self, where: dict[str, str]) -> MagicMock: + return MagicMock(token=where["token"], spend=self._spend) + + +def _prisma_with_key_spend(spend: float) -> MagicMock: + prisma: Final = MagicMock() + prisma.db.litellm_verificationtoken = _VerificationTokenTable(spend) + return prisma + + +@pytest.mark.asyncio +async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypatch): + """The key counter holds the DB spend 10.0, this request's 2.0 reservation and another request's 1.0 hold.""" + counter_cache: Final = ps.DualCache() + counter_cache.in_memory_cache.set_cache(key="spend:key:hashed", value=13.0) + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False} + + await ps.increment_spend_counters( + token="hashed", team_id=None, user_id=None, response_cost=0.5, budget_reservation=reservation + ) + await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.25) + + assert counter_cache.in_memory_cache.get_cache(key="spend:key:hashed") == pytest.approx(11.75) + assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:key:hashed")) == pytest.approx(10.75) + + +@pytest.mark.asyncio +async def test_no_committed_twin_is_kept_without_a_prometheus_logger(monkeypatch): + counter_cache: Final = ps.DualCache() + counter_cache.in_memory_cache.set_cache(key="spend:key:hashed", value=10.0) + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + + await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5) + + assert counter_cache.in_memory_cache.get_cache(key="spend:key:hashed") == pytest.approx(10.5) + assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:key:hashed")) is None + + _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -895,7 +945,7 @@ async def _increment_every_spend_counter() -> None: def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None: assert fake_redis.async_increment.await_count == 0 deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} - assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS) + assert deleted_keys == {*_PENDING_SPEND_COUNTER_KEYS, *committed_spend_keys(_PENDING_SPEND_COUNTER_KEYS)} for key in _PENDING_SPEND_COUNTER_KEYS: assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None @@ -1259,7 +1309,7 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] + assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1277,7 +1327,7 @@ async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_retu await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] + assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1291,7 +1341,7 @@ async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_ await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["spend:key:k", "spend:team:t"] + assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1540,23 +1590,23 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( # --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("counter_key", "deleted"), + [ + ("spend:key:k", ["spend:key:k", "committed_spend:key:k"]), + ("spend:key:k:window:1d", ["spend:key:k:window:1d"]), + ("spend:tag:t", ["spend:tag:t"]), + ], +) @pytest.mark.asyncio -async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch): +async def test_invalidate_spend_counter_deletes_the_counter_and_its_committed_twin(monkeypatch, counter_key, deleted): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._invalidate_spend_counter(counter_key="spend:key:k") + await ps._invalidate_spend_counter(counter_key=counter_key) - observed = { - "in_memory_delete_called": fake_cache.in_memory_cache.delete_cache.called, - "redis_delete_called": fake_cache.redis_cache.async_delete_cache.called, - "delete_args_key": fake_cache.redis_cache.async_delete_cache.call_args.kwargs["key"], - } - assert normalize(observed) == { - "in_memory_delete_called": True, - "redis_delete_called": True, - "delete_args_key": "spend:key:k", - } + assert [c.kwargs["key"] for c in fake_cache.in_memory_cache.delete_cache.call_args_list] == deleted + assert [c.kwargs["key"] for c in fake_cache.redis_cache.async_delete_cache.await_args_list] == deleted @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 2f83ef8b74f..69acc624623 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -14,13 +14,14 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, - carry_live_budget_spend, + carry_committed_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) +from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, KeyBudgetSnapshot, - LiveBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -128,105 +129,92 @@ def test_snapshots_never_reach_the_serialized_token(): user_object=LiteLLM_UserTable(user_id="u1", user_alias="Alice"), ) token.org_budget_snapshot = OrgBudgetSnapshot(spend=1.0, max_budget=2.0) - token.live_budget_spend = LiveBudgetSpend(key=1.0) + metadata_without_committed_spend = dict(carried_budget_metadata(token)) + token.committed_budget_spend = CommittedBudgetSpend(key=1.0, team=2.0, user=3.0, org=4.0, end_user=5.0) dumped = token.model_dump() assert "team_budget_snapshot" not in dumped assert "user_budget_snapshot" not in dumped assert "org_budget_snapshot" not in dumped - assert "live_budget_spend" not in dumped + assert "committed_budget_spend" not in dumped + assert dict(carried_budget_metadata(token)) == metadata_without_committed_spend assert UserAPIKeyAuth(**dumped).team_budget_snapshot is None -def _counter_reader(counters: dict[str, float]): - reads: list[str] = [] +def _committed_reader(committed: CommittedBudgetSpend): + reads: list[EntitySpendCounterKeys] = [] - async def read_spend(counter_key: str, fallback_spend: float) -> float: - reads.append(counter_key) - return counters.get(counter_key, fallback_spend) + async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + reads.append(counter_keys) + return committed - return read_spend, reads - - -def _budgeted_token() -> UserAPIKeyAuth: - token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) - token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) - return token + return read_committed, reads @pytest.mark.asyncio -async def test_live_budget_spend_reads_every_budgeted_counter_and_leaves_the_header_spend_alone(): - """Each counter holds committed spend plus another request's 0.33 reservation, which only the gauges see.""" - token = _budgeted_token() - read_spend, _ = _counter_reader( - { - "spend:key:hashed": 1.63, - "spend:team:t1": 20.63, - "spend:user:u1": 5.63, - "spend:org:o1": 40.63, - "spend:end_user:cust-1": 0.83, - } - ) +async def test_committed_spend_is_read_once_for_every_budgeted_entity(): + token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0, max_budget=10.0) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) + committed = CommittedBudgetSpend(key=1.3, team=20.3, user=5.3, org=40.3, end_user=0.8) + read_committed, reads = _committed_reader(committed) - await carry_live_budget_spend( + await carry_committed_budget_spend( valid_token=token, team_object=LiteLLM_TeamTable(team_id="t1", max_budget=300.0, spend=20.0), user_object=LiteLLM_UserTable(user_id="u1", max_budget=50.0, spend=5.0), end_user_object=LiteLLM_EndUserTable( user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=3.0) ), - read_spend=read_spend, + read_committed=read_committed, ) - expected = LiveBudgetSpend(key=1.63, team=20.63, user=5.63, org=40.63, end_user=0.83) - assert token.live_budget_spend == expected - assert LiveBudgetSpend.from_metadata(carried_budget_metadata(token)) == expected + assert reads == [ + EntitySpendCounterKeys( + key="spend:key:hashed", + team="spend:team:t1", + user="spend:user:u1", + org="spend:org:o1", + end_user="spend:end_user:cust-1", + ) + ] + assert token.committed_budget_spend == committed assert (token.spend, token.team_spend) == (1.0, None) @pytest.mark.asyncio -async def test_live_budget_spend_never_reads_below_the_spend_auth_loaded(): - token = _budgeted_token() - read_spend, _ = _counter_reader({"spend:key:hashed": 0.2, "spend:org:o1": 12.0}) - - await carry_live_budget_spend( - valid_token=token, team_object=None, user_object=None, end_user_object=None, read_spend=read_spend - ) - - assert token.live_budget_spend == LiveBudgetSpend(key=1.0, org=40.0) - - -@pytest.mark.asyncio -async def test_live_budget_spend_skips_entities_without_a_max_budget(): +async def test_committed_spend_skips_entities_without_a_max_budget(): token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) - read_spend, reads = _counter_reader({}) + read_committed, reads = _committed_reader(CommittedBudgetSpend()) - await carry_live_budget_spend( + await carry_committed_budget_spend( valid_token=token, team_object=LiteLLM_TeamTable(team_id="t1", spend=20.0), user_object=LiteLLM_UserTable(user_id="u1", spend=5.0), - end_user_object=LiteLLM_EndUserTable(user_id="cust-1", blocked=False, spend=0.5), - read_spend=read_spend, + end_user_object=LiteLLM_EndUserTable( + user_id="cust-1", blocked=False, spend=0.5, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None) + ), + read_committed=read_committed, ) - assert reads == [] - assert token.live_budget_spend == LiveBudgetSpend() + assert reads == [EntitySpendCounterKeys()] + assert token.committed_budget_spend == CommittedBudgetSpend() @pytest.mark.asyncio -async def test_live_budget_spend_reads_the_team_org_when_the_key_has_none(): +async def test_committed_spend_reads_the_team_org_when_the_key_has_none(): token = UserAPIKeyAuth(token="hashed", team_id="t1") token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=500.0) - read_spend, _ = _counter_reader({"spend:org:team-org": 41.0}) + read_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0)) - await carry_live_budget_spend( + await carry_committed_budget_spend( valid_token=token, team_object=LiteLLM_TeamTable(team_id="t1", organization_id="team-org"), user_object=None, end_user_object=None, - read_spend=read_spend, + read_committed=read_committed, ) - assert token.live_budget_spend == LiveBudgetSpend(org=41.0) + assert reads == [EntitySpendCounterKeys(org="spend:org:team-org")] + assert token.committed_budget_spend == CommittedBudgetSpend(org=41.0) diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py new file mode 100644 index 00000000000..ea76a61047d --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -0,0 +1,183 @@ +"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock + +import pytest +from fakeredis import FakeAsyncRedis, FakeServer + +from litellm._service_logger import ServiceLogging +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker +from litellm.proxy.spend_tracking.committed_spend import ( + EntitySpendCounterKeys, + committed_spend_key, + read_committed_budget_spend, + seed_committed_spend, +) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend + + +class _FakeRedisSettingMax(FakeAsyncRedis): + """fakeredis runs no Lua without lupa, so this stands in for the GET/compare/SET script of ``async_set_max``.""" + + async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> bytes: + key, value, _ttl = keys_and_args + current: Final = await self.get(key) + if current is not None and float(current) >= float(value): + return current + await self.set(key, value) + return value.encode() + + +class _RedisCacheOnClient(RedisCache): + def __init__(self, client: FakeAsyncRedis) -> None: + self._injected_client: Final = client + self.namespace = None + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self.service_logger_obj = ServiceLogging() + self.default_ttl = None + + def init_async_client(self) -> FakeAsyncRedis: + return self._injected_client + + +class _FindUniqueTable: + def __init__(self, row: SimpleNamespace | None) -> None: + self._row: Final = row + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + return self._row + + +def _prisma_with_spend(spend: float | None) -> SimpleNamespace: + row: Final = None if spend is None else SimpleNamespace(spend=spend) + return SimpleNamespace( + db=SimpleNamespace( + litellm_verificationtoken=_FindUniqueTable(row), + litellm_teamtable=_FindUniqueTable(row), + litellm_endusertable=_FindUniqueTable(row), + ) + ) + + +@pytest.fixture(params=["in_memory", "redis"]) +def spend_counter_cache(request: pytest.FixtureRequest) -> DualCache: + if request.param == "in_memory": + return DualCache(in_memory_cache=InMemoryCache()) + return DualCache( + in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(_FakeRedisSettingMax(server=FakeServer())) + ) + + +async def _store(spend_counter_cache: DualCache, key: str, value: float) -> None: + if spend_counter_cache.redis_cache is None: + spend_counter_cache.in_memory_cache.set_cache(key=key, value=value) + return + await spend_counter_cache.redis_cache.async_set_cache(key=key, value=value) + + +@pytest.mark.asyncio +async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(spend_counter_cache: DualCache): + await _store(spend_counter_cache, committed_spend_key("spend:key:hashed"), 1.5) + await _store(spend_counter_cache, committed_spend_key("spend:end_user:cust-1"), 0.25) + await _store(spend_counter_cache, "spend:team:t1", 99.0) + + committed: Final = await read_committed_budget_spend( + spend_counter_cache, + EntitySpendCounterKeys.of(token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1"), + ) + + assert committed == CommittedBudgetSpend(key=1.5, end_user=0.25) + + +@pytest.mark.asyncio +async def test_read_of_a_twin_holding_garbage_reports_no_committed_spend(): + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache()) + spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key("spend:key:hashed"), value="not-a-number") + spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key("spend:team:t1"), value=4.0) + + committed: Final = await read_committed_budget_spend( + spend_counter_cache, + EntitySpendCounterKeys.of(token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None), + ) + + assert committed == CommittedBudgetSpend() + + +@pytest.mark.parametrize( + ("counter_keys", "expected"), + [ + (EntitySpendCounterKeys(key="spend:key:hashed"), CommittedBudgetSpend(key=12.0)), + (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), CommittedBudgetSpend(end_user=12.0)), + ], +) +@pytest.mark.asyncio +async def test_seed_writes_the_db_spend_when_the_twin_is_absent( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys, expected: CommittedBudgetSpend +): + (counter_key,) = counter_keys.tracked() + + seeded: Final = await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), counter_key) + + assert seeded is True + assert await read_committed_budget_spend(spend_counter_cache, counter_keys) == expected + + +@pytest.mark.asyncio +async def test_seed_never_lowers_a_twin_another_replica_already_advanced(spend_counter_cache: DualCache): + await _store(spend_counter_cache, committed_spend_key("spend:team:t1"), 20.0) + + seeded: Final = await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend( + spend_counter_cache, EntitySpendCounterKeys(team="spend:team:t1") + ) + assert seeded is True + assert committed.team == 20.0 + + +@pytest.mark.asyncio +async def test_read_without_any_counter_key_sends_redis_nothing(): + redis_cache: Final = AsyncMock(spec=RedisCache) + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + committed: Final = await read_committed_budget_spend(spend_counter_cache, EntitySpendCounterKeys()) + + assert committed == CommittedBudgetSpend() + redis_cache.async_batch_get_cache.assert_not_awaited() + + +class _UnreachableRedis(FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> bytes: + raise ConnectionError("redis unreachable") + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_redis_rejects_the_write(): + spend_counter_cache: Final = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(_UnreachableRedis(server=FakeServer())) + ) + + seeded: Final = await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), "spend:team:t1") + + committed: Final = await read_committed_budget_spend( + spend_counter_cache, EntitySpendCounterKeys(team="spend:team:t1") + ) + assert seeded is False + assert committed == CommittedBudgetSpend() + + +@pytest.mark.asyncio +async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(spend_counter_cache: DualCache): + seeded: Final = await seed_committed_spend(spend_counter_cache, _prisma_with_spend(None), "spend:key:hashed") + + committed: Final = await read_committed_budget_spend( + spend_counter_cache, EntitySpendCounterKeys(key="spend:key:hashed") + ) + assert seeded is False + assert committed == CommittedBudgetSpend() diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 8ba79a45ff3..041a278acb9 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -13,6 +13,7 @@ import litellm.proxy.proxy_server as ps from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys from litellm.proxy.spend_tracking.spend_counter_batch import ( SpendCounterBatch, active_spend_counter_batch, @@ -539,8 +540,10 @@ async def test_a_failed_post_call_pipeline_invalidates_every_counter_it_carried_ ) assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands - assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == RESERVED_KEYS | { - "spend:user:user" + invalidated = RESERVED_KEYS | {"spend:user:user"} + assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == { + *invalidated, + *committed_spend_keys(invalidated), } assert all("applied_adjustment" not in entry for entry in reservation["entries"]) assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.0 for key in POST_CALL_KEYS} @@ -631,6 +634,8 @@ async def test_a_failed_reservation_pipeline_drops_every_counter_and_reserves_no assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == { "spend:key:hashed", "spend:team:team", + "committed_spend:key:hashed", + "committed_spend:team:team", } assert redis.store == {"spend:key:hashed": 1.0, "spend:team:team": 2.0} diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 55d1587669a..449516e8ba5 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -9403,7 +9403,10 @@ async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure( ) assert counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None - fake_redis.async_delete_cache.assert_awaited_once_with(key="spend:team:redis-fail") + assert [call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list] == [ + "spend:team:redis-fail", + "committed_spend:team:redis-fail", + ] finally: ps.spend_counter_cache = orig_counter From 632067f20d6c851b8c97fa05205dc671827a0373 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 06:57:31 -0700 Subject: [PATCH 4/9] fix(prometheus): count a charge the cached customer row already holds once --- litellm/integrations/prometheus.py | 10 ++++++- .../test_prometheus_user_team_metrics.py | 29 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 9f6b26021c5..935a9ffdaed 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -373,6 +373,12 @@ def _max_known_spend(first: float | None, second: float | None) -> float | None: return first if second is None else max(first, second) +def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float: + if committed_spend is None: + return cached_spend + response_cost + return max(committed_spend + response_cost, cached_spend) + + def _committed_spend_at_auth(metadata: Mapping[str, object]) -> CommittedBudgetSpend: auth: Final = metadata.get("user_api_key_auth") if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: @@ -4707,7 +4713,9 @@ class PrometheusLogger(CustomLogger): budget_table: Final = cached_customer.litellm_budget_table self._set_customer_budget_metrics( end_user_id=end_user_id, - spend=_committed_floor(cached_customer.spend, committed_spend) + response_cost, + spend=_customer_spend_after_request( + cached_spend=cached_customer.spend, committed_spend=committed_spend, response_cost=response_cost + ), max_budget=budget_table.max_budget if budget_table is not None else None, budget_reset_at=None, ) diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index 3c67fd5364d..01eed51fc4e 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1109,6 +1109,35 @@ async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cache assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-live") == pytest.approx(0.5) +@pytest.mark.asyncio +async def test_customer_gauge_counts_a_charge_the_cached_customer_row_already_holds_once( + prometheus_logger, customer_metrics_enabled +): + import sys + + from litellm.models.budget import LiteLLM_BudgetTable + from litellm.models.end_user import LiteLLM_EndUserTable + + mock_proxy_server = MagicMock() + mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( + return_value=LiteLLM_EndUserTable( + user_id="cust-charged", + blocked=False, + spend=0.5, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b", max_budget=3.0), + ) + ) + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): + await prometheus_logger._set_customer_budget_metrics_after_api_request( + end_user_id="cust-charged", + response_cost=0.5, + committed_spend=0.0, + ) + + assert _customer_sample("litellm_remaining_customer_budget_metric", "cust-charged") == pytest.approx(2.5) + + @pytest.mark.asyncio async def test_set_customer_budget_metrics_after_api_request_without_budget_only_emits_remaining( prometheus_logger, customer_metrics_enabled From aec15b56b6fb775c1b448659096c7d7ad6785a80 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:03:50 -0700 Subject: [PATCH 5/9] fix(prometheus): keep committed spend only for budgeted entities and drop it on spend resets Auth now reads committed spend through the admission MGET and names the counters a charge may move, so the master key and unbudgeted entities get no twin. Organization spend edits and the global spend reset drop their twins, and a long request no longer reads as charged. --- litellm/integrations/prometheus.py | 59 +++----- litellm/proxy/auth/user_api_key_auth.py | 16 +- litellm/proxy/db/spend_counter_reseed.py | 3 +- .../proxy/hooks/proxy_track_cost_callback.py | 7 + .../internal_user_endpoints.py | 5 +- .../organization_endpoints.py | 6 + litellm/proxy/proxy_server.py | 27 ++-- .../spend_tracking/carried_budget_state.py | 10 +- .../proxy/spend_tracking/committed_spend.py | 113 +++++++------- .../spend_management_endpoints.py | 33 +++++ litellm/types/proxy/auth/user_api_key_auth.py | 7 + litellm/types/proxy/carried_budget_state.py | 44 +++++- .../test_request_redis_batch_post_call.py | 2 +- .../test_prometheus_carried_budget_state.py | 76 +++++++++- .../test_prometheus_user_team_metrics.py | 6 +- .../test_user_api_key_auth_request_flow.py | 10 +- .../hooks/test_proxy_track_cost_callback.py | 43 ++++++ .../test_internal_user_endpoints.py | 5 +- .../test_organization_endpoints.py | 22 +++ .../proxy/proxy_server/test_spend_counters.py | 69 ++++++--- .../test_carried_budget_state.py | 15 +- .../spend_tracking/test_committed_spend.py | 138 +++++++++++++++--- .../test_spend_counter_batch.py | 8 +- .../test_spend_management_endpoints.py | 59 ++++++++ ...test_proxy_server_endpoints_and_startup.py | 1 - 25 files changed, 614 insertions(+), 170 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 935a9ffdaed..1f3b208a460 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -74,9 +74,10 @@ from litellm.types.integrations.prometheus import ( sanitize_prometheus_label_value, validate_prometheus_deployment_and_latency_caller_identity, ) -from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth, committed_budget_spend_from_metadata from litellm.types.proxy.carried_budget_state import ( CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -363,41 +364,29 @@ def _stored_spend(value: object) -> float: return float(value) if isinstance(value, (int, float)) else 0.0 +_OWN_CHARGE_MATCH_RATIO: Final = 1 - 1e-9 + + def _committed_floor(stored: float, committed: float | None) -> float: return stored if committed is None else max(stored, committed) -def _max_known_spend(first: float | None, second: float | None) -> float | None: - if first is None: - return second - return first if second is None else max(first, second) - - def _customer_spend_after_request(cached_spend: float, committed_spend: float | None, response_cost: float) -> float: if committed_spend is None: return cached_spend + response_cost return max(committed_spend + response_cost, cached_spend) -def _committed_spend_at_auth(metadata: Mapping[str, object]) -> CommittedBudgetSpend: - auth: Final = metadata.get("user_api_key_auth") - if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: - return auth.committed_budget_spend - return CommittedBudgetSpend() - - -async def _committed_spend_now( - token: str | None, team_id: str | None, user_id: str | None, org_id: str | None, end_user_id: str | None -) -> CommittedBudgetSpend: +async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + if not counter_keys.tracked(): + return CommittedBudgetSpend() try: from litellm.proxy.proxy_server import spend_counter_cache - from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys, read_committed_budget_spend + from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend - return await read_committed_budget_spend( - spend_counter_cache, - EntitySpendCounterKeys.of( - token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id - ), + return await asyncio.wait_for( + read_committed_budget_spend(spend_counter_cache, counter_keys), + timeout=_get_budget_metrics_per_request_timeout(), ) except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend verbose_logger.debug("[Non-Blocking] Prometheus: committed spend read failed: %s", e) @@ -408,7 +397,12 @@ def _committed_spend_before_request( at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float ) -> CommittedBudgetSpend: def before(at_auth_spend: float | None, now_spend: float | None) -> float | None: - return _max_known_spend(at_auth_spend, None if now_spend is None else now_spend - response_cost) + if now_spend is None: + return at_auth_spend + if at_auth_spend is None: + return now_spend - response_cost + own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO + return max(at_auth_spend, now_spend - response_cost) if own_charge_counted else now_spend return CommittedBudgetSpend( key=before(at_auth.key, now.key), @@ -2283,7 +2277,7 @@ class PrometheusLogger(CustomLogger): _user_email: Final = _metadata.get("user_api_key_user_email", None) _org_alias: Final = _metadata.get("user_api_key_org_alias", None) _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) - _committed_at_auth: Final = _committed_spend_at_auth(_carried_metadata) + _committed_at_auth: Final = committed_budget_spend_from_metadata(_carried_metadata) # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2297,13 +2291,7 @@ class PrometheusLogger(CustomLogger): async def emit_budget_metrics() -> tuple[object, ...]: committed: Final = _committed_spend_before_request( at_auth=_committed_at_auth, - now=await _committed_spend_now( - token=user_api_key, - team_id=user_api_team, - user_id=user_id, - org_id=user_api_key_org_id, - end_user_id=end_user_id, - ), + now=await _committed_spend_now(_committed_at_auth.counter_keys), response_cost=response_cost, ) return await asyncio.gather( @@ -2874,11 +2862,10 @@ class PrometheusLogger(CustomLogger): get_litellm_metadata_from_kwargs(kwargs) ) failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") + failure_at_auth: Final = committed_budget_spend_from_metadata(failure_metadata) failure_committed: Final = _committed_spend_before_request( - at_auth=_committed_spend_at_auth(failure_metadata), - now=await _committed_spend_now( - token=None, team_id=None, user_id=None, org_id=user_api_key_org_id, end_user_id=None - ), + at_auth=failure_at_auth, + now=await _committed_spend_now(EntitySpendCounterKeys(org=failure_at_auth.counter_keys.org)), response_cost=0, ) await self._set_org_budget_metrics_after_api_request( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 8c8b4e9f774..adaf9da0dfa 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -173,7 +173,11 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_committed_budget_spend, carry_team_and_user_budget_state, ) -from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, read_committed_budget_spend +from litellm.proxy.spend_tracking.committed_spend import ( + bind_committed_admission_counters, + committed_spend_tracked, + read_admission_committed_spend, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( bind_admission_counter_keys, release_spend_counter_batch, @@ -3895,6 +3899,7 @@ async def run_centralized_common_checks( ) bind_admission_counter_keys(user_api_key_auth_obj, end_user_id=end_user_id) + bind_committed_admission_counters(user_api_key_auth_obj, end_user_id=end_user_id) try: _ = await common_checks( request=request, @@ -3970,13 +3975,17 @@ async def run_centralized_common_checks( user_obj=user_object, ) - if committed_spend_tracked() and not route_skips_budget_checks(route=route): + if committed_spend_tracked(): await carry_committed_budget_spend( valid_token=user_api_key_auth_obj, team_object=typed_team_object, user_object=user_object, end_user_object=end_user_object, - read_committed=partial(read_committed_budget_spend, spend_counter_cache), + read_committed=( + None + if route_skips_budget_checks(route=route) + else partial(read_admission_committed_spend, spend_counter_cache) + ), ) await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, @@ -4294,6 +4303,7 @@ async def _prefetch_referenced_auth_objects( """Warm every object and spend counter the checks below will read, in one MGET each (one DB query when cold). Runs after the key's model access check so a denied request costs no more than it did before.""" bind_admission_counter_keys(valid_token, end_user_id=end_user_id or None) + bind_committed_admission_counters(valid_token, end_user_id=end_user_id or None) await prefetch_auth_objects( refs=AuthObjectRefs.from_token(valid_token), user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..f30b351eb5a 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -43,6 +43,7 @@ from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) +from litellm.types.proxy.carried_budget_state import END_USER_COUNTER_PREFIX if TYPE_CHECKING: from prisma.types import LiteLLM_EndUserTableWhereUniqueInput @@ -58,8 +59,6 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType( } ) -END_USER_COUNTER_PREFIX: Final = "spend:end_user:" - _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType( { "Key": "api_key", diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c11b59a809a..97018b74b95 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -67,6 +67,8 @@ from litellm.proxy.utils import ( BilledFailedAttempt, ProxyUpdateSpend, ) +from litellm.types.proxy.auth.user_api_key_auth import committed_budget_spend_from_metadata +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from litellm.types.router import Deployment from litellm.types.utils import ( CallTypes, @@ -1015,6 +1017,8 @@ class _IncrementSpendCounters(Protocol): tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> None: ... @@ -1149,6 +1153,9 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_budget_spend_from_metadata( + _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs)) + ).counter_keys, ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index bb901f22649..0d2038cf781 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1410,8 +1410,11 @@ async def _invalidate_user_spend_counter_if_changed( """ if non_default_values.get("spend") is not None: from litellm.proxy.proxy_server import invalidate_spend_counter + from litellm.proxy.spend_tracking.committed_spend import committed_spend_key - await invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}") + counter_key: Final = f"spend:user:{non_default_values['user_id']}" + await invalidate_spend_counter(counter_key=counter_key) + await invalidate_spend_counter(counter_key=committed_spend_key(counter_key)) def _clears_object_permission(user_request: UpdateUserRequest) -> bool: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..855335c6524 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -66,6 +66,7 @@ from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, management_endpoint_wrapper, ) +from litellm.proxy.spend_tracking.committed_spend import drop_committed_spend from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in @@ -817,6 +818,11 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import spend_counter_cache + + await drop_committed_spend(spend_counter_cache, (f"spend:org:{data.organization_id}",)) + return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4f2f30716a4..e4d12e43cfc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -909,10 +909,8 @@ from litellm.proxy.spend_tracking.budget_reservation import ( stamp_budget_reservation_actual_cost, ) from litellm.proxy.spend_tracking.committed_spend import ( - EntitySpendCounterKeys, committed_spend_key, committed_spend_keys, - committed_spend_tracked, seed_committed_spend, ) from litellm.proxy.spend_tracking.spend_capture_rate import ( @@ -1024,6 +1022,7 @@ from litellm.types.llms.openai import ( ResponseInProgressEvent, ResponsesAPIResponse, ) +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, @@ -3472,6 +3471,7 @@ async def increment_spend_counters( request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ): """ Atomically increment spend counters for budget enforcement. @@ -3508,6 +3508,7 @@ async def increment_spend_counters( request_started_at=request_started_at, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, ) @@ -3532,18 +3533,13 @@ async def _increment_spend_counters_batched( request_started_at: datetime | None, model_access_groups: Sequence[str] | None, project_id: str | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ): """Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET, and the reconcile adjustments go out in the same INCRBYFLOAT pipeline as the counter increments.""" committed_counter_keys: Final = ( - EntitySpendCounterKeys.of( - token=_spend_counter_token(token) if token is not None else None, - team_id=team_id, - user_id=user_id, - org_id=org_id, - end_user_id=end_user_id, - ).tracked() - if response_cost is not None and response_cost > 0 and committed_spend_tracked() + committed_spend_scope.tracked() + if committed_spend_scope is not None and response_cost is not None and response_cost > 0 else () ) bind_spend_counter_keys(frozenset(committed_spend_keys(committed_counter_keys))) @@ -4224,12 +4220,11 @@ _increment_spend_counter_cache: Final = increment_spend_counter_cache @with_service_target(SPEND_COUNTERS_TARGET) async def invalidate_spend_counter(counter_key: str) -> None: - for stale_key in (counter_key, *committed_spend_keys((counter_key,))): - forget_spend_counter(stale_key) - spend_counter_cache.in_memory_cache.delete_cache(key=stale_key) - if spend_counter_cache.redis_cache is not None: - with contextlib.suppress(Exception): - await spend_counter_cache.redis_cache.async_delete_cache(key=stale_key) + forget_spend_counter(counter_key) + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + if spend_counter_cache.redis_cache is not None: + with contextlib.suppress(Exception): + await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) _invalidate_spend_counter: Final = invalidate_spend_counter diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index e86ba782eef..150ddd1efe9 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -9,9 +9,9 @@ from litellm.models.organization import LiteLLM_OrganizationTable from litellm.models.team import LiteLLM_TeamTable from litellm.models.user import LiteLLM_UserTable from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys from litellm.types.proxy.carried_budget_state import ( CommittedBudgetSpend, + EntitySpendCounterKeys, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -56,7 +56,7 @@ async def carry_committed_budget_spend( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, end_user_object: LiteLLM_EndUserTable | None, - read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]], + read_committed: Callable[[EntitySpendCounterKeys], Awaitable[CommittedBudgetSpend]] | None, ) -> None: counter_keys: Final = EntitySpendCounterKeys.of( token=valid_token.token if valid_token.max_budget is not None else None, @@ -65,7 +65,11 @@ async def carry_committed_budget_spend( org_id=_budgeted_org_id(valid_token=valid_token, team_object=team_object), end_user_id=_budgeted_end_user_id(end_user_object=end_user_object), ) - valid_token.committed_budget_spend = await read_committed(counter_keys) # rebind-ok: request-scoped credential + valid_token.committed_budget_spend = ( # rebind-ok: request-scoped credential + CommittedBudgetSpend(counter_keys=counter_keys) + if read_committed is None + else await read_committed(counter_keys) + ) def carry_organization_budget_state(valid_token: UserAPIKeyAuth, org_table: LiteLLM_OrganizationTable) -> None: diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 45bad9e2eed..671662fa8c0 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -1,13 +1,13 @@ """Committed spend counters: twins of the entity spend counters that only ever take charged cost, never budget holds. The remaining budget gauges read them, so a request's worst-case reservation never shows as spent. They are kept -only while a ``PrometheusLogger`` is registered. +only for entities with a budget, while a ``PrometheusLogger`` is registered. """ from __future__ import annotations -from collections.abc import Iterable, Mapping -from dataclasses import dataclass +import asyncio +from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter @@ -16,9 +16,19 @@ import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed -from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET -from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed +from litellm.proxy.spend_tracking.spend_counter_batch import ( + SPEND_COUNTERS_TARGET, + admission_counter_keys, + bind_spend_counter_keys, + read_batched_spend_counter, +) +from litellm.types.proxy.carried_budget_state import ( + END_USER_COUNTER_PREFIX, + CommittedBudgetSpend, + EntitySpendCounterKeys, +) if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -41,42 +51,15 @@ def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if _has_committed_twin(counter_key)) -@dataclass(frozen=True, slots=True) -class EntitySpendCounterKeys: - key: str | None = None - team: str | None = None - user: str | None = None - org: str | None = None - end_user: str | None = None - - @classmethod - def of( - cls, - token: str | None, - team_id: str | None, - user_id: str | None, - org_id: str | None, - end_user_id: str | None, - ) -> EntitySpendCounterKeys: - return cls( - key=f"spend:key:{token}" if token else None, - team=f"spend:team:{team_id}" if team_id else None, - user=f"spend:user:{user_id}" if user_id else None, - org=f"spend:org:{org_id}" if org_id else None, - end_user=f"{END_USER_COUNTER_PREFIX}{end_user_id}" if end_user_id else None, - ) - - def tracked(self) -> tuple[str, ...]: - counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) - return tuple( - counter_key for counter_key in counter_keys if counter_key is not None and _has_committed_twin(counter_key) - ) - - def committed_spend_tracked() -> bool: return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) +def bind_committed_admission_counters(token: UserAPIKeyAuth, end_user_id: str | None) -> None: + if committed_spend_tracked(): + bind_spend_counter_keys(frozenset(committed_spend_keys(admission_counter_keys(token, end_user_id)))) + + def _in_memory_committed_value(spend_counter_cache: DualCache, twin_key: str) -> float | None: return _STORED_VALUE.validate_python( spend_counter_cache.in_memory_cache.get_cache(key=twin_key) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped cache API @@ -94,23 +77,14 @@ async def _read_committed_values(spend_counter_cache: DualCache, twin_keys: tupl return {twin_key: value for twin_key, value in fetched.items() if value is not None} -@with_service_target(SPEND_COUNTERS_TARGET) -async def read_committed_budget_spend( - spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +def _committed_budget_spend( + counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float] ) -> CommittedBudgetSpend: - twin_keys: Final = committed_spend_keys(counter_keys.tracked()) - if not twin_keys: - return CommittedBudgetSpend() - try: - committed: Final = await _read_committed_values(spend_counter_cache=spend_counter_cache, twin_keys=twin_keys) - except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend - verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) - return CommittedBudgetSpend() - def spend_of(counter_key: str | None) -> float | None: return committed.get(committed_spend_key(counter_key)) if counter_key is not None else None return CommittedBudgetSpend( + counter_keys=counter_keys, key=spend_of(counter_keys.key), team=spend_of(counter_keys.team), user=spend_of(counter_keys.user), @@ -119,6 +93,45 @@ async def read_committed_budget_spend( ) +@with_service_target(SPEND_COUNTERS_TARGET) +async def read_committed_budget_spend( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + if not twin_keys: + return CommittedBudgetSpend(counter_keys=counter_keys) + try: + committed: Final = await _read_committed_values(spend_counter_cache=spend_counter_cache, twin_keys=twin_keys) + except Exception as e: # noqa: BLE001 # a gauge read never fails the request; the gauges keep the stored spend + verbose_proxy_logger.debug("committed spend read failed for %s: %s", twin_keys, e) + return CommittedBudgetSpend(counter_keys=counter_keys) + return _committed_budget_spend(counter_keys=counter_keys, committed=committed) + + +async def read_admission_committed_spend( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +) -> CommittedBudgetSpend: + """Rides the admission MGET: a counter the batch cannot answer stays unknown here and the emit-time read has it.""" + if spend_counter_cache.redis_cache is None: + return await read_committed_budget_spend(spend_counter_cache, counter_keys) + twin_keys: Final = committed_spend_keys(counter_keys.tracked()) + batched: Final = await asyncio.gather(*(read_batched_spend_counter(twin_key) for twin_key in twin_keys)) + return _committed_budget_spend( + counter_keys=counter_keys, + committed={ + twin_key: read[0] for twin_key, read in zip(twin_keys, batched) if read is not None and read[0] is not None + }, + ) + + +async def drop_committed_spend(spend_counter_cache: DualCache, counter_keys: Sequence[str]) -> None: + """After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again.""" + try: + await spend_counter_cache.async_delete_cache_keys(committed_spend_keys(counter_keys)) + except Exception as e: # noqa: BLE001 # the write already committed; a kept twin only lasts until its TTL + verbose_proxy_logger.warning("committed spend drop failed: %s", e) + + @with_service_target(SPEND_COUNTERS_TARGET) async def seed_committed_spend( spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 3c65fc0bf06..543773f7e98 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.auth.authorization_dependencies import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, drop_committed_spend from litellm.proxy.spend_tracking.spend_capture_rate import ( ProviderBillingCredentialMissing, ProviderBillingRequestFailed, @@ -3692,6 +3693,36 @@ async def view_spend_logs( ) +class _TokenRow(TypedDict): + token: ReadOnly[str] + + +class _TeamIdRow(TypedDict): + team_id: ReadOnly[str] + + +async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None: + from litellm.proxy.proxy_server import spend_counter_cache + + try: + token_rows: Final[Sequence[_TokenRow]] = await _query_raw( + prisma_client, 'SELECT token FROM "LiteLLM_VerificationToken"' + ) + team_rows: Final[Sequence[_TeamIdRow]] = await _query_raw( + prisma_client, 'SELECT team_id FROM "LiteLLM_TeamTable"' + ) + except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin only lasts until its TTL + verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e) + return + await drop_committed_spend( + spend_counter_cache, + ( + *(f"spend:key:{row['token']}" for row in token_rows), + *(f"spend:team:{row['team_id']}" for row in team_rows), + ), + ) + + @router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], @@ -3720,6 +3751,8 @@ async def global_spend_reset(): await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + if committed_spend_tracked(): + await _drop_all_key_and_team_committed_spend(prisma_client) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/types/proxy/auth/user_api_key_auth.py b/litellm/types/proxy/auth/user_api_key_auth.py index 317a5ab7a19..ad16e26e117 100644 --- a/litellm/types/proxy/auth/user_api_key_auth.py +++ b/litellm/types/proxy/auth/user_api_key_auth.py @@ -326,6 +326,13 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob ) +def committed_budget_spend_from_metadata(metadata: Mapping[str, object]) -> CommittedBudgetSpend: + auth: Final = metadata.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth) and auth.committed_budget_spend is not None: + return auth.committed_budget_spend + return CommittedBudgetSpend() + + def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool: """Return True if the caller's role grants unscoped read access to all tenant resources (managed files, batches, vector stores, spend rows, etc). diff --git a/litellm/types/proxy/carried_budget_state.py b/litellm/types/proxy/carried_budget_state.py index 255f7af779c..bb4f77953e6 100644 --- a/litellm/types/proxy/carried_budget_state.py +++ b/litellm/types/proxy/carried_budget_state.py @@ -9,12 +9,15 @@ reads them back with ``from_metadata``. ``None`` means this request never carrie from collections.abc import Mapping from datetime import datetime from types import MappingProxyType +from typing import Final from pydantic import ConfigDict, Field, ValidationError from typing_extensions import Self from litellm.types.llms.base import LiteLLMBaseModel +END_USER_COUNTER_PREFIX: Final = "spend:end_user:" + class _BudgetSnapshot(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, populate_by_name=True, extra="ignore") @@ -66,12 +69,47 @@ class OrgBudgetSnapshot(_BudgetSnapshot): ) -class CommittedBudgetSpend(LiteLLMBaseModel): - """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the - remaining budget gauges only. ``None`` when the entity has no committed counter.""" +class EntitySpendCounterKeys(LiteLLMBaseModel): + """Spend counter of each entity whose remaining budget a gauge reports. ``None`` for an entity with no budget.""" model_config = ConfigDict(frozen=True) + key: str | None = None + team: str | None = None + user: str | None = None + org: str | None = None + end_user: str | None = None + + @classmethod + def of( + cls, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, + ) -> Self: + return cls( + key=f"spend:key:{token}" if token else None, + team=f"spend:team:{team_id}" if team_id else None, + user=f"spend:user:{user_id}" if user_id else None, + org=f"spend:org:{org_id}" if org_id else None, + end_user=f"{END_USER_COUNTER_PREFIX}{end_user_id}" if end_user_id else None, + ) + + def tracked(self) -> tuple[str, ...]: + counter_keys: Final = (self.key, self.team, self.user, self.org, self.end_user) + return tuple(counter_key for counter_key in counter_keys if counter_key is not None) + + +class CommittedBudgetSpend(LiteLLMBaseModel): + """Spend each entity was charged, budget holds excluded, read from the committed spend counters for the + remaining budget gauges only. ``None`` when the entity has no committed counter. ``counter_keys`` names the + budgeted entities whose committed counters this request charges and reads.""" + + model_config = ConfigDict(frozen=True) + + counter_keys: EntitySpendCounterKeys = EntitySpendCounterKeys() key: float | None = None team: float | None = None user: float | None = None diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 2f6c8ebfe6d..4b73bb6dd34 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -488,7 +488,7 @@ async def test_a_spend_counter_whose_increment_failed_is_invalidated_not_trusted await flush_post_call_redis_batches() assert spend_cache.in_memory_cache.get_cache("spend:key:k1") is None - assert redis_cache.alone == [("DEL", "spend:key:k1"), ("DEL", "committed_spend:key:k1")] + assert redis_cache.alone == [("DEL", "spend:key:k1")] assert spend_cache.in_memory_cache.get_cache("spend:team:t1") == 7.5 diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 8753062bd3c..b7738a62a00 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -4,6 +4,7 @@ from request metadata. get_*_object only runs when that state is missing (custom auth, SDK callers, failure paths) """ +import asyncio from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -11,6 +12,8 @@ import pytest from prometheus_client import REGISTRY from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ( LiteLLM_BudgetTable, @@ -28,6 +31,7 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.types.proxy.carried_budget_state import ( CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -242,7 +246,15 @@ def _remaining(prometheus_logger: PrometheusLogger) -> dict[str, float]: def _token_committed_at_auth() -> UserAPIKeyAuth: token = _authed_token() - token.committed_budget_spend = CommittedBudgetSpend(key=3.0, team=30.0, user=8.0, org=60.0) + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id="u1", org_id="o1", end_user_id=None + ), + key=3.0, + team=30.0, + user=8.0, + org=60.0, + ) return token @@ -264,6 +276,68 @@ async def test_gauges_use_the_committed_spend_at_emit_when_other_requests_finish } +@pytest.mark.asyncio +async def test_gauges_add_the_cost_to_another_requests_charge_while_the_own_charge_is_not_in_yet( + prometheus_logger, getters, spend_counter_cache +): + """1.0 from another request reached the counters during this long request, this request's 2.0 has not.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=4.0, team=31.0, user=9.0, org=61.0) + + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(4.0), + "team": pytest.approx(267.0), + "user": pytest.approx(39.0), + "org": pytest.approx(437.0), + } + + +@pytest.mark.asyncio +async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=9.0, team=99.0, user=19.0, org=99.0) + + await _emit(prometheus_logger, _request_metadata(_authed_token())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + +@pytest.mark.asyncio +async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( + prometheus_logger, getters, monkeypatch +): + async def never_answers(key_list: list[str]) -> dict[str, float]: + await asyncio.sleep(30) + return {} + + redis_cache = AsyncMock(spec=RedisCache) + redis_cache.async_batch_get_cache = never_answers + monkeypatch.setattr( + "litellm.proxy.proxy_server.spend_counter_cache", + DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache), + ) + monkeypatch.setenv("PROMETHEUS_BUDGET_METRICS_PER_REQUEST_TIMEOUT", "0.01") + _stub_gauges(prometheus_logger) + + await asyncio.wait_for( + prometheus_logger.async_log_failure_event( + _failed_request(_request_metadata(_token_committed_at_auth())), None, KEY_RESET_AT, KEY_RESET_AT + ), + timeout=5, + ) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0)) + + @pytest.mark.asyncio async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( prometheus_logger, getters, spend_counter_cache diff --git a/tests/unit/integrations/test_prometheus_user_team_metrics.py b/tests/unit/integrations/test_prometheus_user_team_metrics.py index 01eed51fc4e..210e2aa9ecc 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1077,7 +1077,7 @@ async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cache from litellm.models.end_user import LiteLLM_EndUserTable from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.spend_tracking.committed_spend import committed_spend_key - from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys mock_proxy_server = MagicMock() mock_proxy_server.user_api_key_cache.async_get_cache = AsyncMock( @@ -1093,7 +1093,9 @@ async def test_customer_gauge_subtracts_the_committed_spend_instead_of_the_cache key=committed_spend_key("spend:end_user:cust-live"), value=2.5 ) token = UserAPIKeyAuth(token="hashed") - token.committed_budget_spend = CommittedBudgetSpend(end_user=1.5) + token.committed_budget_spend = CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(end_user="spend:end_user:cust-live"), end_user=1.5 + ) with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_proxy_server}): await prometheus_logger._increment_remaining_budget_metrics( diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index af7a181c83b..afe2d372b36 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5096,7 +5096,7 @@ async def test_centralized_common_checks_carry_the_committed_spend_beside_the_st from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.spend_tracking.committed_spend import committed_spend_key from litellm.proxy.utils import ProxyLogging - from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend + from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys for callback_list in ("success_callback", "failure_callback", "_async_success_callback", "_async_failure_callback"): monkeypatch.setattr(litellm, callback_list, []) @@ -5135,7 +5135,13 @@ async def test_centralized_common_checks_carry_the_committed_spend_beside_the_st setattr(_proxy_server_mod, attribute, value) assert token.committed_budget_spend == ( - CommittedBudgetSpend(key=0.37, team=1.25) if prometheus_registered else None + CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key=f"spend:key:{token.token}", team=f"spend:team:{token.team_id}"), + key=0.37, + team=1.25, + ) + if prometheus_registered + else None ) assert (token.spend, token.team_spend) == (0.0, 0.0) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 4b5957a1237..c3f4c4e9d97 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -38,6 +38,7 @@ from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyUpdateSpend, ) +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.types.router import Deployment, LiteLLM_Params from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage @@ -1028,6 +1029,47 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "expected_scope"), + [ + ( + { + "user_api_key_auth": UserAPIKeyAuth( + committed_budget_spend=CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:test_api_key") + ) + ) + }, + EntitySpendCounterKeys(key="spend:key:test_api_key"), + ), + ({}, EntitySpendCounterKeys()), + ], +) +async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(metadata, expected_scope): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + increment_spend_counters = AsyncMock() + + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": metadata}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() @@ -1156,6 +1198,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda request_started_at=start_time, model_access_groups=("premium",), project_id=None, + committed_spend_scope=EntitySpendCounterKeys(), ) diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d1bbf444d2c..ca0c6fb6afa 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4054,7 +4054,10 @@ async def test_admin_user_update_spend_invalidates_counter(mocker): admin_caller = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) await _update_single_user_helper(user_request=user_request, user_api_key_dict=admin_caller) - mock_invalidate.assert_awaited_once_with(counter_key="spend:user:target-user") + assert [call.kwargs["counter_key"] for call in mock_invalidate.await_args_list] == [ + "spend:user:target-user", + "committed_spend:user:target-user", + ] @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 5e06756f3a9..2336c8c352e 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1187,6 +1187,28 @@ async def _run_legacy_update_organization( return mock_prisma_client +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "twin_after"), + [ + ({"organization_id": "org-1", "spend": 0.0}, None), + ({"organization_id": "org-1", "organization_alias": "renamed"}, 7.0), + ], +) +async def test_legacy_update_drops_the_committed_org_spend_only_when_it_writes_the_stored_spend( + monkeypatch, body, twin_after +): + from litellm.caching.dual_cache import DualCache + + spend_counter_cache = DualCache() + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:org:org-1", value=7.0) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + + await _run_legacy_update_organization(monkeypatch, body=body, existing_budget_id="budget-1") + + assert spend_counter_cache.in_memory_cache.get_cache(key="committed_spend:org:org-1") == twin_after + + @pytest.mark.asyncio @pytest.mark.parametrize( "body", diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 1d3255a9be8..9e4f76cf993 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -51,7 +51,8 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_loca from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team -from litellm.proxy.spend_tracking.committed_spend import committed_spend_key, committed_spend_keys +from litellm.proxy.spend_tracking.committed_spend import committed_spend_key +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from .conftest import normalize @@ -856,8 +857,10 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( class _VerificationTokenTable: def __init__(self, spend: float) -> None: self._spend: Final = spend + self.lookups: list[dict[str, str]] = [] async def find_unique(self, where: dict[str, str]) -> MagicMock: + self.lookups.append(where) return MagicMock(token=where["token"], spend=self._spend) @@ -875,30 +878,60 @@ async def test_committed_twin_takes_the_charge_but_never_a_budget_hold(monkeypat monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) - monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) reservation: Final = {"reserved_cost": 2.0, "entries": [{"counter_key": "spend:key:hashed"}], "finalized": False} + scope: Final = EntitySpendCounterKeys(key="spend:key:hashed") await ps.increment_spend_counters( - token="hashed", team_id=None, user_id=None, response_cost=0.5, budget_reservation=reservation + token="hashed", + team_id=None, + user_id=None, + response_cost=0.5, + budget_reservation=reservation, + committed_spend_scope=scope, + ) + await ps.increment_spend_counters( + token="hashed", team_id=None, user_id=None, response_cost=0.25, committed_spend_scope=scope ) - await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.25) assert counter_cache.in_memory_cache.get_cache(key="spend:key:hashed") == pytest.approx(11.75) assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:key:hashed")) == pytest.approx(10.75) @pytest.mark.asyncio -async def test_no_committed_twin_is_kept_without_a_prometheus_logger(monkeypatch): +async def test_a_charge_outside_the_committed_scope_keeps_no_twin_and_reads_no_stored_spend(monkeypatch): counter_cache: Final = ps.DualCache() counter_cache.in_memory_cache.set_cache(key="spend:key:hashed", value=10.0) + prisma: Final = _prisma_with_key_spend(10.0) monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) - monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + monkeypatch.setattr(ps, "prisma_client", prisma) await ps.increment_spend_counters(token="hashed", team_id=None, user_id=None, response_cost=0.5) assert counter_cache.in_memory_cache.get_cache(key="spend:key:hashed") == pytest.approx(10.5) assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:key:hashed")) is None + assert prisma.db.litellm_verificationtoken.lookups == [] + + +@pytest.mark.asyncio +async def test_only_the_entities_in_the_committed_scope_take_the_charge_on_their_twin(monkeypatch): + counter_cache: Final = ps.DualCache() + counter_cache.in_memory_cache.set_cache(key="spend:key:hashed", value=10.0) + counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=10.0) + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "user_api_key_cache", ps.DualCache()) + monkeypatch.setattr(ps, "prisma_client", _prisma_with_key_spend(10.0)) + + await ps.increment_spend_counters( + token="hashed", + team_id="t1", + user_id=None, + response_cost=0.5, + committed_spend_scope=EntitySpendCounterKeys(key="spend:key:hashed"), + ) + + assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:key:hashed")) == pytest.approx(10.5) + assert counter_cache.in_memory_cache.get_cache(key=committed_spend_key("spend:team:t1")) is None _PENDING_SPEND_COUNTER_KEYS: tuple[str, ...] = ( @@ -945,7 +978,7 @@ async def _increment_every_spend_counter() -> None: def _assert_every_pending_counter_invalidated(fake_redis: AsyncMock) -> None: assert fake_redis.async_increment.await_count == 0 deleted_keys = {call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list} - assert deleted_keys == {*_PENDING_SPEND_COUNTER_KEYS, *committed_spend_keys(_PENDING_SPEND_COUNTER_KEYS)} + assert deleted_keys == set(_PENDING_SPEND_COUNTER_KEYS) for key in _PENDING_SPEND_COUNTER_KEYS: assert ps.spend_counter_cache.in_memory_cache.get_cache(key=key) is None @@ -1309,7 +1342,7 @@ async def test_apply_spend_counter_increments_open_breaker_invalidates_and_retur await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] + assert deleted_keys == ["spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1327,7 +1360,7 @@ async def test_apply_spend_counter_increments_redis_timeout_invalidates_and_retu await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] + assert deleted_keys == ["spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1341,7 +1374,7 @@ async def test_apply_spend_counter_increments_other_redis_error_invalidates_and_ await ps._apply_spend_counter_increments(_two_pending_increments()) deleted_keys = sorted(call.kwargs["key"] for call in fake_cache.in_memory_cache.delete_cache.call_args_list) - assert deleted_keys == ["committed_spend:key:k", "committed_spend:team:t", "spend:key:k", "spend:team:t"] + assert deleted_keys == ["spend:key:k", "spend:team:t"] fake_cache.in_memory_cache.set_cache.assert_not_called() @@ -1590,23 +1623,15 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( # --------------------------------------------------------------------------- -@pytest.mark.parametrize( - ("counter_key", "deleted"), - [ - ("spend:key:k", ["spend:key:k", "committed_spend:key:k"]), - ("spend:key:k:window:1d", ["spend:key:k:window:1d"]), - ("spend:tag:t", ["spend:tag:t"]), - ], -) @pytest.mark.asyncio -async def test_invalidate_spend_counter_deletes_the_counter_and_its_committed_twin(monkeypatch, counter_key, deleted): +async def test_invalidate_spend_counter_keeps_the_committed_twin_a_redis_failure_never_changed(monkeypatch): fake_cache = _make_spend_counter_cache() monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - await ps._invalidate_spend_counter(counter_key=counter_key) + await ps._invalidate_spend_counter(counter_key="spend:key:k") - assert [c.kwargs["key"] for c in fake_cache.in_memory_cache.delete_cache.call_args_list] == deleted - assert [c.kwargs["key"] for c in fake_cache.redis_cache.async_delete_cache.await_args_list] == deleted + assert [c.kwargs["key"] for c in fake_cache.in_memory_cache.delete_cache.call_args_list] == ["spend:key:k"] + assert [c.kwargs["key"] for c in fake_cache.redis_cache.async_delete_cache.await_args_list] == ["spend:key:k"] @pytest.mark.asyncio diff --git a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py index 69acc624623..75b9e422330 100644 --- a/tests/unit/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/unit/proxy/spend_tracking/test_carried_budget_state.py @@ -18,9 +18,9 @@ from litellm.proxy.spend_tracking.carried_budget_state import ( carry_organization_budget_state, carry_team_and_user_budget_state, ) -from litellm.proxy.spend_tracking.committed_spend import EntitySpendCounterKeys from litellm.types.proxy.carried_budget_state import ( CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -182,6 +182,19 @@ async def test_committed_spend_is_read_once_for_every_budgeted_entity(): assert (token.spend, token.team_spend) == (1.0, None) +@pytest.mark.asyncio +async def test_route_without_a_budget_check_carries_the_counters_it_charges_and_reads_nothing(): + token = UserAPIKeyAuth(token="hashed", max_budget=10.0) + + await carry_committed_budget_spend( + valid_token=token, team_object=None, user_object=None, end_user_object=None, read_committed=None + ) + + assert token.committed_budget_spend == CommittedBudgetSpend( + counter_keys=EntitySpendCounterKeys(key="spend:key:hashed") + ) + + @pytest.mark.asyncio async def test_committed_spend_skips_entities_without_a_max_budget(): token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1", org_id="o1", spend=1.0) diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py index ea76a61047d..a6beb006d0c 100644 --- a/tests/unit/proxy/spend_tracking/test_committed_spend.py +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -4,22 +4,32 @@ from __future__ import annotations from types import SimpleNamespace from typing import Final -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from fakeredis import FakeAsyncRedis, FakeServer +import litellm from litellm._service_logger import ServiceLogging from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.spend_tracking.committed_spend import ( - EntitySpendCounterKeys, + bind_committed_admission_counters, committed_spend_key, + drop_committed_spend, + read_admission_committed_spend, read_committed_budget_spend, seed_committed_spend, ) -from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend +from litellm.proxy.spend_tracking.spend_counter_batch import ( + active_spend_counter_batch, + read_batched_spend_counter, + spend_counter_batch_scope, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys class _FakeRedisSettingMax(FakeAsyncRedis): @@ -87,12 +97,13 @@ async def test_read_returns_each_committed_twin_and_none_for_one_never_charged(s await _store(spend_counter_cache, committed_spend_key("spend:end_user:cust-1"), 0.25) await _store(spend_counter_cache, "spend:team:t1", 99.0) - committed: Final = await read_committed_budget_spend( - spend_counter_cache, - EntitySpendCounterKeys.of(token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1"), + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1" ) - assert committed == CommittedBudgetSpend(key=1.5, end_user=0.25) + committed: Final = await read_committed_budget_spend(spend_counter_cache, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5, end_user=0.25) @pytest.mark.asyncio @@ -101,31 +112,34 @@ async def test_read_of_a_twin_holding_garbage_reports_no_committed_spend(): spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key("spend:key:hashed"), value="not-a-number") spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key("spend:team:t1"), value=4.0) - committed: Final = await read_committed_budget_spend( - spend_counter_cache, - EntitySpendCounterKeys.of(token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None), + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None ) - assert committed == CommittedBudgetSpend() + committed: Final = await read_committed_budget_spend(spend_counter_cache, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys) @pytest.mark.parametrize( - ("counter_keys", "expected"), + ("counter_keys", "expected_spend"), [ - (EntitySpendCounterKeys(key="spend:key:hashed"), CommittedBudgetSpend(key=12.0)), - (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), CommittedBudgetSpend(end_user=12.0)), + (EntitySpendCounterKeys(key="spend:key:hashed"), {"key": 12.0}), + (EntitySpendCounterKeys(end_user="spend:end_user:cust-1"), {"end_user": 12.0}), ], ) @pytest.mark.asyncio async def test_seed_writes_the_db_spend_when_the_twin_is_absent( - spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys, expected: CommittedBudgetSpend + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys, expected_spend: dict[str, float] ): (counter_key,) = counter_keys.tracked() seeded: Final = await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), counter_key) assert seeded is True - assert await read_committed_budget_spend(spend_counter_cache, counter_keys) == expected + assert await read_committed_budget_spend(spend_counter_cache, counter_keys) == CommittedBudgetSpend( + counter_keys=counter_keys, **expected_spend + ) @pytest.mark.asyncio @@ -169,7 +183,7 @@ async def test_seed_reports_no_twin_when_redis_rejects_the_write(): spend_counter_cache, EntitySpendCounterKeys(team="spend:team:t1") ) assert seeded is False - assert committed == CommittedBudgetSpend() + assert committed.team is None @pytest.mark.asyncio @@ -180,4 +194,92 @@ async def test_seed_creates_no_twin_when_the_db_has_no_spend_for_the_entity(spen spend_counter_cache, EntitySpendCounterKeys(key="spend:key:hashed") ) assert seeded is False - assert committed == CommittedBudgetSpend() + assert committed.key is None + + +class _CountingMgetRedis(_RedisCacheOnClient): + def __init__(self, client: FakeAsyncRedis) -> None: + super().__init__(client) + self.mgets: list[frozenset[str]] = [] + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + self.mgets.append(frozenset(key_list)) + return await super().async_batch_get_cache(key_list=key_list, **kwargs) + + +@pytest.mark.asyncio +async def test_admission_read_rides_the_admission_mget_with_the_enforcement_counters( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + redis_cache: Final = _CountingMgetRedis(FakeAsyncRedis(server=FakeServer())) + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + await redis_cache.async_set_cache(key=committed_spend_key("spend:key:hashed"), value=1.5) + token: Final = UserAPIKeyAuth(token="hashed", team_id="t1") + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed", team="spend:team:t1") + + with spend_counter_batch_scope(redis_cache, counter_keys=frozenset(("spend:key:hashed",))): + bind_committed_admission_counters(token, end_user_id=None) + await read_batched_spend_counter("spend:key:hashed") + committed: Final = await read_admission_committed_spend(spend_counter_cache, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys, key=1.5) + assert len(redis_cache.mgets) == 1 + assert {committed_spend_key("spend:key:hashed"), committed_spend_key("spend:team:t1")} <= redis_cache.mgets[0] + + +@pytest.mark.asyncio +async def test_admission_read_outside_a_batch_leaves_redis_alone_and_the_spend_unknown(): + redis_cache: Final = _CountingMgetRedis(FakeAsyncRedis(server=FakeServer())) + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + await redis_cache.async_set_cache(key=committed_spend_key("spend:key:hashed"), value=1.5) + counter_keys: Final = EntitySpendCounterKeys(key="spend:key:hashed") + + committed: Final = await read_admission_committed_spend(spend_counter_cache, counter_keys) + + assert committed == CommittedBudgetSpend(counter_keys=counter_keys) + assert redis_cache.mgets == [] + + +@pytest.mark.asyncio +async def test_admission_read_without_redis_reads_the_in_memory_twin(): + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache()) + spend_counter_cache.in_memory_cache.set_cache(key=committed_spend_key("spend:key:hashed"), value=1.5) + + committed: Final = await read_admission_committed_spend( + spend_counter_cache, EntitySpendCounterKeys(key="spend:key:hashed") + ) + + assert committed.key == 1.5 + + +def test_no_committed_counter_is_bound_without_a_prometheus_logger(): + redis_cache: Final = _CountingMgetRedis(FakeAsyncRedis(server=FakeServer())) + + with spend_counter_batch_scope(redis_cache): + bind_committed_admission_counters(UserAPIKeyAuth(token="hashed"), end_user_id=None) + bound: Final = active_spend_counter_batch().counter_keys + + assert bound == frozenset() + + +@pytest.mark.asyncio +async def test_drop_removes_the_twin_and_leaves_the_enforcement_counter(spend_counter_cache: DualCache): + await _store(spend_counter_cache, committed_spend_key("spend:org:o1"), 5.0) + await _store(spend_counter_cache, "spend:org:o1", 5.0) + counter_keys: Final = EntitySpendCounterKeys(org="spend:org:o1") + + await drop_committed_spend(spend_counter_cache, ("spend:org:o1",)) + + assert (await read_committed_budget_spend(spend_counter_cache, counter_keys)).org is None + assert await spend_counter_cache.async_get_cache(key="spend:org:o1") == 5.0 + + +@pytest.mark.asyncio +async def test_a_failed_drop_never_fails_the_write_that_already_committed(): + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache()) + spend_counter_cache.async_delete_cache_keys = AsyncMock(side_effect=ConnectionError("redis unreachable")) + + await drop_committed_spend(spend_counter_cache, ("spend:org:o1",)) + + spend_counter_cache.async_delete_cache_keys.assert_awaited_once() diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 041a278acb9..f32fc3634a5 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -13,7 +13,6 @@ import litellm.proxy.proxy_server as ps from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed -from litellm.proxy.spend_tracking.committed_spend import committed_spend_keys from litellm.proxy.spend_tracking.spend_counter_batch import ( SpendCounterBatch, active_spend_counter_batch, @@ -541,10 +540,7 @@ async def test_a_failed_post_call_pipeline_invalidates_every_counter_it_carried_ assert [c.split()[0] for c in redis.commands] == ["MGET"], redis.commands invalidated = RESERVED_KEYS | {"spend:user:user"} - assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == { - *invalidated, - *committed_spend_keys(invalidated), - } + assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == invalidated assert all("applied_adjustment" not in entry for entry in reservation["entries"]) assert {key: redis.store[key] for key in POST_CALL_KEYS} == {key: 1.0 for key in POST_CALL_KEYS} @@ -634,8 +630,6 @@ async def test_a_failed_reservation_pipeline_drops_every_counter_and_reserves_no assert {call.kwargs["key"] for call in redis.async_delete_cache.await_args_list} == { "spend:key:hashed", "spend:team:team", - "committed_spend:key:hashed", - "committed_spend:team:team", } assert redis.store == {"spend:key:hashed": 1.0, "spend:team:team": 2.0} diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 8d53817f51e..d2b9343b461 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7,6 +7,7 @@ import re import sqlite3 from collections.abc import Mapping from datetime import timezone +from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -17,9 +18,12 @@ from fastapi.testclient import TestClient import litellm import litellm.proxy.proxy_server as ps +from litellm.caching.dual_cache import DualCache +from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.spend_tracking.spend_management_endpoints import global_spend_reset def _default_date_range(): @@ -8268,3 +8272,58 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + + +class _ResetTableRows: + def __init__(self) -> None: + self.update_many: Final = AsyncMock() + + +class _ResetDb: + def __init__(self) -> None: + self.litellm_verificationtoken: Final = _ResetTableRows() + self.litellm_teamtable: Final = _ResetTableRows() + + async def query_raw(self, sql_query: str) -> list[dict[str, str]]: + return [{"token": "hashed"}] if "VerificationToken" in sql_query else [{"team_id": "t1"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("registered_callbacks", "twin_after"), + [ + ([MagicMock(spec=PrometheusLogger)], None), + ([], 7.0), + ], +) +async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gauges_read( + monkeypatch, registered_callbacks, twin_after +): + spend_counter_cache: Final = DualCache() + twins: Final = ("committed_spend:key:hashed", "committed_spend:team:t1") + for twin in twins: + spend_counter_cache.in_memory_cache.set_cache(key=twin, value=7.0) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=_ResetDb())) + monkeypatch.setattr(litellm, "callbacks", registered_callbacks) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert [spend_counter_cache.in_memory_cache.get_cache(key=twin) for twin in twins] == [twin_after, twin_after] + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch): + db: Final = _ResetDb() + db.query_raw = AsyncMock(side_effect=ConnectionError("db unreachable")) + spend_counter_cache: Final = DualCache() + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:team:t1", value=7.0) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + db.litellm_teamtable.update_many.assert_awaited_once() diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 449516e8ba5..6127d6b6abc 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -9405,7 +9405,6 @@ async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure( assert counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None assert [call.kwargs["key"] for call in fake_redis.async_delete_cache.await_args_list] == [ "spend:team:redis-fail", - "committed_spend:team:redis-fail", ] finally: ps.spend_counter_cache = orig_counter From 0ce16cd68d55480125930f39c115a91ff5b07041 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:47:51 -0700 Subject: [PATCH 6/9] fix(prometheus): read the committed org spend on failed requests without the auth object --- litellm/integrations/prometheus.py | 14 ++++++++++---- .../test_prometheus_carried_budget_state.py | 14 ++++++++++++++ 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 1f3b208a460..2636d92a7cf 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -377,6 +377,12 @@ def _customer_spend_after_request(cached_spend: float, committed_spend: float | return max(committed_spend + response_cost, cached_spend) +def _budgeted_org_counter(org_id: str | None, org: OrgBudgetSnapshot | None) -> EntitySpendCounterKeys: + if org is None or org.max_budget is None: + return EntitySpendCounterKeys() + return EntitySpendCounterKeys.of(token=None, team_id=None, user_id=None, org_id=org_id, end_user_id=None) + + async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: if not counter_keys.tracked(): return CommittedBudgetSpend() @@ -2862,16 +2868,16 @@ class PrometheusLogger(CustomLogger): get_litellm_metadata_from_kwargs(kwargs) ) failure_org_alias: Final = failure_metadata.get("user_api_key_org_alias") - failure_at_auth: Final = committed_budget_spend_from_metadata(failure_metadata) + failure_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata) failure_committed: Final = _committed_spend_before_request( - at_auth=failure_at_auth, - now=await _committed_spend_now(EntitySpendCounterKeys(org=failure_at_auth.counter_keys.org)), + at_auth=committed_budget_spend_from_metadata(failure_metadata), + now=await _committed_spend_now(_budgeted_org_counter(user_api_key_org_id, failure_org)), response_cost=0, ) await self._set_org_budget_metrics_after_api_request( org_id=user_api_key_org_id, response_cost=0, - carried=OrgBudgetSnapshot.from_metadata(failure_metadata), + carried=failure_org, org_alias=failure_org_alias if isinstance(failure_org_alias, str) else None, committed_spend=failure_committed.org, ) diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index b7738a62a00..2886ff39586 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -434,6 +434,20 @@ async def test_failed_request_sets_the_org_gauge_from_the_carried_org_and_commit prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) +@pytest.mark.asyncio +async def test_failed_request_without_the_auth_object_still_reads_the_committed_org_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=65.0) + metadata = _request_metadata(_token_committed_at_auth()) + del metadata["user_api_key_auth"] + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) + + @pytest.mark.asyncio async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): _stub_gauges(prometheus_logger) From 3bfbd6a002227a870cfc7142acceac9da70017a7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:24:59 -0700 Subject: [PATCH 7/9] test(prometheus): cover a failed request with no carried org --- .../test_prometheus_carried_budget_state.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 2886ff39586..74525da4c3e 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -448,6 +448,20 @@ async def test_failed_request_without_the_auth_object_still_reads_the_committed_ prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(435.0)) +@pytest.mark.asyncio +async def test_failed_request_without_a_carried_org_reads_no_committed_org_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=3.0, team=30.0, user=8.0, org=65.0) + token = _token_committed_at_auth().model_copy(update={"org_budget_snapshot": None}) + metadata = {key: value for key, value in _request_metadata(token).items() if key != "user_api_key_auth"} + + await prometheus_logger.async_log_failure_event(_failed_request(metadata), None, KEY_RESET_AT, KEY_RESET_AT) + + prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(460.0)) + + @pytest.mark.asyncio async def test_ordinary_metadata_keys_named_like_snapshot_fields_change_no_gauge(prometheus_logger, getters): _stub_gauges(prometheus_logger) From f43b4db2c1cf33d9482c5ba1ce68942a380ed806 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:53:02 -0700 Subject: [PATCH 8/9] fix(prometheus): move the committed counters for charges with no auth on the row A billed background response is charged from its creator's stored attribution, after the request's auth is gone. Its charge moved the live counters but not the committed ones, so every gauge read one charge behind. It now moves the committed counters that already exist, and its own gauge adds its cost to them. --- litellm/integrations/prometheus.py | 16 ++++++- .../proxy/hooks/proxy_track_cost_callback.py | 37 ++++++++++++-- .../proxy/spend_tracking/committed_spend.py | 14 ++++++ .../test_prometheus_carried_budget_state.py | 38 +++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 48 +++++++++++++++++++ 5 files changed, 148 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2636d92a7cf..d71d0f8d403 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2284,6 +2284,7 @@ class PrometheusLogger(CustomLogger): _org_alias: Final = _metadata.get("user_api_key_org_alias", None) _carried_metadata: Final = _REQUEST_METADATA_ADAPTER.validate_python(_metadata) _committed_at_auth: Final = committed_budget_spend_from_metadata(_carried_metadata) + _row_carries_auth: Final = "user_api_key_auth" in _carried_metadata # Bound the per-request budget-metric emission so that slow Redis/DB # lookups under load cannot consume the whole LoggingWorker watchdog @@ -2295,9 +2296,20 @@ class PrometheusLogger(CustomLogger): budget_metrics_timeout: Final = _get_budget_metrics_per_request_timeout() async def emit_budget_metrics() -> tuple[object, ...]: + committed_now: Final = await _committed_spend_now( + _committed_at_auth.counter_keys + if _row_carries_auth + else EntitySpendCounterKeys.of( + token=user_api_key, + team_id=user_api_team, + user_id=user_id, + org_id=user_api_key_org_id, + end_user_id=end_user_id, + ) + ) committed: Final = _committed_spend_before_request( - at_auth=_committed_at_auth, - now=await _committed_spend_now(_committed_at_auth.counter_keys), + at_auth=_committed_at_auth if _row_carries_auth else committed_now, + now=committed_now, response_cost=response_cost, ) return await asyncio.gather( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 97018b74b95..7f95121ea23 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -39,6 +39,7 @@ from litellm.proxy.db.db_spend_update_writer import ( from litellm.proxy.hooks.search_free_tier import apply_search_free_tier from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.spend_tracking.committed_spend import committed_spend_tracked, warm_committed_counter_keys from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope from litellm.proxy.spend_tracking.spend_event import ( ObjectMapping, @@ -1022,6 +1023,30 @@ class _IncrementSpendCounters(Protocol): ) -> None: ... +async def _committed_spend_scope_of_row( + kwargs: Mapping[str, object], + response_cost: float, + token: str | None, + team_id: str | None, + user_id: str | None, + org_id: str | None, + end_user_id: str | None, +) -> EntitySpendCounterKeys: + metadata: Final = _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=dict(kwargs))) + if "user_api_key_auth" in metadata: + return committed_budget_spend_from_metadata(metadata).counter_keys + if response_cost <= 0 or not committed_spend_tracked(): + return EntitySpendCounterKeys() + from litellm.proxy.proxy_server import spend_counter_cache + + return await warm_committed_counter_keys( + spend_counter_cache, + EntitySpendCounterKeys.of( + token=token, team_id=team_id, user_id=user_id, org_id=org_id, end_user_id=end_user_id + ), + ) + + async def _update_database_and_spend_counters( proxy_logging_obj: "ProxyLogging", increment_spend_counters: _IncrementSpendCounters, @@ -1153,9 +1178,15 @@ async def _update_database_and_spend_counters_in_batch( request_started_at=start_time, model_access_groups=model_access_groups, project_id=project_id, - committed_spend_scope=committed_budget_spend_from_metadata( - _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs)) - ).counter_keys, + committed_spend_scope=await _committed_spend_scope_of_row( + kwargs=kwargs, + response_cost=response_cost, + token=user_api_key, + team_id=team_id, + user_id=user_id, + org_id=org_id, + end_user_id=end_user_id, + ), ) except Exception: if budget_reservation is not None: diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 671662fa8c0..0f542040d8e 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -108,6 +108,20 @@ async def read_committed_budget_spend( return _committed_budget_spend(counter_keys=counter_keys, committed=committed) +async def warm_committed_counter_keys( + spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys +) -> EntitySpendCounterKeys: + """For a charge with no auth to name the budgeted entities: only a budgeted entity's twin exists.""" + committed: Final = await read_committed_budget_spend(spend_counter_cache, counter_keys) + return EntitySpendCounterKeys( + key=counter_keys.key if committed.key is not None else None, + team=counter_keys.team if committed.team is not None else None, + user=counter_keys.user if committed.user is not None else None, + org=counter_keys.org if committed.org is not None else None, + end_user=counter_keys.end_user if committed.end_user is not None else None, + ) + + async def read_admission_committed_spend( spend_counter_cache: DualCache, counter_keys: EntitySpendCounterKeys ) -> CommittedBudgetSpend: diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 74525da4c3e..1a746de588e 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -311,6 +311,44 @@ async def test_gauges_read_no_committed_counter_auth_did_not_name_as_budgeted( } +def _billed_row_metadata() -> dict: + return {key: value for key, value in _request_metadata(_authed_token()).items() if key != "user_api_key_auth"} + + +@pytest.mark.asyncio +async def test_row_with_no_auth_adds_its_cost_to_the_committed_spend_of_the_entities_it_is_charged_to( + prometheus_logger, getters, spend_counter_cache +): + """Counter increments flush after the success callbacks, so the row's own charge is not in the counters yet.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=6.0, team=33.0, user=11.0, org=63.0) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(2.0), + "team": pytest.approx(265.0), + "user": pytest.approx(37.0), + "org": 435.0, + } + + +@pytest.mark.asyncio +async def test_row_with_no_auth_and_no_committed_counters_keeps_the_stored_spend( + prometheus_logger, getters, spend_counter_cache +): + _stub_gauges(prometheus_logger) + + await _emit(prometheus_logger, _billed_row_metadata()) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(7.0), + "team": pytest.approx(278.0), + "user": pytest.approx(43.0), + "org": 458.0, + } + + @pytest.mark.asyncio async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_auth_time_spend( prometheus_logger, getters, monkeypatch diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c3f4c4e9d97..c736f92da6d 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1070,6 +1070,54 @@ async def test_charge_moves_only_the_committed_counters_auth_named_as_budgeted(m assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == expected_scope +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tracked", "expected_key", "expected_org"), + [(True, pytest.approx(1.2), pytest.approx(2.2)), (False, 1.0, 2.0)], +) +async def test_charge_with_no_auth_on_its_row_moves_the_committed_counters_that_exist( + tracked, expected_key, expected_org +): + """A background response is billed from its creator's stored attribution, after the request's auth is gone.""" + from litellm.proxy.proxy_server import increment_spend_counters + + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock() + spend_counter_cache = DualCache() + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:key:test_api_key", value=1.0) + spend_counter_cache.in_memory_cache.set_cache(key="committed_spend:org:test_org_id", value=2.0) + + with ( + patch("litellm.proxy.hooks.proxy_track_cost_callback.committed_spend_tracked", return_value=tracked), + patch("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache), + patch( + "litellm.proxy.spend_tracking.committed_spend.SpendCounterReseed.from_db", + AsyncMock(return_value=50.0), + ), + ): + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={"litellm_params": {"metadata": {"user_api_key_hash": "test_api_key"}}}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=None, + ) + + cached = spend_counter_cache.in_memory_cache.get_cache + assert cached(key="committed_spend:key:test_api_key") == expected_key + assert cached(key="committed_spend:org:test_org_id") == expected_org + assert cached(key="committed_spend:team:test_team_id") is None + assert cached(key="committed_spend:user:test_user_id") is None + + @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() From 6046f1b0151286f4b52d6f188b2fc2210a6da36c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:53:21 -0700 Subject: [PATCH 9/9] fix(prometheus): keep the gauges right for pass-through rows, overlapping charges and slow DB spend writes Pass-through rows carry the auth object, so a cold committed counter gets seeded there too. The gauges read a charge still waiting in the post-call pipeline as not counted, whatever other charges reached the counter. Committed counters outlive two DB spend writes. Cold seeds read the DB once per worker, and the global spend reset reads keys and teams in bounded pages. --- litellm/caching/redis_batch.py | 13 +++ litellm/integrations/prometheus.py | 9 ++ litellm/proxy/db/spend_counter_reseed.py | 4 + litellm/proxy/litellm_pre_call_utils.py | 10 +- .../pass_through_endpoints.py | 1 + litellm/proxy/proxy_server.py | 32 +++++- .../proxy/spend_tracking/committed_spend.py | 29 ++++- .../spend_management_endpoints.py | 49 ++++---- .../test_request_redis_batch_post_call.py | 14 +++ .../test_prometheus_carried_budget_state.py | 106 +++++++++++++++++- .../test_pass_through_endpoints.py | 26 +++++ .../spend_tracking/test_committed_spend.py | 81 ++++++++++++- .../test_spend_management_endpoints.py | 37 +++++- 13 files changed, 375 insertions(+), 36 deletions(-) diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..87245023633 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -461,6 +461,7 @@ class RequestRedisBatches: "_deadline", "_deadline_flush", "_post_call", + "_post_call_sent", "post_call_deadline", "prefetched", ) @@ -468,6 +469,7 @@ class RequestRedisBatches: def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None: self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend + self._post_call_sent: bool = False self.post_call_deadline: Final = post_call_deadline self._deadline: asyncio.TimerHandle | None = None self._deadline_flush: asyncio.Task[None] | None = None @@ -508,10 +510,15 @@ class RequestRedisBatches: if self._deadline is not None: self._deadline.cancel() self._deadline = None + self._post_call_sent = self._post_call_sent or any(batch.pending for batch in self._post_call.values()) await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending)) if not any(batch.pending for batch in self._post_call.values()): _open_post_call.discard(self) + @property + def post_call_sent(self) -> bool: + return self._post_call_sent + @property def batches(self) -> tuple[RedisBatch, ...]: return tuple(self._batches.values()) @@ -542,6 +549,12 @@ def active_post_call_redis_batch(redis_cache: RedisCache) -> RedisBatch | None: return batches.post_call(redis_cache) +def post_call_writes_unsent() -> bool: + """Inside a request scope that has sent no post-call pipeline yet, none of its post-call writes reached Redis.""" + batches: Final = _active_request_batches.get() + return batches is not None and not batches.post_call_sent + + async def flush_post_call_redis_batches() -> None: """Called where the success and failure callbacks of a request have all run.""" batches: Final = _active_request_batches.get() diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 28f505370e2..9e0895b4b51 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._internal_context import with_service_target from litellm._logging import print_verbose, verbose_logger +from litellm.caching.redis_batch import post_call_writes_unsent from litellm.constants import ( AUTH_OBJECTS_TARGET, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -405,12 +406,20 @@ async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> Committe return CommittedBudgetSpend() +def _own_charge_waits_in_post_call_pipeline() -> bool: + from litellm.proxy.proxy_server import spend_counter_cache + + return spend_counter_cache.redis_cache is not None and post_call_writes_unsent() + + def _committed_spend_before_request( at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float ) -> CommittedBudgetSpend: def before(at_auth_spend: float | None, now_spend: float | None) -> float | None: if now_spend is None: return at_auth_spend + if _own_charge_waits_in_post_call_pipeline(): + return now_spend if at_auth_spend is None: return now_spend - response_cost own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index f30b351eb5a..d3f26f17371 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -96,6 +96,10 @@ class SpendCounterReseed: _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict() _registry_lock: ClassVar[asyncio.Lock | None] = None + @staticmethod + async def counter_lock(counter_key: str) -> asyncio.Lock: + return await SpendCounterReseed._get_lock(counter_key) + @staticmethod async def _get_lock(counter_key: str) -> asyncio.Lock: if SpendCounterReseed._registry_lock is None: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index edaacd53d22..91ffc3d2ef6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2136,7 +2136,14 @@ class LiteLLMProxyRequestSetup: user_api_key_dict.matched_model_access_groups ) # UserAPIKeyAuth object for MCP server access control - data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy( + data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata( + user_api_key_dict + ) + return data + + @staticmethod + def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth: + return user_api_key_dict.model_copy( update={ "metadata": strip_callback_config(user_api_key_dict.metadata), "team_metadata": strip_callback_config(user_api_key_dict.team_metadata), @@ -2144,7 +2151,6 @@ class LiteLLMProxyRequestSetup: "organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata), } ) - return data @staticmethod def add_management_endpoint_metadata_to_request_metadata( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cfac28db1a9..bbfb8319789 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -707,6 +707,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None) _metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span + _metadata["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata(user_api_key_dict) _metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation _metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups # The per-model budget counters are keyed off these. get_sanitized_user_information_from_key diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dc1133cf596..850e1498608 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -927,6 +927,7 @@ from litellm.proxy.spend_tracking.budget_reservation import ( from litellm.proxy.spend_tracking.committed_spend import ( committed_spend_key, committed_spend_keys, + is_committed_spend_key, seed_committed_spend, ) from litellm.proxy.spend_tracking.spend_capture_rate import ( @@ -3099,6 +3100,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL proxy_batch_write_at = PROXY_BATCH_WRITE_AT +_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5 proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS litellm_master_key_hash = None disable_spend_logs = False @@ -3874,9 +3876,25 @@ def _spend_counter_token(key_token: str) -> str: return hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token +def _committed_spend_ttl() -> int | None: + """A twin outlives two DB spend writes, or the next charge seeds it from a total still missing spend.""" + redis_cache: Final = spend_counter_cache.redis_cache + redis_ttl: Final = redis_cache.default_ttl if redis_cache is not None else None + if redis_ttl is None: + return None + return max(redis_ttl, 2 * (proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS)) + + +def _spend_counter_ttl(counter_key: str, redis_ttl: int | None) -> int | None: + return _committed_spend_ttl() if is_committed_spend_key(counter_key) else redis_ttl + + async def _committed_spend_seeded(counter_key: str) -> bool: return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend( - spend_counter_cache=spend_counter_cache, prisma_client=prisma_client, counter_key=counter_key + spend_counter_cache=spend_counter_cache, + prisma_client=prisma_client, + counter_key=counter_key, + ttl=_committed_spend_ttl(), ) @@ -4301,7 +4319,9 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> ttl: Final = redis_cache.get_ttl() with service_target(SPEND_COUNTERS_TARGET): for item in pending: - batch.increment(item.counter_key, item.increment, ttl).on_settled(_settle_spend_counter_increment(item)) + batch.increment(item.counter_key, item.increment, _spend_counter_ttl(item.counter_key, ttl)).on_settled( + _settle_spend_counter_increment(item) + ) return True @@ -4354,7 +4374,11 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) - ) ttl: Final = redis_cache.get_ttl() increment_list: Final = [ - RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) + RedisPipelineIncrementOperation( + key=item.counter_key, + increment_value=item.increment, + ttl=_spend_counter_ttl(item.counter_key, ttl), + ) for item in pending ] results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list) @@ -11367,7 +11391,7 @@ class ProxyStartupEvent: ) # Ensure minimum interval of 30 seconds for batch writing to prevent memory issues - batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5) + batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS) ### PROXY WORKER HEARTBEAT ### worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client) diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 0f542040d8e..2b69d9cf4a5 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -34,13 +34,14 @@ if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache from litellm.proxy.utils import PrismaClient +_COMMITTED_KEY_PREFIX: Final = "committed_" _BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX) _FETCHED_VALUES: Final = TypeAdapter(dict[str, float | None]) _STORED_VALUE: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) def committed_spend_key(counter_key: str) -> str: - return f"committed_{counter_key}" + return f"{_COMMITTED_KEY_PREFIX}{counter_key}" def _has_committed_twin(counter_key: str) -> bool: @@ -51,6 +52,10 @@ def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]: return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if _has_committed_twin(counter_key)) +def is_committed_spend_key(counter_key: str) -> bool: + return counter_key.startswith(_COMMITTED_KEY_PREFIX) + + def committed_spend_tracked() -> bool: return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger) @@ -142,13 +147,29 @@ async def drop_committed_spend(spend_counter_cache: DualCache, counter_keys: Seq """After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again.""" try: await spend_counter_cache.async_delete_cache_keys(committed_spend_keys(counter_keys)) - except Exception as e: # noqa: BLE001 # the write already committed; a kept twin only lasts until its TTL + except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL verbose_proxy_logger.warning("committed spend drop failed: %s", e) @with_service_target(SPEND_COUNTERS_TARGET) async def seed_committed_spend( - spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str + spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str, ttl: int | None = None +) -> bool: + """One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote.""" + twin_key: Final = committed_spend_key(counter_key) + async with await SpendCounterReseed.counter_lock(twin_key): + try: + stored: Final = await _read_committed_values(spend_counter_cache=spend_counter_cache, twin_keys=(twin_key,)) + except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again + verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e) + return False + return stored.get(twin_key) is not None or await _seed_committed_spend_from_db( + spend_counter_cache, prisma_client, counter_key, ttl + ) + + +async def _seed_committed_spend_from_db( + spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str, ttl: int | None ) -> bool: db_spend: Final = await ( SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key) @@ -164,7 +185,7 @@ async def seed_committed_spend( spend_counter_cache.in_memory_cache.set_cache(key=twin_key, value=db_spend) # pyright: ignore[reportUnknownMemberType] # untyped cache API return True try: - await redis_cache.async_set_max(key=twin_key, value=db_spend) + await redis_cache.async_set_max(key=twin_key, value=db_spend, ttl=ttl) except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e) return False diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 6c3272c74bb..8a89df55cec 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3696,34 +3696,41 @@ async def view_spend_logs( ) -class _TokenRow(TypedDict): - token: ReadOnly[str] +class _EntityIdRow(TypedDict): + id: ReadOnly[str] -class _TeamIdRow(TypedDict): - team_id: ReadOnly[str] +_COMMITTED_DROP_PAGE_SIZE: Final = 1000 +_TOKEN_PAGE_QUERY: Final = ( + 'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2' +) +_TEAM_PAGE_QUERY: Final = 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2' + + +async def _drop_committed_spend_page( + prisma_client: PrismaClient, page_query: str, counter_prefix: str, after_id: str +) -> str | None: + from litellm.proxy.proxy_server import spend_counter_cache + + rows: Final[Sequence[_EntityIdRow]] = await _query_raw( + prisma_client, page_query, after_id, _COMMITTED_DROP_PAGE_SIZE + ) + await drop_committed_spend(spend_counter_cache, tuple(f"{counter_prefix}{row['id']}" for row in rows)) + return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None + + +async def _drop_committed_spend_pages(prisma_client: PrismaClient, page_query: str, counter_prefix: str) -> None: + after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time + while after_id is not None: + after_id = await _drop_committed_spend_page(prisma_client, page_query, counter_prefix, after_id) async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None: - from litellm.proxy.proxy_server import spend_counter_cache - try: - token_rows: Final[Sequence[_TokenRow]] = await _query_raw( - prisma_client, 'SELECT token FROM "LiteLLM_VerificationToken"' - ) - team_rows: Final[Sequence[_TeamIdRow]] = await _query_raw( - prisma_client, 'SELECT team_id FROM "LiteLLM_TeamTable"' - ) - except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin only lasts until its TTL + await _drop_committed_spend_pages(prisma_client, _TOKEN_PAGE_QUERY, "spend:key:") + await _drop_committed_spend_pages(prisma_client, _TEAM_PAGE_QUERY, "spend:team:") + except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e) - return - await drop_committed_spend( - spend_counter_cache, - ( - *(f"spend:key:{row['token']}" for row in token_rows), - *(f"spend:team:{row['team_id']}" for row in team_rows), - ), - ) @router.post( diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 4b73bb6dd34..bbc32b3fe9d 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -21,6 +21,7 @@ from litellm.caching.redis_batch import ( active_post_call_redis_batch, drain_post_call_redis_batches, flush_post_call_redis_batches, + post_call_writes_unsent, request_redis_batch_scope, ) from litellm.integrations.custom_logger import CustomLogger @@ -514,3 +515,16 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it" assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend" assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded" + + +@pytest.mark.asyncio +async def test_post_call_writes_count_as_unsent_until_a_pipeline_carried_them(): + assert post_call_writes_unsent() is False + redis_cache = PostCallFakeRedisCache(FakeClient(_ok_replies)) + with request_redis_batch_scope(): + await flush_post_call_redis_batches() + assert post_call_writes_unsent() is True + active_post_call_redis_batch(redis_cache).increment("spend:key:k", 1.0, None) + assert post_call_writes_unsent() is True + await flush_post_call_redis_batches() + assert post_call_writes_unsent() is False diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 1a746de588e..351d5be1c8f 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -9,11 +9,18 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeAsyncRedis, FakeServer from prometheus_client import REGISTRY from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache +from litellm.caching.redis_batch import ( + active_post_call_redis_batch, + flush_post_call_redis_batches, + request_redis_batch_scope, +) +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker +from litellm._service_logger import ServiceLogging from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ( LiteLLM_BudgetTable, @@ -376,6 +383,103 @@ async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_ prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0)) +class _RedisCacheOnClient(RedisCache): + def __init__(self, client: FakeAsyncRedis) -> None: + self._injected_client = client + self.namespace = None + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self.service_logger_obj = ServiceLogging() + self.default_ttl = None + + def init_async_client(self) -> FakeAsyncRedis: + return self._injected_client + + +_TWIN_KEYS = tuple( + committed_spend_key(counter_key) + for counter_key in ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1") +) + + +@pytest.fixture +def redis_spend_counter_cache(): + redis_cache = _RedisCacheOnClient(FakeAsyncRedis(server=FakeServer())) + cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam + "litellm.proxy.proxy_server.spend_counter_cache", cache + ): + yield redis_cache + + +async def _store_auth_time_spend_plus_a_two_dollar_charge(redis_cache: _RedisCacheOnClient) -> None: + for twin_key, spend in zip(_TWIN_KEYS, (5.0, 32.0, 10.0, 62.0)): + await redis_cache.init_async_client().set(twin_key, spend) + + +@pytest.mark.asyncio +async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline( + prometheus_logger, getters, redis_spend_counter_cache +): + """Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count.""" + _stub_gauges(prometheus_logger) + await _store_auth_time_spend_plus_a_two_dollar_charge(redis_spend_counter_cache) + + with request_redis_batch_scope(): + batch = active_post_call_redis_batch(redis_spend_counter_cache) + for twin_key in _TWIN_KEYS: + batch.increment(twin_key, 2.0, None) + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + await flush_post_call_redis_batches() + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(3.0), + "team": pytest.approx(266.0), + "user": pytest.approx(38.0), + "org": pytest.approx(436.0), + } + + +@pytest.mark.asyncio +async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out( + prometheus_logger, getters, redis_spend_counter_cache +): + """The 2.0 over the auth time spend is this request's own charge, sent before the gauges read it.""" + _stub_gauges(prometheus_logger) + await _store_auth_time_spend_plus_a_two_dollar_charge(redis_spend_counter_cache) + + with request_redis_batch_scope(): + batch = active_post_call_redis_batch(redis_spend_counter_cache) + batch.increment("unrelated-write", 1.0, None) + await flush_post_call_redis_batches() + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + +@pytest.mark.asyncio +async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own( + prometheus_logger, getters, spend_counter_cache +): + """Without Redis the charge increments in memory at once, so a request scope says nothing about it.""" + _stub_gauges(prometheus_logger) + _store_committed(spend_counter_cache, key=5.0, team=32.0, user=10.0, org=62.0) + + with request_redis_batch_scope(): + await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth())) + + assert _remaining(prometheus_logger) == { + "key": pytest.approx(5.0), + "team": pytest.approx(268.0), + "user": pytest.approx(40.0), + "org": pytest.approx(438.0), + } + + @pytest.mark.asyncio async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters( prometheus_logger, getters, spend_counter_cache diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a9098ac3d59..738a7c709c3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -39,6 +39,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, @@ -7829,6 +7830,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +@pytest.mark.asyncio +async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted(): + """A forged body value stays out, and the row's auth keeps no callback credentials for its loggers.""" + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None + ) + user_api_key_dict: Final = UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys), + metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"}, + ) + + kwargs: Final = _passthrough_kwargs_for_reservation( + user_api_key_dict, + parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}}, + ) + carried: Final = kwargs["litellm_params"]["metadata"]["user_api_key_auth"] + increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs) + + assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys + assert carried.metadata == {"kept": "yes"} + assert user_api_key_dict.metadata["logging"] == [{"callback_name": "langfuse"}] + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """ diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py index a6beb006d0c..1d7c8fdf304 100644 --- a/tests/unit/proxy/spend_tracking/test_committed_spend.py +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -2,9 +2,11 @@ from __future__ import annotations +import asyncio from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 import pytest from fakeredis import FakeAsyncRedis, FakeServer @@ -13,6 +15,7 @@ import litellm from litellm._service_logger import ServiceLogging from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_batch import flush_post_call_redis_batches, request_redis_batch_scope from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.spend_tracking.committed_spend import ( @@ -23,7 +26,13 @@ from litellm.proxy.spend_tracking.committed_spend import ( read_committed_budget_spend, seed_committed_spend, ) +from litellm.proxy.proxy_server import ( + _defer_spend_counter_increments, + _prepare_committed_spend_increments, + run_spend_counter_pipeline, +) from litellm.proxy.spend_tracking.spend_counter_batch import ( + PendingSpendIncrement, active_spend_counter_batch, read_batched_spend_counter, spend_counter_batch_scope, @@ -36,11 +45,11 @@ class _FakeRedisSettingMax(FakeAsyncRedis): """fakeredis runs no Lua without lupa, so this stands in for the GET/compare/SET script of ``async_set_max``.""" async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> bytes: - key, value, _ttl = keys_and_args + key, value, ttl = keys_and_args current: Final = await self.get(key) if current is not None and float(current) >= float(value): return current - await self.set(key, value) + await self.set(key, value, ex=int(ttl) if int(ttl) > 0 else None) return value.encode() @@ -59,8 +68,11 @@ class _RedisCacheOnClient(RedisCache): class _FindUniqueTable: def __init__(self, row: SimpleNamespace | None) -> None: self._row: Final = row + self.reads = 0 async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + self.reads += 1 + await asyncio.sleep(0) return self._row @@ -142,6 +154,71 @@ async def test_seed_writes_the_db_spend_when_the_twin_is_absent( ) +@pytest.mark.asyncio +async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(spend_counter_cache: DualCache): + prisma_client: Final = _prisma_with_spend(12.0) + counter_key: Final = f"spend:key:{uuid4()}" + + seeded: Final = await asyncio.gather( + *(seed_committed_spend(spend_counter_cache, prisma_client, counter_key) for _ in range(5)) + ) + + assert seeded == [True] * 5 + assert prisma_client.db.litellm_verificationtoken.reads == 1 + + +@pytest.mark.asyncio +async def test_seed_reports_no_twin_when_redis_cannot_be_read(): + redis_cache: Final = AsyncMock(spec=RedisCache) + redis_cache.async_batch_get_cache.side_effect = ConnectionError("redis unreachable") + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + assert await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), "spend:team:t1") is False + + +@pytest.mark.parametrize( + ("batch_write_seconds", "expected_twin_ttl"), + [ + (10, 60), + (120, 250), + ], +) +@pytest.mark.asyncio +async def test_twin_outlives_two_db_spend_writes_on_every_write_path( + monkeypatch: pytest.MonkeyPatch, batch_write_seconds: int, expected_twin_ttl: int +): + """A twin expiring before the DB write would be seeded from a total still missing its charges.""" + client: Final = _FakeRedisSettingMax(server=FakeServer()) + redis_cache: Final = _RedisCacheOnClient(client) + redis_cache.default_ttl = 60 + spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_spend(12.0)) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_batch_write_at", batch_write_seconds) + seeded_twin: Final = committed_spend_key("spend:key:seeded") + piped_twin: Final = committed_spend_key("spend:key:piped") + batched_twin: Final = committed_spend_key("spend:key:batched") + + await _prepare_committed_spend_increments(("spend:key:seeded",), 1.0) + await run_spend_counter_pipeline( + ( + PendingSpendIncrement(counter_key=piped_twin, increment=1.0), + PendingSpendIncrement(counter_key="spend:key:piped", increment=1.0), + ) + ) + with request_redis_batch_scope(): + assert _defer_spend_counter_increments( + ( + PendingSpendIncrement(counter_key=batched_twin, increment=1.0), + PendingSpendIncrement(counter_key="spend:key:batched", increment=1.0), + ) + ) + await flush_post_call_redis_batches() + + assert [await client.ttl(twin) for twin in (seeded_twin, piped_twin, batched_twin)] == [expected_twin_ttl] * 3 + assert [await client.ttl(counter) for counter in ("spend:key:piped", "spend:key:batched")] == [60, 60] + + @pytest.mark.asyncio async def test_seed_never_lowers_a_twin_another_replica_already_advanced(spend_counter_cache: DualCache): await _store(spend_counter_cache, committed_spend_key("spend:team:t1"), 20.0) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c82dcb57bfd..057f512ab0d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8284,8 +8284,9 @@ class _ResetDb: self.litellm_verificationtoken: Final = _ResetTableRows() self.litellm_teamtable: Final = _ResetTableRows() - async def query_raw(self, sql_query: str) -> list[dict[str, str]]: - return [{"token": "hashed"}] if "VerificationToken" in sql_query else [{"team_id": "t1"}] + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"] + return [{"id": row} for row in rows if row > after_id][:limit] @pytest.mark.asyncio @@ -8313,6 +8314,38 @@ async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gau assert [spend_counter_cache.in_memory_cache.get_cache(key=twin) for twin in twins] == [twin_after, twin_after] +class _PagedResetDb(_ResetDb): + def __init__(self, token_count: int) -> None: + super().__init__() + self._tokens: Final = sorted(f"token-{index}" for index in range(token_count)) + self.largest_page = 0 + + async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]: + rows: Final = self._tokens if "VerificationToken" in sql_query else [] + page: Final = [{"id": row} for row in rows if row > after_id][:limit] + self.largest_page = max(self.largest_page, len(page)) + return page + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500]) +async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count): + db: Final = _PagedResetDb(token_count) + spend_counter_cache: Final = DualCache() + twins: Final = tuple(f"committed_spend:key:token-{index}" for index in range(token_count)) + for twin in twins: + spend_counter_cache.in_memory_cache.set_cache(key=twin, value=7.0) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=db)) + monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)]) + + response: Final = await global_spend_reset() + + assert response["status"] == "success" + assert [twin for twin in twins if spend_counter_cache.in_memory_cache.get_cache(key=twin) is not None] == [] + assert db.largest_page <= 1000 + + @pytest.mark.asyncio async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch): db: Final = _ResetDb()