scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(proxy): keep the spend already counted when a budget hold is far over the budget

Browse source
This commit is contained in:
mateo-berri 2026-10-29 14:00:25 -07:00
parent a443e4ab9f
commit 0c66c95f46
3 changed files with 39 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -180,7 +180,7 @@ async def _apply_over_budget_reservation_policy(
entry: dict[str, float | str],
applied_entries: list[dict[str, float | str]],
reservation_cost: float,
current_spend: float,
spend_before_reservation: float,
fail_closed_budget_enforcement: bool = False,
) -> float:
"""
@ -196,13 +196,13 @@ async def _apply_over_budget_reservation_policy(
applied_entries.remove(entry)
return reservation_cost
remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost)
remaining_before_reservation: Final = counter.max_budget - spend_before_reservation
if remaining_before_reservation <= 1e-12:
_raise_counter_budget_exceeded(counter=counter, current_cost=current_spend)
if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12:
_raise_counter_budget_exceeded(counter=counter, current_cost=spend_before_reservation + reservation_cost)
if fail_closed_budget_enforcement and reservation_cost - remaining_before_reservation > 1e-12:
_raise_counter_budget_exceeded(
counter=counter,
current_cost=current_spend - reservation_cost,
current_cost=spend_before_reservation,
estimated_cost=reservation_cost,
)
await _resize_applied_reservation(
@ -992,7 +992,7 @@ async def _reserve_reservable_counters(
current_spend_by_counter_key=current_spend_by_counter_key,
reservation_cost=reservation_cost,
):
charged_cost = reservation_cost
charged_cost = min(reservation_cost, max(counter.max_budget for counter in group))
entries = tuple(_counter_to_reservation_entry(counter=counter, reserved_cost=charged_cost) for counter in group)
applied_entries.extend(entries)
reserved_values = await _reserve_counters(counters=group, entries=entries, reservation_cost=charged_cost)
@ -1005,18 +1005,19 @@ async def _reserve_reservable_counters(
for counter, entry, reserved_value in zip(group, entries, reserved_values):
if entry not in applied_entries:
continue
if reserved_value is not None:
current_spend = reserved_value - (charged_cost - reservation_cost)
else:
current_spend = current_spend_by_counter_key[counter.counter_key] + reservation_cost
if current_spend > counter.max_budget:
spend_before_reservation = (
current_spend_by_counter_key[counter.counter_key]
if reserved_value is None
else reserved_value - charged_cost
)
if spend_before_reservation + reservation_cost > counter.max_budget:
reservation_cost = await _apply_over_budget_reservation_policy(
counter=counter,
valid_token=valid_token,
entry=entry,
applied_entries=applied_entries,
reservation_cost=reservation_cost,
current_spend=current_spend,
spend_before_reservation=spend_before_reservation,
fail_closed_budget_enforcement=fail_closed_budget_enforcement,
)
return reservation_cost

25
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -357,6 +357,31 @@ def test_an_n_too_large_to_price_holds_one_choice(count: object) -> None:
assert _chat_hold({**GEMINI_CHAT_BODY, "n": count}) == one_choice_hold
HOLDS_FAR_OVER_A_TINY_BUDGET: Final = (
(GEMINI_ROUTE, {"contents": [GEMINI_TEXT_CONTENT], "generationConfig": {"candidateCount": 1e300}}),
("/v1/chat/completions", {**GEMINI_CHAT_BODY, "n": 1e300}),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("route", "request_body"), HOLDS_FAR_OVER_A_TINY_BUDGET)
async def test_a_hold_far_over_the_budget_keeps_the_spend_already_counted(
spend_counter_cache: DualCache, route: str, request_body: dict[str, object]
) -> None:
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
counted_spend: Final = 0.004
hold_estimate: Final = estimate_request_max_cost(request_body=request_body, route=route, llm_router=None)
assert hold_estimate is not None and 1e20 < hold_estimate < math.inf
await spend_counter_cache.async_set_cache(key=counter_key, value=counted_spend)
reservation: Final = await _reserve_for_tiny_budget_key(route, request_body)
assert reservation is not None
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.01)
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(counted_spend)
def test_a_negative_max_tokens_holds_the_default_output() -> None:
uncapped_hold: Final = _chat_hold(GEMINI_CHAT_BODY)
assert uncapped_hold is not None

2
tests/unit/proxy/test_budget_reservation.py
View file

@ -207,7 +207,7 @@ async def test_over_budget_window_counter_tags_clean_entity_id():
entry={"counter_key": counter.counter_key},
applied_entries=[],
reservation_cost=0.5,
current_spend=2.0,
spend_before_reservation=1.5,
)
assert exc_info.value.entity_type == "key"
assert exc_info.value.entity_id == "test-token"