From d9f13e9075c640783d3b257fbfb663ba4e1e6b66 Mon Sep 17 00:00:00 2001 From: Fernando Jordan Silva Date: Tue, 20 Oct 2026 14:17:12 +0200 Subject: [PATCH] feat(proxy): support YAML model group budgets and deployment attribution --- cookbook/model_access_group_budgets/README.md | 71 +++++++ litellm/proxy/auth/auth_checks.py | 9 + litellm/proxy/db/db_spend_update_writer.py | 32 +-- .../proxy/hooks/proxy_track_cost_callback.py | 45 ++++ ...model_access_group_management_endpoints.py | 133 ++++++++---- litellm/proxy/management_helpers/utils.py | 5 + litellm/proxy/proxy_server.py | 5 +- .../types/proxy/model_access_group_budget.py | 5 + litellm/types/router.py | 1 + .../authorization/test_access_group_budget.py | 192 ++++++++++++++++++ .../auth/test_model_access_group_budgets.py | 2 +- .../proxy/db/test_model_access_group_spend.py | 44 ++++ .../test_access_group_management.py | 159 +++++++++++---- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 14 files changed, 613 insertions(+), 96 deletions(-) create mode 100644 cookbook/model_access_group_budgets/README.md diff --git a/cookbook/model_access_group_budgets/README.md b/cookbook/model_access_group_budgets/README.md new file mode 100644 index 00000000000..6bfecb4ff6b --- /dev/null +++ b/cookbook/model_access_group_budgets/README.md @@ -0,0 +1,71 @@ +# Track a shared monthly model budget + +This guide configures a USD 300 calendar-month soft budget on existing Model Access Groups. It requires PostgreSQL and a proxy with the deployment attribution option. It does not filter by provider or grant model access. + +## Configure the deployments + +Add the same group and `access_group_spend_attribution: deployment` to each deployment whose spend should count. Include every deployment of a model, including fallback deployments. Keep your existing provider parameters. + +```yaml +model_list: + - model_name: budgeted-chat + litellm_params: + model: openai/gpt-6.1-sol + api_key: os.environ/OPENAI_API_KEY + model_info: + access_groups: [shared-pool] + access_group_spend_attribution: deployment + +litellm_settings: + timezone: Europe/Madrid + budget_reset_time: "00:00" + +general_settings: + database_url: os.environ/DATABASE_URL + alerting: [webhook] + alert_types: [budget_alerts] +``` + +Set `WEBHOOK_URL` to your alert receiver and reload the configuration using your usual deployment process. The timezone and reset time are proxy-wide settings, so check their effect on other budgets first. + +Set this option on the models selected by a smart router. A router's classifier or embedding call counts only if its own deployment belongs to the group and enables deployment attribution. Setting the option only on the outer router is insufficient. + +Leave the option unset, or set it to `permission`, to retain the existing behavior. In that mode, only groups explicitly granted by an applicable caller allowlist are charged. Deployment attribution changes accounting, not authorization or hard-budget enforcement. Existing permission-based hard limits still apply. + +## Create the budget + +Use an administrative key against your own proxy. + +```bash +curl --fail-with-body -X PUT "$LITELLM_URL/access_group/shared-pool/budget" \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" \ + -H 'Content-Type: application/json' \ + -d '{"soft_budget":300,"budget_duration":"1mo"}' +``` + +For a new pool, the response has `max_budget: null` and a `budget_reset_at` on the first of the next month in the configured timezone. A soft budget sends alerts without rejecting requests. Other key, team, organization or deployment limits remain effective. If reusing a pool that already has a hard budget, remove its old budget first or explicitly clear that limit through the budget API. Do not assume setting a soft limit clears an existing hard limit. + +The group can come from YAML, PostgreSQL, or both. PostgreSQL remains required for budgets and recorded spend. This does not move YAML models into PostgreSQL or rewrite the YAML file. + +## Verify the result + +```bash +curl --fail-with-body "$LITELLM_URL/access_group/list" \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" +curl --fail-with-body "$LITELLM_URL/access_group/shared-pool/budget" \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" +``` + +Call a configured model using an unrestricted or wildcard key, then read the budget again after the spend queue flushes. Confirm that the increment matches the priced request. Multiple applicable groups each receive the full cost, but duplicate membership in one group is counted once. There is no extra charge for a permission grant to the same group. + +Crossing the soft limit triggers the existing budget alert delivery mechanism after the cost callback updates the counters. Configure `budget_alerts` and a supported destination. Delivery uses the existing alert TTL and is best effort. Concurrent workers can send duplicate alerts, and delivery is not an exactly-once notification service. + +## Changes and operational limits + +Budget rows and recorded spend survive proxy restarts. Removing the last YAML member does not erase the budget. The group remains listed with no deployments, and its budget can still be read or deleted. Manage YAML membership in YAML. The group membership mutation endpoints still manage PostgreSQL deployments and reject replacement or deletion of groups containing YAML members. Budget deletion remains available. + +Attribution resolves the actual deployment ID against the serving router at cost-callback time. The resolved groups are shared by the counter and queued database update. A later YAML reload cannot change that queued charge. A reload *before* the callback can change membership, and removing a deployment before its callback can lose deployment-only attribution. Drain in-flight requests before removing or changing those deployments. Keep configuration consistent across replicas and spend workers. This feature does not retroactively reclassify historical spend. + +Only costs reported by the existing LiteLLM cost pipeline can be counted. Streaming uses its final or recovered usage event. Image generation needs supported pricing. Asynchronous video, batches, background responses and provider pass-through require their existing completion-cost tracking and a resolvable deployment ID. Status polls without a charge must not be treated as new inference. This example does not establish coverage for every provider, endpoint or plugin. Inspect spend logs and validate the specific route before relying on its totals. + +The existing spend queue, atomic database increments, event deduplication and reset job remain in use. Queue flushes and resets are asynchronous. This is an operational budget, not an invoice ledger with event-time month boundaries, guaranteed delivery during crashes, or retroactive recovery of missing provider costs. Use the existing spend-worker durability controls and reconcile against provider billing where needed. diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a04fba2d905..24b6b45c2e6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,6 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from functools import partial from itertools import chain from types import MappingProxyType @@ -356,6 +357,12 @@ def _end_user_table(repo: _PrismaTableHolder[_PrismaEndUserRow]) -> _PrismaAuthT class _PrismaMaxBudgetRow(Protocol): + @property + def soft_budget(self) -> float | None: ... + + @property + def budget_reset_at(self) -> datetime | None: ... + @property def max_budget(self) -> float | None: ... @@ -2415,6 +2422,8 @@ def _model_access_group_budget(row: _PrismaModelAccessGroupBudgetRow) -> ModelAc access_group_name=row.access_group_name, spend=row.spend or 0.0, max_budget=None if budget_table is None else budget_table.max_budget, + soft_budget=None if budget_table is None else budget_table.soft_budget, + budget_reset_at=None if budget_table is None else budget_table.budget_reset_at, ) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 166a42adfc0..f170cd317f4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -458,15 +458,13 @@ def debitable_model_access_groups( served_model_id: str | None, router: _DeploymentLookup | None, ) -> tuple[str, ...]: - """Groups to debit: the set attributed at auth time, narrowed to those the served model belongs to. - - The router may fall back to a model outside the pool auth reserved against, so the - attributed set is the hard upper bound: a group absent from it is never debited. - """ + """Use the served deployment for opt-in accounting, otherwise narrow the auth grant.""" ordered: Final = coerce_model_access_groups(attributed) - if not ordered: - return () + deployment: Final = router.get_model_info(id=served_model_id) if router is not None and served_model_id else None + model_info: Final = deployment.get("model_info") if deployment is not None else None served: Final = _served_model_access_groups(router=router, served_model_id=served_model_id) + if isinstance(model_info, Mapping) and model_info.get("access_group_spend_attribution") == "deployment": + return tuple(sorted(served or ())) if served is None: return ordered return tuple(group for group in ordered if group in served) @@ -513,6 +511,7 @@ class DBSpendUpdateWriter: end_time: datetime, response_cost: float | None, project_id: str | None = None, + resolved_model_access_groups: Sequence[str] | None = None, ) -> bool: """Record the request's spend, answering whether its cost still needs charging. @@ -604,6 +603,7 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, request_model_access_groups=get_request_model_access_groups(kwargs), + resolved_model_access_groups=resolved_model_access_groups, ) ) @@ -1023,6 +1023,7 @@ class DBSpendUpdateWriter: payload: SpendLogsPayload, request_model_access_groups: Sequence[str] = (), project_id: str | None = None, + resolved_model_access_groups: Sequence[str] | None = None, ): """ Runs all 13 spend-update helpers sequentially inside a single asyncio task. @@ -1116,6 +1117,7 @@ class DBSpendUpdateWriter: served_model_id=payload_copy.get("model_id"), prisma_client=prisma_client, router=get_llm_router(), + resolved_model_access_groups=resolved_model_access_groups, ) _agent_id_for_spend: Final = payload_copy.get("agent_id") @@ -1479,6 +1481,7 @@ class DBSpendUpdateWriter: served_model_id: str | None, prisma_client: PrismaClient | None, router: _DeploymentLookup | None = None, + resolved_model_access_groups: Sequence[str] | None = None, ) -> None: """ Update spend for every model access group this request is billed against. @@ -1494,11 +1497,16 @@ class DBSpendUpdateWriter: if prisma_client is None: return - for model_access_group in debitable_model_access_groups( - attributed=request_model_access_groups, - served_model_id=served_model_id, - router=router, - ): + groups: Final = ( + tuple(dict.fromkeys(resolved_model_access_groups)) + if resolved_model_access_groups is not None + else debitable_model_access_groups( + attributed=request_model_access_groups, + served_model_id=served_model_id, + router=router, + ) + ) + for model_access_group in groups: await self.spend_update_queue.add_update( update=SpendUpdateQueueItem( entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3195394b9aa..50c12a1cb5b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -24,9 +24,11 @@ from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, + get_model_access_group_budgets_batch, get_team_object, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( DBSpendUpdateWriter, @@ -56,6 +58,8 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # should_store_prompts_and_responses_in_spend_logs, ) from litellm.proxy.utils import ProxyUpdateSpend +from litellm.types.integrations.slack_alerting import CallInfo, Litellm_EntityType +from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget from litellm.types.utils import ( CallTypes, LiteLLMBatch, @@ -465,6 +469,8 @@ class _ProxyDBLogger(CustomLogger): if not charged: return + await alert_model_access_group_soft_budgets(model_access_groups, proxy_logging_obj) + # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) asyncio.create_task( @@ -864,6 +870,7 @@ async def _update_database_and_spend_counters_in_batch( end_time=end_time, org_id=org_id, project_id=project_id, + resolved_model_access_groups=model_access_groups, ) except Exception: if budget_reservation is not None: @@ -963,3 +970,41 @@ async def _invalidate_budget_reservation_counters( await invalidate_budget_reservation_counters( budget_reservation=budget_reservation, ) + + +async def alert_model_access_group_soft_budgets(groups: Sequence[str], proxy_logging_obj: "ProxyLogging") -> None: + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if not groups or prisma_client is None or not proxy_logging_obj.alerting: + return + budgets: Final = await get_model_access_group_budgets_batch( + access_group_names=groups, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache + ) + await asyncio.gather( + *(_alert_model_access_group_soft_budget(budget, proxy_logging_obj) for budget in budgets.values()), + return_exceptions=True, + ) + + +async def _alert_model_access_group_soft_budget( + budget: ModelAccessGroupBudget, proxy_logging_obj: "ProxyLogging" +) -> None: + from litellm.proxy.proxy_server import get_current_spend + + if budget.soft_budget is None: + return + spend: Final = await get_current_spend( + counter_key=model_access_group_spend_counter_key(budget.access_group_name), fallback_spend=budget.spend + ) + if spend < budget.soft_budget: + return + await proxy_logging_obj.budget_alerts( + type="soft_budget", + user_info=CallInfo( + spend=spend, + soft_budget=budget.soft_budget, + token=f"model_access_group:{budget.access_group_name}:{budget.budget_reset_at}", + key_alias=budget.access_group_name, + event_group=Litellm_EntityType.MODEL_ACCESS_GROUP, + ), + ) diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index af469376289..f94ce77e2fa 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -19,9 +19,11 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel, ConfigDict from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.user_api_key_cache import ( @@ -273,9 +275,67 @@ async def _delete_model_access_group_budget_row( return row is not None +class _GroupDeployment(BaseModel): + model_config = ConfigDict(frozen=True) + model_name: str + model_info: Mapping[str, object] + + +def access_groups_from_deployments(deployments: Sequence[_GroupDeployment]) -> Mapping[str, AccessGroupInfo]: + memberships: Final = tuple(_deployment_groups(deployment) for deployment in deployments) + names: Final = frozenset[str]().union(*(groups for _, groups in memberships)) + return { + group: AccessGroupInfo( + access_group=group, + model_names=sorted({name for name, groups in memberships if group in groups}), + deployment_count=sum(group in groups for _, groups in memberships), + ) + for group in names + } + + +def _deployment_groups(deployment: _GroupDeployment) -> tuple[str, frozenset[str]]: + declared: Final = deployment.model_info.get("access_groups") + return deployment.model_name, frozenset(coerce_model_access_groups(declared)) + + +async def get_all_access_groups(prisma_client: PrismaClient) -> Mapping[str, AccessGroupInfo]: + from litellm.proxy.proxy_server import llm_router + + rows: Final = await _model_table(prisma_client).find_many() + database_ids: Final = frozenset(row.model_id for row in rows) + database: Final = tuple( + _GroupDeployment(model_name=row.model_name, model_info=dict(model_info_as_mapping(row.model_info) or {})) + for row in rows + ) + live: Final = llm_router.get_model_list() or () if llm_router is not None else () + configured: Final = tuple(_GroupDeployment.model_validate(deployment) for deployment in live) + config_only: Final = tuple( + deployment + for deployment in configured + if deployment.model_info.get("id") not in database_ids and deployment.model_info.get("db_model") is not True + ) + return access_groups_from_deployments((*database, *config_only)) + + +async def _reject_config_membership_mutation( + access_group: str, database_groups: Mapping[str, AccessGroupInfo], prisma_client: PrismaClient +) -> None: + all_groups: Final = await get_all_access_groups(prisma_client) + discovered: Final = all_groups.get(access_group) + database: Final = database_groups.get(access_group) + if discovered is not None and discovered.deployment_count > (database.deployment_count if database else 0): + raise _http_error( + 400, "This group contains config deployments. Edit their access_groups in config.yaml instead." + ) + + async def _raise_404_if_model_access_group_missing(access_group: str, prisma_client: PrismaClient) -> None: - access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) - if access_group not in access_groups_map: + access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client) + if ( + access_group not in access_groups_map + and await _model_access_group_budget_row(access_group, prisma_client) is None + ): raise _http_error(404, f"Access group '{access_group}' not found") @@ -529,41 +589,14 @@ def remove_access_group_from_deployment(model_info: dict[str, Any], access_group async def get_all_access_groups_from_db( prisma_client: PrismaClient, -) -> dict[str, AccessGroupInfo]: - """ - Get all access groups from the database. - - Returns: - Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info - """ - # Get all deployments - deployments: Final = await ModelRepository(prisma_client).table.find_many() - - # Build access group map - model_names_by_group: Final[dict[str, list[str]]] = {} - - for deployment in deployments: - model_info = deployment.model_info or {} - access_groups = model_info.get( # pyright: ignore[reportAttributeAccessIssue] # Json reads back as a dict - "access_groups", [] +) -> Mapping[str, AccessGroupInfo]: + deployments: Final = await _model_table(prisma_client).find_many() + return access_groups_from_deployments( + tuple( + _GroupDeployment(model_name=row.model_name, model_info=dict(model_info_as_mapping(row.model_info) or {})) + for row in deployments ) - model_name = deployment.model_name - - for access_group in access_groups: - if access_group not in model_names_by_group: - model_names_by_group[access_group] = [] - - model_names_by_group[access_group].append(model_name) - - # Convert to AccessGroupInfo objects - return { - access_group: AccessGroupInfo( - access_group=access_group, - model_names=sorted(frozenset(model_names)), - deployment_count=len(model_names), - ) - for access_group, model_names in model_names_by_group.items() - } + ) @router.post( @@ -653,7 +686,7 @@ async def create_model_group( try: # Check if access group already exists - existing_access_groups: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + existing_access_groups: Final = await get_all_access_groups(prisma_client=prisma_client) if data.access_group in existing_access_groups: raise HTTPException( @@ -744,11 +777,20 @@ async def list_access_groups( ) try: - access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client) budget_rows: Final = await _model_access_group_budget_rows(prisma_client) access_groups_list: Final = sorted( - (_with_budget(info, budget_rows.get(info.access_group)) for info in access_groups_map.values()), + ( + _with_budget(info, budget_rows.get(info.access_group)) + for info in { + **{ + name: AccessGroupInfo(access_group=name, model_names=[], deployment_count=0) + for name in budget_rows + }, + **access_groups_map, + }.values() + ), key=lambda group: group.access_group, ) @@ -798,17 +840,20 @@ async def get_access_group_info( ) try: - access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + access_groups_map: Final = await get_all_access_groups(prisma_client=prisma_client) - if access_group not in access_groups_map: + budget_row: Final = await _model_access_group_budget_row(access_group, prisma_client) + if access_group not in access_groups_map and budget_row is None: raise HTTPException( status_code=404, detail={"error": f"Access group '{access_group}' not found"}, ) return _with_budget( - access_groups_map[access_group], - await _model_access_group_budget_row(access_group, prisma_client), + access_groups_map.get( + access_group, AccessGroupInfo(access_group=access_group, model_names=[], deployment_count=0) + ), + budget_row, ) except HTTPException: @@ -884,6 +929,7 @@ async def update_access_group( # Validation: Check if access group exists try: access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + await _reject_config_membership_mutation(access_group, access_groups_map, prisma_client) if access_group not in access_groups_map: raise HTTPException( status_code=404, @@ -1018,6 +1064,7 @@ async def delete_access_group( # Validation: Check if access group exists try: access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client) + await _reject_config_membership_mutation(access_group, access_groups_map, prisma_client) if access_group not in access_groups_map: raise HTTPException( status_code=404, diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 30aaaacad04..2ab953e795d 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -246,6 +246,11 @@ async def handle_budget_for_entity( _budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create( data={ **new_budget_data, + **( + {"budget_reset_at": get_budget_reset_time(budget_row.budget_duration)} + if budget_row.budget_duration is not None + else {} + ), "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 660d7d529cb..c59db9db6df 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3069,8 +3069,9 @@ def cost_tracking(): from litellm.integrations.shadow_eval_logger import ShadowEvalLogger spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event) - litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer)) - litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer)) + spend_logger: Final = ProxyDBLogger(spend_event_producer) + litellm.logging_callback_manager.add_litellm_callback(spend_logger) + litellm.logging_callback_manager.add_litellm_async_success_callback(spend_logger) litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger()) diff --git a/litellm/types/proxy/model_access_group_budget.py b/litellm/types/proxy/model_access_group_budget.py index 08ea04f44f3..5ed9a147281 100644 --- a/litellm/types/proxy/model_access_group_budget.py +++ b/litellm/types/proxy/model_access_group_budget.py @@ -2,6 +2,8 @@ from __future__ import annotations +from datetime import datetime + from litellm.types.llms.base import LiteLLMBaseModel @@ -17,3 +19,6 @@ class ModelAccessGroupBudget(LiteLLMBaseModel): access_group_name: str spend: float = 0.0 max_budget: float | None = None + + soft_budget: float | None = None + budget_reset_at: datetime | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 68dbc701c9e..a1938972ca8 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -275,6 +275,7 @@ class ModelInfo(MirroredPricingParams): # the model_name that can be used by the team when making LLM calls team_public_model_name: str | None = None member_auto_router: bool = False + access_group_spend_attribution: Literal["permission", "deployment"] = "permission" # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked blocked: bool | None = None diff --git a/tests/integration/authorization/test_access_group_budget.py b/tests/integration/authorization/test_access_group_budget.py index 71302be2f37..962b3237f51 100644 --- a/tests/integration/authorization/test_access_group_budget.py +++ b/tests/integration/authorization/test_access_group_budget.py @@ -429,3 +429,195 @@ def test_reservation_refuses_a_concurrent_call_while_an_in_flight_call_holds_the wire, f"Budget has been exceeded! Model access group={group.name} Current cost: {_MAX_BUDGET}, Max budget: {_MAX_BUDGET}", ) + + +def test_yaml_deployment_accounting_includes_router_subcalls_and_survives_restart(tmp_path: Path) -> None: + group: Final = f"yaml-pool-{uuid.uuid4().hex}" + target: Final = f"target-{uuid.uuid4().hex}" + judge_name: Final = f"judge-{uuid.uuid4().hex}" + smart: Final = f"smart-{uuid.uuid4().hex}" + + def classifier(request: Request) -> Reply: + response: Final = _billed(request) + if request.method == "GET": + return response + body: Final = json.loads(response.body) + return Reply( + body=json.dumps( + { + **body, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": '{"tier":"COMPLEX"}'}, + "finish_reason": "stop", + } + ], + } + ).encode() + ) + + def receive_alert(request: Request) -> Reply: + assert request.method == "POST" + return Reply(body=b"{}") + + with ( + wire_server(_billed) as upstream, + wire_server(classifier) as judge, + wire_server(receive_alert) as alerts, + httpx.Client() as client, + ): + seed: Final = Gateway(client, "sk-local-group-budget-test", upstream.url) + deployments: Final = [ + { + "model_name": name, + "litellm_params": { + "model": f"openai/{_PROVIDER_MODEL}", + "api_base": f"{url}/v1", + "api_key": "synthetic", + "input_cost_per_token": _INPUT_RATE, + "output_cost_per_token": _OUTPUT_RATE, + }, + "model_info": {"id": name, "access_groups": [group], "access_group_spend_attribution": "deployment"}, + } + for name, url in ((target, upstream.url), (judge_name, judge.url)) + ] + config: Final = { + "model_list": [ + *deployments, + { + "model_name": smart, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "llm", + "classifier_llm_config": {"model": judge_name, "timeout_ms": 20000}, + "tiers": {"SIMPLE": target, "MEDIUM": target, "COMPLEX": target, "REASONING": target}, + }, + }, + }, + ], + "general_settings": { + "database_url": "os.environ/DATABASE_URL", + "master_key": seed.key, + "store_model_in_db": True, + "alerting": ["webhook"], + "alert_types": ["budget_alerts"], + }, + "litellm_settings": {"timezone": "UTC"}, + } + path: Final = tmp_path / "yaml-groups.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(seed, tmp_path, {"WEBHOOK_URL": alerts.url}, config=path) as candidate: + created: Final = candidate.request( + "PUT", f"/access_group/{group}/budget", {"soft_budget": _CALL_COST / 2, "budget_duration": "1mo"} + ) + assert created.status_code == 200, created.text + budget: Final = created.json()["budget"] + assert budget["max_budget"] is None + assert budget["budget_reset_at"][8:10] == "01" + with candidate.scenario() as scenario: + key: Final = scenario.key(models=["*"]) + for stream in (False, True): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": smart, + "stream": stream, + "messages": [{"role": "user", "content": f"route {uuid.uuid4().hex}"}], + }, + key=key, + ) + assert response.status_code == 200, response.text + assert "billed" in response.text + + def call(index: int) -> int: + result: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": target, + "messages": [{"role": "user", "content": str(index)}], + }, + key=key, + ) + assert result.status_code == 200, result.text + return result.status_code + + with ThreadPoolExecutor(max_workers=4) as executor: + assert tuple(executor.map(call, range(4))) == (200, 200, 200, 200) + assert len(_provider_calls(upstream)) == 6 + assert len(_provider_calls(judge)) == 2 + spent: Final = eventually( + lambda: candidate.request("GET", f"/access_group/{group}/budget").json()["spend"], + lambda value: value >= 8 * _CALL_COST - 1e-10, + seconds=70, + ) + assert spent == pytest.approx(8 * _CALL_COST) + notifications: Final = tuple(json.loads(request.body) for request in alerts.drain()) + assert any( + event.get("event") == "soft_budget_crossed" + and event.get("key_alias") == group + and event.get("event_group") == "model_access_group" + for event in notifications + ), notifications + with owned_proxy(seed, tmp_path, {}, config=path) as restarted: + persisted: Final = restarted.request("GET", f"/access_group/{group}/budget") + assert persisted.status_code == 200, persisted.text + assert persisted.json()["spend"] == pytest.approx(8 * _CALL_COST) + assert persisted.json()["budget"]["budget_id"] == budget["budget_id"] + assert restarted.request("DELETE", f"/access_group/{group}/budget").status_code == 200 + assert restarted.request("POST", "/budget/delete", {"id": budget["budget_id"]}).status_code == 200 + + +def test_yaml_image_group_counts_identical_generations_once_each(tmp_path: Path) -> None: + group: Final = f"image-pool-{uuid.uuid4().hex}" + name: Final = f"images-{uuid.uuid4().hex}" + price: Final = 0.25 + + def image(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=b'{"object":"list","data":[]}') + assert request.target == "/images/generations" + body: Final = json.loads(request.body) + assert "access_group_spend_attribution" not in body + return Reply(body=b'{"created":1700000000,"data":[{"b64_json":"aW1n"}]}') + + with wire_server(image) as upstream, httpx.Client() as client: + seed: Final = Gateway(client, "sk-local-image-budget-test", upstream.url) + config: Final = { + "model_list": [ + { + "model_name": name, + "litellm_params": { + "model": "openai/dall-e-3", + "api_base": upstream.url, + "api_key": "synthetic", + "output_cost_per_image": price, + }, + "model_info": {"access_groups": [group], "access_group_spend_attribution": "deployment"}, + } + ], + "general_settings": {"database_url": "os.environ/DATABASE_URL", "master_key": seed.key}, + } + path: Final = tmp_path / "image-group.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(seed, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario: + created: Final = candidate.request("PUT", f"/access_group/{group}/budget", {"soft_budget": 0.1}) + assert created.status_code == 200, created.text + budget_id: Final = created.json()["budget"]["budget_id"] + scenario.cleanups.callback(scenario.delete_budget, budget_id) + scenario.cleanups.callback(candidate.request, "DELETE", f"/access_group/{group}/budget") + key: Final = scenario.key(models=["*"]) + body: Final = {"model": name, "prompt": "a scripted sea otter", "size": "1024x1024", "n": 1} + for _ in range(2): + response: Final = candidate.request("POST", "/v1/images/generations", body, key=key) + assert response.status_code == 200, response.text + spent: Final = eventually( + lambda: candidate.request("GET", f"/access_group/{group}/budget").json()["spend"], + lambda value: value >= 2 * price, + seconds=70, + ) + assert spent == pytest.approx(2 * price) + assert len(tuple(request for request in upstream.drain() if request.method == "POST")) == 2 diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7506bd031d9..d95650c1d3d 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -311,7 +311,7 @@ class _MagBudgetRow: def __init__(self, access_group_name: str, spend: float = 0.0, max_budget: float | None = None) -> None: self.access_group_name = access_group_name self.spend = spend - self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget) + self.litellm_budget_table = None if max_budget is None else SimpleNamespace(max_budget=max_budget, soft_budget=None, budget_reset_at=None) class _RecordingPrismaClient: diff --git a/tests/unit/proxy/db/test_model_access_group_spend.py b/tests/unit/proxy/db/test_model_access_group_spend.py index 0ee86e783ea..e427b1f8113 100644 --- a/tests/unit/proxy/db/test_model_access_group_spend.py +++ b/tests/unit/proxy/db/test_model_access_group_spend.py @@ -511,3 +511,47 @@ async def test_batch_database_updates_enqueues_nothing_without_access_groups(): updates = await _drain(writer.spend_update_queue) assert [update for update in updates if update["entity_type"] is Litellm_EntityType.MODEL_ACCESS_GROUP] == [] + + +@pytest.mark.parametrize("attributed", [(), ("pool",), ("wrong-pool",), ("pool", "pool")]) +def test_deployment_accounting_charges_actual_groups_independently_of_grants(attributed): + from typing import Final + from litellm import Router + + router: Final = Router( + model_list=[ + { + "model_name": "selected-by-smart-router", + "litellm_params": {"model": "openai/group-budget-probe", "api_key": "synthetic"}, + "model_info": { + "id": "actual", + "access_groups": ["pool", "pool", "other"], + "access_group_spend_attribution": "deployment", + }, + } + ] + ) + assert debitable_model_access_groups(attributed, "actual", router) == ("other", "pool") + assert debitable_model_access_groups((), "unknown", router) == () + + +@pytest.mark.asyncio +async def test_resolved_group_charge_is_not_recomputed_after_yaml_reload(): + from typing import Final + + writer: Final = DBSpendUpdateWriter() + await writer._update_model_access_group_db( + response_cost=0.42, + request_model_access_groups=(), + served_model_id="deployment-1", + prisma_client=object(), + router=_FakeRouter({"deployment-1": ["replacement-group"]}), + resolved_model_access_groups=("original-group", "original-group"), + ) + assert await _drain(writer.spend_update_queue) == [ + SpendUpdateQueueItem( + entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP, + entity_id="original-group", + response_cost=0.42, + ) + ] diff --git a/tests/unit/proxy/management_endpoints/test_access_group_management.py b/tests/unit/proxy/management_endpoints/test_access_group_management.py index 59c2921e0d0..d59f951f093 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_management.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_management.py @@ -74,7 +74,6 @@ async def test_create_duplicate_access_group_fails(): patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), ): - # Should raise 409 Conflict with pytest.raises(HTTPException) as exc_info: await create_model_group(data=request_data, user_api_key_dict=mock_user) @@ -103,9 +102,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments mock_prisma = MagicMock() mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=deploy_a - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() mock_user = UserAPIKeyAuth( @@ -126,15 +123,11 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): - response = await create_model_group( - data=request_data, user_api_key_dict=mock_user - ) + response = await create_model_group(data=request_data, user_api_key_dict=mock_user) assert response.models_updated == 1 assert response.model_ids == ["deploy-A"] - mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with( - where={"model_id": "deploy-A"} - ) + mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "deploy-A"}) assert mock_prisma.db.litellm_proxymodeltable.update.call_count == 1 update_call = mock_prisma.db.litellm_proxymodeltable.update.call_args assert update_call.kwargs["where"] == {"model_id": "deploy-A"} @@ -169,9 +162,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments(): ) mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( - side_effect=[[], [deploy_a, deploy_b, deploy_c]] - ) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[], [deploy_a, deploy_b, deploy_c]]) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() mock_user = UserAPIKeyAuth( @@ -179,9 +170,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments(): user_role=LitellmUserRoles.PROXY_ADMIN, ) - request_data = NewModelGroupRequest( - access_group="production-models", model_names=["gpt-4o"] - ) + request_data = NewModelGroupRequest(access_group="production-models", model_names=["gpt-4o"]) with ( patch("litellm.proxy.proxy_server.llm_router", mock_router), @@ -191,9 +180,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments(): new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): - response = await create_model_group( - data=request_data, user_api_key_dict=mock_user - ) + response = await create_model_group(data=request_data, user_api_key_dict=mock_user) assert response.models_updated == 3 assert response.model_names == ["gpt-4o"] @@ -217,9 +204,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names(): mock_prisma = MagicMock() mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) - mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock( - return_value=deploy_a - ) + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock() mock_user = UserAPIKeyAuth( @@ -241,14 +226,10 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names(): new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), ), ): - response = await create_model_group( - data=request_data, user_api_key_dict=mock_user - ) + response = await create_model_group(data=request_data, user_api_key_dict=mock_user) assert response.models_updated == 1 - mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with( - where={"model_id": "deploy-A"} - ) + mock_prisma.db.litellm_proxymodeltable.find_unique.assert_called_once_with(where={"model_id": "deploy-A"}) @pytest.mark.asyncio @@ -367,7 +348,6 @@ async def test_create_access_group_surfaces_dropped_models(): assert "deploy-A" in str(exc_info.value.detail) - @pytest.mark.asyncio async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read(): """A concurrent reconcile sampled after the lock is released must not make this @@ -396,9 +376,7 @@ async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_r patch( "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache", new=AsyncMock( - return_value=ReconcileOutcome( - still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"}) - ) + return_value=ReconcileOutcome(still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"})) ), ), ): @@ -575,6 +553,7 @@ async def test_create_access_group_model_missing_everywhere_still_400s(): assert exc_info.value.status_code == 400 assert model_name in str(exc_info.value.detail) + @dataclass class _FakeBudgetRow: budget_id: str @@ -619,6 +598,7 @@ class _FakeBudgetTable: max_budget=data.get("max_budget"), soft_budget=data.get("soft_budget"), budget_duration=data.get("budget_duration"), + budget_reset_at=data.get("budget_reset_at"), ) self.rows[budget_id] = row self.create_calls.append(dict(data)) @@ -1128,9 +1108,7 @@ async def test_deleting_the_access_group_strips_deployments_before_dropping_the_ cache = _FakeAuthCache() with _proxy_with_stubbed_reload(prisma): - response = await delete_access_group( - access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache - ) + response = await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) assert response.models_updated == 1 assert prisma.access_group_budget_table.rows == {} @@ -1282,10 +1260,115 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts(): cache = _FakeAuthCache(journal) with _proxy_with_stubbed_reload(prisma): - response = await delete_access_group( - access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache - ) + response = await delete_access_group(access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache) assert response.models_updated == 1 assert prisma.access_group_budget_table.rows == {} _assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models") + + +@pytest.mark.asyncio +async def test_yaml_group_budget_survives_router_replacement_and_removal(): + from typing import Final + + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group_budget, + get_access_group_budget, + list_access_groups, + set_access_group_budget, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import AccessGroupBudgetRequest + + prisma: Final = _FakePrismaClient([]) + cache: Final = _FakeAuthCache([]) + models: Final = [ + { + "model_name": "yaml-model", + "litellm_params": {"model": "openai/group-budget-probe", "api_key": "synthetic"}, + "model_info": {"id": "yaml-id", "access_groups": ["yaml-pool", "yaml-pool"]}, + } + ] + with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=models)): + created: Final = await set_access_group_budget( + "yaml-pool", AccessGroupBudgetRequest(soft_budget=300, budget_duration="1mo"), _admin(), cache + ) + listed: Final = await list_access_groups(_admin()) + assert [(g.access_group, g.model_names, g.deployment_count) for g in listed.access_groups] == [ + ("yaml-pool", ["yaml-model"], 1) + ] + assert created.budget.soft_budget == 300 + assert created.budget.max_budget is None + with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=models)): + assert await get_access_group_budget("yaml-pool") == created + with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", Router(model_list=[])): + assert await get_access_group_budget("yaml-pool") == created + remaining: Final = await list_access_groups(_admin()) + assert [(g.access_group, g.model_names, g.deployment_count) for g in remaining.access_groups] == [ + ("yaml-pool", [], 0) + ] + assert (await delete_access_group_budget("yaml-pool", cache)).budget_deleted is True + + +@pytest.mark.asyncio +async def test_discovery_counts_database_deployment_once_and_uses_current_database_membership(): + from typing import Final + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import get_all_access_groups + + prisma: Final = _FakePrismaClient([], deployments=[_deployment("db-id", "same-name", ("pool",))]) + router: Final = Router( + model_list=[ + { + "model_name": "same-name", + "litellm_params": {"model": "openai/probe", "api_key": "synthetic"}, + "model_info": {"id": "db-id", "db_model": True, "access_groups": ["stale"]}, + }, + { + "model_name": "same-name", + "litellm_params": {"model": "openai/probe", "api_key": "synthetic"}, + "model_info": {"id": "yaml-id", "access_groups": ["pool"]}, + }, + ] + ) + with patch("litellm.proxy.proxy_server.llm_router", router): + groups: Final = await get_all_access_groups(prisma) + assert tuple(groups) == ("pool",) + assert groups["pool"].model_names == ["same-name"] + assert groups["pool"].deployment_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["update", "delete"]) +@pytest.mark.parametrize("with_database_member", [False, True]) +async def test_membership_mutations_preserve_yaml_groups_and_their_budget(operation, with_database_member): + from typing import Final + + from fastapi import HTTPException + from litellm.proxy.management_endpoints.model_access_group_management_endpoints import ( + delete_access_group, + update_access_group, + ) + from litellm.types.proxy.management_endpoints.model_management_endpoints import UpdateModelGroupRequest + + prisma: Final = _FakePrismaClient([], deployments=[_deployment()] if with_database_member else []) + _seed_budget(prisma, "prod-models", spend=12.0, soft_budget=300.0) + router: Final = Router( + model_list=[ + { + "model_name": "yaml-model", + "litellm_params": {"model": "openai/probe", "api_key": "synthetic"}, + "model_info": {"id": "yaml-id", "access_groups": ["prod-models"]}, + } + ] + ) + with _proxy(prisma), patch("litellm.proxy.proxy_server.llm_router", router): + mutation: Final = ( + delete_access_group("prod-models", _admin(), _FakeAuthCache()) + if operation == "delete" + else update_access_group("prod-models", UpdateModelGroupRequest(model_names=["yaml-model"]), _admin()) + ) + with pytest.raises(HTTPException) as error: + await mutation + assert error.value.status_code == 400 + assert "config.yaml" in str(error.value.detail) + assert prisma.access_group_budget_table.rows["prod-models"].spend == 12.0 + assert prisma.model_table.updates == [] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index daffc3c0c41..4e7346557d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50021,6 +50021,12 @@ export interface components { }; /** ModelInfo */ litellm__types__router__ModelInfo: { + /** + * Access Group Spend Attribution + * @default permission + * @enum {string} + */ + access_group_spend_attribution: "permission" | "deployment"; /** Access Windows */ access_windows?: components["schemas"]["ModelAccessWindow"][] | null; /** Allow Fail Open */