diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index cf912ddaa25..14b2b4e3299 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -177,15 +177,15 @@ quota_management... behavior : ratelimit | budget | spend_tracking variant : rpm | tpm | priority_generous | priority_strict 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 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] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index 42a075681e0..d0afcaca848 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -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"} diff --git a/tests/e2e/quota_management/budgets/budget_client.py b/tests/e2e/quota_management/budgets/budget_client.py index 543d5f959e0..087dc8ca522 100644 --- a/tests/e2e/quota_management/budgets/budget_client.py +++ b/tests/e2e/quota_management/budgets/budget_client.py @@ -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( diff --git a/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py new file mode 100644 index 00000000000..9c927a31216 --- /dev/null +++ b/tests/e2e/quota_management/budgets/test_model_access_group_budget_e2e.py @@ -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" + )