scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

test(e2e): cover model access group budgets against a live proxy

Browse source 
Four cases in tests/e2e/quota_management/budgets, driving real OpenAI calls
through a group whose shared pool is drained to exhaustion: the spender key
stays blocked, a key that spent nothing of its own is blocked by the same
pool, a sibling group with no budget keeps serving, and the budget read
reports the spend drawn against the group.

Adds set/get/delete access group budget to BudgetClient and the four
matching rows to the coverage registry.
This commit is contained in:
ryan-crabbe-berri 2026-08-29 15:46:49 -07:00
parent 40a7fe9221
commit e263c09e4f
4 changed files with 234 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
tests/e2e/CLAUDE.md
View file

@ -177,15 +177,15 @@ quota_management.<behavior>.<variant>.<assertion>
behavior : ratelimit | budget | spend_tracking
variant : <ratelimit> rpm | tpm | priority_generous | priority_strict
<budget> key | internal_user | end_user | organization | team | team_member | tag
| model_max | soft | key_multi_window | team_multi_window
| fallback | spend_counter
| model_access_group | model_max | soft | key_multi_window
| team_multi_window | fallback | spend_counter
<spend_tracking> chat_completions | stream | messages_bridge | embeddings
| cache_hit | key_rollup | concurrent_burst | tags | end_user
| per_model | failure | spend_calculate | pagination
assertion : blocks_over_limit | resets_after_window | headers_report_remaining | picks_under_tpm
| blocks_then_resets | resets_windows_independently | alerts_without_blocking
| isolates_per_model | isolates_per_member | enforced_across_keys | routes_to_fallback
| reseed_matches_db | logs_cost | zero_cost
| isolates_per_model | isolates_per_member | isolates_per_group | enforced_across_keys
| routes_to_fallback | reseed_matches_db | reports_spend | logs_cost | zero_cost
| matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
| writes_failure_row | returns_cost | keeps_total
e.g. quota_management.ratelimit.rpm.blocks_over_limit exercised_on=[chat_completions, messages]

