diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index db419fde1e6..8bca681ea12 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -15,7 +15,18 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, cast from fastapi import APIRouter, Depends, HTTPException, Request -from pydantic import TypeAdapter +from litellm_enterprise.proxy.management_endpoints.project_model_quota import ( + ENFORCE_PROJECT_MODEL_QUOTA_SETTING, + REQUIRE_BOTH_LIMITS, + ProjectModelQuotaPolicy, + ProvidersOf, + describe_missing_quota, + missing_project_model_quota, + no_providers, + parse_project_model_quota_setting, + router_deployment_providers, +) +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -265,26 +276,6 @@ def _check_team_project_limits( ) -def _project_models_missing_positive_quota( - models: list[str] | None, - rpm_limits: Mapping[str, object] | None, - tpm_limits: Mapping[str, object] | None, -) -> list[str]: - """Return the models that lack a positive `rpm` AND `tpm` quota. - - A valid quota is a positive integer; null, zero, and negative are rejected - because downstream rate limiters treat a non-positive limit as immediately - exhausted (every request blocked). - """ - - def _is_positive(value: object) -> bool: - return isinstance(value, int) and not isinstance(value, bool) and value > 0 - - rpm = rpm_limits or {} - tpm = tpm_limits or {} - return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))] - - def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]: return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset() @@ -311,64 +302,79 @@ def _raise_on_project_models_expanding_at_request_time( raise HTTPException( status_code=400, detail={ - "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead." + "error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while '{ENFORCE_PROJECT_MODEL_QUOTA_SETTING}' is enabled. List concrete model names instead." }, ) +def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ProjectModelQuotaPolicy | None: + try: + return parse_project_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING)) + except ValidationError as e: + raise HTTPException( + status_code=500, + detail={"error": f"general_settings.{ENFORCE_PROJECT_MODEL_QUOTA_SETTING} is invalid: {e}"}, + ) from e + + def _raise_on_missing_project_model_quota( - data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset() + data: NewProjectRequest | UpdateProjectRequest, + access_group_names: frozenset[str] = frozenset(), + policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + providers_of: ProvidersOf = no_providers, ) -> None: - """Require a positive `rpm`/`tpm` quota for every model on project CREATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE. `model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request model's `set_model_info` validator, so they are read from there. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. """ _raise_on_project_models_expanding_at_request_time(data.models, access_group_names) metadata = data.metadata or {} - missing = _project_models_missing_positive_quota( - data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + data.models or (), metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit"), policy, providers_of ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models added to project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) def _raise_on_missing_project_model_quota_on_update( - data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset() + data: UpdateProjectRequest, + existing_project: object, + access_group_names: frozenset[str] = frozenset(), + policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS, + providers_of: ProvidersOf = no_providers, ) -> None: - """Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE. + """Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE. `/project/update` replaces `models` and `metadata` when they are provided, so the check runs on what the project WILL look like: a partial update that doesn't touch models/quota keeps the existing values, while one that adds a model or clears a - model's quota must leave every resulting model with a positive limit. - - Only invoked when `general_settings.enforce_project_model_quota` is enabled - (default off), so it is opt-in and does not change behavior for existing users. + model's quota must leave every resulting model with the limits it needs. """ resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or []) resulting_metadata = ( data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {}) ) _raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names) - missing = _project_models_missing_positive_quota( - resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit") + missing: Final = missing_project_model_quota( + resulting_models, + resulting_metadata.get("model_rpm_limit"), + resulting_metadata.get("model_tpm_limit"), + policy, + providers_of, ) if not missing: return raise HTTPException( status_code=400, detail={ - "error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model." + "error": f"models would be left on the project without a positive rpm/tpm quota: {describe_missing_quota(missing)}. Set a positive value for each listed limit." }, ) @@ -567,9 +573,14 @@ async def new_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model added to the project. - if general_settings.get("enforce_project_model_quota", False): - _raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router)) + quota_policy: Final = _project_model_quota_policy(general_settings) + if quota_policy is not None: + _raise_on_missing_project_model_quota( + data, + _router_access_group_names(llm_router), + quota_policy, + router_deployment_providers(llm_router, data.team_id), + ) # Check if user has permission to create projects for this team # only team admins can create projects for their team @@ -817,10 +828,14 @@ async def update_project( data=data, ) - # Opt-in (default off): require rpm/tpm for every model the update would leave on the project. - if general_settings.get("enforce_project_model_quota", False): + quota_policy: Final = _project_model_quota_policy(general_settings) + if quota_policy is not None: _raise_on_missing_project_model_quota_on_update( - data, existing_project, _router_access_group_names(llm_router) + data, + existing_project, + _router_access_group_names(llm_router), + quota_policy, + router_deployment_providers(llm_router, target_team_id), ) # Prepare update data diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py new file mode 100644 index 00000000000..5cdfe158fd1 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_model_quota.py @@ -0,0 +1,125 @@ +import re +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter + +from litellm.router_utils.common_utils import provider_for_generic_call +from litellm.router_utils.pattern_match_deployments import PatternUtils + +if TYPE_CHECKING: + from litellm import Router + +QuotaLimit: TypeAlias = Literal["rpm", "tpm"] +ProvidersOf: TypeAlias = Callable[[str], frozenset[str]] + +ENFORCE_PROJECT_MODEL_QUOTA_SETTING: Final = "enforce_project_model_quota" +_BOTH_LIMITS: Final[frozenset[QuotaLimit]] = frozenset(("rpm", "tpm")) +_NO_LIMITS_SET: Final[Mapping[str, object]] = MappingProxyType({}) + + +class ProjectModelQuotaPolicy(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + default: frozenset[QuotaLimit] = _BOTH_LIMITS + providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict) + + +REQUIRE_BOTH_LIMITS: Final = ProjectModelQuotaPolicy() + +_QUOTA_SETTING: Final = TypeAdapter(bool | ProjectModelQuotaPolicy | None) + + +@dataclass(frozen=True, slots=True) +class MissingModelQuota: + model: str + limits: tuple[QuotaLimit, ...] + + +def parse_project_model_quota_setting(setting: object) -> ProjectModelQuotaPolicy | None: + parsed: Final = _QUOTA_SETTING.validate_python(setting) + match parsed: + case ProjectModelQuotaPolicy(): + return parsed + case True: + return REQUIRE_BOTH_LIMITS + case _: + return None + + +def _wildcard_matches(pattern: str, model: str) -> bool: + return re.fullmatch(re.escape(pattern).replace(r"\*", ".*"), model) is not None + + +def _most_specific_model_pattern(model: str, entries: Mapping[str, frozenset[QuotaLimit]]) -> str | None: + matching: Final = tuple(pattern for pattern in entries if "*" in pattern and _wildcard_matches(pattern, model)) + return max(matching, key=PatternUtils.calculate_pattern_specificity, default=None) + + +def required_limits(policy: ProjectModelQuotaPolicy, model: str, providers: frozenset[str]) -> frozenset[QuotaLimit]: + """A model entry (exact name, then the most specific `*` pattern) wins over the providers of the + model's deployments, which win over the default. A model served by several providers needs every + limit any of them requires.""" + exact: Final = policy.models.get(model) + if exact is not None: + return exact + pattern: Final = _most_specific_model_pattern(model, policy.models) + if pattern is not None: + return policy.models[pattern] + if not providers: + return policy.default + return frozenset(chain.from_iterable(policy.providers.get(provider, policy.default) for provider in providers)) + + +def _is_positive(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value > 0 + + +def _unset_limits( + model: str, required: frozenset[QuotaLimit], set_limits: Mapping[QuotaLimit, Mapping[str, object]] +) -> tuple[QuotaLimit, ...]: + return tuple(sorted(limit for limit in required if not _is_positive(set_limits[limit].get(model)))) + + +def missing_project_model_quota( + models: Sequence[str], + rpm_limits: Mapping[str, object] | None, + tpm_limits: Mapping[str, object] | None, + policy: ProjectModelQuotaPolicy, + providers_of: ProvidersOf, +) -> tuple[MissingModelQuota, ...]: + set_limits: Final[Mapping[QuotaLimit, Mapping[str, object]]] = MappingProxyType( + {"rpm": rpm_limits or _NO_LIMITS_SET, "tpm": tpm_limits or _NO_LIMITS_SET} + ) + gaps: Final = tuple( + MissingModelQuota(model, _unset_limits(model, required_limits(policy, model, providers_of(model)), set_limits)) + for model in models + ) + return tuple(gap for gap in gaps if gap.limits) + + +def describe_missing_quota(missing: Sequence[MissingModelQuota]) -> str: + return ", ".join(f"{gap.model} ({' and '.join(f'model_{limit}_limit' for limit in gap.limits)})" for gap in missing) + + +def no_providers(_model: str) -> frozenset[str]: + return frozenset() + + +def router_deployment_providers(llm_router: "Router | None", team_id: str | None) -> ProvidersOf: + if llm_router is None: + return no_providers + + def providers_of(model: str) -> frozenset[str]: + deployments: Final = llm_router.get_model_list(model_name=model, team_id=team_id) or () + return frozenset( + provider + for deployment in deployments + if (provider := provider_for_generic_call(deployment["litellm_params"])) is not None + ) + + return providers_of