diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 5c481e33259..5a6d6fc4cdf 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -3,17 +3,19 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting """ import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict import litellm -from litellm import ModelResponse, Router +from litellm import Router from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache +from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, @@ -28,8 +30,12 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l PROXY_MaxParallelRequestsHandler_v3, RateLimitDescriptor, RateLimitDescriptorRateLimitObject, + RateLimitedModel, RateLimitResponse, + RateLimitStatus, _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export + _socket_turn_tpm_windows, # pyright: ignore[reportPrivateUsage] # same per-turn TPM window the regular limiter checks + borrowing_team_model_limits, claim_request_stash_for_data, get_or_create_request_stash, ) @@ -40,12 +46,56 @@ from litellm.proxy.hooks.rate_limiter_utils import ( from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET, CounterIncrement, RateLimitCounter, RateLimitScope from litellm.types.router import ModelGroupInfo -from litellm.types.utils import CallTypesLiteral, LLMResponseTypes +from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, Usage if TYPE_CHECKING: from litellm.types.utils import PriorityReservationSettings +class _SettlementMetadata(BaseModel): + """The two fields of a request's standard logging metadata that token settlement reads.""" + + model_config = ConfigDict(frozen=True) + + user_api_key_auth_metadata: Mapping[str, object] | None = None + user_api_key_team_id: str | None = None + + @property + def priority(self) -> str | None: + value: Final = (self.user_api_key_auth_metadata or {}).get("priority") + return value if isinstance(value, str) else None + + +def _usage_of(response_obj: object) -> Usage | None: + usage: Final = getattr(response_obj, "usage", None) + return usage if isinstance(usage, Usage) else get_usage_object(response_obj) + + +TEAM_MODEL_LIMIT_AT_SATURATION: Final = "model_per_team_saturated" + + +def _team_limit_name(rate_limit_type: str) -> str: + return "model_rpm_limit" if rate_limit_type == "requests" else "model_tpm_limit" + + +def _team_model_token_increment( + descriptor: RateLimitDescriptor | None, total_tokens: int, ttl: int +) -> tuple[CounterIncrement, ...]: + """Charge the team counter the pre-call hook tracked, and only when that team limit has a TPM share.""" + if descriptor is None: + return () + rate_limit: Final = descriptor["rate_limit"] + if rate_limit is None or rate_limit.get("tokens_per_unit") is None: + return () + return ( + CounterIncrement( + counter=RateLimitScope(descriptor["key"], descriptor["value"]).counter("tokens"), + amount=total_tokens, + ttl=ttl, + ), + ) + + def _get_priority_settings() -> "PriorityReservationSettings": """ Get the priority reservation settings, guaranteed to be non-None. @@ -342,6 +392,87 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return descriptors + def _create_borrowing_team_model_limit_descriptor( + self, model: str, user_api_key_dict: UserAPIKeyAuth + ) -> RateLimitDescriptor | None: + """One counter per team and resolved model group, so an alias and its target share the team's share.""" + rate_limited: Final = self.v3_limiter.rate_limited_model(model) or RateLimitedModel( + requested=model, group=model + ) + limits: Final = borrowing_team_model_limits(user_api_key_dict, rate_limited) + if limits.tpm is None and limits.rpm is None: + return None + return RateLimitDescriptor( + key=TEAM_MODEL_LIMIT_AT_SATURATION, + value=f"{user_api_key_dict.team_id}:{rate_limited.group}", + rate_limit={ + "tokens_per_unit": limits.tpm, + "requests_per_unit": limits.rpm, + "window_size": self.v3_limiter.window_size, + }, + ) + + def _team_limit_error( + self, + model: str, + status: RateLimitStatus, + saturation: float, + priority: str | None, + user_api_key_dict: UserAPIKeyAuth, + ) -> ProxyRateLimitError: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) + return ProxyRateLimitError( + detail={ + "error": f"Team {_team_limit_name(status['rate_limit_type'])} reached for {model} " + "and the model is saturated, so there is no idle capacity to borrow. " + f"Team: {user_api_key_dict.team_id}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Team {_team_limit_name(status['rate_limit_type'])}: {status['current_limit']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "rate_limited_by": TEAM_MODEL_LIMIT_AT_SATURATION, + "x-litellm-priority": priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]), + model=resolved_model, + llm_provider=llm_provider, + ) + + async def enforce_token_limits_for_frame( + self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None + ) -> None: + """Hold a Responses WebSocket turn to the team's share once the model saturates, like the regular limiter does.""" + if requested_model is None: + return + team_limit_descriptor: Final = self._create_borrowing_team_model_limit_descriptor( + model=requested_model, user_api_key_dict=user_api_key_dict + ) + model_group_info: Final = self.llm_router.registry().group(requested_model) + if team_limit_descriptor is None or model_group_info is None: + return + saturation: Final = await self._check_model_saturation(requested_model, model_group_info) + if saturation < _get_priority_settings().saturation_threshold: + return + tpm_windows: Final = _socket_turn_tpm_windows((team_limit_descriptor,)) + response: Final = await self.v3_limiter.should_rate_limit( + tpm_windows, parent_otel_span=user_api_key_dict.parent_otel_span, read_only=True, skip_parallel_check=True + ) + over_limit: Final = next((s for s in response["statuses"] if s["code"] == "OVER_LIMIT"), None) + if over_limit is None: + return + raise self._team_limit_error( + requested_model, + over_limit, + saturation, + self._get_priority_from_user_api_key_dict(user_api_key_dict), + user_api_key_dict, + ) + def _create_model_tracking_descriptor( self, model: str, @@ -438,9 +569,18 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # asyncio.Lock + in-memory fallback for single-process deployments. # All-or-nothing: if any enforced descriptor would exceed its limit, # no counter is modified and the response carries "OVER_LIMIT". + team_limit_descriptor: Final = self._create_borrowing_team_model_limit_descriptor( + model=model, user_api_key_dict=user_api_key_dict + ) + get_or_create_request_stash().borrowing_team_model_limit = team_limit_descriptor + if team_limit_descriptor is not None: + descriptors_to_check.append(team_limit_descriptor) + enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor] if priority_descriptors and should_enforce_priority: enforced_descriptors.extend(priority_descriptors) + if team_limit_descriptor is not None and should_enforce_priority: + enforced_descriptors.append(team_limit_descriptor) per_request_increment: Final[dict[Literal["requests", "tokens"], int]] = { "requests": 1, @@ -505,6 +645,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): llm_provider=llm_provider, ) + if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION: + raise self._team_limit_error(model, status, saturation, priority, user_api_key_dict) + # Fail-closed guard: overall_code says OVER_LIMIT but no status # matched a descriptor key we know how to translate into a 429. # Refuse the request rather than silently fall through and let an @@ -544,6 +687,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # priority allocation has accurate usage when it later becomes # enforced. The increment-then-check semantics of should_rate_limit # are fine because we ignore the OVER_LIMIT response. + if team_limit_descriptor is not None and not should_enforce_priority: + await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=[team_limit_descriptor], + increments=[per_request_increment], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) if priority_descriptors and not should_enforce_priority: priority_tracking_response: Final = await self.v3_limiter.should_rate_limit( descriptors=priority_descriptors, @@ -678,7 +827,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return response @with_service_target(RATE_LIMITS_TARGET) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: """ Update token usage for priority-based rate limiting after successful API calls. @@ -692,7 +841,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) - from litellm.types.utils import Usage try: verbose_proxy_logger.debug("INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING") @@ -701,7 +849,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Get metadata from standard_logging_object standard_logging_object: Final = kwargs.get("standard_logging_object") or {} - standard_logging_metadata: Final = standard_logging_object.get("metadata") or {} + settlement_metadata: Final = _SettlementMetadata.model_validate( + standard_logging_object.get("metadata") or {} + ) # Get model and priority model_group: Final = get_model_group_from_litellm_kwargs(kwargs) @@ -710,22 +860,20 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Get priority from user_api_key_auth_metadata in standard_logging_metadata # This is where user_api_key_dict.metadata is stored during pre-call - user_api_key_auth_metadata: Final = standard_logging_metadata.get("user_api_key_auth_metadata") or {} - key_priority: Final[str | None] = user_api_key_auth_metadata.get("priority") + key_priority: Final = settlement_metadata.priority # Get total tokens from response total_tokens = 0 rate_limit_type: Final = self.v3_limiter.get_rate_limit_type() - if isinstance(response_obj, ModelResponse): - _usage: Final = getattr(response_obj, "usage", None) - if _usage and isinstance(_usage, Usage): - if rate_limit_type == "output": - total_tokens = _usage.completion_tokens - elif rate_limit_type == "input": - total_tokens = _usage.prompt_tokens - elif rate_limit_type == "total": - total_tokens = _usage.total_tokens + _usage: Final = _usage_of(response_obj) + if _usage is not None: + if rate_limit_type == "output": + total_tokens = _usage.completion_tokens + elif rate_limit_type == "input": + total_tokens = _usage.prompt_tokens + elif rate_limit_type == "total": + total_tokens = _usage.total_tokens if total_tokens == 0: return @@ -750,6 +898,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): else: priority_key = f"{model_group}:default_pool" + team_model_tokens: Final = _team_model_token_increment( + get_or_create_request_stash().borrowing_team_model_limit, total_tokens, self.v3_limiter.window_size + ) pipeline_operations: Final = ( model_tokens, CounterIncrement( @@ -757,6 +908,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): amount=total_tokens, ttl=self.v3_limiter.window_size, ), + *team_model_tokens, ) # Execute token increments with TTL preservation diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 24c35d78c2e..7f266e67b41 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -31,6 +31,7 @@ from pydantic import TypeAdapter, ValidationError from starlette.status import HTTP_400_BAD_REQUEST, HTTP_503_SERVICE_UNAVAILABLE from typing_extensions import NotRequired, ReadOnly +import litellm from litellm import DualCache from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -166,6 +167,66 @@ class RateLimitedModel: return requested_limit if requested_limit is not None else limits.get(self.group) +def dynamic_rate_limiter_covers(model_group: str) -> bool: + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import PROXY_DynamicRateLimitHandlerV3 + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + return False + group: Final = llm_router.registry().group(model_group) + if group is None or (group.tpm is None and group.rpm is None): + return False + return any( + callback == "dynamic_rate_limiter_v3" or isinstance(callback, PROXY_DynamicRateLimitHandlerV3) + for callback in litellm.callbacks + ) + + +def inherited_team_model_limit( + user_api_key_dict: UserAPIKeyAuth, + model: RateLimitedModel, + rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], +) -> int | None: + """The team's per-model limit a key inherits, or None when the key sets its own limit for that model.""" + 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 model.limit_in(get_key_own_model_rate_limit(user_api_key_dict, rate_limit_key)) is not None + ): + return None + return team_limit + + +@dataclass(frozen=True, slots=True) +class BorrowingTeamModelLimits: + """Inherited team per-model limits the dynamic rate limiter enforces at saturation instead of the regular limiter.""" + + tpm: int | None + rpm: int | None + + +NO_BORROWING: Final = BorrowingTeamModelLimits(tpm=None, rpm=None) + + +def borrowing_team_model_limits(user_api_key_dict: UserAPIKeyAuth, model: RateLimitedModel) -> BorrowingTeamModelLimits: + settings: Final = litellm.priority_reservation_settings + if ( + settings is None + or not settings.team_model_limit_borrows_idle_capacity_from_other_teams + or user_api_key_dict.team_id is None + ): + return NO_BORROWING + limits: Final = BorrowingTeamModelLimits( + tpm=inherited_team_model_limit(user_api_key_dict, model, "model_tpm_limit"), + rpm=inherited_team_model_limit(user_api_key_dict, model, "model_rpm_limit"), + ) + if limits == NO_BORROWING or not dynamic_rate_limiter_covers(model.group): + return NO_BORROWING + return limits + + def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]: if router is None: return (model,) @@ -512,6 +573,7 @@ class RequestRateLimiterStash: otpm_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset) batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None batch_tpd_refund_ops: tuple[ReservationOperation, ...] = () + borrowing_team_model_limit: RateLimitDescriptor | None = None reservation_released: bool = False io_reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) @@ -2440,6 +2502,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if response["overall_code"] == "OVER_LIMIT": self._handle_rate_limit_error(response, tpm_windows, requested_model) + def rate_limited_model(self, requested_model: str | None) -> RateLimitedModel | None: + """The requested name and the model group it resolves to, so limits keyed by either name are found.""" + return self._rate_limited_model(requested_model) + def _rate_limited_model(self, requested_model: str | None) -> RateLimitedModel | None: if not requested_model: return None @@ -2551,11 +2617,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): model: Final = self._rate_limited_model(requested_model) if model is None: return - model_specific_tpm_limit: Final = model.limit_in( - get_key_model_tpm_limit(user_api_key_dict, model_name=model.group) + borrowing: Final = borrowing_team_model_limits(user_api_key_dict, model) + model_specific_tpm_limit: Final = ( + None + if borrowing.tpm is not None + else model.limit_in(get_key_model_tpm_limit(user_api_key_dict, model_name=model.group)) ) - model_specific_rpm_limit: Final = model.limit_in( - get_key_model_rpm_limit(user_api_key_dict, model_name=model.group) + model_specific_rpm_limit: Final = ( + None + if borrowing.rpm is not None + else 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 @@ -3134,12 +3205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): model: RateLimitedModel, rate_limit_key: Literal["model_rpm_limit", "model_tpm_limit"], ) -> int | None: - 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 + return inherited_team_model_limit(user_api_key_dict, model, rate_limit_key) def _key_owns_model_tpm_limit_from_request_metadata( self, @@ -3165,8 +3231,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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, model, "model_rpm_limit") - team_tpm_limit: Final = self._inherited_team_model_limit(user_api_key_dict, model, "model_tpm_limit") + borrowing: Final = borrowing_team_model_limits(user_api_key_dict, model) + team_rpm_limit: Final = ( + None + if borrowing.rpm is not None + else inherited_team_model_limit(user_api_key_dict, model, "model_rpm_limit") + ) + team_tpm_limit: Final = ( + None + if borrowing.tpm is not None + else 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( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5990135a5cb..9121c016ca1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -5389,6 +5389,10 @@ class PriorityReservationSettings(LiteLLMBaseModel): default=0.50, description="Saturation threshold (0.0-1.0) at which strict priority enforcement begins. Below this threshold, generous mode allows priority borrowing. Above this threshold, strict mode enforces normalized priority limits.", ) + team_model_limit_borrows_idle_capacity_from_other_teams: bool = Field( + default=False, + description="When true, a team's model_tpm_limit and model_rpm_limit on a model this limiter covers stop being ceilings. The team may run past them on capacity other teams are not using while the model is below saturation_threshold, and is held to them once the model saturates.", + ) saturation_check_cache_ttl: int = Field( default=60, diff --git a/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py b/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py new file mode 100644 index 00000000000..03951dfca21 --- /dev/null +++ b/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py @@ -0,0 +1,117 @@ +import json +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +OPENAI_MODEL: Final = "gpt-4.1-mini" +MODEL_TPM: Final = 1000 +MODEL_RPM: Final = 1000 +TEAM_TPM: Final = 100 +TEAM_RPM: Final = 3 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl_team_model_limit_borrowing", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "borrowed"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + } +).encode() + + +def _respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions", request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + return Reply(body=UPSTREAM_REPLY) + + +def _proxy(tmp_path_factory: pytest.TempPathFactory, saturation_threshold: float) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp(f"team-borrowing-{saturation_threshold}") + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["litellm_settings"] = { + **configuration["litellm_settings"], + "callbacks": ["dynamic_rate_limiter_v3"], + "priority_reservation": {"interactive": 0.6, "batch": 0.4}, + "priority_reservation_settings": { + "default_priority": 1.0, + "saturation_threshold": saturation_threshold, + "team_model_limit_borrows_idle_capacity_from_other_teams": True, + }, + } + path: Final = directory / "team_model_limit_borrowing.yaml" + path.write_text(yaml.safe_dump(configuration)) + with ( + gateway_from_environment() as rig, + wire_server(_respond) as wire, + owned_proxy(rig, directory, {}, config=path) as owned, + ): + yield owned, wire.url + + +@pytest.fixture(scope="module") +def idle_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + yield from _proxy(tmp_path_factory, saturation_threshold=0.5) + + +@pytest.fixture(scope="module") +def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + yield from _proxy(tmp_path_factory, saturation_threshold=0.0) + + +def _team_calls( + owned: Gateway, upstream: str, calls: int, limit: str = "model_tpm_limit" +) -> tuple[httpx.Response, ...]: + with owned.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", + api_base=f"{upstream}/v1", + api_key="synthetic-openai-key", + tpm=MODEL_TPM, + rpm=MODEL_RPM, + ) + key: Final = scenario.key( + team_id=scenario.team(metadata={limit: {model: TEAM_TPM if limit == "model_tpm_limit" else TEAM_RPM}}) + ) + return tuple( + owned.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"team limit borrowing {uuid.uuid4().hex}"}]}, + key=key, + ) + for _ in range(calls) + ) + + +@pytest.mark.parametrize("limit", ["model_tpm_limit", "model_rpm_limit"]) +def test_team_borrows_idle_model_capacity_past_its_model_limit(idle_model: tuple[Gateway, str], limit: str) -> None: + owned, upstream = idle_model + responses: Final = _team_calls(owned, upstream, calls=8, limit=limit) + assert [response.status_code for response in responses] == [200] * 8, [r.text for r in responses] + + +@pytest.mark.parametrize( + "limit, rate_limit_type", [("model_tpm_limit", "tokens"), ("model_rpm_limit", "requests")], ids=["tpm", "rpm"] +) +def test_team_is_held_to_its_model_limit_once_the_model_saturates( + saturated_model: tuple[Gateway, str], limit: str, rate_limit_type: str +) -> None: + owned, upstream = saturated_model + responses: Final = _team_calls(owned, upstream, calls=8, limit=limit) + blocked: Final = tuple(response for response in responses if response.status_code == 429) + assert blocked, [r.text for r in responses] + assert responses[0].status_code == 200, responses[0].text + first: Final = blocked[0] + assert first.headers.get("rate_limited_by") == "model_per_team_saturated", dict(first.headers) + assert first.headers.get("rate_limit_type") == rate_limit_type, dict(first.headers) + assert f"Team {limit} reached" in first.json()["error"]["message"], first.text + assert "no idle capacity to borrow" in first.json()["error"]["message"], first.text diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 0b1cb9bb425..4977e8d5737 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -5,17 +5,20 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of """ import asyncio -import os import time from datetime import datetime, timedelta +from typing import Final from unittest.mock import AsyncMock, patch import pytest - import litellm from litellm import DualCache, Router from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + PROXY_DynamicRateLimitHandlerV3, +) from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler, ) @@ -23,6 +26,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3, ) from litellm.redis.stores.rate_limits import RateLimitScope +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits @@ -353,9 +357,9 @@ async def test_100_concurrent_priority_requests(time_controller, monkeypatch): assert result["rpm"] == 50, f"Low priority user {result['user_id']} got {result['rpm']} RPM, expected 50" print(f"✅ Successfully processed 100 concurrent requests in {total_duration:.2f}s") - print(f" - 70 high priority users: 900 TPM, 450 RPM each") - print(f" - 30 low priority users: 100 TPM, 50 RPM each") - print(f" - Priority ratio maintained: 9:1 (TPM) and 9:1 (RPM)") + print(" - 70 high priority users: 900 TPM, 450 RPM each") + print(" - 30 low priority users: 100 TPM, 50 RPM each") + print(" - Priority ratio maintained: 9:1 (TPM) and 9:1 (RPM)") @pytest.mark.asyncio @@ -572,7 +576,7 @@ async def test_concurrent_pre_call_hooks_stress(monkeypatch): print(f" - Premium users: {premium_success}/{len(premium_results)} success ({premium_success_rate:.1%})") print(f" - Standard users: {standard_success}/{len(standard_results)} success ({standard_success_rate:.1%})") print(f" - Total successful: {successful_count}/50 ({successful_count / 50:.1%})") - print(f" - Priority system working: Premium > Standard success rates") + print(" - Priority system working: Premium > Standard success rates") # These tests make actual async_pre_call_hook calls to simulate real traffic @@ -676,7 +680,7 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(monkeypatch): total_successful = successful_requests["key_a"] + successful_requests["key_b"] total_rate_limited = rate_limited_requests["key_a"] + rate_limited_requests["key_b"] - print(f"Test Case 1 - Saturation-Aware Rate Limiting:") + print("Test Case 1 - Saturation-Aware Rate Limiting:") print(f" - Duration: {end_time - start_time:.2f}s") print(f" - Key A: {successful_requests['key_a']}/1 successful (reserved 75 RPM)") print(f" - Key B: {successful_requests['key_b']}/100 successful (reserved 25 RPM)") @@ -796,7 +800,7 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(monkeypatch): key_a_success_rate = successful_requests["key_a"] / 200 key_b_success_rate = successful_requests["key_b"] / 200 - print(f"Test Case 2 - Priority Queue Behavior During Saturation:") + print("Test Case 2 - Priority Queue Behavior During Saturation:") print(f" - Duration: {end_time - start_time:.2f}s") print(f" - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") print(f" - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") @@ -929,7 +933,7 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(monkeypatch): # Analyze results total_successful = sum(successful_requests.values()) - print(f"Test Case 3 - Spillover Capacity for Default Keys:") + print("Test Case 3 - Spillover Capacity for Default Keys:") print(f" - Duration: {end_time - start_time:.2f}s") print(f" - Key A: {successful_requests['key_a']}/150 successful") print(f" - Key B: {successful_requests['key_b']}/150 successful (default)") @@ -1047,7 +1051,7 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(monkeypatch): key_a_success_rate = successful_requests["key_a"] / 200 key_b_success_rate = successful_requests["key_b"] / 200 - print(f"Test Case 4 - Over-Allocated Priority Reservations with Normalization:") + print("Test Case 4 - Over-Allocated Priority Reservations with Normalization:") print(f" - Duration: {end_time - start_time:.2f}s") print(f" - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") print(f" - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") @@ -1184,7 +1188,7 @@ async def test_fake_calls_case_5_default_value_priority_reservation(monkeypatch) # Analyze results total_successful = sum(successful_requests.values()) - print(f"Test Case 5 - Default value for priority reservation:") + print("Test Case 5 - Default value for priority reservation:") print(f" - Duration: {end_time - start_time:.2f}s") print(f" - Key A (0.50): {successful_requests['key_a']}/150 successful") print(f" - Key B (0.20): {successful_requests['key_b']}/150 successful") @@ -1430,7 +1434,6 @@ async def test_saturation_check_cache_ttl_configuration(monkeypatch): class _RecordingRateLimits(FakeRateLimits): async def get(self, entry, *, backfill_ttl_seconds=None): get_cache_calls.append({"entry": entry, "ttl": backfill_ttl_seconds}) - return None # Simulate cache miss handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( handler.internal_usage_cache, @@ -1819,3 +1822,261 @@ async def test_post_call_success_hook_priority_header_is_always_http_encodable(t http_response = Response(headers={key: str(value) for key, value in additional_headers.items()}) assert http_response.headers.get("x-litellm-priority") == expected_priority_header assert http_response.headers["x-litellm-rate-limiter-version"] == "v3" + + +def _borrowing_handler( + monkeypatch: pytest.MonkeyPatch, + borrows: bool, + saturation_threshold: float = 0.5, + model_group_alias: dict[str, str] | None = None, +) -> PROXY_DynamicRateLimitHandlerV3: + from litellm.types.utils import PriorityReservationSettings + + monkeypatch.setenv("LITELLM_LICENSE", "test-license-key") + monkeypatch.setattr(litellm, "priority_reservation", {"interactive": 0.6, "batch": 0.4}) + monkeypatch.setattr( + litellm, + "priority_reservation_settings", + PriorityReservationSettings( + saturation_threshold=saturation_threshold, team_model_limit_borrows_idle_capacity_from_other_teams=borrows + ), + ) + monkeypatch.setattr(litellm, "callbacks", ["dynamic_rate_limiter_v3"]) + router: Final = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000, "rpm": 100}, + } + ], + model_group_alias=model_group_alias or {}, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + handler = PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=DualCache()) + handler.update_variables(llm_router=router) + return handler + + +async def _admit_then_settle(handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None: + await _check(handler, auth, saturation=0.0) + await _settle_tokens(handler, tokens) + + +async def _settle_tokens(handler: PROXY_DynamicRateLimitHandlerV3, tokens: int, response: object | None = None) -> None: + await handler.async_log_success_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": "shared-model"}}, + "standard_logging_object": { + "metadata": {"user_api_key_team_id": "team-a", "user_api_key_auth_metadata": {}} + }, + }, + response_obj=response + if response is not None + else litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=tokens, completion_tokens=0, total_tokens=tokens)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + +def _team_key( + team_id: str = "team-a", + team_limits: dict[str, dict[str, int]] | None = None, + **extra: dict[str, dict[str, int]], +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-borrow", + team_id=team_id, + team_metadata=team_limits if team_limits is not None else {"model_tpm_limit": {"shared-model": 100}}, + **extra, + ) + + +async def _check( + handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, saturation: float, model: str = "shared-model" +) -> None: + group: Final = handler.llm_router.registry().group(model) + assert group is not None + await handler._check_rate_limits( + model=model, model_group_info=group, user_api_key_dict=auth, priority=None, saturation=saturation + ) + + +def _held_to_team_limit(exc: ProxyRateLimitError, rate_limit_type: str, limit_sentence: str) -> None: + assert exc.status_code == 429 + assert exc.headers is not None + assert exc.headers["rate_limited_by"] == "model_per_team_saturated" + assert exc.headers["rate_limit_type"] == rate_limit_type + assert str(exc.detail).startswith( + "{'error': 'Team " + limit_sentence + " reached for shared-model and the model is saturated, " + "so there is no idle capacity to borrow. Team: team-a, Rate limit type: " + rate_limit_type + ", " + ), exc.detail + + +@pytest.mark.asyncio +async def test_team_past_its_model_tpm_limit_borrows_while_the_model_is_below_the_saturation_threshold( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _team_key(), 150) + await _check(handler, _team_key(), saturation=0.49) + + +@pytest.mark.asyncio +async def test_team_past_its_model_tpm_limit_is_held_to_it_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _team_key(), 150) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + assert "Team model_tpm_limit: 100, Remaining: 0, Model saturation: 50.0%" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_team_past_its_model_rpm_limit_is_held_to_it_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + auth: Final = _team_key(team_limits={"model_rpm_limit": {"shared-model": 2}}) + await _check(handler, auth, saturation=0.0) + await _check(handler, auth, saturation=0.0) + await _check(handler, auth, saturation=0.49) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, auth, saturation=0.5) + _held_to_team_limit(exc.value, "requests", "model_rpm_limit") + assert "Team model_rpm_limit: 2, Remaining: 0, Model saturation: 50.0%" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_requests_admitted_below_saturation_count_only_their_settled_tokens(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + for _ in range(20): + await _admit_then_settle(handler, _team_key(), 4) + await _check(handler, _team_key(), saturation=0.9) + + +@pytest.mark.asyncio +async def test_an_embedding_responses_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.0) + await _settle_tokens( + handler, + 150, + response=litellm.EmbeddingResponse( + usage=litellm.Usage(prompt_tokens=150, completion_tokens=0, total_tokens=150) + ), + ) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_responses_api_calls_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.0) + await _settle_tokens( + handler, + 150, + response=ResponsesAPIResponse( + id="resp", + created_at=0, + model="m", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150), + ), + ) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_null_auth_metadata_still_settles_the_teams_tokens(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.0) + await handler.async_log_success_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": "shared-model"}}, + "standard_logging_object": { + "metadata": {"user_api_key_team_id": "team-a", "user_api_key_auth_metadata": None} + }, + }, + response_obj=litellm.ModelResponse( + usage=litellm.Usage(prompt_tokens=150, completion_tokens=0, total_tokens=150) + ), + start_time=None, + end_time=None, + ) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_another_teams_usage_does_not_count_against_a_teams_limit(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _team_key("team-b"), 150) + await _check(handler, _team_key("team-a"), saturation=0.9) + + +@pytest.mark.asyncio +async def test_borrowing_is_off_by_default_so_the_dynamic_limiter_adds_no_team_limit(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=False) + await _admit_then_settle(handler, _team_key(), 150) + await _check(handler, _team_key(), saturation=0.9) + + +@pytest.mark.asyncio +async def test_a_key_that_owns_its_own_model_tpm_limit_is_not_held_to_the_team_limit(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + owner: Final = _team_key(metadata={"model_tpm_limit": {"shared-model": 5000}}) + await _admit_then_settle(handler, owner, 150) + await _check(handler, owner, saturation=0.9) + + +@pytest.mark.asyncio +async def test_a_key_that_owns_a_limit_on_another_model_is_still_held_to_the_team_limit( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + other_owner: Final = _team_key(metadata={"model_tpm_limit": {"other-model": 5000}}) + await _admit_then_settle(handler, other_owner, 150) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, other_owner, saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_an_exempt_keys_tokens_are_not_charged_to_its_teammates(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + owner: Final = _team_key(metadata={"model_tpm_limit": {"shared-model": 5000}}) + await _admit_then_settle(handler, owner, 150) + await _check(handler, _team_key(), saturation=0.9) + + +@pytest.mark.asyncio +async def test_an_alias_request_counts_against_the_same_team_share_as_its_target(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, model_group_alias={"alias-model": "shared-model"}) + await _check(handler, _team_key(), saturation=0.0, model="alias-model") + await _settle_tokens(handler, 150) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_websocket_turn_is_held_to_the_team_limit_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + await _admit_then_settle(handler, _team_key(), 150) + with pytest.raises(ProxyRateLimitError) as exc: + await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_websocket_turn_still_borrows_while_the_model_is_idle(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) + await _admit_then_settle(handler, _team_key(), 150) + await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") 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 c6e25dbeec3..40144736ebd 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5,13 +5,12 @@ Unit Tests for the max parallel request limiter v3 for the proxy import asyncio import contextvars import logging -import sys import time from collections.abc import AsyncGenerator, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from functools import partial -from typing import Any, Dict, Final, List, Optional +from typing import Any, Final from unittest.mock import patch import httpx @@ -29,8 +28,8 @@ 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.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired, total_deadline from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth @@ -67,8 +66,6 @@ from litellm.redis.stores.rate_limits import ( RateLimitStore, ) from litellm.responses.streaming_iterator import _record_usage_for_failure -from tests._support.clock_loops import AutojumpClockLoop -from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits, scope_key from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -83,6 +80,8 @@ from litellm.types.utils import ( TextCompletionResponse, Usage, ) +from tests._support.clock_loops import AutojumpClockLoop +from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits, scope_key class TimeController: @@ -144,7 +143,8 @@ def reservation() -> bool: @pytest.fixture def collision_limiter( - reservation: bool, monkeypatch: pytest.MonkeyPatch, + 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" @@ -169,23 +169,31 @@ def collision_limiter( (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")), + ) + + (("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, + 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", + 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", + auth, + cache, + {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, + "acompletion", ) assert exc.value.status_code == 429 @@ -196,19 +204,27 @@ async def test_latest_alias_collision_keeps_the_stricter_limit( @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, + 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", + 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 + "key": auth.api_key, + "team": auth.team_id, + "project": auth.project_id, + "organization": auth.org_id, }[scope] rpm_key: Final = _counter(f"model_per_{scope}", f"{entity}:{target}", "requests") tpm_key: Final = _counter(f"model_per_{scope}", f"{entity}:{target}", "tokens") @@ -221,8 +237,10 @@ async def test_latest_alias_collision_charges_and_settles_one_counter_per_metric "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, + "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}}, @@ -245,11 +263,13 @@ async def test_latest_alias_collision_charges_and_settles_one_counter_per_metric @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, + 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", + 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( @@ -258,7 +278,8 @@ async def test_latest_alias_collision_project_io_reservations_settle_once( 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), + ("model_per_project_itpm", stash.itpm_reserved_tokens), + ("model_per_project_otpm", stash.otpm_reserved_tokens), ): assert reserved > 0 counter: Final = _counter(scope, f"{auth.project_id}:{target}", "tokens") @@ -288,12 +309,15 @@ async def test_latest_alias_collision_project_io_reservations_settle_once( @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, + 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}}, + 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" @@ -323,7 +347,8 @@ async def test_latest_alias_collision_project_frames_charge_once_and_enforce_the ) -> None: cache, handler, alias, target = collision_limiter auth: Final = UserAPIKeyAuth( - api_key="collision-frame-key", project_id="collision-project", + 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) @@ -369,7 +394,9 @@ async def test_latest_alias_is_blocked_by_either_rate_limit(scope: str, metric: @pytest.mark.parametrize("scope", ("key", "team", "project", "organization")) @pytest.mark.parametrize("metric", ("model_rpm_limit", "model_tpm_limit")) @pytest.mark.parametrize("model", ("anthropic/claude-opus-5-5", "anthropic/unknown-model", None)) -async def test_latest_alias_rate_limit_does_not_govern_other_requests(scope: str, metric: str, model: str | None) -> None: +async def test_latest_alias_rate_limit_does_not_govern_other_requests( + scope: str, metric: str, model: str | None +) -> None: cache: Final = DualCache() handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) limits: Final = {metric: {"anthropic/claude-opus-latest": 0}} @@ -408,7 +435,10 @@ async def test_latest_alias_rpm_draws_from_both_model_counters(scope: str) -> No assert await handler.async_pre_call_hook(auth, cache, {"model": alias}, "acompletion") is None descriptor: Final = f"model_per_{scope}" entity: Final = { - "key": auth.api_key, "team": auth.team_id, "project": auth.project_id, "organization": auth.org_id + "key": auth.api_key, + "team": auth.team_id, + "project": auth.project_id, + "organization": auth.org_id, }[scope] for model in (alias, target): counter: Final = _counter(descriptor, f"{entity}:{model}", "requests") @@ -455,11 +485,12 @@ async def test_latest_alias_tpm_settles_both_counters( assert stash is not None descriptor: Final = f"model_per_{scope}" entity: Final = { - "key": auth.api_key, "team": auth.team_id, "project": auth.project_id, "organization": auth.org_id + "key": auth.api_key, + "team": auth.team_id, + "project": auth.project_id, + "organization": auth.org_id, }[scope] - keys: Final = tuple( - _counter(descriptor, f"{entity}:{model}", "tokens") for model in (alias, target) - ) + keys: Final = tuple(_counter(descriptor, f"{entity}:{model}", "tokens") for model in (alias, target)) for key in keys: assert int(await handler.rate_limits.get(key) or 0) == (stash.reserved_tokens if reservation else 1) kwargs: Final = { @@ -580,7 +611,10 @@ async def test_latest_alias_project_io_limits_block_either_name(metric: str, lim ) 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" + auth, + cache, + {"model": alias, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 1}, + "acompletion", ) assert exc.value.status_code == 429 @@ -627,7 +661,9 @@ async def test_latest_alias_project_io_reservations_settle_both_counters(event: for model in (alias, target): counter: Final = _counter(scope, f"{auth.project_id}:{model}", "tokens") assert await handler.rate_limits.get(counter) == reserved - usage: Final = Usage(prompt_tokens=3, completion_tokens=2, total_tokens=5, prompt_tokens_details={"cached_tokens": 1}) + 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, @@ -845,8 +881,8 @@ async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): ) # Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction - window_starts: Dict[str, int] = {} - request_counts: Dict[str, int] = {} + window_starts: dict[str, int] = {} + request_counts: dict[str, int] = {} async def mock_batch_rate_limiter(*args, **kwargs): keys = kwargs.get("keys") if kwargs else args[0] @@ -929,8 +965,8 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller ) # Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction - window_starts: Dict[str, int] = {} - request_counts: Dict[str, int] = {} + window_starts: dict[str, int] = {} + request_counts: dict[str, int] = {} async def mock_batch_rate_limiter(*args, **kwargs): keys = kwargs.get("keys") if kwargs else args[0] @@ -1071,8 +1107,8 @@ async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object, time_co ) # Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction - window_starts: Dict[str, int] = {} - request_counts: Dict[str, int] = {} + window_starts: dict[str, int] = {} + request_counts: dict[str, int] = {} async def mock_batch_rate_limiter(*args, **kwargs): print(f"args: {args}, kwargs: {kwargs}") @@ -1458,7 +1494,7 @@ async def test_rejected_request_does_not_consume_parallel_slot_v3(): _api_key = hash_token("sk-98765") user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) - admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + admitted_data: dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -1524,7 +1560,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): fakes.parallel_acquire_script = fake_acquire - data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + data: dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -1700,10 +1736,6 @@ async def test_model_specific_rate_limits_only_called_when_configured_v3(): """ Test that model-specific rate limits only trigger should_rate_limit when actually configured for the requested model. """ - from litellm.proxy.auth.auth_utils import ( - get_key_model_rpm_limit, - get_key_model_tpm_limit, - ) _api_key = "sk-98765" _api_key = hash_token(_api_key) @@ -2369,7 +2401,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return empty dict (no token_rate_limit_type set) - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server original_settings = getattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "general_settings", {}) @@ -2390,7 +2422,7 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): parallel_request_handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Mock general_settings to return an invalid token_rate_limit_type - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server original_settings = getattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "general_settings", {"token_rate_limit_type": "invalid_type"}) @@ -3022,8 +3054,8 @@ async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): rpm_limit=2, ) - window_starts: Dict[str, int] = {} - request_counts: Dict[str, int] = {} + window_starts: dict[str, int] = {} + request_counts: dict[str, int] = {} async def mock_batch_rate_limiter(*args, **kwargs): keys = kwargs.get("keys") if kwargs else args[0] @@ -3466,7 +3498,7 @@ async def test_pre_call_hook_keeps_internal_stash_out_of_request_body(): parallel_request_handler.should_rate_limit = mock_should_rate_limit parallel_request_handler.reserve_tpm_tokens = mock_reserve_tpm_tokens - data: Dict[str, Any] = { + data: dict[str, Any] = { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10, @@ -3509,7 +3541,7 @@ async def test_responses_route_body_untouched_by_pre_call_hook(caller_metadata): internal_usage_cache=InternalUsageCache(local_cache), ) - data: Dict[str, Any] = { + data: dict[str, Any] = { "model": "gpt-4o-mini", "input": "hello", } @@ -3608,7 +3640,7 @@ async def test_pre_call_hook_ignores_caller_supplied_stash_values(): "_litellm_tpm_reserved_tokens": 9999, "_litellm_rate_limit_descriptors": victim_descriptors, } - data: Dict[str, Any] = { + data: dict[str, Any] = { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "metadata": dict(injected), @@ -3976,7 +4008,7 @@ async def test_a_pre_call_after_a_failure_leaves_the_success_charging_project_io ], ) async def test_a_nested_mcp_pre_call_settles_each_project_io_hold_once_v3( - steps: tuple[str, ...], key_tpm_limit: Optional[int], expected_counters: tuple[int, int, int, int] + steps: tuple[str, ...], key_tpm_limit: int | None, expected_counters: tuple[int, int, int, int] ) -> None: api_key: Final = hash_token("sk-nested-mcp-hold-flags") model: Final = "gpt-4o-mini" @@ -4626,10 +4658,7 @@ async def test_a_reservation_made_after_a_released_one_in_the_same_context_settl "itpm": ("model_per_project_itpm", "proj-second:gpt-4o-mini"), "otpm": ("model_per_project_otpm", "proj-second:gpt-4o-mini"), } - settled: Final = { - name: int(_local(handler, _counter(key, value)) or 0) - for name, (key, value) in counters.items() - } + settled: Final = {name: int(_local(handler, _counter(key, value)) or 0) for name, (key, value) in counters.items()} assert settled == {"tokens": 27, "itpm": 20, "otpm": 7} @@ -5453,7 +5482,7 @@ async def test_post_call_failure_hook_releases_parallel_slot_v3(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = RateLimitScope("api_key", _api_key).gauge - admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + admitted_data: dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -5501,7 +5530,7 @@ async def test_success_event_releases_parallel_slot_v3(monkeypatch): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1) counter_key = RateLimitScope("api_key", _api_key).gauge - admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + admitted_data: dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -5833,7 +5862,7 @@ async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): assert exc_info.value.status_code == 429 handler.rate_limits.set_local(counter_key, 4, None) - admitted_data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} + admitted_data: dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -5938,7 +5967,7 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect gap that let a disconnect leak the slot in the default limiter-only config. A mid-stream disconnect must still refund the pre-call +1. """ - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() proxy_logging_obj = proxy_server.proxy_logging_obj @@ -5988,8 +6017,8 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): from the outer generator: the counter returns to 0 (not -1), proving the nested hook does not also refund and there is no double decrement. """ - import litellm.proxy.proxy_server as proxy_server from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy import proxy_server class _PassthroughIteratorOverride(CustomLogger): async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): @@ -6054,14 +6083,14 @@ async def test_pre_call_hook_reserves_tpm_when_enabled(monkeypatch): user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) - should_rate_limit_calls: List[Dict[str, Any]] = [] + should_rate_limit_calls: list[dict[str, Any]] = [] original_should_rate_limit = handler.should_rate_limit async def spy_should_rate_limit(*args, **kwargs): should_rate_limit_calls.append(kwargs) return await original_should_rate_limit(*args, **kwargs) - reserve_calls: List[int] = [] + reserve_calls: list[int] = [] original_reserve = handler.reserve_tpm_tokens async def spy_reserve(*args, **kwargs): @@ -6094,14 +6123,14 @@ async def test_pre_call_hook_skips_reservation_when_disabled(monkeypatch): user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-tpm"), tpm_limit=10_000) - should_rate_limit_calls: List[Dict[str, Any]] = [] + should_rate_limit_calls: list[dict[str, Any]] = [] original_should_rate_limit = handler.should_rate_limit async def spy_should_rate_limit(*args, **kwargs): should_rate_limit_calls.append(kwargs) return await original_should_rate_limit(*args, **kwargs) - reserve_calls: List[Any] = [] + reserve_calls: list[Any] = [] async def spy_reserve(*args, **kwargs): reserve_calls.append(kwargs) @@ -6428,7 +6457,7 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) # Real pre-call: populates data and stashes the response into metadata # so the success callback can find it via litellm_params.metadata. - data: Dict[str, Any] = { + data: dict[str, Any] = { "model": "gpt-4o-mini", "metadata": {}, "stream": True, @@ -6451,7 +6480,7 @@ async def test_streaming_end_to_end_populates_slp_ratelimit_headers(monkeypatch) usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), choices=[], ) - mock_kwargs: Dict[str, Any] = { + mock_kwargs: dict[str, Any] = { "standard_logging_object": { "metadata": { "user_api_key_hash": _api_key, @@ -6503,7 +6532,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - data: Dict[str, Any] = {"model": "gpt-4o-mini", "metadata": {}, "stream": True} + data: dict[str, Any] = {"model": "gpt-4o-mini", "metadata": {}, "stream": True} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -6520,7 +6549,7 @@ async def test_streaming_populates_model_per_key_ratelimit_headers(monkeypatch): choices=[], ) - mock_kwargs: Dict[str, Any] = { + mock_kwargs: dict[str, Any] = { "standard_logging_object": { "metadata": { "user_api_key_hash": _api_key, @@ -6570,7 +6599,7 @@ async def test_async_log_success_event_no_mirror_when_no_snapshot(monkeypatch): choices=[], ) - mock_kwargs: Dict[str, Any] = { + mock_kwargs: dict[str, Any] = { "standard_logging_object": { "metadata": { "user_api_key_hash": _api_key, @@ -6615,7 +6644,7 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) # Drive pre-call once so both paths have the same authoritative snapshot. - data: Dict[str, Any] = {"model": "gpt-4o-mini", "metadata": {}} + data: dict[str, Any] = {"model": "gpt-4o-mini", "metadata": {}} await handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, @@ -6641,7 +6670,7 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): non_stream_headers = non_stream_response._hidden_params.get("additional_headers", {}) # Streaming path: async_logging_hook mirrors into standard_logging_object. - stream_kwargs: Dict[str, Any] = { + stream_kwargs: dict[str, Any] = { "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, "litellm_params": {"metadata": data["metadata"]}, "model": "gpt-4o-mini", @@ -6661,7 +6690,7 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): ) stream_slp_headers = stream_kwargs["standard_logging_object"].get("hidden_params", {}).get("additional_headers", {}) - def _rl_only(headers: Dict[str, Any]) -> Dict[str, Any]: + def _rl_only(headers: dict[str, Any]) -> dict[str, Any]: return {k: v for k, v in headers.items() if k.startswith("x-ratelimit-")} assert _rl_only(stream_slp_headers) == _rl_only(non_stream_headers), ( @@ -6794,7 +6823,7 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): value="tpm-only-model", rate_limit={"tokens_per_unit": 100, "window_size": 60}, ) - zero_token_increment: Dict[str, int] = {"requests": 1, "tokens": 0} + zero_token_increment: dict[str, int] = {"requests": 1, "tokens": 0} under_limit = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], @@ -6822,7 +6851,7 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): assert _local(handler, counter_key) == 100 - negative_increment: Dict[str, int] = {"requests": -1, "tokens": -50} + negative_increment: dict[str, int] = {"requests": -1, "tokens": -50} refund_attempt = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], increments=[negative_increment], @@ -7194,7 +7223,7 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) local_cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + raised_data: dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} raised_reserved = await _reserved_tokens_for( handler, local_cache, @@ -7208,7 +7237,7 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) assert raised_data["max_tokens"] == 900 assert raised_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 900 - lowered_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + lowered_data: dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} lowered_reserved = await _reserved_tokens_for( handler, local_cache, @@ -7222,7 +7251,7 @@ async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch) assert lowered_data["max_tokens"] == 500 assert lowered_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 500 - unconfigured_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + unconfigured_data: dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} unconfigured_reserved = await _reserved_tokens_for( handler, local_cache, @@ -7432,7 +7461,7 @@ def test_internal_call_origin_success_ops_are_skipped(origin, rate_limit_type): choices=[], ) - def _kwargs(metadata: Dict[str, Any]) -> Dict[str, Any]: + def _kwargs(metadata: dict[str, Any]) -> dict[str, Any]: return { "standard_logging_object": { "metadata": { @@ -7458,7 +7487,7 @@ def test_internal_call_origin_success_ops_are_skipped(origin, rate_limit_type): assert skipped == (charged if origin == "autorouter_compaction" else []) -def _conflicting_budget_bodies() -> Dict[str, Dict[str, object]]: +def _conflicting_budget_bodies() -> dict[str, dict[str, object]]: """The same request, three ways of declaring the output budget.""" base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]} return { @@ -8054,10 +8083,7 @@ async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired( assert redis.released_slots == [ ( - [ - scope_key(gauge, "max_parallel_requests") - for gauge in acquired["gauges"] - ], + [scope_key(gauge, "max_parallel_requests") for gauge in acquired["gauges"]], [acquired["slot_id"]], ) ] @@ -8066,7 +8092,7 @@ async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired( @pytest.mark.asyncio async def test_fail_closed_rate_limit_enforcement_is_read_from_general_settings(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-settings"), rpm_limit=2) @@ -8088,7 +8114,7 @@ async def test_fail_closed_rate_limit_enforcement_is_read_from_general_settings( async def test_fail_closed_rate_limit_enforcement_coerces_the_general_settings_value( monkeypatch, configured_value, rejects ): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-coerced"), rpm_limit=2) monkeypatch.setitem(proxy_server.general_settings, "fail_closed_rate_limit_enforcement", configured_value) @@ -8283,7 +8309,7 @@ async def test_request_capacity_enforces_shared_rpm_scopes(limits, request_data, @pytest.mark.asyncio async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server router = Router( model_list=[ @@ -8659,7 +8685,7 @@ async def _rpm_request( 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 + from litellm.proxy import proxy_server router: Final = Router( model_list=[ @@ -9947,7 +9973,6 @@ async def test_end_user_model_rpm_limit_is_shared_across_keys_and_isolated_per_c await _rpm_request(handler, cache, acme_via_key_one, "other-model") - @pytest.mark.asyncio async def test_success_tpm_accounting_moves_the_charge_to_the_model_group_that_served_a_fallback() -> None: cache: Final = DualCache() @@ -10047,7 +10072,6 @@ async def test_success_tpm_accounting_settles_an_alias_to_target_fallback_once_p assert [inc.amount for inc in increments if inc.counter == alias_counter] == [] - @pytest.mark.asyncio async def test_success_tpm_accounting_ignores_a_client_supplied_model_group_without_a_router_fallback() -> None: cache: Final = DualCache() @@ -10086,3 +10110,71 @@ async def test_success_tpm_accounting_ignores_a_client_supplied_model_group_with charged: Final = {_increment_of(op).counter: _increment_of(op).amount for op in ops} assert charged[_counter("model_per_key", f"{auth.api_key}:primary")] == 150 - stash.reserved_tokens assert not any(":uncapped" in counter.scope.descriptor_value for counter in charged) + + +@pytest.mark.parametrize( + "rate_limit_type, team_limits, data", + [ + ( + "tokens", + {"model_tpm_limit": {"test-model": 500}}, + {"model": "test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 200}, + ), + ("requests", {"model_rpm_limit": {"test-model": 2}}, {"model": "test-model"}), + ], + ids=["model_tpm_limit", "model_rpm_limit"], +) +@pytest.mark.parametrize( + "borrows, third_request_passes", + [(False, False), (True, True)], + ids=["by_default_the_team_limit_stays_a_ceiling", "with_borrowing_on_the_dynamic_limiter_takes_over_the_limit"], +) +@pytest.mark.asyncio +async def test_team_model_limits_stop_being_hard_caps_when_borrowing_is_on( + monkeypatch: pytest.MonkeyPatch, + borrows: bool, + third_request_passes: bool, + rate_limit_type: str, + team_limits: dict[str, dict[str, int]], + data: dict[str, object], +): + from litellm.types.utils import PriorityReservationSettings + + monkeypatch.setattr( + litellm, + "priority_reservation_settings", + PriorityReservationSettings(team_model_limit_borrows_idle_capacity_from_other_teams=borrows), + ) + monkeypatch.setattr(litellm, "callbacks", ["dynamic_rate_limiter_v3"]) + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k", "api_base": "b", "tpm": 1000}, + } + ] + ), + ) + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + auth = UserAPIKeyAuth( + api_key=hash_token("sk-borrow"), + team_id="t", + team_metadata=team_limits, + ) + + async def request(): + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=dict(data), call_type="acompletion") + + await request() + await request() + if third_request_passes: + await request() + return + with pytest.raises(HTTPException) as exc: + await request() + assert exc.value.status_code == 429 + assert exc.value.headers is not None + assert exc.value.headers["rate_limit_type"] == rate_limit_type