From 54286fca60471fdef2fe1d98112a36a656a59e4a Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Sat, 31 Jan 2026 13:45:33 -0800 Subject: [PATCH] Revert "fix: make cache updates synchronous for budget enforcement" This reverts commit d0383412e8a48f7fe0296aa781943ba03200b440. --- litellm/proxy/auth/auth_checks.py | 8 +- .../proxy/hooks/proxy_track_cost_callback.py | 20 +++-- litellm/proxy/proxy_server.py | 84 ++++++++++--------- 3 files changed, 56 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index faca49ccc6..e0b056d450 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2295,13 +2295,7 @@ async def _check_team_member_budget( and team_membership.litellm_budget_table.max_budget is not None ): team_member_budget = team_membership.litellm_budget_table.max_budget - # Prefer valid_token.team_member_spend (from token cache) over team_membership.spend - # The token cache is updated synchronously after each request, while the team membership - # cache may have stale data since DB updates are batched - if valid_token.team_member_spend is not None: - team_member_spend = valid_token.team_member_spend - else: - team_member_spend = team_membership.spend or 0.0 + team_member_spend = team_membership.spend or 0.0 if team_member_spend >= team_member_budget: raise litellm.BudgetExceededError( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 8f681c727e..dab5fb1bfd 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -181,15 +181,17 @@ class _ProxyDBLogger(CustomLogger): org_id=org_id, ) - # update cache - await to ensure budget checks see updated spend - await update_cache( - token=user_api_key, - user_id=user_id, - end_user_id=end_user_id, - response_cost=response_cost, - team_id=team_id, - parent_otel_span=parent_otel_span, - tags=tags, + # update cache + asyncio.create_task( + update_cache( + token=user_api_key, + user_id=user_id, + end_user_id=end_user_id, + response_cost=response_cost, + team_id=team_id, + parent_otel_span=parent_otel_span, + tags=tags, + ) ) await proxy_logging_obj.slack_alerting_instance.customer_spend_alert( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91094c49c9..f12d4d6ab4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,5 +1,6 @@ import asyncio import copy +import enum import inspect import io import os @@ -12,7 +13,6 @@ import time import traceback import warnings from datetime import datetime, timedelta, timezone -import enum from typing import ( TYPE_CHECKING, Any, @@ -29,41 +29,9 @@ from typing import ( get_origin, get_type_hints, ) + from pydantic import BaseModel, Json -from litellm.proxy._types import ( - ProxyException, - UserAPIKeyAuth, - LiteLLM_UserTable, - CommonProxyErrors, - LitellmUserRoles, - ConfigList, - ConfigYAML, - ConfigFieldUpdate, - ConfigGeneralSettings, - ConfigFieldInfo, - PassThroughGenericEndpoint, - FieldDetail, - ConfigFieldDelete, - CallbackDelete, - InvitationClaim, - InvitationModel, - InvitationNew, - InvitationUpdate, - InvitationDelete, - CallInfo, - Litellm_EntityType, - TeamDefaultSettings, - RoleBasedPermissions, - SupportedDBObjectType, - ProxyErrorTypes, - EnterpriseLicenseData, - LiteLLM_JWTAuth, - TokenCountRequest, - TransformRequestBody, - LiteLLM_TeamTable, - SpecialModelNames, -) from litellm._uuid import uuid from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT, @@ -84,6 +52,39 @@ from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import ( + CallbackDelete, + CallInfo, + CommonProxyErrors, + ConfigFieldDelete, + ConfigFieldInfo, + ConfigFieldUpdate, + ConfigGeneralSettings, + ConfigList, + ConfigYAML, + EnterpriseLicenseData, + FieldDetail, + InvitationClaim, + InvitationDelete, + InvitationModel, + InvitationNew, + InvitationUpdate, + Litellm_EntityType, + LiteLLM_JWTAuth, + LiteLLM_TeamTable, + LiteLLM_UserTable, + LitellmUserRoles, + PassThroughGenericEndpoint, + ProxyErrorTypes, + ProxyException, + RoleBasedPermissions, + SpecialModelNames, + SupportedDBObjectType, + TeamDefaultSettings, + TokenCountRequest, + TransformRequestBody, + UserAPIKeyAuth, +) from litellm.proxy.common_utils.callback_utils import ( normalize_callback_names, process_callback, @@ -1703,11 +1704,12 @@ async def update_cache( # noqa: PLR0915 if tags is not None: await _update_tag_cache() - # Await cache update to ensure budget checks see updated spend values - await user_api_key_cache.async_set_cache_pipeline( - cache_list=values_to_update_in_cache, - ttl=60, - litellm_parent_otel_span=parent_otel_span, + asyncio.create_task( + user_api_key_cache.async_set_cache_pipeline( + cache_list=values_to_update_in_cache, + ttl=60, + litellm_parent_otel_span=parent_otel_span, + ) ) @@ -3643,7 +3645,9 @@ class ProxyConfig: ) else: # Interval-based scheduling (existing behavior) - from litellm.litellm_core_utils.duration_parser import duration_in_seconds + from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, + ) retention_interval = general_settings.get( "maximum_spend_logs_retention_interval", "1d"