4
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -20,6 +20,10 @@
- {id: quota_management.budget.organization.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "An organization's max_budget blocks keys under its teams"}
- {id: quota_management.budget.team_member.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A member's per-team budget blocks independently of the team budget"}
- {id: quota_management.budget.team_member.isolates_per_member, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [isolates_per_member], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "One team member's exhausted per-team budget does not block a different member on the same team"}
- {id: quota_management.budget.model_access_group.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A model access group's shared max_budget blocks further calls to deployments in the group once the pool is spent"}
- {id: quota_management.budget.model_access_group.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "The pool is shared, so a key that spent nothing of its own is blocked once another key granted the same group drained it"}
- {id: quota_management.budget.model_access_group.isolates_per_group, module: quota_management, tier: P1, behavior: budget, variant: model_access_group, assertions: [isolates_per_group], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A request is charged only to the granted groups that serve the model it called, so an exhausted group never blocks a sibling group"}
- {id: quota_management.budget.model_access_group.reports_spend, module: quota_management, tier: P2, behavior: budget, variant: model_access_group, assertions: [reports_spend], exercised_on: [chat_completions], source: "proxy/management_endpoints/model_access_group_management_endpoints.py", rationale: "GET /access_group/{name}/budget reports the pool and the spend drawn against it, so an admin can see why calls are being refused"}
- {id: quota_management.budget.tag.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: tag, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "router_strategy/budget_limiter.py", rationale: "Proxy-level tag budgets block tagged requests at the cap"}
- {id: quota_management.budget.end_user_model_max.blocks_over_limit, module: quota_management, tier: P1, behavior: budget, variant: end_user_model_max, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "budget_management_endpoints.py", fail_before_fix: proven, rationale: "A per-model rpm_limit on an end-user budget is accepted and stored but never enforced; only key-attached budgets honour it"}
- {id: quota_management.budget.model_max.isolates_per_model, module: quota_management, tier: P1, behavior: budget, variant: model_max, assertions: [isolates_per_model], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "model_max_budget caps one model without touching a sibling's budget"}

65
tests/e2e/quota_management/budgets/budget_client.py
View file

@ -151,6 +151,28 @@ class TagDeleteBody(BaseModel):
name: str
class AccessGroupBudgetBody(BaseModel):
max_budget: float | None = None
soft_budget: float | None = None
budget_duration: str | None = None
class AccessGroupBudgetView(BaseModel):
budget_id: str
max_budget: float | None = None
soft_budget: float | None = None
budget_duration: str | None = None
class AccessGroupBudgetResponse(BaseModel):
"""GET/PUT /access_group/{name}/budget: the group's shared pool and the spend
every key that can reach the group has drawn against it."""
access_group: str
spend: float
budget: AccessGroupBudgetView | None = None
class BudgetNewBody(BaseModel):
max_budget: float | None = None
soft_budget: float | None = None
@ -514,6 +536,49 @@ class BudgetClient:
response_type=NoBody,
)
# ---- model access group ---------------------------------------------
def set_access_group_budget(
self,
access_group: str,
*,
max_budget: float | None = None,
soft_budget: float | None = None,
budget_duration: str | None = None,
) -> AccessGroupBudgetResponse:
"""Give a model access group one shared budget. Every key that can reach a
deployment in the group draws from it."""
return unwrap(
self.proxy.transport.put(
f"/access_group/{access_group}/budget",
headers=self.proxy.transport.master,
json=AccessGroupBudgetBody(
max_budget=max_budget,
soft_budget=soft_budget,
budget_duration=budget_duration,
),
response_type=AccessGroupBudgetResponse,
)
)
def access_group_budget(self, access_group: str) -> AccessGroupBudgetResponse:
return unwrap(
self.proxy.transport.get(
f"/access_group/{access_group}/budget",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=AccessGroupBudgetResponse,
)
)
def delete_access_group_budget(self, access_group: str) -> None:
_ = self.proxy.transport.delete(
f"/access_group/{access_group}/budget",
headers=self.proxy.transport.master,
json=NoBody(),
response_type=NoBody,
)
# ---- budget table ---------------------------------------------------
def create_budget(

161
tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py Normal file
View file

@ -0,0 +1,161 @@
"""Live e2e: one shared budget across every key that can reach a model access group.
A model access group is a free-text label on a deployment (`model_info.access_groups`),
and a key is granted the group by name. The budget hangs off the group, not the key, so
the interesting behaviors are the ones a per-key budget cannot produce: a key that has
spent nothing of its own is refused once somebody else drained the pool, and draining one
group leaves a second group untouched, because a request is only charged to the groups
the caller was granted that also serve the model being called.
"""
from __future__ import annotations
import os
import time
from collections.abc import Iterator
from dataclasses import dataclass
from typing import Final
import pytest
from budget_client import BudgetClient, is_budget_block
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call
from lifecycle import ResourceManager
from models import KeyGenerateBody, LiteLLMParamsBody, ModelInfoBody, ModelNewBody
pytestmark = pytest.mark.e2e
BACKEND: Final = "openai/gpt-5.4-nano"
TINY_BUDGET: Final = 5e-6
MAX_TOKENS: Final = 16
DRAIN_TIMEOUT_SECONDS: Final = 180
@dataclass(frozen=True, slots=True)
class DrainedPool:
"""A model access group whose shared budget has been spent to exhaustion, the
deployment inside it, the key that did the spending, and a second group holding
its own deployment that was never given a budget at all."""
access_group: str
model: str
spender_key: str
free_access_group: str
free_model: str
def _provider_key(env_var: str) -> str:
return os.environ.get(env_var) or f"os.environ/{env_var}"
def _grouped_model(model_name: str, access_group: str) -> ModelNewBody:
return ModelNewBody(
model_name=model_name,
litellm_params=LiteLLMParamsBody(model=BACKEND, api_key=_provider_key("OPENAI_API_KEY")),
model_info=ModelInfoBody(access_groups=[access_group]),
)
def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse:
return client.chat(key, model, f"hi {unique_marker()}", max_tokens=MAX_TOKENS)
def _drain(client: BudgetClient, key: str, model: str, access_group: str) -> None:
"""Spend the group's pool until the proxy refuses the next request. The first call
lands under the cap and the block comes from the spend it recorded, so this needs at
least one round trip through the spend writer, not just one request."""
deadline: Final = time.monotonic() + DRAIN_TIMEOUT_SECONDS
while time.monotonic() < deadline:
result = _call(client, key, model)
if is_budget_block(result):
return
require_successful_call(result)
time.sleep(1)
pytest.fail(f"budget on model access group {access_group!r} never blocked a request")
@pytest.fixture(scope="module")
def drained(client: BudgetClient) -> Iterator[DrainedPool]:
marker: Final = unique_marker()
pool: Final = DrainedPool(
access_group=f"e2e-mag-budget-{marker}",
model=f"e2e-mag-budgeted-{marker}",
spender_key=client.proxy.generate_key(KeyGenerateBody(models=[f"e2e-mag-budget-{marker}"])),
free_access_group=f"e2e-mag-free-{marker}",
free_model=f"e2e-mag-unbudgeted-{marker}",
)
created: Final = (
client.proxy.register_model(_grouped_model(pool.model, pool.access_group)),
client.proxy.register_model(_grouped_model(pool.free_model, pool.free_access_group)),
)
try:
client.set_access_group_budget(pool.access_group, max_budget=TINY_BUDGET)
_drain(client, pool.spender_key, pool.model, pool.access_group)
yield pool
finally:
client.delete_access_group_budget(pool.access_group)
client.proxy.delete_key(pool.spender_key)
for model_id in created:
client.proxy.delete_model(model_id)
class TestModelAccessGroupBudget:
@pytest.mark.covers("quota_management.budget.model_access_group.blocks_over_limit")
def test_the_key_that_drained_the_pool_stays_blocked(
self, client: BudgetClient, drained: DrainedPool
) -> None:
result = _call(client, drained.spender_key, drained.model)
assert is_budget_block(result), (
f"an exhausted pool served {drained.model!r} again: {result.status_code} {result.body[:300]}"
)
assert drained.access_group in result.body, (
f"the block did not name the group that caused it: {result.body[:300]}"
)
@pytest.mark.covers("quota_management.budget.model_access_group.enforced_across_keys")
def test_a_key_that_spent_nothing_is_blocked_by_the_shared_pool(
self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool
) -> None:
newcomer = resources.key(models=[drained.access_group])
result = _call(client, newcomer, drained.model)
assert is_budget_block(result), (
"a freshly minted key with no spend of its own was served by an exhausted "
f"shared pool: {result.status_code} {result.body[:300]}"
)
@pytest.mark.covers("quota_management.budget.model_access_group.isolates_per_group")
def test_a_drained_group_does_not_block_a_different_group(
self, client: BudgetClient, resources: ResourceManager, drained: DrainedPool
) -> None:
other = resources.key(models=[drained.free_access_group])
result = _call(client, other, drained.free_model)
assert not is_budget_block(result), (
f"{drained.free_access_group!r} has no budget of its own but was blocked by "
f"{drained.access_group!r}'s exhausted pool: {result.body[:300]}"
)
require_successful_call(result)
@pytest.mark.covers("quota_management.budget.model_access_group.reports_spend")
def test_the_budget_read_reports_the_spend_drawn_against_the_pool(
self, client: BudgetClient, drained: DrainedPool
) -> None:
"""Enforcement runs off a live counter while the group's row is written by the
batched spend writer, so the recorded spend an admin reads lands a beat after the
block. Poll for it: what matters is that it arrives and matches the pool."""
deadline = time.monotonic() + client.proxy.poll_timeout
reported = client.access_group_budget(drained.access_group)
while reported.spend < TINY_BUDGET and time.monotonic() < deadline:
time.sleep(client.proxy.poll_interval)
reported = client.access_group_budget(drained.access_group)
assert reported.budget is not None, "the group lost the budget that just blocked it"
assert reported.budget.max_budget == TINY_BUDGET
assert reported.spend >= TINY_BUDGET, (
f"the pool blocked at {TINY_BUDGET} but only {reported.spend} was ever recorded "
f"against the group within {client.proxy.poll_timeout}s"
)