diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ce6c9330620..6a2d5a84c3d 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -103,35 +103,7 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: return counter_key == f"spend:key:{valid_token.token}" and should_throttle_budget_exceeded(valid_token) -async def _apply_over_budget_reservation_policy( - counter: _BudgetCounter, - valid_token: UserAPIKeyAuth | None, - entry: dict[str, float | str], - applied_entries: list[dict[str, float | str]], - reservation_cost: float, - current_spend: float, -) -> float: - """ - Decide what to do when a counter is over budget, and return the reservation - cost to carry into the next counter. Three outcomes: an over-budget key that - opted into throttling releases its own reservation (the rate limiter slows - it) and keeps the cost; a partially-remaining budget resizes the reservation - down to what is left; anything else hard-blocks by raising. - """ - if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): - await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) - applied_entries.remove(entry) - return reservation_cost - - remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation > 1e-12: - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, - ) - return remaining_before_reservation - +def _raise_counter_budget_exceeded(counter: _BudgetCounter, current_spend: float) -> NoReturn: raise litellm.BudgetExceededError( current_cost=current_spend, max_budget=counter.max_budget, @@ -146,6 +118,49 @@ async def _apply_over_budget_reservation_policy( ) +async def _apply_over_budget_reservation_policy( + counter: _BudgetCounter, + valid_token: UserAPIKeyAuth | None, + entry: dict[str, float | str], + applied_entries: list[dict[str, float | str]], + reservation_cost: float, + current_spend: float, + fail_closed_budget_enforcement: bool, +) -> float: + """ + Decide what to do when a counter is over budget, and return the reservation + cost to carry into the next counter. An over-budget key that opted into + throttling releases its own reservation (the rate limiter slows it) and + keeps the cost. Otherwise strict enforcement blocks the request, since the + estimated cost is already known not to fit, while soft mode resizes the + reservation down to the budget that was left and admits the request. + """ + if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): + await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) + applied_entries.remove(entry) + return reservation_cost + + if fail_closed_budget_enforcement: + verbose_proxy_logger.warning( + "fail_closed_budget_enforcement: rejecting request, estimated cost %s does not fit the budget " + "remaining on %s", + reservation_cost, + counter.counter_key, + ) + _raise_counter_budget_exceeded(counter=counter, current_spend=current_spend) + + remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) + if remaining_before_reservation > 1e-12: + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + return remaining_before_reservation + + _raise_counter_budget_exceeded(counter=counter, current_spend=current_spend) + + async def reserve_budget_for_request( request_body: dict, route: str, @@ -241,6 +256,7 @@ async def reserve_budget_for_request( applied_entries=applied_entries, reservation_cost=reservation_cost, current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) continue except Exception: diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 133b53bb18d..987a7a83287 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -185,6 +185,7 @@ async def test_over_budget_window_counter_tags_clean_entity_id(): applied_entries=[], reservation_cost=0.5, current_spend=2.0, + fail_closed_budget_enforcement=False, ) assert exc_info.value.entity_type == "key" assert exc_info.value.entity_id == "test-token" @@ -802,6 +803,184 @@ async def test_should_cap_known_estimate_to_remaining_budget( ) == pytest.approx(0.9) +@pytest.mark.asyncio +async def test_should_reject_known_estimate_over_remaining_budget_when_fail_closed( + spend_counter_state, +): + """Capping a known estimate to the leftover budget admits a request whose full + cost is already known not to fit, and reconciliation then pushes the counter + past max_budget. Strict mode must block it and roll the counter back.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-known-estimate-strict", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-known-estimate-strict", + value=0.9, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.entity_type == "key" + assert exc_info.value.entity_id == "key-budget-known-estimate-strict" + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-strict" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_should_admit_exact_boundary_estimate_when_fail_closed( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-strict-boundary", + spend=0.4, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-strict-boundary", + value=0.4, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-strict-boundary" + ) == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_strict_rejection_releases_counters_reserved_before_the_full_one( + spend_counter_state, +): + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-strict-multi", + spend=0.0, + max_budget=10.0, + team_id="team-budget-strict-multi", + ) + team_object = LiteLLM_TeamTable( + team_id="team-budget-strict-multi", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:team:team-budget-strict-multi", + value=0.9, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.entity_type == "team" + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-strict-multi" + ) == pytest.approx(0.0) + assert counter_cache.in_memory_cache.get_cache( + key="spend:team:team-budget-strict-multi" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_strict_mode_keeps_throttled_key_exemption( + spend_counter_state, monkeypatch +): + """A key that opted into budget-exceeded throttling is slowed by the rate + limiter, so strict enforcement must not turn its own counter into a block.""" + monkeypatch.setattr(litellm, "budget_exceeded_throttle_percentage", 0.1) + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-strict-throttle", + spend=0.9, + max_budget=1.0, + tpm_limit=1000, + rpm_limit=100, + metadata={"throttle_on_budget_exceeded": True}, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-strict-throttle", + value=0.9, + ) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert reservation is None + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-strict-throttle" + ) == pytest.approx(0.9) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_default_when_output_cap_missing( spend_counter_state,