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:
ryan-crabbe-berri 2026-08-29 15:20:10 -07:00
parent d7c0bc1e6d
commit 40a7fe9221
2 changed files with 44 additions and 7 deletions

View file

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

View file

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