mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): apply the quota policy to model add, follow router wildcard rules, keep old setting values
This commit is contained in:
mateo-berri
parent
232bbea175
commit
7f04cbe781
6 changed files with 326 additions and 151 deletions
|
|
@ -17,14 +17,8 @@ from typing import TYPE_CHECKING, Final, cast
|
|||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
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
|
||||
|
||||
|
|
@ -38,6 +32,14 @@ from litellm.proxy.management.teams.dependencies import get_team_access
|
|||
from litellm.proxy.management.users.service import org_admin_org_ids
|
||||
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
|
||||
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
|
||||
from litellm.proxy.management_helpers.model_quota_policy import (
|
||||
REQUIRE_BOTH_LIMITS,
|
||||
ModelQuotaPolicy,
|
||||
ProvidersOf,
|
||||
no_providers,
|
||||
parse_model_quota_setting,
|
||||
router_deployment_providers,
|
||||
)
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
management_endpoint_wrapper,
|
||||
)
|
||||
|
|
@ -307,9 +309,9 @@ def _raise_on_project_models_expanding_at_request_time(
|
|||
)
|
||||
|
||||
|
||||
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ProjectModelQuotaPolicy | None:
|
||||
def _project_model_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None:
|
||||
try:
|
||||
return parse_project_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
|
||||
return parse_model_quota_setting(general_settings.get(ENFORCE_PROJECT_MODEL_QUOTA_SETTING))
|
||||
except ValidationError as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
@ -320,7 +322,7 @@ def _project_model_quota_policy(general_settings: Mapping[str, object]) -> Proje
|
|||
def _raise_on_missing_project_model_quota(
|
||||
data: NewProjectRequest | UpdateProjectRequest,
|
||||
access_group_names: frozenset[str] = frozenset(),
|
||||
policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
providers_of: ProvidersOf = no_providers,
|
||||
) -> None:
|
||||
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model on project CREATE.
|
||||
|
|
@ -347,7 +349,7 @@ def _raise_on_missing_project_model_quota_on_update(
|
|||
data: UpdateProjectRequest,
|
||||
existing_project: object,
|
||||
access_group_names: frozenset[str] = frozenset(),
|
||||
policy: ProjectModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
policy: ModelQuotaPolicy = REQUIRE_BOTH_LIMITS,
|
||||
providers_of: ProvidersOf = no_providers,
|
||||
) -> None:
|
||||
"""Require the positive `rpm`/`tpm` quotas the policy asks of each model over the RESULTING state on project UPDATE.
|
||||
|
|
|
|||
|
|
@ -1,103 +1,46 @@
|
|||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
from typing import Final
|
||||
|
||||
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]]
|
||||
from litellm.proxy.management_helpers.model_quota_policy import (
|
||||
ModelQuotaPolicy,
|
||||
ProvidersOf,
|
||||
QuotaLimit,
|
||||
required_limits,
|
||||
unmet_limits,
|
||||
)
|
||||
|
||||
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]] = 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[QuotaLimit]().union(*(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 _limits_set_for(
|
||||
model: str, rpm_set: Mapping[str, object], tpm_set: Mapping[str, object]
|
||||
) -> Mapping[QuotaLimit, object]:
|
||||
return MappingProxyType({"rpm": rpm_set.get(model), "tpm": tpm_set.get(model)})
|
||||
|
||||
|
||||
def missing_project_model_quota(
|
||||
models: Sequence[str],
|
||||
rpm_limits: Mapping[str, object] | None,
|
||||
tpm_limits: Mapping[str, object] | None,
|
||||
policy: ProjectModelQuotaPolicy,
|
||||
policy: ModelQuotaPolicy,
|
||||
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}
|
||||
)
|
||||
rpm_set: Final = rpm_limits or _NO_LIMITS_SET
|
||||
tpm_set: Final = tpm_limits or _NO_LIMITS_SET
|
||||
gaps: Final = tuple(
|
||||
MissingModelQuota(model, _unset_limits(model, required_limits(policy, model, providers_of(model)), set_limits))
|
||||
MissingModelQuota(
|
||||
model,
|
||||
unmet_limits(required_limits(policy, model, providers_of(model)), _limits_set_for(model, rpm_set, tpm_set)),
|
||||
)
|
||||
for model in models
|
||||
)
|
||||
return tuple(gap for gap in gaps if gap.limits)
|
||||
|
|
@ -105,22 +48,3 @@ def missing_project_model_quota(
|
|||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -102,6 +102,13 @@ from litellm.proxy.management_helpers.auto_router_permissions import (
|
|||
authorize_member_auto_router_write,
|
||||
)
|
||||
from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model
|
||||
from litellm.proxy.management_helpers.model_quota_policy import (
|
||||
ModelQuotaPolicy,
|
||||
deployment_provider,
|
||||
parse_model_quota_setting,
|
||||
required_limits,
|
||||
unmet_limits,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
@ -151,6 +158,7 @@ from litellm.types.router import (
|
|||
SPECIAL_MODEL_INFO_PARAMS,
|
||||
Deployment,
|
||||
GenericLiteLLMParams,
|
||||
LiteLLM_Params,
|
||||
ModelInfo,
|
||||
updateDeployment,
|
||||
updateLiteLLMParams,
|
||||
|
|
@ -760,26 +768,39 @@ async def _auto_router_capability_slot(
|
|||
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
|
||||
|
||||
|
||||
def _raise_if_rate_limits_required_but_missing(*, litellm_params: GenericLiteLLMParams, enforced: bool) -> None:
|
||||
"""Require both rpm and tpm (each a positive value) when the operator opts in via config.yaml.
|
||||
def _model_add_quota_policy(general_settings: Mapping[str, object]) -> ModelQuotaPolicy | None:
|
||||
try:
|
||||
return parse_model_quota_setting(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING))
|
||||
except ValidationError as e:
|
||||
raise ProxyException(
|
||||
message=f"general_settings.{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING} is invalid: {e}",
|
||||
type=ProxyErrorTypes.internal_server_error.value,
|
||||
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
param=None,
|
||||
) from e
|
||||
|
||||
Off by default, so deployments keep adding models without limits. When
|
||||
``enforce_rpm_tpm_on_model_add: true`` is set under general_settings, a model added
|
||||
without both rpm and tpm set to a positive value is rejected rather than stored
|
||||
unbounded (or effectively excluded from routing by a zero/negative limit).
|
||||
"""
|
||||
if not enforced:
|
||||
|
||||
def _raise_if_rate_limits_required_but_missing(
|
||||
*, model_name: str, litellm_params: LiteLLM_Params, policy: ModelQuotaPolicy | None
|
||||
) -> None:
|
||||
"""Off by default, so deployments keep adding models without limits. Once the operator opts in under
|
||||
general_settings, the policy decides per model name and provider which of rpm and tpm a new model needs,
|
||||
and a model missing one, or carrying a zero or negative one, is rejected rather than stored unbounded
|
||||
(or effectively excluded from routing)."""
|
||||
if policy is None:
|
||||
return
|
||||
missing: Final = tuple(
|
||||
field
|
||||
for field, value in (("rpm", litellm_params.rpm), ("tpm", litellm_params.tpm))
|
||||
if value is None or value <= 0
|
||||
provider: Final = deployment_provider(
|
||||
MappingProxyType({"model": litellm_params.model, "custom_llm_provider": litellm_params.custom_llm_provider})
|
||||
)
|
||||
missing: Final = unmet_limits(
|
||||
required_limits(policy, model_name, frozenset() if provider is None else frozenset((provider,))),
|
||||
MappingProxyType({"rpm": litellm_params.rpm, "tpm": litellm_params.tpm}),
|
||||
)
|
||||
if not missing:
|
||||
return
|
||||
raise ProxyException(
|
||||
message=(
|
||||
f"{' and '.join(missing)} must be set to a positive value when "
|
||||
f"{' and '.join(missing)} must be set to a positive value for '{model_name}' when "
|
||||
f"'{ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING}' is enabled in general_settings"
|
||||
),
|
||||
type=ProxyErrorTypes.validation_error.value,
|
||||
|
|
@ -2717,8 +2738,9 @@ async def add_new_model(
|
|||
)
|
||||
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name=model_params.model_name,
|
||||
litellm_params=model_params.litellm_params,
|
||||
enforced=bool(general_settings.get(ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING, False)),
|
||||
policy=_model_add_quota_policy(general_settings),
|
||||
)
|
||||
|
||||
clean_model_info: Final = ModelInfo.model_validate(
|
||||
|
|
|
|||
110
litellm/proxy/management_helpers/model_quota_policy.py
Normal file
110
litellm/proxy/management_helpers/model_quota_policy.py
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.router_utils.common_utils import provider_for_generic_call
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
|
||||
QuotaLimit: TypeAlias = Literal["rpm", "tpm"]
|
||||
ProvidersOf: TypeAlias = Callable[[str], frozenset[str]]
|
||||
|
||||
_LIMITS: Final[tuple[QuotaLimit, ...]] = ("rpm", "tpm")
|
||||
_ROUTER_PATTERNS: Final = PatternMatchRouter()
|
||||
|
||||
|
||||
class ModelQuotaPolicy(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
default: frozenset[QuotaLimit] = frozenset(_LIMITS)
|
||||
providers: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
models: Mapping[str, frozenset[QuotaLimit]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
REQUIRE_BOTH_LIMITS: Final = ModelQuotaPolicy()
|
||||
|
||||
|
||||
def parse_model_quota_setting(setting: object) -> ModelQuotaPolicy | None:
|
||||
"""Any falsy value is off and any other non-mapping value requires both limits, as the boolean
|
||||
setting always behaved, so only a mapping is read as a per-provider and per-model policy."""
|
||||
if not setting:
|
||||
return None
|
||||
if is_str_keyed_mapping(setting):
|
||||
return ModelQuotaPolicy.model_validate(setting)
|
||||
return REQUIRE_BOTH_LIMITS
|
||||
|
||||
|
||||
def _wildcard_matches(pattern: str, model: str) -> bool:
|
||||
return re.fullmatch(_ROUTER_PATTERNS.pattern_to_regex(pattern), model) is not None
|
||||
|
||||
|
||||
def _router_specificity(pattern: str) -> tuple[int, int]:
|
||||
return PatternUtils.calculate_pattern_specificity(_ROUTER_PATTERNS.pattern_to_regex(pattern))
|
||||
|
||||
|
||||
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=_router_specificity, default=None)
|
||||
|
||||
|
||||
def required_limits(policy: ModelQuotaPolicy, 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[QuotaLimit]().union(*(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 unmet_limits(required: frozenset[QuotaLimit], values: Mapping[QuotaLimit, object]) -> tuple[QuotaLimit, ...]:
|
||||
return tuple(
|
||||
limit
|
||||
for limit in _LIMITS
|
||||
if (limit in required or values.get(limit) is not None) and not _is_positive(values.get(limit))
|
||||
)
|
||||
|
||||
|
||||
def no_providers(_model: str) -> frozenset[str]:
|
||||
return frozenset()
|
||||
|
||||
|
||||
def deployment_provider(litellm_params: object) -> str | None:
|
||||
return provider_for_generic_call(litellm_params) if is_str_keyed_mapping(litellm_params) else None
|
||||
|
||||
|
||||
def _routed_deployments(llm_router: "Router", model: str, team_id: str | None) -> Sequence[Mapping[str, object]]:
|
||||
named: Final[tuple[Mapping[str, object], ...]] = (
|
||||
*llm_router.get_all_deployments(model_name=model, team_id=team_id),
|
||||
*llm_router.get_model_list_from_model_alias(model_name=model),
|
||||
*llm_router.get_model_list_from_routing_groups(model_name=model),
|
||||
)
|
||||
return named or llm_router.wildcard_route_deployments(model, team_id)
|
||||
|
||||
|
||||
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]:
|
||||
return frozenset(
|
||||
provider
|
||||
for deployment in _routed_deployments(llm_router, model, team_id)
|
||||
if (provider := deployment_provider(deployment.get("litellm_params"))) is not None
|
||||
)
|
||||
|
||||
return providers_of
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import os
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
|
@ -1627,22 +1627,20 @@ def _quota_error(
|
|||
rpm: dict[str, int] | None = None,
|
||||
tpm: dict[str, int] | None = None,
|
||||
setting: object = _CUSTOMER_QUOTA_POLICY,
|
||||
providers_of=_customer_providers_of,
|
||||
providers_of: Callable[[str], frozenset[str]] = _customer_providers_of,
|
||||
) -> str | None:
|
||||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_raise_on_missing_project_model_quota,
|
||||
)
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
|
||||
|
||||
policy = parse_project_model_quota_setting(setting)
|
||||
policy: Final = parse_model_quota_setting(setting)
|
||||
assert policy is not None
|
||||
data = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm)
|
||||
data: Final = NewProjectRequest(team_id="test-team", models=models, model_rpm_limit=rpm, model_tpm_limit=tpm)
|
||||
try:
|
||||
_raise_on_missing_project_model_quota(data, frozenset(), policy, providers_of)
|
||||
except HTTPException as e:
|
||||
rejection = e
|
||||
rejection: Final = e
|
||||
else:
|
||||
return None
|
||||
assert rejection.status_code == 400
|
||||
|
|
@ -1664,6 +1662,27 @@ def test_quota_policy_model_pattern_overrides_its_provider_per_limit():
|
|||
assert "model_tpm_limit" not in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "rpm", "tpm", "expected_missing"),
|
||||
[
|
||||
("runway-gen4-turbo", {"runway-gen4-turbo": 0}, None, "runway-gen4-turbo (model_rpm_limit)"),
|
||||
("jev-latest", None, {"jev-latest": -5}, "jev-latest (model_tpm_limit)"),
|
||||
(
|
||||
"veo-3.1-lite-generate-001",
|
||||
{"veo-3.1-lite-generate-001": 5},
|
||||
{"veo-3.1-lite-generate-001": 0},
|
||||
"veo-3.1-lite-generate-001 (model_tpm_limit)",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_quota_policy_rejects_a_non_positive_limit_set_on_an_exempt_model(
|
||||
model: str, rpm: dict[str, int] | None, tpm: dict[str, int] | None, expected_missing: str
|
||||
):
|
||||
error: Final = _quota_error([model], rpm=rpm, tpm=tpm)
|
||||
assert error is not None
|
||||
assert expected_missing in error
|
||||
|
||||
|
||||
def test_quota_policy_error_names_each_model_with_the_limits_it_misses():
|
||||
error = _quota_error(["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo"])
|
||||
assert error is not None
|
||||
|
|
@ -1684,10 +1703,13 @@ def test_quota_policy_error_names_each_model_with_the_limits_it_misses():
|
|||
),
|
||||
({"models": {"runway-*": [], "runway-gen4-turbo": ["rpm"]}}, "runway-gen4-turbo", "(model_rpm_limit)"),
|
||||
({"models": {"veo-*": [], "veo-3.1-lite-*": ["rpm"]}}, "veo-3.1-lite-generate-001", "(model_rpm_limit)"),
|
||||
({"models": {"gpt-5-mini*": [], "gpt-*-*-*": ["rpm"]}}, "gpt-5-mini-fast", "(model_rpm_limit)"),
|
||||
],
|
||||
)
|
||||
def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(setting, model, expected_missing):
|
||||
error = _quota_error([model], setting=setting)
|
||||
def test_quota_policy_precedence_is_model_then_pattern_then_provider_then_default(
|
||||
setting: object, model: str, expected_missing: str
|
||||
):
|
||||
error: Final = _quota_error([model], setting=setting)
|
||||
assert error is not None
|
||||
assert expected_missing in error
|
||||
|
||||
|
|
@ -1708,20 +1730,18 @@ def test_quota_policy_model_on_several_providers_needs_every_required_limit():
|
|||
assert "video-pool (model_tpm_limit)" in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [True, "true", {}])
|
||||
def test_quota_setting_true_or_empty_mapping_still_requires_both_limits(setting):
|
||||
error = _quota_error(["runway-gen4-turbo"], setting=setting)
|
||||
@pytest.mark.parametrize("setting", [True, "true", "false", 1, ["rpm"]])
|
||||
def test_quota_setting_truthy_non_mapping_still_requires_both_limits(setting: object):
|
||||
error: Final = _quota_error(["runway-gen4-turbo"], setting=setting)
|
||||
assert error is not None
|
||||
assert "runway-gen4-turbo (model_rpm_limit and model_tpm_limit)" in error
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [False, None])
|
||||
def test_quota_setting_off_enforces_nothing(setting):
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
@pytest.mark.parametrize("setting", [False, None, {}, "", 0, []])
|
||||
def test_quota_setting_off_enforces_nothing(setting: object):
|
||||
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
|
||||
|
||||
assert parse_project_model_quota_setting(setting) is None
|
||||
assert parse_model_quota_setting(setting) is None
|
||||
|
||||
|
||||
def test_update_quota_policy_resolves_providers_for_the_resulting_models():
|
||||
|
|
@ -1730,11 +1750,9 @@ def test_update_quota_policy_resolves_providers_for_the_resulting_models():
|
|||
from litellm_enterprise.proxy.management_endpoints.project_endpoints import (
|
||||
_raise_on_missing_project_model_quota_on_update,
|
||||
)
|
||||
from litellm_enterprise.proxy.management_endpoints.project_model_quota import (
|
||||
parse_project_model_quota_setting,
|
||||
)
|
||||
from litellm.proxy.management_helpers.model_quota_policy import parse_model_quota_setting
|
||||
|
||||
policy = parse_project_model_quota_setting(_CUSTOMER_QUOTA_POLICY)
|
||||
policy = parse_model_quota_setting(_CUSTOMER_QUOTA_POLICY)
|
||||
assert policy is not None
|
||||
existing = types.SimpleNamespace(models=["runway-gen4-turbo"], metadata={})
|
||||
assert (
|
||||
|
|
@ -1777,6 +1795,25 @@ def _customer_router() -> "litellm.Router":
|
|||
)
|
||||
|
||||
|
||||
def test_quota_providers_come_from_the_one_wildcard_route_requests_use():
|
||||
from litellm.proxy.management_helpers.model_quota_policy import router_deployment_providers
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "gen4*", "litellm_params": {"model": "runwayml/gen4*", "api_key": "qa-runway-key"}},
|
||||
{
|
||||
"model_name": "media-team-gen",
|
||||
"litellm_params": {"model": "vertex_ai/*", "vertex_project": "qa-project"},
|
||||
"model_info": {"team_id": "media-team", "team_public_model_name": "gen*"},
|
||||
},
|
||||
]
|
||||
)
|
||||
providers_of: Final = router_deployment_providers(router, "media-team")
|
||||
|
||||
assert providers_of("gen4_turbo") == frozenset({"runwayml"})
|
||||
assert providers_of("gen3a_turbo") == frozenset({"vertex_ai"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_project_quota_policy_reads_providers_from_router_deployments(monkeypatch):
|
||||
team_models = ["gemini-3.8-flash", "veo-3.1-lite-generate-001", "runway-gen4-turbo", "jev-latest"]
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|||
patch_model,
|
||||
update_model,
|
||||
)
|
||||
from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIMITS, ModelQuotaPolicy
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams
|
||||
|
|
@ -7352,12 +7353,19 @@ class TestAutoRouterClassifierDefaultPrompt:
|
|||
assert response.system_prompt == classification_system_prompt(5)
|
||||
|
||||
|
||||
_CUSTOMER_MODEL_ADD_POLICY: Final = ModelQuotaPolicy.model_validate(
|
||||
{
|
||||
"providers": {"vertex_ai": ["rpm", "tpm"], "runwayml": []},
|
||||
"models": {"veo-*": ["rpm"], "jev-latest": []},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TestEnforceRpmTpmOnModelAdd:
|
||||
def test_passes_when_disabled_even_without_limits(self):
|
||||
assert (
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2"),
|
||||
enforced=False,
|
||||
model_name="gpt-5.2", litellm_params=LiteLLM_Params(model="azure/gpt-5.2"), policy=None
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
|
@ -7365,8 +7373,9 @@ class TestEnforceRpmTpmOnModelAdd:
|
|||
def test_passes_when_enabled_and_both_set(self):
|
||||
assert (
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name="gpt-5.2",
|
||||
litellm_params=LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=1000),
|
||||
enforced=True,
|
||||
policy=REQUIRE_BOTH_LIMITS,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
|
@ -7381,14 +7390,85 @@ class TestEnforceRpmTpmOnModelAdd:
|
|||
(LiteLLM_Params(model="azure/gpt-5.2", rpm=10, tpm=-1), "tpm"),
|
||||
],
|
||||
)
|
||||
def test_raises_when_enabled_and_missing(self, params, expected_missing):
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
def test_raises_when_enabled_and_missing(self, params: LiteLLM_Params, expected_missing: str):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
_raise_if_rate_limits_required_but_missing(litellm_params=params, enforced=True)
|
||||
assert expected_missing in str(exc_info.value.message)
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name="gpt-5.2", litellm_params=params, policy=REQUIRE_BOTH_LIMITS
|
||||
)
|
||||
assert exc_info.value.message.startswith(f"{expected_missing} must be set")
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "params"),
|
||||
[
|
||||
("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo")),
|
||||
("my-runway", LiteLLM_Params(model="gen4_turbo", custom_llm_provider="runwayml")),
|
||||
("jev-latest", LiteLLM_Params(model="openai/jev-latest")),
|
||||
("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", rpm=10)),
|
||||
("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10, tpm=1000)),
|
||||
],
|
||||
)
|
||||
def test_policy_asks_each_model_only_for_the_limits_its_entry_requires(
|
||||
self, model_name: str, params: LiteLLM_Params
|
||||
):
|
||||
assert (
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "params", "expected_missing"),
|
||||
[
|
||||
("veo-3.1-generate-001", LiteLLM_Params(model="vertex_ai/veo-3.1-generate-001", tpm=1000), "rpm"),
|
||||
("gemini-3-pro", LiteLLM_Params(model="vertex_ai/gemini-3-pro", rpm=10), "tpm"),
|
||||
("openai-gpt", LiteLLM_Params(model="openai/gpt-5.2"), "rpm and tpm"),
|
||||
("runway-gen4", LiteLLM_Params(model="runwayml/gen4_turbo", rpm=0), "rpm"),
|
||||
],
|
||||
)
|
||||
def test_policy_rejects_a_model_missing_a_required_or_carrying_a_non_positive_limit(
|
||||
self, model_name: str, params: LiteLLM_Params, expected_missing: str
|
||||
):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
_raise_if_rate_limits_required_but_missing(
|
||||
model_name=model_name, litellm_params=params, policy=_CUSTOMER_MODEL_ADD_POLICY
|
||||
)
|
||||
assert exc_info.value.message.startswith(f"{expected_missing} must be set to a positive value for '{model_name}'")
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("setting", "expected_code", "expected_message"),
|
||||
[
|
||||
({"default": [], "models": {"veo-*": ["rpm"]}}, "400", "rpm must be set to a positive value for 'veo-3.1'"),
|
||||
("yes", "400", "rpm and tpm must be set to a positive value for 'veo-3.1'"),
|
||||
(
|
||||
{"providers": {"vertex_ai": ["qps"]}},
|
||||
"500",
|
||||
"general_settings.enforce_rpm_tpm_on_model_add is invalid",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_model_add_resolves_the_quota_policy_for_the_public_model_name(
|
||||
monkeypatch: pytest.MonkeyPatch, setting: object, expected_code: str, expected_message: str
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
|
||||
|
||||
prisma: Final = MagicMock()
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"enforce_rpm_tpm_on_model_add": setting})
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await add_new_model(
|
||||
Deployment(model_name="veo-3.1", litellm_params={"model": "vertex_ai/veo-3.1"}),
|
||||
UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
assert error.value.code == expected_code
|
||||
assert error.value.message.startswith(expected_message)
|
||||
prisma.db.litellm_proxymodeltable.create.assert_not_called()
|
||||
|
||||
|
||||
class TestBlockModelResponseSerialization:
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue