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 42fe0837c29..852b2b68e76 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, @@ -74,8 +75,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, @@ -364,6 +367,73 @@ 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 + + +_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 _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 _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() + try: + from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.spend_tracking.committed_spend import read_committed_budget_spend + + 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) + 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 + 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), + 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 @@ -2228,6 +2298,9 @@ 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) + _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 @@ -2237,45 +2310,66 @@ 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_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 if _row_carries_auth else committed_now, + now=committed_now, response_cost=response_cost, - key_max_budget=_api_key_max_budget, - key_spend=_api_key_spend, - carried=KeyBudgetSnapshot.from_metadata(_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_max_budget=_team_max_budget, - response_cost=response_cost, - carried=TeamBudgetSnapshot.from_metadata(_metadata), - ), - self._set_user_budget_metrics_after_api_request( - user_id=user_id, - user_spend=_user_spend, - user_max_budget=_user_max_budget, - response_cost=response_cost, - carried=UserBudgetSnapshot.from_metadata(_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), - org_alias=_org_alias if isinstance(_org_alias, str) else None, - ), - self._set_customer_budget_metrics_after_api_request( - end_user_id=end_user_id, - response_cost=response_cost, - ), - 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 " @@ -2798,9 +2892,22 @@ 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_org: Final = OrgBudgetSnapshot.from_metadata(failure_metadata) + failure_committed: Final = _committed_spend_before_request( + 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=failure_org, + 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) @@ -4511,6 +4618,7 @@ class PrometheusLogger(CustomLogger): response_cost: float, carried: OrgBudgetSnapshot | None = None, org_alias: str | None = None, + committed_spend: float | None = None, ): """ Set org budget metrics after an LLM API request @@ -4528,7 +4636,7 @@ class PrometheusLogger(CustomLogger): self._set_org_budget_metrics( org_id=org_id, org_alias=org_alias or "", - spend=carried.spend + response_cost, + spend=_committed_floor(carried.spend, committed_spend) + response_cost, max_budget=carried.max_budget, budget_reset_at=None, ) @@ -4554,7 +4662,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 @@ -4621,6 +4729,7 @@ class PrometheusLogger(CustomLogger): self, end_user_id: str | None, response_cost: float, + committed_spend: float | None = None, ): if self._customer_budget_gauges_are_noop() or not _customer_budget_metrics_enabled(): return @@ -4646,7 +4755,9 @@ 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=_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/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2b0522c09b5..84d2ba64674 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,7 +170,15 @@ 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_committed_budget_spend, + carry_team_and_user_budget_state, +) +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, @@ -3571,6 +3580,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, ) @@ -3881,6 +3891,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, @@ -3956,6 +3967,18 @@ async def run_centralized_common_checks( user_obj=user_object, ) + 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=( + 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, request=request, @@ -4272,6 +4295,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/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/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..d3f26f17371 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", @@ -97,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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 2951bc02d8c..b64abf24d67 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -41,6 +41,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_request 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, @@ -69,6 +70,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, @@ -238,6 +241,11 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _committed_spend_scope(user_api_key_dict: UserAPIKeyAuth) -> EntitySpendCounterKeys: + committed: Final = user_api_key_dict.committed_budget_spend + return committed.counter_keys if committed is not None else EntitySpendCounterKeys() + + def _proxy_spend_counters() -> "_IncrementSpendCounters": from litellm.proxy.proxy_server import increment_spend_counters @@ -622,6 +630,7 @@ class _ProxyDBLogger(CustomLogger): team_id=user_api_key_dict.team_id, org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), kwargs=dict(row_kwargs), completion_response=original_exception, start_time=start_time, @@ -739,6 +748,7 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), project_id=user_api_key_dict.project_id, + committed_spend_scope=_committed_spend_scope(user_api_key_dict), ) except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") @@ -1201,9 +1211,34 @@ class _IncrementSpendCounters(Protocol): 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: ... +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, @@ -1222,6 +1257,7 @@ async def _update_database_and_spend_counters( model_access_groups: Sequence[str] | None = None, project_id: str | None = None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: """The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then spans the database write and the counter update, so the post-call counters are read with a single MGET after the @@ -1264,6 +1300,7 @@ async def _update_database_and_spend_counters( request_tags=request_tags, model_access_groups=model_access_groups, project_id=project_id, + committed_spend_scope=committed_spend_scope, ) @@ -1285,6 +1322,7 @@ async def _update_database_and_spend_counters_in_batch( model_access_groups: Sequence[str] | None, project_id: str | None, spend_writer: DBSpendUpdateWriter | None = None, + committed_spend_scope: EntitySpendCounterKeys | None = None, ) -> bool: try: writer: Final = spend_writer if spend_writer is not None else proxy_logging_obj.db_spend_update_writer @@ -1338,6 +1376,19 @@ 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_spend_scope + if committed_spend_scope is not None + else 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/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..27d1ecc1653 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/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 54a2d4062d2..743041cfdaa 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1443,8 +1443,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/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 61da5599886..d99e0d6b1d6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -157,6 +157,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 @@ -4263,7 +4264,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) @@ -4278,6 +4279,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( @@ -6964,6 +6966,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, @@ -7018,6 +7021,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/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/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 a767f5423ed..ac513e88ba7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -924,6 +924,12 @@ 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 ( + committed_spend_key, + committed_spend_keys, + is_committed_spend_key, + seed_committed_spend, +) from litellm.proxy.spend_tracking.spend_capture_rate import ( run_scheduled_spend_capture_rate_check, ) @@ -931,6 +937,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, @@ -1033,6 +1040,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, @@ -3092,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 @@ -3511,6 +3520,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. @@ -3547,6 +3557,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, ) @@ -3571,9 +3582,16 @@ 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 = ( + 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))) reservation_update: Final = await _reconcile_budget_reservation_for_counter_update( budget_reservation=budget_reservation, response_cost=response_cost, @@ -3596,9 +3614,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, ...]] = ( () @@ -3778,6 +3794,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 ) @@ -3853,6 +3872,43 @@ 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 + + +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, + ttl=_committed_spend_ttl(), + ) + + +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, @@ -4263,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 @@ -4316,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) @@ -11325,7 +11387,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/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index 0dfb38272b0..150ddd1efe9 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -1,14 +1,17 @@ """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 +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 ( + CommittedBudgetSpend, + EntitySpendCounterKeys, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -34,6 +37,41 @@ def carry_team_and_user_budget_state( ) +def _budgeted_org_id(valid_token: UserAPIKeyAuth, team_object: LiteLLM_TeamTable | None) -> str | None: + org_snapshot: Final = valid_token.org_budget_snapshot + if org_snapshot is None or org_snapshot.max_budget is None: + return None + return valid_token.org_id or (team_object.organization_id if team_object is not None else 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 end_user_object.user_id + + +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_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, + 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 = ( # 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: budget_table: Final = org_table.litellm_budget_table valid_token.organization_alias = org_table.organization_alias diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py new file mode 100644 index 00000000000..2b69d9cf4a5 --- /dev/null +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -0,0 +1,192 @@ +"""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 for entities with a budget, while a ``PrometheusLogger`` is registered. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable, Mapping, Sequence +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._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 + 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_KEY_PREFIX}{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)) + + +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) + + +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 + ) + + +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} + + +def _committed_budget_spend( + counter_keys: EntitySpendCounterKeys, committed: Mapping[str, float] +) -> 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), + org=spend_of(counter_keys.org), + end_user=spend_of(counter_keys.end_user), + ) + + +@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 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: + """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 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, 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) + 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, 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 + return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 89200792c0b..8a89df55cec 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -52,6 +52,7 @@ from litellm.proxy.auth.authorization_dependencies import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.openai_error_payload import error_status_code, openai_error_param, openai_error_type 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, @@ -3695,6 +3696,43 @@ async def view_spend_logs( ) +class _EntityIdRow(TypedDict): + 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: + try: + 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) + + @router.post( "/global/spend/reset", tags=["Budget & Spend Tracking"], @@ -3723,6 +3761,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 1183c36eda7..b9eb81e98d5 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 ( + CommittedBudgetSpend, OrgBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -206,6 +207,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) + 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 @@ -329,6 +331,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 1a0c55a38ff..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") @@ -25,7 +28,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 @@ -64,3 +67,51 @@ 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 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 + 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..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 e49e6baa8b0..351d5be1c8f 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -4,12 +4,23 @@ 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 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_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, @@ -24,7 +35,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, + EntitySpendCounterKeys, KeyBudgetSnapshot, TeamBudgetSnapshot, UserBudgetSnapshot, @@ -55,6 +69,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.""" @@ -137,6 +160,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, } @@ -198,6 +222,411 @@ 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() +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) + + +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], + } + + +def _token_committed_at_auth() -> UserAPIKeyAuth: + token = _authed_token() + 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 + + +@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_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, + } + + +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 +): + 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)) + + +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 +): + _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), + "org": pytest.approx(438.0), + } + + +@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_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_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) + 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 ab0fb67b52f..210e2aa9ecc 100644 --- a/tests/unit/integrations/test_prometheus_user_team_metrics.py +++ b/tests/unit/integrations/test_prometheus_user_team_metrics.py @@ -1065,6 +1065,81 @@ 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_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, EntitySpendCounterKeys + + 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), + ) + ) + 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( + 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( + user_api_team=None, + user_api_team_alias=None, + user_api_key=None, + user_api_key_alias=None, + 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(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 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..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 @@ -5084,6 +5084,68 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +@pytest.mark.asyncio +@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, EntitySpendCounterKeys + + 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() + 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": None, + "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.committed_budget_spend == ( + 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) + + 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} 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/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index e482b6552a4..42057c0d198 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -40,6 +40,7 @@ from litellm.proxy.utils import ( BilledFailedAttempt, 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 @@ -1080,6 +1081,95 @@ 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 +@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() @@ -1208,6 +1298,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(), ) @@ -2713,7 +2804,9 @@ _LIVE_COUNTERS: Final = { } -async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: +async def _fail_request( + logger: ProxyDBLogger, failure: Mapping[str, object], committed: CommittedBudgetSpend | None = None +) -> None: await logger.async_post_call_failure_hook( request_data={ "model": "gpt-4.1", @@ -2730,6 +2823,7 @@ async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> end_user_id="end-user-1", org_id="org-1", project_id="project-1", + committed_budget_spend=committed, ), ) @@ -2768,6 +2862,43 @@ async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ({"combined_usage_object": Usage(prompt_tokens=30), "response_cost": 0.125}, 0.125), + (_failure_with_billed_attempts(), 1.25), + ({"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}, 0.125), + ], + ids=["recovered_stream_cost", "billed_retries", "completed_cost"], +) +async def test_a_failed_charge_moves_only_the_committed_counters_auth_named_as_budgeted( + live_spend_counters: DualCache, failure: Mapping[str, object], charged: float +): + """The remaining budget gauges read the committed counters, so a charge missing from them shows as budget the + request never left behind, until the next reseed from the database.""" + budgeted: Final = EntitySpendCounterKeys.of( + token="hashed-key", team_id="team-1", user_id=None, org_id=None, end_user_id=None + ) + starting: Final = { + "committed_spend:key:hashed-key": 0.2, + "committed_spend:team:team-1": 0.3, + "committed_spend:user:user-1": 0.4, + } + for key, spend in starting.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), failure, CommittedBudgetSpend(counter_keys=budgeted)) + + assert {key: await live_spend_counters.async_get_cache(key) for key in starting} == pytest.approx( + { + "committed_spend:key:hashed-key": 0.2 + charged, + "committed_spend:team:team-1": 0.3 + charged, + "committed_spend:user:user-1": 0.4, + } + ) + + @pytest.mark.asyncio async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache): """A failure whose cost was already finalized settles through the success path's write, which raises the 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 39586a1a204..be18a1e1b07 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_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 0198f582be7..cb2c564dbc3 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -8672,6 +8672,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 @@ -8988,6 +8991,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/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/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/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..9e4f76cf993 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,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 +from litellm.types.proxy.carried_budget_state import EntitySpendCounterKeys from .conftest import normalize @@ -851,6 +854,86 @@ 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 + 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) + + +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)) + 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, + 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 + ) + + 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_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) + + 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, ...] = ( "spend:key:hashed", "spend:team:team-1", @@ -1541,22 +1624,14 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates( @pytest.mark.asyncio -async def test_invalidate_spend_counter_deletes_in_memory_and_redis(monkeypatch): +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="spend:key:k") - 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] == ["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 0bdf43b396c..75b9e422330 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,10 +14,13 @@ from litellm.proxy._types import ( ) from litellm.proxy.spend_tracking.carried_budget_state import ( carried_budget_metadata, + carry_committed_budget_spend, carry_organization_budget_state, carry_team_and_user_budget_state, ) from litellm.types.proxy.carried_budget_state import ( + CommittedBudgetSpend, + EntitySpendCounterKeys, KeyBudgetSnapshot, OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -123,10 +129,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) + 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 "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 _committed_reader(committed: CommittedBudgetSpend): + reads: list[EntitySpendCounterKeys] = [] + + async def read_committed(counter_keys: EntitySpendCounterKeys) -> CommittedBudgetSpend: + reads.append(counter_keys) + return committed + + return read_committed, reads + + +@pytest.mark.asyncio +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_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_committed=read_committed, + ) + + 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_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) + token.org_budget_snapshot = OrgBudgetSnapshot(spend=40.0, max_budget=None) + read_committed, reads = _committed_reader(CommittedBudgetSpend()) + + 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, litellm_budget_table=LiteLLM_BudgetTable(max_budget=None) + ), + read_committed=read_committed, + ) + + assert reads == [EntitySpendCounterKeys()] + assert token.committed_budget_spend == CommittedBudgetSpend() + + +@pytest.mark.asyncio +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_committed, reads = _committed_reader(CommittedBudgetSpend(org=41.0)) + + 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_committed=read_committed, + ) + + 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..1d7c8fdf304 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -0,0 +1,362 @@ +"""Committed spend twins hold only charged cost, so the remaining budget gauges never count a budget hold.""" + +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 + +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 ( + bind_committed_admission_counters, + committed_spend_key, + drop_committed_spend, + read_admission_committed_spend, + 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, +) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys + + +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, ex=int(ttl) if int(ttl) > 0 else None) + 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 + self.reads = 0 + + async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None: + self.reads += 1 + await asyncio.sleep(0) + 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) + + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id="cust-1" + ) + + 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 +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) + + counter_keys: Final = EntitySpendCounterKeys.of( + token="hashed", team_id="t1", user_id=None, org_id=None, end_user_id=None + ) + + 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_spend"), + [ + (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_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) == CommittedBudgetSpend( + counter_keys=counter_keys, **expected_spend + ) + + +@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) + + 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.team is None + + +@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.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 8ba79a45ff3..f32fc3634a5 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -539,9 +539,8 @@ 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 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} 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 6598b94e60f..057f512ab0d 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(): @@ -8270,6 +8274,94 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] +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, 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 +@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] + + +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() + 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() + + def test_spend_calculate_rejection_body_is_openai_shaped(client): app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user", api_key="hashed-k" 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 3647146eb4e..921c79d67bc 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -9693,7 +9693,9 @@ 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", + ] finally: ps.spend_counter_cache = orig_counter