mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge branch 'litellm_project_model_quota_per_provider' into litellm_project_model_quota_policy
This commit is contained in:
mateo-berri
commit
dbb898dc51
2 changed files with 185 additions and 45 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue