diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9d35178891c..33744906b13 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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"] diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index f6b38243896..8a0edf4a6ca 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -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"} diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 03167692727..02a131f60ae 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -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 diff --git a/tests/e2e/quota_management/ratelimit/test_model_group_alias_rate_limit_e2e.py b/tests/e2e/quota_management/ratelimit/test_model_group_alias_rate_limit_e2e.py new file mode 100644 index 00000000000..1c3fff47b78 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_model_group_alias_rate_limit_e2e.py @@ -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) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 0f19675edb9..8c0dcd3383c 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index bdaca9ffc2d..e6795bb22f3 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -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 = {