scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 455

feat(rate-limit): team_model_limit_borrows_idle_capacity_from_other_teams lets a team's model TPM and RPM limits borrow idle capacity in the dynamic rate limiter (#51328)

Browse source 
* feat(rate-limit): opt-in team_model_limit_is_floor turns a team's model_tpm_limit into a guaranteed share that borrows idle capacity

With priority_reservation_settings.team_model_limit_is_floor: true, a team's inherited model_tpm_limit on a model the dynamic rate limiter covers stops being a hard cap in the regular limiter. The dynamic limiter tracks it as a team_model_floor descriptor, lets the team past it while the model is below saturation_threshold, and holds the team to it once the model saturates, with rate_limited_by: team_model_floor on the 429. Default false keeps today's behavior. Key-owned model_tpm_limit, model_rpm_limit, the model-wide tpm and customer limits stay hard.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(rate-limit): name the setting team_model_limit_mode with hard and soft values

soft replaces the team_model_limit_is_floor flag. The 429 header reads rate_limited_by: team_model_soft_limit and the message says the team soft limit was reached.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(rate-limit): team_model_limit_borrows_idle_capacity_from_other_teams covers team model TPM and RPM

Replaces team_model_limit_mode. A team's inherited model_tpm_limit and model_rpm_limit on a model the dynamic
rate limiter covers are enforced by that limiter once the model saturates, so the team borrows idle capacity below
the threshold. The 429 names the limit and reads rate_limited_by: model_per_team_saturated.

Folds in the bot findings: key ownership is checked per requested model, team limits resolve through the model
alias like the regular limiter, the team counter is charged only for requests the pre-call hook tracked, every
response type that carries usage is settled, and tracking below saturation no longer adds a token per request.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rate-limit): keep the type gate flat in the dynamic limiter's team limit paths

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rate-limit): settle responses usage and null auth metadata, resolve aliases for the borrowing team limit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rate-limit): key the borrowing team counter by resolved model group and read usage through get_usage_object

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rate-limit): hold Responses WebSocket turns to the borrowing team limit once the model saturates

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-11 20:58:34 -08:00 • committed by GitHub
parent 6a98d09261
commit 19ffe3a5f2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 831 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

186
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -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

99
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -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(

4
litellm/types/utils.py
View file

@ -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,

117
tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py Normal file
View file

@ -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

285
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -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")

270
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -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