diff --git a/litellm/proxy/auth/fallback_rate_limit.py b/litellm/proxy/auth/fallback_rate_limit.py index 1d4a33b2994..d2f6cf9498a 100644 --- a/litellm/proxy/auth/fallback_rate_limit.py +++ b/litellm/proxy/auth/fallback_rate_limit.py @@ -13,10 +13,12 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Final +import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.fallback_budget import ( _user_api_key_auth_from_request, # pyright: ignore[reportPrivateUsage] # shared auth lookup for fallback checks ) +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import PROXY_DynamicRateLimitHandlerV3 from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # the proxy's limiter type ) @@ -42,7 +44,18 @@ class RouterFallbackRateLimitCheck: if valid_token is None: return True try: - return await limiter.async_fallback_target_allowed(valid_token, model, request_kwargs) + dynamic_limiter: Final = next( + (cb for cb in litellm.callbacks if isinstance(cb, PROXY_DynamicRateLimitHandlerV3)), None + ) + if dynamic_limiter is not None and not await dynamic_limiter.async_fallback_target_allowed( + valid_token, model, request_kwargs, commit=False + ): + return False + if not await limiter.async_fallback_target_allowed(valid_token, model, request_kwargs): + return False + return dynamic_limiter is None or await dynamic_limiter.async_fallback_target_allowed( + valid_token, model, request_kwargs + ) except Exception as e: # noqa: BLE001 # fail closed: a limiter failure must not hand out unchecked capacity verbose_proxy_logger.warning("Skipping fallback to model=%s: rate limit check failed: %s", model, e) return False diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 059efe440ee..89818c21850 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -7,6 +7,7 @@ import os from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias from fastapi import HTTPException @@ -20,18 +21,22 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.cost_calculator import get_usage_object +from litellm.exceptions import ServiceUnavailableError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, ) from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, map_v3_rate_limit_type, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # legacy module exports + RATE_LIMIT_UNVERIFIABLE_MESSAGE, + BorrowingTeamModelHold, BorrowingTeamModelLimits, PROXY_MaxParallelRequestsHandler_v3, RateLimitDescriptor, @@ -39,15 +44,20 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l RateLimitedModel, RateLimitResponse, RateLimitStatus, - _call_id_from_callback_kwargs, # pyright: ignore[reportPrivateUsage] # same call-ownership check the regular limiter uses - _is_socket_turn, # pyright: ignore[reportPrivateUsage] # same socket-turn check the regular limiter uses + RateLimitUnverifiableError, + RequestRateLimiterStash, + ReservationOperation, + ReservationWindowIdentity, + WindowGuardedIncrement, _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, + callback_kwargs_mapping, claim_request_stash_for_data, get_or_create_request_stash, get_request_stash, get_request_stash_for_call, + request_stash_for_callback, ) from litellm.proxy.hooks.priority_class_share import ( DEFAULT_CLASS, @@ -61,6 +71,7 @@ from litellm.proxy.hooks.priority_class_share import ( from litellm.proxy.hooks.rate_limiter_utils import ( convert_priority_to_percent, resolve_llm_provider_for_rate_limit, + socket_session_usage, ) from litellm.proxy.utils import InternalUsageCache, check_and_merge_model_level_guardrails from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET, CounterIncrement, RateLimitCounter, RateLimitScope @@ -349,9 +360,36 @@ class _SettlementMetadata(BaseModel): return value if isinstance(value, str) else None -def _usage_of(response_obj: object) -> Usage | None: +class _SettlementLoggingObject(BaseModel): + """The slice of a request's standard logging object that token settlement reads.""" + + model_config = ConfigDict(frozen=True) + + metadata: _SettlementMetadata | None = None + + +def _usage_of(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: + """Usage of a modeled response, a Responses WebSocket turn's combined usage, or its raw event list.""" usage: Final = getattr(response_obj, "usage", None) - return usage if isinstance(usage, Usage) else get_usage_object(response_obj) + if isinstance(usage, Usage): + return usage + combined_usage: Final = kwargs.get("combined_usage_object") + if isinstance(combined_usage, Usage): + return combined_usage + socket_usage: Final = socket_session_usage(kwargs, response_obj) + return socket_usage if socket_usage is not None else get_usage_object(response_obj) + + +def _settled_tokens(usage: Usage | None, rate_limit_type: str) -> int: + if usage is None: + return 0 + match rate_limit_type: + case "output": + return usage.completion_tokens + case "input": + return usage.prompt_tokens + case _: + return usage.total_tokens TEAM_MODEL_LIMIT_AT_SATURATION: Final = "model_per_team_saturated" @@ -376,10 +414,10 @@ def _team_limit_name(rate_limit_type: str) -> str: def _team_model_token_increment( - descriptor: RateLimitDescriptor | None, total_tokens: int, ttl: int + descriptor: RateLimitDescriptor | None, amount: 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: + """Move the team counter the pre-call hook tracked by ``amount``, and only when that team limit has a TPM share.""" + if descriptor is None or amount == 0: return () rate_limit: Final = descriptor["rate_limit"] if rate_limit is None or rate_limit.get("tokens_per_unit") is None: @@ -387,12 +425,60 @@ def _team_model_token_increment( return ( CounterIncrement( counter=RateLimitScope(descriptor["key"], descriptor["value"]).counter("tokens"), - amount=total_tokens, + amount=amount, ttl=ttl, ), ) +def _window_guarded( + operations: tuple[CounterIncrement, ...], window: ReservationWindowIdentity | None +) -> tuple[ReservationOperation, ...]: + """Apply a reservation's settlement only while the window it was held in is still open, never to a newer window.""" + if window is None: + return operations + return tuple( + WindowGuardedIncrement(increment=op, expected_window_start=window[1], reservation_backend=window[2]) + if op.counter == window[0] + else op + for op in operations + ) + + +def _team_model_settlement( + stash: RequestRateLimiterStash | None, model_group: str | None, total_tokens: int, ttl: int +) -> tuple[ReservationOperation, ...]: + """Settle the counter that served the call to actual usage and return the admitting reservation it did not use. + + The admitting team counter serves unless a fallback group did, in which case that fallback's hold serves and + the admitting reservation is returned. A hold on a fallback target that did not serve stays for that attempt's + own failure callback, which may still be reporting the usage it recorded. + """ + if stash is None: + return () + served: Final = stash.borrowing_team_model_fallback_limits.get(model_group) if model_group is not None else None + admitted: Final = stash.borrowing_team_model_limit + reserved_tokens: Final = stash.borrowing_team_model_reserved_tokens + window: Final = stash.borrowing_team_model_reserved_window + if served is None: + return _window_guarded(_team_model_token_increment(admitted, total_tokens - reserved_tokens, ttl), window) + return ( + *_window_guarded(_team_model_token_increment(admitted, -reserved_tokens, ttl), window), + *_window_guarded( + _team_model_token_increment(served.descriptor, total_tokens - served.reserved_tokens, ttl), served.window + ), + ) + + +def _reserved_window( + response: RateLimitResponse, descriptor: RateLimitDescriptor | None +) -> ReservationWindowIdentity | None: + if descriptor is None: + return None + counter: Final = RateLimitScope(descriptor["key"], descriptor["value"]).counter("tokens") + return next((w for w in response.get("reservation_windows", frozenset()) if w[0] == counter), None) + + def _get_priority_settings() -> "PriorityReservationSettings": """ Get the priority reservation settings, guaranteed to be non-None. @@ -918,9 +1004,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if not descriptors: return tpm_windows: Final = _socket_turn_tpm_windows(descriptors) - 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 - ) + try: + 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, + ) + except RateLimitUnverifiableError as unverifiable: + raise ServiceUnavailableError( + message=RATE_LIMIT_UNVERIFIABLE_MESSAGE, llm_provider="litellm_proxy", model=requested_model + ) from unverifiable over_limit: Final = next((s for s in response["statuses"] if s["code"] == "OVER_LIMIT"), None) if over_limit is None: return @@ -966,6 +1060,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): user_api_key_dict: UserAPIKeyAuth, priority: str | None, saturation: float, + data: Mapping[str, object] = MappingProxyType({}), + call_type: str | None = None, ) -> None: """ Check rate limits using THREE-PHASE approach to prevent partial increments. @@ -992,6 +1088,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): user_api_key_dict: User authentication info priority: User's priority level saturation: Current saturation level + data: Request body, used to size the token reservation a saturated team limit holds + call_type: The endpoint's call type, used to read the body's prompt fields Raises: HTTPException: If any limit is exceeded @@ -1051,9 +1149,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): "requests": 1, "tokens": 0, } + team_reservation: Final = ( + await self._team_token_reservation(enforced_team_share, data, model, call_type, user_api_key_dict) + if enforced_team_share is not None + else 0 + ) atomic_response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( descriptors=enforced_descriptors, - increments=[per_request_increment for _ in enforced_descriptors], + increments=[ + {"requests": 1, "tokens": team_reservation} + if descriptor is enforced_team_share + else per_request_increment + for descriptor in enforced_descriptors + ], parent_otel_span=user_api_key_dict.parent_otel_span, ) @@ -1154,6 +1262,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. + stash: Final = get_or_create_request_stash() + stash.borrowing_team_model_reserved_tokens = team_reservation + stash.borrowing_team_model_reserved_window = ( + _reserved_window(atomic_response, enforced_team_share) if team_reservation else None + ) + self._cap_admitted_output(enforced_team_share, team_reservation, data, call_type) if team_share is not None and enforced_team_share is None: await self.v3_limiter.atomic_check_and_increment_by_n( descriptors=[team_share.share], @@ -1188,6 +1302,124 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): else: get_or_create_request_stash().rate_limit_response = atomic_response + def _cap_admitted_output( + self, + enforced_team_share: RateLimitDescriptor | None, + team_reservation: int, + data: Mapping[str, object], + call_type: str | None, + ) -> None: + """Write the output cap the hold was estimated with onto the admitted request, unless only input tokens count.""" + tpm_share: Final = ( + None if enforced_team_share is None else (enforced_team_share["rate_limit"] or {}).get("tokens_per_unit") + ) + if team_reservation and tpm_share is not None and self.v3_limiter.get_rate_limit_type() != "input": + self.v3_limiter.bound_output_for_limit(data, tpm_share, call_type) + + async def _team_token_reservation( + self, + descriptor: RateLimitDescriptor, + data: Mapping[str, object], + model: str, + call_type: str | None, + user_api_key_dict: UserAPIKeyAuth, + ) -> int: + """Tokens to hold against the team's share before the call, as the regular limiter does for a hard limit.""" + tpm_share: Final = (descriptor["rate_limit"] or {}).get("tokens_per_unit") + if tpm_share is None or not data: + return 0 + as_dispatched: Final = dict(data, model=model) + if self.v3_limiter.get_rate_limit_type() != "input": + self.v3_limiter.bound_output_for_limit(as_dispatched, tpm_share, call_type) + try: + estimate: Final = await offload_token_count(self.v3_limiter.estimate_request_tokens)( + data=as_dispatched, + model=model, + min_configured_tpm_limit=tpm_share, + call_type=call_type, + team_id=user_api_key_dict.team_id, + rate_limit_type=self.v3_limiter.get_rate_limit_type(), + ) + except Exception as e: # noqa: BLE001 # an unreadable prompt still holds the minimum so admission runs + verbose_proxy_logger.warning( + "Dynamic rate limiter: token estimate failed for %s, holding one token: %s", model, e + ) + return 1 + return max(estimate, 1) + + async def async_fallback_target_allowed( + self, user_api_key_dict: UserAPIKeyAuth, model: str, request_data: Mapping[str, object], *, commit: bool = True + ) -> bool: + """Hold a router fallback target to the team's share on it once it saturates, and to its borrow ceiling before. + + With ``commit`` the attempt is counted toward the target's saturation and the team's counter, holding the + prompt's estimated tokens when the share is enforced, otherwise only checked. + """ + team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict) + model_group_info: Final = self.llm_router.registry().group(model) + if team_share is None or model_group_info is None: + return True + saturation: Final = await self._check_model_saturation(model, model_group_info) + enforced: Final = team_share.enforced(saturation >= _get_priority_settings().saturation_threshold) + reservation: Final = ( + await self._team_token_reservation( + enforced, + {**request_data, "model": model}, + model, + get_or_create_request_stash().call_type, + user_api_key_dict, + ) + if enforced is not None + else 0 + ) + if not commit: + return enforced is None or await self._share_has_room(enforced, reservation, user_api_key_dict) + model_wide_descriptor: Final = self._create_model_tracking_descriptor( + model=model, model_group_info=model_group_info, high_limit_multiplier=1 + ) + response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=[model_wide_descriptor, *(() if enforced is None else (enforced,))], + increments=[ + {"requests": 1, "tokens": 0}, + *(() if enforced is None else ({"requests": 1, "tokens": reservation},)), + ], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + if response["overall_code"] == "OVER_LIMIT": + return False + if enforced is None: + await self.v3_limiter.atomic_check_and_increment_by_n( + descriptors=[team_share.share], + increments=[{"requests": 1, "tokens": reservation}], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + self._hold_fallback_share(model, team_share.share, reservation, _reserved_window(response, enforced)) + return True + + async def _share_has_room( + self, descriptor: RateLimitDescriptor, tokens: int, user_api_key_dict: UserAPIKeyAuth + ) -> bool: + checked: Final = await self.v3_limiter.should_rate_limit( + [descriptor], parent_otel_span=user_api_key_dict.parent_otel_span, read_only=True + ) + return all( + status["code"] != "OVER_LIMIT" + and (status["rate_limit_type"] != "tokens" or (status["limit_remaining"] or 0) >= tokens) + for status in checked["statuses"] + ) + + def _hold_fallback_share( + self, model: str, share: RateLimitDescriptor, reserved_tokens: int, window: ReservationWindowIdentity | None + ) -> None: + """Remember the hold under the requested name and the resolved group, so settlement finds it either way.""" + stash: Final = get_or_create_request_stash() + earlier: Final = stash.borrowing_team_model_fallback_limits.get(model) + carried: Final = earlier.reserved_tokens if earlier is not None and earlier.window == window else 0 + hold: Final = BorrowingTeamModelHold(descriptor=share, reserved_tokens=reserved_tokens + carried, window=window) + stash.borrowing_team_model_fallback_limits = MappingProxyType( + {**stash.borrowing_team_model_fallback_limits, model: hold, share["value"].rpartition(":")[2]: hold} + ) + @with_service_target(RATE_LIMITS_TARGET) async def async_pre_call_hook( self, @@ -1247,7 +1479,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): steered: Final = await self._steer_to_alternate(model, requested_group_info, user_api_key_dict, data, call_type) try: admitted: Final = await self._admit_steered_or_preferred( - model, requested_group_info, steered, user_api_key_dict + model, requested_group_info, steered, user_api_key_dict, data, call_type ) except HTTPException: raise @@ -1269,21 +1501,30 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): requested_group_info: ModelGroupInfo, steered: "SteeredModel | None", user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str | None, ) -> "SteeredModel | None": """The alternate once it admits the request, else None after the preferred model's own check.""" if steered is not None: try: - await self._admit(steered.model, steered.group_info, user_api_key_dict) + await self._admit(steered.model, steered.group_info, user_api_key_dict, data, call_type) except Exception: # noqa: BLE001 # the alternate's refusal or error keeps the request on its own model verbose_proxy_logger.info( "Alternate %s refused team %s, staying on %s", steered.model, user_api_key_dict.team_id, model ) else: return steered - await self._admit(model, requested_group_info, user_api_key_dict) + await self._admit(model, requested_group_info, user_api_key_dict, data, call_type) return None - async def _admit(self, model: str, model_group_info: ModelGroupInfo, user_api_key_dict: UserAPIKeyAuth) -> None: + async def _admit( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str | None, + ) -> None: """Read saturation, then the three-phase rate limit check, raising the 429 when ``model`` is out of capacity.""" priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) saturation: Final = await self._check_model_saturation(model, model_group_info) @@ -1298,6 +1539,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): user_api_key_dict=user_api_key_dict, priority=priority, saturation=saturation, + data=data, + call_type=call_type, ) async def _steer_to_alternate( @@ -1450,7 +1693,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return {STEERED_FROM_HEADER: stash.steered_from_model} @with_service_target(RATE_LIMITS_TARGET) - async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: """ Update token usage for priority-based rate limiting after successful API calls. @@ -1469,36 +1714,41 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): verbose_proxy_logger.debug("INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING") litellm_parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - - # Get metadata from standard_logging_object - standard_logging_object: Final = kwargs.get("standard_logging_object") or {} - settlement_metadata: Final = _SettlementMetadata.model_validate( - standard_logging_object.get("metadata") or {} + callback_kwargs: Final = callback_kwargs_mapping(kwargs) + settlement_metadata: Final = ( + _SettlementLoggingObject.model_validate(callback_kwargs.get("standard_logging_object") or {}).metadata + or _SettlementMetadata() ) - # Get model and priority - model_group: Final = get_model_group_from_litellm_kwargs(kwargs) + model_group: Final = get_model_group_from_litellm_kwargs(dict(kwargs)) + total_tokens: Final = _settled_tokens( + _usage_of(callback_kwargs, response_obj), self.v3_limiter.get_rate_limit_type() + ) + stash: Final = request_stash_for_callback(callback_kwargs) + team_model_tokens: Final = _team_model_settlement( + stash, model_group, total_tokens, self.v3_limiter.window_size + ) + if stash is not None: + served: Final = ( + stash.borrowing_team_model_fallback_limits.get(model_group) if model_group is not None else None + ) + stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_fallback_limits = MappingProxyType( + {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not served} + ) if not model_group: + if team_model_tokens: + await self.v3_limiter.async_increment_reservation_aware_tokens( + pipeline_operations=team_model_tokens, parent_otel_span=litellm_parent_otel_span + ) return # 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 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() - - _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: + if total_tokens == 0 and not team_model_tokens: return # Model-wide token tracking (model_saturation_check) @@ -1521,13 +1771,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): else: priority_key = f"{model_group}:default_pool" - call_stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - stash: Final = ( - call_stash if call_stash is not None else (get_request_stash() if _is_socket_turn(kwargs) else None) - ) - team_model_tokens: Final = _team_model_token_increment( - None if stash is None else stash.borrowing_team_model_limit, total_tokens, self.v3_limiter.window_size - ) class_share: Final = None if stash is None else stash.priority_class_share settled_team_id: Final = settlement_metadata.user_api_key_team_id priority_scope: Final = ( @@ -1538,19 +1781,23 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): class_share["value"] if settled_team_id is None else counter_value(model_group, settled_team_id), ) ) - pipeline_operations: Final = ( - model_tokens, - CounterIncrement( - counter=priority_scope.counter("tokens"), - amount=total_tokens, - ttl=self.v3_limiter.window_size, - ), - *team_model_tokens, + usage_operations: Final = ( + () + if total_tokens == 0 + else ( + model_tokens, + CounterIncrement( + counter=priority_scope.counter("tokens"), + amount=total_tokens, + ttl=self.v3_limiter.window_size, + ), + ) ) + pipeline_operations: Final[tuple[ReservationOperation, ...]] = (*usage_operations, *team_model_tokens) # Execute token increments with TTL preservation if pipeline_operations: - await self.v3_limiter.async_increment_tokens_with_ttl_preservation( + await self.v3_limiter.async_increment_reservation_aware_tokens( pipeline_operations=pipeline_operations, parent_otel_span=litellm_parent_otel_span, ) @@ -1568,5 +1815,55 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): except Exception as e: verbose_proxy_logger.exception("Error in dynamic rate limiter success event: %s", e) + @with_service_target(RATE_LIMITS_TARGET) + async def async_log_failure_event( + self, kwargs: object, response_obj: object, start_time: object, end_time: object + ) -> None: + """Settle the tokens held for a failed attempt to the usage it recorded, or give them all back when it recorded + none, on the admitting team counter or the fallback target that failed, leaving holds on other fallback + targets to the attempts still running.""" + from litellm.proxy.common_utils.callback_utils import get_model_group_from_litellm_kwargs + + callback_kwargs: Final = callback_kwargs_mapping(kwargs) + stash: Final = request_stash_for_callback(callback_kwargs) + if stash is None: + return + failed_group: Final = get_model_group_from_litellm_kwargs(dict(callback_kwargs)) + failed_hold: Final = ( + stash.borrowing_team_model_fallback_limits.get(failed_group) if failed_group is not None else None + ) + used: Final = _settled_tokens( + _usage_of(callback_kwargs, None if isinstance(response_obj, BaseException) else response_obj), + self.v3_limiter.get_rate_limit_type(), + ) + refund: Final = ( + *_window_guarded( + _team_model_token_increment( + stash.borrowing_team_model_limit, + (used if failed_hold is None else 0) - stash.borrowing_team_model_reserved_tokens, + self.v3_limiter.window_size, + ), + stash.borrowing_team_model_reserved_window, + ), + *( + () + if failed_hold is None + else _window_guarded( + _team_model_token_increment( + failed_hold.descriptor, used - failed_hold.reserved_tokens, self.v3_limiter.window_size + ), + failed_hold.window, + ) + ), + ) + stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_reserved_window = None + if failed_hold is not None: + stash.borrowing_team_model_fallback_limits = MappingProxyType( + {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not failed_hold} + ) + if refund: + await self.v3_limiter.async_increment_reservation_aware_tokens(pipeline_operations=refund) + PROXY_DynamicRateLimitHandlerV3: Final = _PROXY_DynamicRateLimitHandlerV3 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 16e12eb1e03..39409002058 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -570,6 +570,15 @@ class CallTypeRateLimiter(Protocol): ) -> Exception | str | dict[str, object] | None: ... +@dataclass(frozen=True, slots=True) +class BorrowingTeamModelHold: + """A router fallback target's team share with the tokens held on it, settled to actual usage or given back.""" + + descriptor: RateLimitDescriptor + reserved_tokens: int = 0 + window: ReservationWindowIdentity | None = None + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -594,6 +603,7 @@ class RequestRateLimiterStash: """ owner_litellm_call_id: str | None = None + call_type: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None parallel_slot_reused_at: datetime | None = None @@ -615,6 +625,11 @@ class RequestRateLimiterStash: batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None batch_tpd_refund_ops: tuple[ReservationOperation, ...] = () borrowing_team_model_limit: RateLimitDescriptor | None = None + borrowing_team_model_reserved_tokens: int = 0 + borrowing_team_model_reserved_window: ReservationWindowIdentity | None = None + borrowing_team_model_fallback_limits: Mapping[str, BorrowingTeamModelHold] = field( + default_factory=lambda: MappingProxyType({}) + ) priority_class_share: RateLimitDescriptor | None = None reservation_released: bool = False io_reservation_released: bool = False @@ -763,6 +778,23 @@ def _is_socket_turn(kwargs: Mapping[str, object]) -> bool: ) +def callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]: + """A logging callback's kwargs as a typed mapping, empty when they are not a mapping.""" + return _callback_kwargs_mapping(kwargs) + + +def request_stash_for_callback(kwargs: Mapping[str, object]) -> RequestRateLimiterStash | None: + """The stash of the request a logging callback settles, or the connection's for a Responses WebSocket turn. + + A nested call (an LLM judge, a silent experiment) mints its own call id and gets None. + """ + call_id: Final = kwargs.get("litellm_call_id") + call_stash: Final = get_request_stash_for_call(call_id if isinstance(call_id, str) else None) + if call_stash is not None: + return call_stash + return get_request_stash() if _is_socket_turn(kwargs) else None + + _AGENT_DESCRIPTOR_KEYS: Final = frozenset({"agent", "agent_session"}) _SOCKET_TURN_UNCOUNTED_DESCRIPTOR_KEYS: Final = frozenset( @@ -1060,6 +1092,37 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if existing_cap is None or effective_cap < existing_cap: data[cap_field] = effective_cap # rebind-ok: downstream routing requires the bounded output cap + def estimate_request_tokens( + self, + data: object, + model: str, + min_configured_tpm_limit: int, + call_type: str | None, + team_id: str | None, + rate_limit_type: str = "total", + ) -> int: + """The reservation this limiter would hold for the request, for another limiter settling the same call.""" + if rate_limit_type == "total": + return self._estimate_tokens_for_request( + data=data, + model=model, + min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, + team_id=team_id, + ) + estimated_input, output_budget = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, + team_id=team_id, + media_model_resolver=self._media_model_resolver, + ) + return output_budget if rate_limit_type == "output" else estimated_input + + def bound_output_for_limit(self, data: object, min_configured_limit: int, call_type: str | None) -> None: + """Cap a request without an explicit output cap the way this limiter does for its own small TPM limits.""" + self._apply_implicit_output_cap(data=data, min_configured_limit=min_configured_limit, call_type=call_type) + def _estimate_tokens_for_request( self, data: object, @@ -3988,6 +4051,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + stash.call_type = call_type reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None if reuses_request_slot: stash.parallel_slot_reused_at = self._get_current_time() 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 index b9f853d7da8..ee0fcbea41b 100644 --- a/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py +++ b/tests/integration/routing/test_team_model_limit_borrows_idle_capacity.py @@ -30,8 +30,10 @@ UPSTREAM_REPLY: Final = json.dumps( def _respond(request: Request) -> Reply: - assert request.method == "POST" and request.target == "/v1/chat/completions", request.target + assert request.method == "POST" and request.target.endswith("/v1/chat/completions"), request.target assert request.headers["authorization"] == "Bearer synthetic-openai-key" + if request.target.startswith("/down/"): + return Reply(status=500, body=b'{"error": {"message": "upstream down"}}') return Reply(body=UPSTREAM_REPLY) @@ -123,6 +125,40 @@ def test_team_is_held_to_its_model_limit_once_the_model_saturates( assert "no idle capacity to borrow" in first.json()["error"]["message"], first.text +def test_a_fallback_target_is_held_to_the_teams_share_on_it_once_it_saturates( + saturated_model: tuple[Gateway, str], +) -> None: + owned, upstream = saturated_model + with owned.scenario() as scenario: + source: Final = scenario.model( + name=f"down-{uuid.uuid4().hex[:8]}", + model=f"openai/{OPENAI_MODEL}", + api_base=f"{upstream}/down/v1", + api_key="synthetic-openai-key", + ) + target: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", + api_base=f"{upstream}/v1", + api_key="synthetic-openai-key", + tpm=MODEL_TPM, + rpm=MODEL_RPM, + ) + owned.post("/config/update", {"router_settings": {"num_retries": 0, "fallbacks": [{source: [target]}]}}) + key: Final = scenario.key(team_id=scenario.team(metadata={"model_rpm_limit": {target: 1}})) + responses: Final = tuple( + owned.request( + "POST", + "/v1/chat/completions", + {"model": source, "messages": [{"role": "user", "content": f"fallback borrowing {uuid.uuid4().hex}"}]}, + key=key, + ) + for _ in range(2) + ) + assert responses[0].status_code == 200, responses[0].text + assert responses[0].headers.get("x-litellm-attempted-fallbacks") == "1", dict(responses[0].headers) + assert responses[1].status_code != 200, responses[1].text + + @pytest.mark.parametrize( "limit, rate_limit_type", [("model_tpm_limit", "tokens"), ("model_rpm_limit", "requests")], ids=["tpm", "rpm"] ) 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 22e6d5e3146..efd2e9e7e18 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1340,7 +1340,7 @@ async def test_async_log_success_event_increments_by_actual_tokens(monkeypatch): } ) - handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + handler.v3_limiter.async_increment_reservation_aware_tokens = mock_increment # Create mock response with 50 completion tokens mock_response = MagicMock(spec=ModelResponse) @@ -1514,7 +1514,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(mon for op in pipeline_operations: incremented_keys.append(op.counter.scope) - handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment + handler.v3_limiter.async_increment_reservation_aware_tokens = mock_increment # Create mock response mock_response = MagicMock(spec=ModelResponse) @@ -2166,14 +2166,21 @@ async def _admit_then_settle(handler: _PROXY_DynamicRateLimitHandlerV3, auth: Us async def _settle_tokens( - handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int, response: object | None = None + handler: PROXY_DynamicRateLimitHandlerV3, + tokens: int, + response: object | None = None, + call_id: str | None = None, + model_group: str = "shared-model", + **extra_kwargs: object, ) -> None: await handler.async_log_success_event( kwargs={ - "litellm_params": {"metadata": {"model_group": "shared-model"}}, + "litellm_params": {"metadata": {"model_group": model_group}}, "standard_logging_object": { "metadata": {"user_api_key_team_id": "team-a", "user_api_key_auth_metadata": {}} }, + **({"litellm_call_id": call_id} if call_id is not None else {}), + **extra_kwargs, }, response_obj=response if response is not None @@ -2197,19 +2204,45 @@ def _team_key( async def _check( - handler: _PROXY_DynamicRateLimitHandlerV3, + handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, saturation: float, model: str = "shared-model", + data: dict[str, object] | None = None, priority: str | None = None, ) -> None: group: Final = handler.llm_router.registry().group(model) assert group is not None - await handler._check_rate_limits( # pyright: ignore[reportPrivateUsage] # saturation is an input here, the public hook reads it from Redis - model=model, model_group_info=group, user_api_key_dict=auth, priority=priority, saturation=saturation + await handler._check_rate_limits( + model=model, + model_group_info=group, + user_api_key_dict=auth, + priority=priority, + saturation=saturation, + data=data or {}, + call_type="acompletion" if data else None, ) +async def _team_tokens_remaining( + handler: PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth | None = None, model: str = "shared-model" +) -> int: + share: Final = handler._borrowing_team_share(model, auth if auth is not None else _team_key()) + assert share is not None + response: Final = await handler.v3_limiter.should_rate_limit( + [share.share], read_only=True, skip_parallel_check=True + ) + return next(s["limit_remaining"] for s in response["statuses"] if s["rate_limit_type"] == "tokens") + + +def _prompt(words: int, max_tokens: int = 10) -> dict[str, object]: + return { + "model": "shared-model", + "messages": [{"role": "user", "content": "word " * words}], + "max_tokens": max_tokens, + } + + 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 @@ -2391,6 +2424,376 @@ async def test_a_websocket_turn_still_borrows_while_the_model_is_idle(monkeypatc await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") +@pytest.mark.asyncio +async def test_a_nested_calls_tokens_are_not_charged_to_the_requests_team_share(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.hooks.parallel_request_limiter_v3 import claim_request_stash_for_data + + handler = _borrowing_handler(monkeypatch, borrows=True) + claim_request_stash_for_data({"litellm_call_id": "outer-call"}) + await _check(handler, _team_key(), saturation=0.0) + await _settle_tokens(handler, 60, call_id="llm-judge-call") + assert await _team_tokens_remaining(handler) == 100 + await _settle_tokens(handler, 60, call_id="outer-call") + assert await _team_tokens_remaining(handler) == 40 + + +@pytest.mark.asyncio +async def test_a_websocket_turns_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=[], + combined_usage_object=litellm.Usage(prompt_tokens=100, completion_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_websocket_turn_gets_a_503_when_the_team_counter_cannot_be_read(monkeypatch: pytest.MonkeyPatch): + from litellm.exceptions import ServiceUnavailableError + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitUnverifiableError + + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + await _admit_then_settle(handler, _team_key(), 150) + monkeypatch.setattr(handler.v3_limiter, "should_rate_limit", AsyncMock(side_effect=RateLimitUnverifiableError())) + with pytest.raises(ServiceUnavailableError): + await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") + + +@pytest.mark.asyncio +async def test_a_prompt_larger_than_the_team_share_is_refused_before_the_call_once_the_model_saturates( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.0, data=_prompt(words=400)) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5, data=_prompt(words=400)) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + + +@pytest.mark.asyncio +async def test_a_reservation_at_saturation_settles_to_the_tokens_actually_used(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.5, data=_prompt(words=30)) + assert await _team_tokens_remaining(handler) < 70 + await _settle_tokens(handler, 5) + assert await _team_tokens_remaining(handler) == 95 + + +@pytest.mark.asyncio +async def test_a_reservation_at_saturation_is_refunded_when_the_call_fails(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.5, data=_prompt(words=30)) + assert await _team_tokens_remaining(handler) < 70 + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": {"model_group": "shared-model"}}}, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler) == 100 + + +@pytest.mark.asyncio +async def test_a_fallback_target_is_refused_once_it_saturates_and_the_teams_share_on_it_is_used( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + key: Final = _team_key( + team_limits={"model_tpm_limit": {"shared-model": 100}, "model_rpm_limit": {"shared-model": 1}} + ) + await _admit_then_settle(handler, key, 150) + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is False + + +@pytest.mark.asyncio +async def test_a_fallback_target_still_borrows_while_idle_and_its_tokens_count_against_that_teams_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", {}) is True + 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_fallback_attempt_counts_toward_the_targets_saturation(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + group: Final = handler.llm_router.registry().group("shared-model") + assert group is not None + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", {}) is True + assert await handler._check_model_saturation("shared-model", group) > 0 + + +@pytest.mark.asyncio +async def test_a_read_only_fallback_check_does_not_use_the_teams_share(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + key: Final = _team_key( + team_limits={"model_tpm_limit": {"shared-model": 100}, "model_rpm_limit": {"shared-model": 2}} + ) + await _admit_then_settle(handler, key, 50) + assert await handler.async_fallback_target_allowed(key, "shared-model", {}, commit=False) is True + assert await handler.async_fallback_target_allowed(key, "shared-model", {}, commit=False) is True + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is True + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is False + + +async def _model_requests_remaining(handler: PROXY_DynamicRateLimitHandlerV3) -> int: + group: Final = handler.llm_router.registry().group("shared-model") + assert group is not None + descriptor: Final = handler._create_model_tracking_descriptor("shared-model", group, high_limit_multiplier=1) + response: Final = await handler.v3_limiter.should_rate_limit([descriptor], read_only=True, skip_parallel_check=True) + return next(s["limit_remaining"] for s in response["statuses"] if s["rate_limit_type"] == "requests") + + +def _fallback_prompt(words: int) -> dict[str, object]: + return {"messages": [{"role": "user", "content": "word " * words}], "max_tokens": 10} + + +@pytest.mark.asyncio +async def test_idle_fallbacks_past_the_teams_rpm_share_still_count_toward_the_targets_saturation( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + key: Final = _team_key(team_limits={"model_rpm_limit": {"shared-model": 1}}) + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is True + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is True + assert await handler.async_fallback_target_allowed(key, "shared-model", {}) is True + assert await _model_requests_remaining(handler) == 97 + + +@pytest.mark.asyncio +async def test_a_saturated_fallback_target_refuses_a_prompt_larger_than_the_teams_remaining_tokens( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", _fallback_prompt(300)) is False + assert ( + await handler.async_fallback_target_allowed(_team_key(), "shared-model", _fallback_prompt(300), commit=False) + is False + ) + assert await _team_tokens_remaining(handler) == 100 + assert ( + await handler.async_fallback_target_allowed(_team_key(), "shared-model", _fallback_prompt(30), commit=False) + is True + ) + assert await _team_tokens_remaining(handler) == 100 + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", _fallback_prompt(30)) is True + assert await _team_tokens_remaining(handler) < 70 + await _settle_tokens(handler, 5) + assert await _team_tokens_remaining(handler) == 95 + + +@pytest.mark.asyncio +async def test_a_saturated_responses_fallback_holds_the_output_budget_as_well_as_the_input( + monkeypatch: pytest.MonkeyPatch, +): + from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash + + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + get_or_create_request_stash().call_type = "aresponses" + prompt: Final = {"input": "hi", "max_output_tokens": 500} + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", prompt) is False + assert await _team_tokens_remaining(handler) == 100 + + +@pytest.mark.asyncio +async def test_an_output_token_limit_holds_only_the_prompts_output_budget(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "token_rate_limit_type", "output") + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + await _check(handler, _team_key(), saturation=1.0, data=_fallback_prompt(300)) + assert await _team_tokens_remaining(handler) == 90 + + +@pytest.mark.asyncio +async def test_a_request_without_max_tokens_leaves_admission_capped_to_the_teams_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + prompt: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}]} + await _check(handler, _team_key(), saturation=1.0, data=prompt) + assert prompt["max_tokens"] == handler.v3_limiter.no_max_tokens_output_floor(100) + + +@pytest.mark.asyncio +async def test_a_prompt_that_fits_once_its_output_is_capped_is_admitted(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + prompt: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "word " * 60}]} + await _check(handler, _team_key(), saturation=1.0, data=prompt) + held: Final = handler.v3_limiter.estimate_request_tokens( + data={**prompt, "model": "shared-model"}, + model="shared-model", + min_configured_tpm_limit=100, + call_type="acompletion", + team_id="team-a", + ) + assert prompt["max_tokens"] == handler.v3_limiter.no_max_tokens_output_floor(100) + assert await _team_tokens_remaining(handler) == 100 - held + + +@pytest.mark.asyncio +async def test_a_refused_admission_leaves_the_requests_output_uncapped(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + prompt: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "word " * 300}]} + with pytest.raises(ProxyRateLimitError): + await _check(handler, _team_key(), saturation=1.0, data=prompt) + assert "max_tokens" not in prompt + + +@pytest.mark.asyncio +async def test_an_input_token_limit_leaves_the_reply_length_alone(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "token_rate_limit_type", "input") + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + prompt: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}]} + await _check(handler, _team_key(), saturation=1.0, data=prompt) + assert "max_tokens" not in prompt + + +@pytest.mark.asyncio +async def test_a_late_failure_from_a_fallback_that_did_not_serve_is_charged_to_that_target( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + key: Final = _team_key(team_limits={"model_tpm_limit": {"shared-model": 100, "other-model": 100, "chat": 100}}) + await _check(handler, key, saturation=1.0, data=_fallback_prompt(30)) + assert await handler.async_fallback_target_allowed(key, "other-model", _fallback_prompt(30)) is True + assert await handler.async_fallback_target_allowed(key, "chat", _fallback_prompt(30)) is True + await _settle_tokens(handler, 7, model_group="chat") + await handler.async_log_failure_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": "other-model"}}, + "combined_usage_object": litellm.Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + }, + response_obj=Exception("stream cut"), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler, key, "shared-model") == 100 + assert await _team_tokens_remaining(handler, key, "chat") == 93 + assert await _team_tokens_remaining(handler, key, "other-model") == 75 + + +@pytest.mark.asyncio +async def test_a_served_fallbacks_hold_is_not_refunded_again_under_the_name_the_request_used( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler( + monkeypatch, borrows=True, saturation_threshold=0.0, model_group_alias={"alias-model": "other-model"} + ) + key: Final = _team_key(team_limits={"model_tpm_limit": {"shared-model": 100, "other-model": 100}}) + await _check(handler, key, saturation=1.0, data=_fallback_prompt(30)) + assert await handler.async_fallback_target_allowed(key, "alias-model", _fallback_prompt(30)) is True + await _settle_tokens(handler, 7, model_group="other-model") + assert await _team_tokens_remaining(handler, key, "other-model") == 93 + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": {"model_group": "alias-model"}}}, + response_obj=Exception("late"), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler, key, "other-model") == 93 + + +@pytest.mark.asyncio +async def test_a_steered_requests_images_are_counted_for_the_model_that_serves_it(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + asked: Final = [] # mutable-ok: records the models whose image cost was looked up + + def resolver(model: str, team_id: str | None) -> tuple[str, ...]: + asked.append(model) + return (model,) + + handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( + handler.internal_usage_cache, media_model_resolver=resolver + ) + key: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 10000}}) + pixel: Final = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + steered: Final[dict[str, object]] = { + "model": "shared-model", + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": pixel}}]}], + "max_tokens": 5, + } + await _check(handler, key, saturation=1.0, model="other-model", data=steered) + assert asked and set(asked) == {"other-model"} + + +@pytest.mark.asyncio +async def test_a_failed_attempt_that_recorded_usage_is_charged_that_usage_instead_of_refunded( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + await _check(handler, _team_key(), saturation=1.0, data=_fallback_prompt(30)) + assert await _team_tokens_remaining(handler) < 70 + await handler.async_log_failure_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": "shared-model"}}, + "combined_usage_object": litellm.Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25), + }, + response_obj=Exception("stream cut"), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler) == 75 + + +@pytest.mark.asyncio +async def test_a_fallback_attempt_that_fails_gives_back_the_tokens_it_held(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) + assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", _fallback_prompt(30)) is True + assert await _team_tokens_remaining(handler) < 70 + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": {"model_group": "shared-model"}}}, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler) == 100 + + +@pytest.mark.asyncio +async def test_a_prompt_the_token_estimator_cannot_read_is_still_held_to_the_teams_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + key: Final = _team_key( + team_limits={"model_tpm_limit": {"shared-model": 100}, "model_rpm_limit": {"shared-model": 1}} + ) + unreadable: Final = { + "model": "shared-model", + "messages": [{"role": "user", "content": [{"type": "text", "text": 1}, {"type": "text", "text": "hi"}]}], + } + await _check(handler, key, saturation=0.5, data=unreadable) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, key, saturation=0.5, data=unreadable) + _held_to_team_limit(exc.value, "requests", "model_rpm_limit") + + +@pytest.mark.asyncio +async def test_a_success_without_a_model_group_still_releases_the_reservation(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.5, data=_prompt(words=30)) + assert await _team_tokens_remaining(handler) < 70 + await handler.async_log_success_event( + kwargs={"litellm_params": {"metadata": {}}}, + response_obj=litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=5, completion_tokens=0, total_tokens=5)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert await _team_tokens_remaining(handler) == 95 + + def _class_team(team_id: str, priority: str | None = "interactive") -> UserAPIKeyAuth: return UserAPIKeyAuth(api_key="sk-class", team_id=team_id, team_metadata={}, metadata={"priority": priority})