mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
test(e2e): cover team model_max_budget shared cap and key override
Two new quota_management budget cells: a team's per-model cap blocks a sibling team key on its first call for that model while other models stay open, and a key carrying its own model_max_budget entry keeps serving after the team cap is exhausted. Adds model_max_budget to the e2e team create helper and registers the team_model_max variant and key_override_wins assertion in the registry grammar.
This commit is contained in:
ryan-crabbe-berri
parent
a261722efb
commit
5ca606aa4e
4 changed files with 97 additions and 2 deletions
|
|
@ -150,14 +150,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
|
||||
| model_max | team_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
|
||||
| isolates_per_model | isolates_per_member | enforced_across_keys | key_override_wins
|
||||
| routes_to_fallback
|
||||
| reseed_matches_db | logs_cost | zero_cost
|
||||
| matches_sum_of_logs | loses_no_spend | attributes_spend | writes_own_rows
|
||||
| writes_failure_row | returns_cost | keeps_total
|
||||
|
|
|
|||
|
|
@ -19,6 +19,8 @@
|
|||
- {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.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.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"}
|
||||
- {id: quota_management.budget.team_model_max.enforced_across_keys, module: quota_management, tier: P1, behavior: budget, variant: team_model_max, assertions: [enforced_across_keys], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "A team's model_max_budget entry is one shared cap: exhausting it on one key blocks every other team key for that model while other models stay open"}
|
||||
- {id: quota_management.budget.team_model_max.key_override_wins, module: quota_management, tier: P1, behavior: budget, variant: team_model_max, assertions: [key_override_wins], exercised_on: [chat_completions], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "A key's own model_max_budget entry takes precedence over the team's cap for that model, so the key keeps serving after the team cap is exhausted"}
|
||||
- {id: quota_management.budget.soft.alerts_without_blocking, module: quota_management, tier: P1, behavior: budget, variant: soft, assertions: [alerts_without_blocking], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "soft_budget alerts but never blocks traffic"}
|
||||
- {id: quota_management.budget.key.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: key, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes key spend after the window; a blocked key serves again"}
|
||||
- {id: quota_management.budget.team.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes a team's spend after the window; every key on the team serves again"}
|
||||
|
|
|
|||
|
|
@ -97,6 +97,7 @@ class TeamNewBody(BaseModel):
|
|||
budget_duration: str | None = None
|
||||
organization_id: str | None = None
|
||||
budget_limits: list[BudgetWindow] | None = None
|
||||
model_max_budget: dict[str, ModelBudgetEntry] | None = None
|
||||
|
||||
|
||||
class TeamNewResponse(BaseModel):
|
||||
|
|
@ -385,6 +386,7 @@ class BudgetClient:
|
|||
budget_duration: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
budget_limits: list[BudgetWindow] | None = None,
|
||||
model_max_budget: dict[str, ModelBudgetEntry] | None = None,
|
||||
) -> str:
|
||||
team_id = unwrap(
|
||||
self.proxy.transport.post(
|
||||
|
|
@ -396,6 +398,7 @@ class BudgetClient:
|
|||
budget_duration=budget_duration,
|
||||
organization_id=organization_id,
|
||||
budget_limits=budget_limits,
|
||||
model_max_budget=model_max_budget,
|
||||
),
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,89 @@
|
|||
"""Live e2e: a team's `model_max_budget` is one shared per-model cap for every
|
||||
team key, and a key carrying its own entry for that model is exempt from it."""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient, is_budget_block, model_budget
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import StreamingResponse, require_successful_call
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
CAPPED_MODEL = "claude-haiku-4-5"
|
||||
FREE_MODEL = "gemini-2.5-flash"
|
||||
BLOCK_DEADLINE_SECONDS = 60
|
||||
|
||||
|
||||
def _call(client: BudgetClient, key: str, model: str) -> StreamingResponse:
|
||||
result = client.chat(key, model, f"hi {unique_marker()}", max_tokens=16)
|
||||
if not result.ok and not is_budget_block(result):
|
||||
require_successful_call(result)
|
||||
return result
|
||||
|
||||
|
||||
def _drive_to_team_block(client: BudgetClient, key: str, team_id: str) -> None:
|
||||
deadline = time.monotonic() + BLOCK_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
result = _call(client, key, CAPPED_MODEL)
|
||||
if is_budget_block(result):
|
||||
assert team_id in result.body, result.body
|
||||
return
|
||||
time.sleep(1)
|
||||
pytest.fail(f"team model_max_budget on {CAPPED_MODEL} never enforced")
|
||||
|
||||
|
||||
# User flow (TLDR^2)
|
||||
# 1. Admin caps claude-haiku tiny on the team, leaves gemini roomy
|
||||
# 2. Alice's team key burns through the claude cap
|
||||
# 3. Bob's untouched team key is refused claude on its first try
|
||||
# 4. Bob's same key still gets gemini answers
|
||||
@pytest.mark.covers("quota_management.budget.team_model_max.enforced_across_keys")
|
||||
def test_team_model_cap_blocks_sibling_key(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-model-max-{unique_marker()}",
|
||||
model_max_budget={
|
||||
**model_budget(CAPPED_MODEL, 1e-6),
|
||||
**model_budget(FREE_MODEL, 1000.0),
|
||||
},
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
spender = client.generate_key(team_id=team_id)
|
||||
resources.defer(lambda: client.delete_key(spender))
|
||||
bystander = client.generate_key(team_id=team_id)
|
||||
resources.defer(lambda: client.delete_key(bystander))
|
||||
|
||||
_drive_to_team_block(client, spender, team_id)
|
||||
|
||||
first = _call(client, bystander, CAPPED_MODEL)
|
||||
assert is_budget_block(first), (
|
||||
f"sibling team key served {CAPPED_MODEL} after the team cap was exhausted: "
|
||||
f"HTTP {first.status_code} {first.body}"
|
||||
)
|
||||
assert team_id in first.body, first.body
|
||||
|
||||
require_successful_call(_call(client, bystander, FREE_MODEL))
|
||||
|
||||
|
||||
# User flow (TLDR^2)
|
||||
# 1. Admin caps claude-haiku tiny on the team
|
||||
# 2. Admin issues Carol a team key with its own big claude budget
|
||||
# 3. Dave's plain team key exhausts the team cap and is refused
|
||||
# 4. Carol keeps getting claude answers on the same team
|
||||
@pytest.mark.covers("quota_management.budget.team_model_max.key_override_wins")
|
||||
def test_key_override_exempts_from_team_cap(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-model-max-override-{unique_marker()}",
|
||||
model_max_budget=model_budget(CAPPED_MODEL, 1e-6),
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
inherit = client.generate_key(team_id=team_id)
|
||||
resources.defer(lambda: client.delete_key(inherit))
|
||||
override = client.generate_key(team_id=team_id, model_max_budget=model_budget(CAPPED_MODEL, 1000.0))
|
||||
resources.defer(lambda: client.delete_key(override))
|
||||
|
||||
_drive_to_team_block(client, inherit, team_id)
|
||||
|
||||
require_successful_call(_call(client, override, CAPPED_MODEL))
|
||||
Loading…
Add table
Reference in a new issue