mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
fix(proxy): share model rate-limit buckets between a model_group_alias and its target (#42516)
* fix(proxy): share model rate-limit buckets between a model_group_alias and its target A request sent under a model_group_alias counted in its own per-key, per-team, per-org, and per-project model bucket, so a key could double a deployment's default_api_key_rpm_limit / tpm_limit by alternating the alias and the model group name, and a metadata model_rpm_limit / model_tpm_limit keyed by the model group never applied to alias requests. The limiter now resolves the requested name to its model group before keying any model bucket, looks the limit up by the requested name first and the model group second, and charges post-call tokens to the same bucket. * fix(proxy): charge the model group resolved at admission when reconciling reserved tokens --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
fc0055497c
commit
c6c3881d7f
6 changed files with 389 additions and 116 deletions
|
|
@ -63,6 +63,7 @@ from litellm.router_utils.add_retry_fallback_headers import (
|
|||
ensure_response_additional_headers,
|
||||
response_has_hidden_params,
|
||||
)
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -91,6 +92,26 @@ else:
|
|||
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RateLimitedModel:
|
||||
requested: str
|
||||
group: str
|
||||
|
||||
def limit_in(self, limits: Mapping[str, int] | None) -> int | None:
|
||||
if limits is None:
|
||||
return None
|
||||
requested_limit: Final = limits.get(self.requested)
|
||||
return requested_limit if requested_limit is not None else limits.get(self.group)
|
||||
|
||||
|
||||
def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if llm_router is None:
|
||||
return None
|
||||
return resolve_model_group_alias(llm_router.model_group_alias, model)
|
||||
|
||||
|
||||
def _sibling_counter_keys(window_key: str) -> tuple[str, str]:
|
||||
prefix: Final = window_key.removesuffix(":window")
|
||||
return f"{prefix}:requests", f"{prefix}:tokens"
|
||||
|
|
@ -546,7 +567,7 @@ class RequestRateLimiterStash:
|
|||
parallel_slot: ParallelSlotAcquisition | None = None
|
||||
parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
|
||||
reserved_tokens: int = 0
|
||||
reserved_model: str | None = None
|
||||
reserved_model: RateLimitedModel | None = None
|
||||
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
|
||||
itpm_reserved_tokens: int = 0
|
||||
itpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
|
||||
|
|
@ -626,9 +647,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self,
|
||||
internal_usage_cache: InternalUsageCache,
|
||||
time_provider: Callable[[], datetime] | None = None,
|
||||
model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router,
|
||||
):
|
||||
self.internal_usage_cache = internal_usage_cache
|
||||
self._time_provider = time_provider or datetime.now
|
||||
self._model_group_resolver = model_group_resolver
|
||||
if self.internal_usage_cache.dual_cache.redis_cache is not None:
|
||||
self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
|
||||
BATCH_RATE_LIMITER_SCRIPT
|
||||
|
|
@ -2346,6 +2369,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if response["overall_code"] == "OVER_LIMIT":
|
||||
self._handle_rate_limit_error(response, descriptors, requested_model)
|
||||
|
||||
def _rate_limited_model(self, requested_model: str | None) -> RateLimitedModel | None:
|
||||
if not requested_model:
|
||||
return None
|
||||
return RateLimitedModel(
|
||||
requested=requested_model,
|
||||
group=self._model_group_resolver(requested_model) or requested_model,
|
||||
)
|
||||
|
||||
def create_organization_rate_limit_descriptor(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None = None
|
||||
) -> list[RateLimitDescriptor]:
|
||||
|
|
@ -2367,43 +2398,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
)
|
||||
|
||||
# Model specific org rate limits
|
||||
if (
|
||||
model: Final = self._rate_limited_model(requested_model)
|
||||
if model is None:
|
||||
return descriptors
|
||||
model_specific_tpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_tpm_limit")
|
||||
)
|
||||
model_specific_rpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_rpm_limit")
|
||||
is not None
|
||||
or get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_tpm_limit")
|
||||
is not None
|
||||
):
|
||||
_tpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_tpm_limit") or {}
|
||||
)
|
||||
if model_specific_tpm_limit is None and model_specific_rpm_limit is None:
|
||||
return descriptors
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_organization",
|
||||
value=f"{user_api_key_dict.org_id}:{model.group}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
_rpm_limit_for_team_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "organization_metadata", "model_rpm_limit") or {}
|
||||
)
|
||||
|
||||
should_check_rate_limit = False
|
||||
if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model:
|
||||
should_check_rate_limit = True
|
||||
|
||||
if should_check_rate_limit:
|
||||
model_specific_tpm_limit = None
|
||||
model_specific_rpm_limit = None
|
||||
if requested_model in _tpm_limit_for_team_model:
|
||||
model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model]
|
||||
if requested_model in _rpm_limit_for_team_model:
|
||||
model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model]
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_organization",
|
||||
value=f"{user_api_key_dict.org_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
)
|
||||
return descriptors
|
||||
|
||||
def _add_model_per_key_rate_limit_descriptor(
|
||||
|
|
@ -2425,34 +2441,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
get_key_model_tpm_limit,
|
||||
)
|
||||
|
||||
if not requested_model:
|
||||
model: Final = self._rate_limited_model(requested_model)
|
||||
if model is None:
|
||||
return
|
||||
|
||||
_tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict, model_name=requested_model)
|
||||
_rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict, model_name=requested_model)
|
||||
|
||||
if _tpm_limit_for_key_model is None and _rpm_limit_for_key_model is None:
|
||||
return
|
||||
|
||||
_tpm_limit_for_key_model = _tpm_limit_for_key_model or {}
|
||||
_rpm_limit_for_key_model = _rpm_limit_for_key_model or {}
|
||||
|
||||
# Check if model has any rate limits configured
|
||||
should_check_rate_limit: Final = (
|
||||
requested_model in _tpm_limit_for_key_model or requested_model in _rpm_limit_for_key_model
|
||||
model_specific_tpm_limit: Final = model.limit_in(
|
||||
get_key_model_tpm_limit(user_api_key_dict, model_name=model.group)
|
||||
)
|
||||
|
||||
if not should_check_rate_limit:
|
||||
model_specific_rpm_limit: Final = model.limit_in(
|
||||
get_key_model_rpm_limit(user_api_key_dict, model_name=model.group)
|
||||
)
|
||||
if model_specific_tpm_limit is None and model_specific_rpm_limit is None:
|
||||
return
|
||||
|
||||
# Get model-specific limits
|
||||
model_specific_tpm_limit: Final[int | None] = _tpm_limit_for_key_model.get(requested_model)
|
||||
model_specific_rpm_limit: Final[int | None] = _rpm_limit_for_key_model.get(requested_model)
|
||||
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_key",
|
||||
value=f"{user_api_key_dict.api_key}:{requested_model}",
|
||||
value=f"{user_api_key_dict.api_key}:{model.group}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
|
|
@ -2955,32 +2959,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
def _key_owns_model_limit(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str,
|
||||
model: RateLimitedModel,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> bool:
|
||||
key_own_limits: Final = get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)
|
||||
return key_own_limits is not None and key_own_limits.get(requested_model) is not None
|
||||
return model.limit_in(get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)) is not None
|
||||
|
||||
def _inherited_team_model_limit(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
requested_model: str,
|
||||
model: RateLimitedModel,
|
||||
rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"],
|
||||
) -> int | None:
|
||||
team_limits: Final = get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key)
|
||||
team_limit: Final = team_limits.get(requested_model) if team_limits else None
|
||||
if team_limit is None:
|
||||
return None
|
||||
if self._key_owns_model_limit(user_api_key_dict, requested_model, rate_limit_key):
|
||||
team_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "team_metadata", rate_limit_key)
|
||||
)
|
||||
if team_limit is None or self._key_owns_model_limit(user_api_key_dict, model, rate_limit_key):
|
||||
return None
|
||||
return team_limit
|
||||
|
||||
def _key_owns_model_tpm_limit_from_request_metadata(
|
||||
self,
|
||||
request_metadata: Mapping[str, object],
|
||||
model_group: str | None,
|
||||
model: RateLimitedModel | None,
|
||||
) -> bool:
|
||||
if model_group is None:
|
||||
if model is None:
|
||||
return False
|
||||
key_view: Final = UserAPIKeyAuth.model_validate(
|
||||
{
|
||||
|
|
@ -2988,7 +2990,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
"model_max_budget": request_metadata.get("user_api_key_model_max_budget") or {},
|
||||
}
|
||||
)
|
||||
return self._key_owns_model_limit(key_view, model_group, "model_tpm_limit")
|
||||
return self._key_owns_model_limit(key_view, model, "model_tpm_limit")
|
||||
|
||||
def _add_team_model_rate_limit_descriptor_from_metadata(
|
||||
self,
|
||||
|
|
@ -2996,16 +2998,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
requested_model: str | None,
|
||||
descriptors: list[RateLimitDescriptor],
|
||||
) -> None:
|
||||
if requested_model is None:
|
||||
model: Final = self._rate_limited_model(requested_model)
|
||||
if model is None:
|
||||
return
|
||||
team_rpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_rpm_limit")
|
||||
team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, requested_model, "model_tpm_limit")
|
||||
team_rpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, model, "model_rpm_limit")
|
||||
team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, model, "model_tpm_limit")
|
||||
if team_rpm_limit is None and team_tpm_limit is None:
|
||||
return
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{requested_model}",
|
||||
value=f"{user_api_key_dict.team_id}:{model.group}",
|
||||
rate_limit={
|
||||
"requests_per_unit": team_rpm_limit,
|
||||
"tokens_per_unit": team_tpm_limit,
|
||||
|
|
@ -3021,34 +3024,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
descriptors: list[RateLimitDescriptor],
|
||||
) -> None:
|
||||
"""Add project model rate limit descriptor from project_metadata if applicable."""
|
||||
if (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_rpm_limit") is not None
|
||||
or get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") is not None
|
||||
):
|
||||
_tpm_limit_for_project_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") or {}
|
||||
model: Final = self._rate_limited_model(requested_model)
|
||||
if model is None:
|
||||
return
|
||||
model_specific_tpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit")
|
||||
)
|
||||
model_specific_rpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_rpm_limit")
|
||||
)
|
||||
if model_specific_tpm_limit is None and model_specific_rpm_limit is None:
|
||||
return
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_project",
|
||||
value=f"{user_api_key_dict.project_id}:{model.group}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
_rpm_limit_for_project_model: Final = (
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_rpm_limit") or {}
|
||||
)
|
||||
should_check_rate_limit: Final = (
|
||||
requested_model in _tpm_limit_for_project_model or requested_model in _rpm_limit_for_project_model
|
||||
)
|
||||
|
||||
if should_check_rate_limit and requested_model is not None:
|
||||
model_specific_tpm_limit: Final = _tpm_limit_for_project_model.get(requested_model)
|
||||
model_specific_rpm_limit: Final = _rpm_limit_for_project_model.get(requested_model)
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="model_per_project",
|
||||
value=f"{user_api_key_dict.project_id}:{requested_model}",
|
||||
rate_limit={
|
||||
"requests_per_unit": model_specific_rpm_limit,
|
||||
"tokens_per_unit": model_specific_tpm_limit,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def add_project_io_token_rate_limit_descriptors_from_metadata(
|
||||
self,
|
||||
|
|
@ -3062,25 +3059,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
TPM descriptor above -- these give Bedrock Mantle-style separate input/output
|
||||
token quotas at the project level.
|
||||
"""
|
||||
if requested_model is None or user_api_key_dict.project_id is None:
|
||||
model: Final = self._rate_limited_model(requested_model)
|
||||
if model is None or user_api_key_dict.project_id is None:
|
||||
return
|
||||
|
||||
itpm_limit_for_project_model: Final = (
|
||||
model_itpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit")
|
||||
or {} # mutable-ok: metadata helper returns an optional mapping
|
||||
)
|
||||
otpm_limit_for_project_model: Final = (
|
||||
model_otpm_limit: Final = model.limit_in(
|
||||
get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit")
|
||||
or {} # mutable-ok: metadata helper returns an optional mapping
|
||||
)
|
||||
|
||||
model_itpm_limit: Final = itpm_limit_for_project_model.get(requested_model)
|
||||
model_otpm_limit: Final = otpm_limit_for_project_model.get(requested_model)
|
||||
|
||||
if model_itpm_limit is None and model_otpm_limit is None:
|
||||
return
|
||||
|
||||
descriptor_value: Final = f"{user_api_key_dict.project_id}:{requested_model}"
|
||||
descriptor_value: Final = f"{user_api_key_dict.project_id}:{model.group}"
|
||||
if model_itpm_limit is not None:
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
|
|
@ -3767,7 +3760,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# the (actual - reserved) delta to those — unreserved
|
||||
# scopes get charged the full actual usage instead.
|
||||
stash.reserved_tokens = estimated_tokens
|
||||
stash.reserved_model = requested_model
|
||||
stash.reserved_model = self._rate_limited_model(requested_model)
|
||||
stash.reserved_scopes = frozenset(
|
||||
(d["key"], d["value"])
|
||||
for d in descriptors
|
||||
|
|
@ -4514,9 +4507,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
reserved_scopes: Final[frozenset[tuple[str, str]]] = stash.reserved_scopes if stash is not None else frozenset()
|
||||
# Reconciliation must target the same model-scoped counter that the
|
||||
# pre-call reservation incremented. If a reservation was made,
|
||||
# ``reserved_model`` is authoritative; otherwise fall back to the
|
||||
# router's ``model_group`` (covers the no-reservation charge path).
|
||||
reconcile_model: Final = reserved_model or model_group
|
||||
# ``reserved_model`` (resolved at admission, so an alias map reload
|
||||
# mid-flight cannot move the charge) is authoritative; otherwise fall
|
||||
# back to the router's ``model_group`` (the no-reservation charge path).
|
||||
reconcile_model: Final = reserved_model if reserved_model is not None else self._rate_limited_model(model_group)
|
||||
|
||||
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
|
||||
|
||||
|
|
@ -4534,7 +4528,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
targets: Final = self._collect_tpm_scope_targets(
|
||||
standard_logging_metadata=standard_logging_metadata,
|
||||
kwargs=kwargs,
|
||||
model_group=reconcile_model,
|
||||
model_group=reconcile_model.group if reconcile_model is not None else None,
|
||||
)
|
||||
charged_targets: Final = (
|
||||
[target for target in targets if target[0] != "model_per_team"]
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@
|
|||
- {id: quota_management.ratelimit.redis_backed.blocks_over_limit, module: quota_management, tier: P0, behavior: ratelimit, variant: redis_backed, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "With Redis configured, RPM still enforces 429 across the shared limiter path customers run multi-replica"}
|
||||
- {id: quota_management.ratelimit.rpm.resets_after_window, module: quota_management, tier: P1, behavior: ratelimit, variant: rpm, assertions: [resets_after_window], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py", rationale: "Rate-limit window (LITELLM_RATE_LIMIT_WINDOW_SIZE, 60s default) expires; a blocked key serves again in the next window"}
|
||||
- {id: quota_management.ratelimit.rpm.headers_report_remaining, module: quota_management, tier: P1, behavior: ratelimit, variant: rpm, assertions: [headers_report_remaining], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py async_post_call_success_hook", rationale: "Successful responses carry x-ratelimit-api_key-{limit,remaining}-{requests,tokens} so clients can pace"}
|
||||
- {id: quota_management.ratelimit.model_group_alias.shares_bucket, module: quota_management, tier: P1, behavior: ratelimit, variant: model_group_alias, assertions: [shares_bucket], exercised_on: [chat_completions], source: "parallel_request_limiter_v3.py:_add_model_per_key_rate_limit_descriptor", rationale: "A model_group_alias draws on the same per-key deployment rpm bucket as the resolved model group, so alias plus real-name traffic cannot exceed the configured limit"}
|
||||
- {id: quota_management.ratelimit.priority_generous.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_generous, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:36-52", rationale: "Generous mode (<80% sat) allows priority borrowing"}
|
||||
- {id: quota_management.ratelimit.priority_strict.picks_under_tpm, module: quota_management, tier: P1, behavior: ratelimit, variant: priority_strict, assertions: [picks_under_tpm], exercised_on: [chat_completions, messages], source: "dynamic_rate_limiter_v3.py:53-71", rationale: "Strict mode (>=80% sat) enforces priority fairness"}
|
||||
- {id: quota_management.budget.key.blocks_over_limit, module: quota_management, tier: P0, behavior: budget, variant: key, assertions: [blocks_over_limit], exercised_on: [chat_completions], source: "proxy/auth/auth_checks.py", rationale: "A key's max_budget blocks further paid calls once spend crosses it"}
|
||||
|
|
|
|||
|
|
@ -43,6 +43,8 @@ router_settings:
|
|||
num_retries: 3
|
||||
allowed_fails: 5
|
||||
cooldown_time: 30
|
||||
model_group_alias:
|
||||
e2e-alias-rl-alias: e2e-alias-rl-target
|
||||
|
||||
model_list:
|
||||
- model_name: gpt-5.5
|
||||
|
|
@ -63,6 +65,11 @@ model_list:
|
|||
litellm_params:
|
||||
model: gemini/gemini-2.5-flash
|
||||
api_key: os.environ/GEMINI_API_KEY
|
||||
- model_name: e2e-alias-rl-target
|
||||
litellm_params:
|
||||
model: anthropic/claude-haiku-4-5
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
default_api_key_rpm_limit: 3
|
||||
- model_name: openai-text-embedding-3-small
|
||||
litellm_params:
|
||||
model: openai/text-embedding-3-small
|
||||
|
|
|
|||
|
|
@ -0,0 +1,85 @@
|
|||
"""Live e2e: a model group alias must share its per-key deployment rate-limit
|
||||
bucket with the model group it resolves to.
|
||||
|
||||
Covers quota_management.ratelimit.model_group_alias.shares_bucket: the proxy
|
||||
config declares `e2e-alias-rl-target` (a cheap Anthropic deployment) with
|
||||
`default_api_key_rpm_limit: 3` and `router_settings.model_group_alias` mapping
|
||||
`e2e-alias-rl-alias` -> `e2e-alias-rl-target`. Both spellings must draw on
|
||||
one per-key rpm bucket, so a key that exhausts the limit on one spelling is
|
||||
blocked on the other spelling inside the same window; each test in this file
|
||||
exhausts the budget on one name and asserts the other name 429s.
|
||||
|
||||
All calls of one test must land inside a single window
|
||||
(LITELLM_RATE_LIMIT_WINDOW_SIZE, 60s default), which real chat latency
|
||||
comfortably allows.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import StreamingResponse, require_successful_call
|
||||
from quota_client import QuotaClient
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MODEL_GROUP = "e2e-alias-rl-target"
|
||||
MODEL_ALIAS = "e2e-alias-rl-alias"
|
||||
RPM_LIMIT = 3
|
||||
WINDOW_SECONDS = 60
|
||||
LAST_CALL_LATENCY_MARGIN_SECONDS = 10
|
||||
|
||||
|
||||
def _chat(client: QuotaClient, key: str, model: str) -> StreamingResponse:
|
||||
return client.chat(key, model, f"reply with one word {unique_marker()}")
|
||||
|
||||
|
||||
def _exhaust_rpm(client: QuotaClient, key: str, model: str) -> float:
|
||||
"""Send RPM_LIMIT successful calls on `model`, opening the rate-limit
|
||||
window; returns the send timestamp of the first call as a lower bound on
|
||||
the window start. A fresh key may briefly 401 until the data plane's auth
|
||||
cache picks it up, so retry on 401 to a deadline; a 401 never reaches the
|
||||
rate limiter."""
|
||||
deadline = time.monotonic() + client.proxy.poll_timeout
|
||||
first_sent_at: float | None = None
|
||||
sent = 0
|
||||
while sent < RPM_LIMIT:
|
||||
if first_sent_at is None:
|
||||
first_sent_at = time.monotonic()
|
||||
outcome = _chat(client, key, model)
|
||||
if outcome.status_code == 401 and time.monotonic() < deadline:
|
||||
time.sleep(client.proxy.poll_interval)
|
||||
continue
|
||||
require_successful_call(outcome)
|
||||
sent += 1
|
||||
assert first_sent_at is not None
|
||||
return first_sent_at
|
||||
|
||||
|
||||
def _assert_blocked_inside_window(
|
||||
client: QuotaClient, key: str, model: str, window_opened_at: float
|
||||
) -> StreamingResponse:
|
||||
assert time.monotonic() < window_opened_at + WINDOW_SECONDS - LAST_CALL_LATENCY_MARGIN_SECONDS, (
|
||||
f"the {RPM_LIMIT} exhaust calls took too long; the follow-up call could land in the "
|
||||
"next window and mask a shared-bucket regression"
|
||||
)
|
||||
outcome = _chat(client, key, model)
|
||||
assert outcome.status_code == 429, (
|
||||
f"{model} must share its rpm bucket with the model group/alias that was already "
|
||||
f"exhausted, expected a 429 but got {outcome.status_code}: {outcome.body[:300]}"
|
||||
)
|
||||
return outcome
|
||||
|
||||
|
||||
class TestModelGroupAliasRateLimit:
|
||||
@pytest.mark.covers("quota_management.ratelimit.model_group_alias.shares_bucket")
|
||||
def test_alias_shares_rpm_bucket_with_model_group(self, client: QuotaClient, scoped_key: str) -> None:
|
||||
opened_at = _exhaust_rpm(client, scoped_key, MODEL_GROUP)
|
||||
_assert_blocked_inside_window(client, scoped_key, MODEL_ALIAS, opened_at)
|
||||
|
||||
@pytest.mark.covers("quota_management.ratelimit.model_group_alias.shares_bucket")
|
||||
def test_model_group_shares_rpm_bucket_with_alias(self, client: QuotaClient, scoped_key: str) -> None:
|
||||
opened_at = _exhaust_rpm(client, scoped_key, MODEL_ALIAS)
|
||||
_assert_blocked_inside_window(client, scoped_key, MODEL_GROUP, opened_at)
|
||||
|
|
@ -26,6 +26,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
|||
PARALLEL_REQUEST_SLOT_TTL_SECONDS,
|
||||
ParallelSlotAcquisition,
|
||||
RateLimitDescriptor,
|
||||
RateLimitedModel,
|
||||
RateLimitResponse,
|
||||
RequestRateLimiterStash,
|
||||
_request_stash,
|
||||
|
|
@ -3367,7 +3368,7 @@ async def test_pre_call_hook_keeps_internal_stash_out_of_request_body():
|
|||
stash = get_request_stash()
|
||||
assert stash is not None
|
||||
assert stash.reserved_tokens > 0
|
||||
assert stash.reserved_model == "gpt-4o-mini"
|
||||
assert stash.reserved_model == RateLimitedModel(requested="gpt-4o-mini", group="gpt-4o-mini")
|
||||
assert stash.reserved_scopes == frozenset({("api_key", _api_key)})
|
||||
|
||||
|
||||
|
|
@ -6965,3 +6966,187 @@ def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_t
|
|||
"Rate limit exceeded for api_key: sk-test. Limit type: requests. "
|
||||
f"Current limit: 2, Remaining: 0. Limit resets at: {expected_reset}"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_alias_to_target(model: str) -> str | None:
|
||||
return "target" if model == "alias" else None
|
||||
|
||||
|
||||
async def _rpm_request(handler: _PROXY_MaxParallelRequestsHandler, cache: DualCache, auth: UserAPIKeyAuth, model: str) -> None:
|
||||
await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data={"model": model}, call_type="acompletion")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("first_name, second_name", [("target", "alias"), ("alias", "target")])
|
||||
async def test_model_group_alias_shares_deployment_default_rpm_bucket_with_its_target(
|
||||
monkeypatch: pytest.MonkeyPatch, first_name: str, second_name: str
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "target",
|
||||
"litellm_params": {"model": "openai/gpt-test", "api_key": "test-key", "default_api_key_rpm_limit": 2},
|
||||
"model_info": {"id": "target-deployment"},
|
||||
}
|
||||
],
|
||||
model_group_alias={"alias": "target"},
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
|
||||
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-default"))
|
||||
|
||||
await _rpm_request(handler, cache, auth, first_name)
|
||||
await _rpm_request(handler, cache, auth, first_name)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _rpm_request(handler, cache, auth, second_name)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
assert "model_per_key" in str(exc.value.detail)
|
||||
assert f"{auth.api_key}:target" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("first_name, second_name", [("target", "alias"), ("alias", "target")])
|
||||
@pytest.mark.parametrize(
|
||||
"limits, counter_scope",
|
||||
[
|
||||
({"metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_key"),
|
||||
(
|
||||
{
|
||||
"team_id": "t",
|
||||
"metadata": {"model_rpm_limit": {"other-model": 100}},
|
||||
"team_metadata": {"model_rpm_limit": {"target": 1}},
|
||||
},
|
||||
"model_per_team",
|
||||
),
|
||||
({"org_id": "o", "organization_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_organization"),
|
||||
({"project_id": "p", "project_metadata": {"model_rpm_limit": {"target": 1}}}, "model_per_project"),
|
||||
],
|
||||
ids=["key_metadata", "team_metadata", "organization_metadata", "project_metadata"],
|
||||
)
|
||||
async def test_model_group_alias_shares_metadata_model_rpm_bucket_with_its_target(
|
||||
limits: dict[str, object], counter_scope: str, first_name: str, second_name: str
|
||||
) -> None:
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(cache), model_group_resolver=_resolve_alias_to_target
|
||||
)
|
||||
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-metadata"), **limits)
|
||||
|
||||
await _rpm_request(handler, cache, auth, first_name)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _rpm_request(handler, cache, auth, second_name)
|
||||
|
||||
assert exc.value.status_code == 429
|
||||
assert counter_scope in str(exc.value.detail)
|
||||
assert ":target" in str(exc.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_rpm_limit_keyed_by_the_alias_name_still_limits_alias_requests_only() -> None:
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(cache), model_group_resolver=_resolve_alias_to_target
|
||||
)
|
||||
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-keyed"), metadata={"model_rpm_limit": {"alias": 1}})
|
||||
|
||||
await _rpm_request(handler, cache, auth, "alias")
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await _rpm_request(handler, cache, auth, "alias")
|
||||
assert exc.value.status_code == 429
|
||||
assert "model_per_key" in str(exc.value.detail)
|
||||
|
||||
await _rpm_request(handler, cache, auth, "target")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"key_metadata, charges_team_model_pool",
|
||||
[({}, True), ({"model_tpm_limit": {"target": 500}}, False)],
|
||||
ids=["no_key_override", "key_owns_target_tpm_limit"],
|
||||
)
|
||||
def test_success_tpm_accounting_charges_the_alias_target_bucket(
|
||||
key_metadata: dict[str, object], charges_team_model_pool: bool
|
||||
) -> None:
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache()), model_group_resolver=_resolve_alias_to_target
|
||||
)
|
||||
response: Final = ModelResponse(
|
||||
id="alias-tpm",
|
||||
object="chat.completion",
|
||||
created=int(datetime.now().timestamp()),
|
||||
model="alias",
|
||||
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
|
||||
choices=[],
|
||||
)
|
||||
kwargs: Final = {
|
||||
"standard_logging_object": {
|
||||
"metadata": {"user_api_key_hash": hash_token("sk-alias-tpm"), "user_api_key_team_id": "t"}
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"model_group": "alias",
|
||||
"user_api_key_metadata": key_metadata,
|
||||
"user_api_key_team_metadata": {"model_tpm_limit": {"target": 500}},
|
||||
}
|
||||
},
|
||||
"model": "alias",
|
||||
}
|
||||
|
||||
ops: Final = handler._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs, response_obj=response, rate_limit_type="output"
|
||||
)
|
||||
|
||||
charged_keys: Final = {op["key"] for op in ops}
|
||||
assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-alias-tpm')}:target", "tokens") in charged_keys
|
||||
assert not any(":alias" in key for key in charged_keys)
|
||||
team_pool_key: Final = handler.create_rate_limit_keys("model_per_team", "t:target", "tokens")
|
||||
assert (team_pool_key in charged_keys) is charges_team_model_pool
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_success_tpm_accounting_keeps_the_admission_target_after_an_alias_reload() -> None:
|
||||
alias_map: Final[dict[str, str]] = {"alias": "target-a"}
|
||||
cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(cache), model_group_resolver=alias_map.get
|
||||
)
|
||||
key_metadata: Final = {"model_tpm_limit": {"target-a": 1000, "target-b": 1000}}
|
||||
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-reload"), metadata=key_metadata)
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=auth,
|
||||
cache=cache,
|
||||
data={"model": "alias", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10},
|
||||
call_type="acompletion",
|
||||
)
|
||||
stash: Final = get_request_stash()
|
||||
assert stash is not None
|
||||
assert stash.reserved_model == RateLimitedModel(requested="alias", group="target-a")
|
||||
assert stash.reserved_tokens > 0
|
||||
|
||||
alias_map["alias"] = "target-b"
|
||||
response: Final = ModelResponse(
|
||||
id="alias-reload",
|
||||
object="chat.completion",
|
||||
created=int(datetime.now().timestamp()),
|
||||
model="alias",
|
||||
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
|
||||
choices=[],
|
||||
)
|
||||
kwargs: Final = {
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}},
|
||||
"litellm_params": {"metadata": {"model_group": "alias", "user_api_key_metadata": key_metadata}},
|
||||
"model": "alias",
|
||||
}
|
||||
|
||||
ops: Final = handler._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs, response_obj=response, rate_limit_type="total"
|
||||
)
|
||||
|
||||
admission_bucket: Final = handler.create_rate_limit_keys("model_per_key", f"{auth.api_key}:target-a", "tokens")
|
||||
charged: Final = {op["key"]: op["increment_value"] for op in ops}
|
||||
assert charged[admission_bucket] == 150 - stash.reserved_tokens
|
||||
assert not any(":target-b" in key for key in charged)
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
PROJECT_ITPM_DESCRIPTOR_KEY,
|
||||
PROJECT_OTPM_DESCRIPTOR_KEY,
|
||||
RateLimitedModel,
|
||||
_AUDIO_BYTES_PER_TOKEN,
|
||||
_PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler,
|
||||
)
|
||||
|
|
@ -308,7 +309,7 @@ async def test_model_scope_refund_targets_reserved_model(rate_limiter):
|
|||
|
||||
stash = get_or_create_request_stash()
|
||||
stash.reserved_tokens = 100
|
||||
stash.reserved_model = reserved_model
|
||||
stash.reserved_model = RateLimitedModel(requested=reserved_model, group=reserved_model)
|
||||
stash.reserved_scopes = frozenset({("model_per_team", f"{team_id}:{reserved_model}")})
|
||||
|
||||
mock_kwargs = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue