mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(budgets): make the model access group ceiling exclusive
A pool whose recorded spend has reached max_budget has nothing left to give, so the next request is refused rather than admitted. This departs from the tag check it otherwise mirrors and matches where keys and organizations already draw the line. A non-positive budget now means no budget here too, so the read-time check and the reservation path agree on what counts as unbudgeted.
This commit is contained in:
parent
d7c0bc1e6d
commit
40a7fe9221
2 changed files with 44 additions and 7 deletions
|
|
@ -5620,6 +5620,11 @@ async def _model_access_group_max_budget_check(
|
|||
overshoot the ceiling slightly. The reservation counters are the precise path; this one covers
|
||||
the ``disable_budget_reservation`` case.
|
||||
|
||||
The ceiling is exclusive, unlike the tag check it otherwise mirrors: a pool whose recorded
|
||||
spend has reached ``max_budget`` has nothing left to give, so the next request is refused.
|
||||
Keys and organizations already draw the line there. A non-positive budget means no budget,
|
||||
matching what the reservation path treats as unbudgeted.
|
||||
|
||||
Raises:
|
||||
BudgetExceededError if a matched group is over its max budget.
|
||||
"""
|
||||
|
|
@ -5636,7 +5641,7 @@ async def _model_access_group_max_budget_check(
|
|||
|
||||
for group in matched_model_access_groups:
|
||||
budget = budgets.get(group)
|
||||
if budget is None or budget.max_budget is None:
|
||||
if budget is None or budget.max_budget is None or budget.max_budget <= 0:
|
||||
continue
|
||||
|
||||
group_spend = await get_current_spend(
|
||||
|
|
@ -5645,7 +5650,7 @@ async def _model_access_group_max_budget_check(
|
|||
max_budget=budget.max_budget,
|
||||
fallback_authoritative=True,
|
||||
)
|
||||
if group_spend <= budget.max_budget:
|
||||
if group_spend < budget.max_budget:
|
||||
continue
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=group_spend,
|
||||
|
|
|
|||
|
|
@ -364,19 +364,51 @@ async def test_group_under_its_max_budget_passes():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_exactly_at_its_max_budget_passes():
|
||||
"""The ceiling is inclusive, matching the tag check it mirrors; only spend strictly above it blocks.
|
||||
async def test_group_exactly_at_its_max_budget_blocks_the_request():
|
||||
"""A pool whose spend has reached the ceiling has nothing left, so the next request is refused.
|
||||
|
||||
Asserting the counter was read is what keeps this honest: a group that got skipped entirely,
|
||||
because its row never arrived or carried no budget, would also not raise.
|
||||
This is where the check departs from the tag one it otherwise mirrors, and it matches where
|
||||
keys and organizations already draw the line.
|
||||
"""
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _enforce(
|
||||
("tier-a",),
|
||||
_MagBudgetRow("tier-a", max_budget=10.0),
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0},
|
||||
)
|
||||
|
||||
assert exc_info.value.entity_id == "tier-a"
|
||||
assert exc_info.value.current_cost == 10.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_just_under_its_max_budget_passes():
|
||||
"""Asserting the counter was read is what keeps this honest: a group that got skipped entirely,
|
||||
because its row never arrived or carried no budget, would also not raise."""
|
||||
assert await _enforce(
|
||||
("tier-a",),
|
||||
_MagBudgetRow("tier-a", max_budget=10.0),
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0},
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 9.99},
|
||||
) == [MODEL_ACCESS_GROUP_COUNTER_KEY]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_non_positive_budget_means_no_budget():
|
||||
"""The reservation path treats max_budget <= 0 as unbudgeted, so the read-time check must agree.
|
||||
|
||||
Without this the exclusive ceiling would turn a zero into a total freeze on one path and a
|
||||
no-op on the other.
|
||||
"""
|
||||
assert (
|
||||
await _enforce(
|
||||
("tier-a",),
|
||||
_MagBudgetRow("tier-a", max_budget=0.0),
|
||||
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 5.0},
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_over_its_max_budget_blocks_the_request_and_names_the_group():
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue