scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(budget): reject partial-headroom requests under fail_closed_budget_enforcement

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
milan 2026-08-21 01:42:33 +00:00
parent dad4c1a0fb
commit b664bf5420
2 changed files with 224 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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:

179
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -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,