From 35f8a298aedf15cb5f6580b9d08f6b45a382e06b Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Fri, 6 Nov 2026 04:18:39 +0000 Subject: [PATCH] Coalesce colliding floating alias rate limit counters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hooks/parallel_request_limiter_v3.py | 153 ++++++++------ .../hooks/test_parallel_request_limiter_v3.py | 197 ++++++++++++++++++ 2 files changed, 292 insertions(+), 58 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c81f1f78fd2..aaab4bf70ca 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -626,6 +626,44 @@ class RateLimitDescriptor(TypedDict): rate_limit: RateLimitDescriptorRateLimitObject | None +def _minimum_descriptor_limit( + descriptors: Sequence[RateLimitDescriptor], + metric: Literal["requests_per_unit", "tokens_per_unit", "max_parallel_requests"], +) -> int | None: + return min( + (limit for descriptor in descriptors if (limit := (descriptor["rate_limit"] or {}).get(metric)) is not None), + default=None, + ) + + +def _merge_rate_limit_descriptors(descriptors: Sequence[RateLimitDescriptor]) -> RateLimitDescriptor: + first: Final = descriptors[0] + if len(descriptors) == 1: + return first + return RateLimitDescriptor( + key=first["key"], + value=first["value"], + rate_limit=RateLimitDescriptorRateLimitObject( + requests_per_unit=_minimum_descriptor_limit(descriptors, "requests_per_unit"), + tokens_per_unit=_minimum_descriptor_limit(descriptors, "tokens_per_unit"), + max_parallel_requests=_minimum_descriptor_limit(descriptors, "max_parallel_requests"), + window_size=(first["rate_limit"] or {}).get("window_size"), + ), + ) + + +def _coalesce_rate_limit_descriptors(descriptors: Sequence[RateLimitDescriptor]) -> Sequence[RateLimitDescriptor]: + identities: Final = tuple(dict.fromkeys((descriptor["key"], descriptor["value"]) for descriptor in descriptors)) + if len(identities) == len(descriptors): + return descriptors + return tuple( + _merge_rate_limit_descriptors( + tuple(descriptor for descriptor in descriptors if (descriptor["key"], descriptor["value"]) == identity) + ) + for identity in identities + ) + + class ParallelRequestGauge(TypedDict): counter_key: str limit: int @@ -3127,14 +3165,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - return [ - *descriptors, - *( - descriptor - for model in self._rate_limited_models(self._rate_limited_model(requested_model)) - if (descriptor := self._organization_model_rate_limit_descriptor(user_api_key_dict, model)) is not None - ), - ] + model_descriptors: Final = tuple( + descriptor + for model in self._rate_limited_models(self._rate_limited_model(requested_model)) + if (descriptor := self._organization_model_rate_limit_descriptor(user_api_key_dict, model)) is not None + ) + return [*descriptors, *_coalesce_rate_limit_descriptors(model_descriptors)] def _organization_model_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, model: RateLimitedModel @@ -3685,7 +3721,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) - return descriptors + return descriptors if len(model_candidates) == 1 else list(_coalesce_rate_limit_descriptors(descriptors)) _create_rate_limit_descriptors = create_rate_limit_descriptors @@ -3838,19 +3874,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: str | None, descriptors: _RateLimitDescriptorSink, ) -> None: - for candidate in expand_model_latest_alias(requested_model) if requested_model else (None,): - self._add_single_project_model_rate_limit_descriptor(user_api_key_dict, candidate, descriptors) + candidates: Final = expand_model_latest_alias(requested_model) if requested_model else (None,) + model_descriptors: Final = tuple( + descriptor + for candidate in candidates + if (descriptor := self._project_model_rate_limit_descriptor(user_api_key_dict, candidate)) is not None + ) + for descriptor in _coalesce_rate_limit_descriptors(model_descriptors): + descriptors.append(descriptor) - def _add_single_project_model_rate_limit_descriptor( + def _project_model_rate_limit_descriptor( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None, - descriptors: _RateLimitDescriptorSink, - ) -> None: + ) -> RateLimitDescriptor | None: """Add project model rate limit descriptor from project_metadata if applicable.""" model: Final = self._rate_limited_model(requested_model) if model is None: - return + return None model_specific_tpm_limit: Final = model.limit_in( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_tpm_limit") ) @@ -3858,17 +3899,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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, - }, - ) + return None + return 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, + }, ) def add_project_io_token_rate_limit_descriptors_from_metadata( @@ -3877,15 +3916,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: str | None, descriptors: _RateLimitDescriptorSink, ) -> None: - for candidate in expand_model_latest_alias(requested_model) if requested_model else (None,): - self._add_single_project_io_token_rate_limit_descriptors(user_api_key_dict, candidate, descriptors) + candidates: Final = expand_model_latest_alias(requested_model) if requested_model else (None,) + model_descriptors: Final = tuple( + itertools.chain.from_iterable( + self._project_io_token_rate_limit_descriptors(user_api_key_dict, candidate) for candidate in candidates + ) + ) + for descriptor in _coalesce_rate_limit_descriptors(model_descriptors): + descriptors.append(descriptor) - def _add_single_project_io_token_rate_limit_descriptors( + def _project_io_token_rate_limit_descriptors( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None, - descriptors: _RateLimitDescriptorSink, - ) -> None: + ) -> tuple[RateLimitDescriptor, ...]: """Add project-scoped ITPM/OTPM descriptors from project_metadata. Enforced independently of, and alongside, the combined ``model_per_project`` @@ -3894,7 +3938,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ model: Final = self._rate_limited_model(requested_model) if model is None or user_api_key_dict.project_id is None: - return + return () model_itpm_limit: Final = model.limit_in( get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_itpm_limit") @@ -3903,34 +3947,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): get_model_rate_limit_from_metadata(user_api_key_dict, "project_metadata", "model_otpm_limit") ) - if model_itpm_limit is None and model_otpm_limit is None: - return - descriptor_value: Final = f"{user_api_key_dict.project_id}:{model.group}" - if model_itpm_limit is not None: - descriptors.append( - RateLimitDescriptor( - key=PROJECT_ITPM_DESCRIPTOR_KEY, - value=descriptor_value, - rate_limit={ - "requests_per_unit": None, - "tokens_per_unit": model_itpm_limit, - "window_size": self.window_size, - }, - ) + return tuple( + RateLimitDescriptor( + key=key, + value=descriptor_value, + rate_limit={ + "requests_per_unit": None, + "tokens_per_unit": limit, + "window_size": self.window_size, + }, ) - if model_otpm_limit is not None: - descriptors.append( - RateLimitDescriptor( - key=PROJECT_OTPM_DESCRIPTOR_KEY, - value=descriptor_value, - rate_limit={ - "requests_per_unit": None, - "tokens_per_unit": model_otpm_limit, - "window_size": self.window_size, - }, - ) + for key, limit in ( + (PROJECT_ITPM_DESCRIPTOR_KEY, model_itpm_limit), + (PROJECT_OTPM_DESCRIPTOR_KEY, model_otpm_limit), ) + if limit is not None + ) def _handle_rate_limit_error( self, @@ -5526,7 +5559,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): key_owned_team_scopes: Final = frozenset( ("model_per_team", f"{logging_metadata.get('user_api_key_team_id')}:{model.group}") for model in reconcile_models - if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, model) + if all( + self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, candidate) + for candidate in reconcile_models + if candidate.group == model.group + ) ) charged_targets: Final = [target for target in targets if target not in key_owned_team_scopes] if reserved_tokens > 0 and total_tokens < reserved_tokens: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 412a72cf3c5..d212e91871b 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -29,6 +29,7 @@ from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.litellm_core_utils import ptu_pricing from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth @@ -96,6 +97,202 @@ def _isolated_request_stash(): _request_stash.reset(token) +@pytest.fixture +def reservation() -> bool: + return True + + +@pytest.fixture +def collision_limiter( + reservation: bool, monkeypatch: pytest.MonkeyPatch, +) -> tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str]: + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", str(reservation).lower()) + alias: Final = "anthropic/claude-opus-latest" + target: Final = f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}" + router: Final = Router( + model_list=[{"model_name": "anthropic/*", "litellm_params": {"model": "anthropic/*"}}], + model_group_alias={alias: target}, + ) + cache: Final = DualCache() + handler: Final = PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=InternalUsageCache(cache), + model_group_resolver=partial(resolve_model_group_alias, router.registry().aliases()), + time_provider=lambda: datetime(2026, 1, 1, tzinfo=timezone.utc), + ) + return cache, handler, alias, target + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "scope,metric", + tuple( + (scope, metric) + for scope in ("key", "team", "project", "organization") + for metric in ("model_rpm_limit", "model_tpm_limit") + ) + (("project", "model_itpm_limit"), ("project", "model_otpm_limit")), +) +@pytest.mark.parametrize("limit_alias", (False, True)) +async def test_latest_alias_collision_keeps_the_stricter_limit( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], + scope: str, metric: str, limit_alias: bool, +) -> None: + cache, handler, alias, target = collision_limiter + limits: Final = {metric: {alias: 0 if limit_alias else 100, target: 100 if limit_alias else 0}} + auth: Final = UserAPIKeyAuth( + api_key="collision-key", team_id="collision-team", project_id="collision-project", org_id="collision-org", + **{("metadata" if scope == "key" else f"{scope}_metadata"): limits}, + ) + with pytest.raises(ProxyRateLimitError) as exc: + await handler.async_pre_call_hook( + auth, cache, + {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, "acompletion", + ) + assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ("key", "team", "project", "organization")) +@pytest.mark.parametrize("reservation", (False, True)) +@pytest.mark.parametrize("event", ("success", "failure", "partial")) +async def test_latest_alias_collision_charges_and_settles_one_counter_per_metric( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], + scope: str, reservation: bool, event: str, +) -> None: + cache, handler, alias, target = collision_limiter + limits: Final = {"model_rpm_limit": {alias: 2}, "model_tpm_limit": {target: 10000}} + auth: Final = UserAPIKeyAuth( + api_key="collision-key", team_id="collision-team", project_id="collision-project", org_id="collision-org", + **{("metadata" if scope == "key" else f"{scope}_metadata"): limits}, + ) + await handler.async_pre_call_hook( + auth, cache, {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 64}, "acompletion" + ) + entity: Final = { + "key": auth.api_key, "team": auth.team_id, "project": auth.project_id, "organization": auth.org_id + }[scope] + rpm_key: Final = handler.create_rate_limit_keys(f"model_per_{scope}", f"{entity}:{target}", "requests") + tpm_key: Final = handler.create_rate_limit_keys(f"model_per_{scope}", f"{entity}:{target}", "tokens") + stash: Final = get_request_stash() + assert stash is not None + assert await cache.async_get_cache(rpm_key) == 1 + assert await cache.async_get_cache(tpm_key) == (stash.reserved_tokens if reservation else 1) + usage: Final = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) + kwargs: Final = { + "standard_logging_object": { + "model_group": alias, + "metadata": { + "user_api_key_hash": auth.api_key, "user_api_key_team_id": auth.team_id, + "user_api_key_project_id": auth.project_id, "user_api_key_org_id": auth.org_id, + }, + }, + "litellm_params": {"metadata": {"model_group": alias, "user_api_key_metadata": auth.metadata}}, + **({"combined_usage_object": usage} if event == "partial" else {}), + } + if event == "success": + await handler.async_log_success_event(kwargs, ModelResponse(model=target, usage=usage), None, None) + else: + await handler.async_log_failure_event(kwargs, None, None, None) + assert await cache.async_get_cache(tpm_key) == (0 if event == "failure" else 5) + (0 if reservation else 1) + alias_key: Final = handler.create_rate_limit_keys(f"model_per_{scope}", f"{entity}:{alias}", "tokens") + assert await cache.async_get_cache(alias_key) is None + assert await handler.async_pre_call_hook(auth, cache, {"model": alias}, "acompletion") is None + with pytest.raises(ProxyRateLimitError) as exc: + await handler.async_pre_call_hook(auth, cache, {"model": alias}, "acompletion") + assert exc.value.status_code == 429 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event", ("success", "failure", "partial")) +async def test_latest_alias_collision_project_io_reservations_settle_once( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], event: str, +) -> None: + cache, handler, alias, target = collision_limiter + auth: Final = UserAPIKeyAuth( + api_key="collision-io-key", project_id="collision-project", + project_metadata={metric: {alias: 10000, target: 10000} for metric in ("model_itpm_limit", "model_otpm_limit")}, + ) + await handler.async_pre_call_hook( + auth, cache, {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 64}, "acompletion" + ) + stash: Final = get_request_stash() + assert stash is not None + for scope, reserved in ( + ("model_per_project_itpm", stash.itpm_reserved_tokens), ("model_per_project_otpm", stash.otpm_reserved_tokens), + ): + assert reserved > 0 + counter: Final = handler.create_rate_limit_keys(scope, f"{auth.project_id}:{target}", "tokens") + assert await cache.async_get_cache(counter) == reserved + usage: Final = Usage( + prompt_tokens=3, completion_tokens=2, total_tokens=5, prompt_tokens_details={"cached_tokens": 1} + ) + kwargs: Final = { + "standard_logging_object": { + "model_group": alias, + "metadata": {"user_api_key_hash": auth.api_key, "user_api_key_project_id": auth.project_id}, + }, + "litellm_params": {"metadata": {"model_group": alias}}, + **({"combined_usage_object": usage} if event == "partial" else {}), + } + if event == "success": + await handler.async_log_success_event(kwargs, ModelResponse(model=target, usage=usage), None, None) + else: + await handler.async_log_failure_event(kwargs, None, None, None) + for scope in ("model_per_project_itpm", "model_per_project_otpm"): + counter: Final = handler.create_rate_limit_keys(scope, f"{auth.project_id}:{target}", "tokens") + assert int(await cache.async_get_cache(counter) or 0) == (0 if event == "failure" else 2) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reservation", (False, True)) +@pytest.mark.parametrize("event", ("success", "partial")) +async def test_latest_alias_collision_settles_the_target_team_limit_beside_the_key_alias_limit( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], + reservation: bool, event: str, +) -> None: + cache, handler, alias, target = collision_limiter + auth: Final = UserAPIKeyAuth( + api_key="collision-inherited-key", team_id="collision-team", + metadata={"model_tpm_limit": {alias: 10000}}, team_metadata={"model_tpm_limit": {target: 10000}}, + ) + await handler.async_pre_call_hook( + auth, cache, {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 64}, "acompletion" + ) + usage: Final = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5) + kwargs: Final = { + "standard_logging_object": { + "model_group": alias, + "metadata": {"user_api_key_hash": auth.api_key, "user_api_key_team_id": auth.team_id}, + }, + "litellm_params": {"metadata": {"model_group": alias, "user_api_key_metadata": auth.metadata}}, + **({"combined_usage_object": usage} if event == "partial" else {}), + } + if event == "success": + await handler.async_log_success_event(kwargs, ModelResponse(model=target, usage=usage), None, None) + else: + await handler.async_log_failure_event(kwargs, None, None, None) + for scope, entity in (("model_per_key", auth.api_key), ("model_per_team", auth.team_id)): + counter: Final = handler.create_rate_limit_keys(scope, f"{entity}:{target}", "tokens") + assert await cache.async_get_cache(counter) == 5 + (0 if reservation else 1) + + +@pytest.mark.asyncio +async def test_latest_alias_collision_project_frames_charge_once_and_enforce_the_lower_cap( + collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str], +) -> None: + cache, handler, alias, target = collision_limiter + auth: Final = UserAPIKeyAuth( + api_key="collision-frame-key", project_id="collision-project", + project_metadata={metric: {alias: 5, target: 100} for metric in ("model_itpm_limit", "model_otpm_limit")}, + ) + await handler.enforce_project_io_token_quota_for_frame(auth, alias, 3, 2) + with pytest.raises(ProxyRateLimitError) as exc: + await handler.enforce_project_io_token_quota_for_frame(auth, alias, 3, 2) + assert exc.value.status_code == 429 + for scope, expected in (("model_per_project_itpm", 3), ("model_per_project_otpm", 2)): + counter: Final = handler.create_rate_limit_keys(scope, f"{auth.project_id}:{target}", "tokens") + assert await cache.async_get_cache(counter) == expected + + @pytest.mark.asyncio @pytest.mark.parametrize("scope", ("key", "team", "project", "organization")) @pytest.mark.parametrize("metric", ("model_rpm_limit", "model_tpm_limit"))