From 52d1c4cd80162fbf03d7f92def0cb1631f1edfd4 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 01:10:23 -0800 Subject: [PATCH] feat(dynamic_rate_limiter_v3): steer flexible teams to a permitted alternate when their preferred model is saturated (#51350) * feat(router): steer flexible teams to a permitted alternate model when their preferred model is saturated Teams opt in with metadata.model_preferences: [{"model": "luna", "alternates": ["opus"]}]. The dynamic rate limiter v3 pre-call hook reads the saturation cache for the preferred model and, once it is at or above saturation_threshold, rewrites the request to the first alternate the team may call that is below the threshold. The response carries x-litellm-model-steered-from with the preferred model. Rate limits and spend are charged to the model that served. Teams without model_preferences are untouched. * fix(dynamic_rate_limiter_v3): type the hook request dicts so the steering rewrite passes the lint gates * fix(dynamic_rate_limiter_v3): authorize alternates against the key and its model budget, send the steering header on streams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): enforce member and project allowlists on alternates, skip the steering header for non Latin-1 names * fix(dynamic_rate_limiter_v3): check every budget layer for an alternate, fall back to the preferred model when the alternate refuses, merge the served model's guardrails * fix(dynamic_rate_limiter_v3): never steer agent keys, type the steering unit tests * fix(dynamic_rate_limiter_v3): hold alternates to customer grants, access group budgets, guardrails and the zero-cost waiver An alternate now passes every check a budget fallback passes. Customer model grants are enforced, the access groups that grant the alternate are budget checked and stamped on the token and request metadata so spend charges them, a zero-cost preferred model that waived the aggregate budget checks is never steered to a priced alternate, and an alternate whose deployments add model-level guardrails the request did not run is skipped. * fix(dynamic_rate_limiter_v3): stamp the alternate's groups on every bucket and skip alternates the request cannot safely take The served model's access groups now land on every metadata bucket and the auth copy inside it even when the preferred model had none, so group spend is charged. An alternate is skipped when a budget reservation already holds a different group, when it does not serve the request's endpoint, when unpriced models are blocked and it has no price, or when its guardrails differ from the request's by name or settings. A failing alternate check or alternate admission leaves the request on the preferred model instead of failing it. * fix(dynamic_rate_limiter_v3): keep a reserved budget off pricier alternates While the request holds a budget reservation sized for the preferred model, an alternate that costs more per token is skipped, so fail-closed budget enforcement never admits a call the reservation does not cover. * fix(dynamic_rate_limiter_v3): an unpriced alternate never takes a reserved budget While a budget reservation is held, an alternate with no input or output price is skipped rather than treated as free. --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 472 ++++++++++++++++-- .../hooks/parallel_request_limiter_v3.py | 3 +- .../test_team_model_preference_steering.py | 144 ++++++ .../hooks/test_dynamic_rate_limiter_v3.py | 285 +++++++++++ 4 files changed, 872 insertions(+), 32 deletions(-) create mode 100644 tests/integration/routing/test_team_model_preference_steering.py diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 3654d0b9176..bd20fbf3ab3 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -2,14 +2,15 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting """ +import json import os -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import datetime -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from fastapi import HTTPException -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm import Router @@ -22,6 +23,7 @@ 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.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, @@ -40,12 +42,13 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l borrowing_team_model_limits, claim_request_stash_for_data, get_or_create_request_stash, + get_request_stash_for_call, ) from litellm.proxy.hooks.rate_limiter_utils import ( convert_priority_to_percent, resolve_llm_provider_for_rate_limit, ) -from litellm.proxy.utils import InternalUsageCache +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 from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, Usage @@ -54,6 +57,270 @@ if TYPE_CHECKING: from litellm.types.utils import PriorityReservationSettings +STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from" + + +@dataclass(frozen=True, slots=True) +class AlternateGrant: + """What an authorized alternate carries onto the request: the budgeted access groups that granted it.""" + + access_groups: tuple[str, ...] = () + + +AlternateAuthorizer: TypeAlias = Callable[[UserAPIKeyAuth, str], Awaitable[AlternateGrant | None]] + + +class ModelPreference(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + model: str + alternates: tuple[str, ...] = () + + +MODEL_PREFERENCES: Final = TypeAdapter(tuple[ModelPreference, ...]) + + +@dataclass(frozen=True, slots=True) +class SteeredModel: + model: str + group_info: ModelGroupInfo + access_groups: tuple[str, ...] = () + + +async def alternate_grant_for(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None: + """The grant a budget fallback to ``model`` would get: the key, team, member, project and customer allowlists, + every model budget layer and the budgets of the access groups that authorize it, else None.""" + from litellm.exceptions import BudgetExceededError + from litellm.proxy.auth.auth_checks import ( + _model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract + can_customer_access_model, + collect_matched_model_access_groups, + get_end_user_object, + get_project_object, + get_team_membership, + get_team_object, + ) + from litellm.proxy.auth.user_api_key_auth import ( + _is_fallback_model_authorized, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract + _is_within_every_model_budget, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract + ) + from litellm.proxy.proxy_server import ( + llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries + llm_router, + model_max_budget_limiter, + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + user_id: Final = user_api_key_dict.user_id + team_id: Final = user_api_key_dict.team_id + project_id: Final = user_api_key_dict.project_id + end_user_id: Final = user_api_key_dict.end_user_id + try: + team_object: Final = ( + await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if team_id is not None + else None + ) + team_membership: Final = ( + await get_team_membership( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if user_id is not None and team_id is not None + else None + ) + project_object: Final = ( + await get_project_object( + project_id=project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if project_id is not None + else None + ) + end_user_object: Final = ( + await get_end_user_object( + end_user_id=end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if end_user_id is not None + else None + ) + if end_user_object is not None and end_user_object.models: + can_customer_access_model( + model=model, end_user_object=end_user_object, llm_router=llm_router, valid_token=user_api_key_dict + ) + except Exception: # noqa: BLE001 # a failed team, member, project or customer lookup must not authorize an alternate + return None + if not await _is_fallback_model_authorized( + model=model, + valid_token=user_api_key_dict, + llm_model_list=llm_model_list, # pyright: ignore[reportUnknownArgumentType] # proxy global has untyped deployment entries + llm_router=llm_router, + prisma_client=prisma_client, + team_membership=team_membership, + project_object=project_object, + refuse_agent_keys=True, + team_object=team_object, + ): + return None + if not await _is_within_every_model_budget( + valid_token=user_api_key_dict, model_max_budget_limiter=model_max_budget_limiter, model=model + ): + return None + try: + access_groups: Final = await collect_matched_model_access_groups( + model=model, + valid_token=user_api_key_dict, + team_object=team_object, + project_object=project_object, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + team_membership=team_membership, + team_membership_loaded=True, + ) + await _model_access_group_max_budget_check( + matched_model_access_groups=access_groups, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + except BudgetExceededError: + return None + except Exception: # noqa: BLE001 # an unreadable group budget must not authorize an alternate + return None + return AlternateGrant(access_groups=access_groups) + + +REQUEST_BUCKET: Final = TypeAdapter(dict[str, object]) +REQUEST_ENTRIES: Final = TypeAdapter(tuple[object, ...]) +RESERVATION_ENTRIES: Final = TypeAdapter(tuple[Mapping[str, object], ...]) +MODEL_ACCESS_GROUP_ENTITY: Final = "Model access group" + + +def _metadata_bucket(data: Mapping[str, object], key: str) -> Mapping[str, object] | None: + """The request's ``key`` metadata bucket as a typed dict, or None when it is absent or not a mapping.""" + value: Final = data.get(key) + if not isinstance(value, Mapping): + return None + try: + return REQUEST_BUCKET.validate_python(value) + except ValidationError: + return None + + +def _guardrail_entries(data: Mapping[str, object]) -> frozenset[str]: + """The request's guardrail entries, settings included, in a form two lists can be compared by.""" + metadata: Final = _metadata_bucket(data, "metadata") + guardrails: Final = metadata.get("guardrails") if metadata is not None else None + entries: Final = ( + REQUEST_ENTRIES.validate_python(guardrails) + if isinstance(guardrails, list) + else (guardrails,) + if guardrails + else () + ) + return frozenset( + entry if isinstance(entry, str) else json.dumps(entry, sort_keys=True, default=str) for entry in entries + ) + + +def _auth_copy_with_groups(bucket: Mapping[str, object], access_groups: tuple[str, ...]) -> dict[str, object]: + auth: Final = bucket.get("user_api_key_auth") + if isinstance(auth, UserAPIKeyAuth): + return {"user_api_key_auth": auth.model_copy(update={"matched_model_access_groups": list(access_groups)})} + if isinstance(auth, Mapping): + return { + "user_api_key_auth": { + **REQUEST_BUCKET.validate_python(auth), + "matched_model_access_groups": list(access_groups), + } + } + return {} + + +def _with_access_groups(data: Mapping[str, object], access_groups: tuple[str, ...]) -> dict[str, object]: + """``data`` with the served model's access groups stamped on every metadata bucket, the auth copy in it + included, so spend attribution charges the groups that granted the served model and nothing else.""" + stamped: Final = { + key: { + **bucket, + MODEL_ACCESS_GROUP_METADATA_KEY: list(access_groups), + **_auth_copy_with_groups(bucket, access_groups), + } + for key in ("metadata", "litellm_metadata") + if (bucket := _metadata_bucket(data, key)) is not None + } + return {**data, **stamped} + + +def _priced_within(preferred: ModelGroupInfo, alternate: ModelGroupInfo) -> bool: + """Whether ``alternate`` costs no more per token than ``preferred``, so a budget reserved for the preferred + model's estimate still covers the alternate.""" + if alternate.input_cost_per_token is None or alternate.output_cost_per_token is None: + return False + return alternate.input_cost_per_token <= (preferred.input_cost_per_token or 0.0) and ( + alternate.output_cost_per_token <= (preferred.output_cost_per_token or 0.0) + ) + + +def _reservation_covers( + user_api_key_dict: UserAPIKeyAuth, + access_groups: tuple[str, ...], + preferred: ModelGroupInfo, + alternate: ModelGroupInfo, +) -> bool: + """Whether the budget the request already reserved still fits the alternate: the alternate costs no more per + token than the preferred model the estimate was made for, and every model access group reserved against also + grants the alternate, so settlement never charges a group that did not serve.""" + reservation: Final = user_api_key_dict.budget_reservation + if reservation is None: + return True + if not _priced_within(preferred, alternate): + return False + entries: Final = reservation.get("entries") + if not isinstance(entries, list): + return True + try: + reserved: Final = frozenset( + entry.get("entity_id") + for entry in RESERVATION_ENTRIES.validate_python(entries) + if entry.get("entity_type") == MODEL_ACCESS_GROUP_ENTITY + ) + except ValidationError: + return False + return reserved <= frozenset(access_groups) + + +def alternates_for(user_api_key_dict: UserAPIKeyAuth, model: str) -> tuple[str, ...]: + """The team's opt-in alternates for ``model`` from ``team_metadata.model_preferences``, or none.""" + raw: Final = (user_api_key_dict.team_metadata or {}).get("model_preferences") + if raw is None: + return () + try: + preferences: Final = MODEL_PREFERENCES.validate_python(raw) + except ValidationError as error: + verbose_proxy_logger.warning( + "Ignoring malformed team model_preferences for team %s: %s", user_api_key_dict.team_id, error + ) + return () + return next((preference.alternates for preference in preferences if preference.model == model), ()) + + class _SettlementMetadata(BaseModel): """The two fields of a request's standard logging metadata that token settlement reads.""" @@ -155,8 +422,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): self, internal_usage_cache: DualCache, time_provider: Callable[[], datetime] | None = None, + is_model_authorized: AlternateAuthorizer = alternate_grant_for, ): self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) + self._is_model_authorized = is_model_authorized self.v3_limiter = PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider) def update_variables(self, llm_router: Router) -> None: @@ -781,7 +1050,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, - data: dict, + data: Mapping[str, object], call_type: CallTypesLiteral, ) -> Exception | str | dict | None: """ @@ -823,38 +1092,20 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): claim_request_stash_for_data(data) model: Final = data["model"] - priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) + if not isinstance(model, str): + return None # Get model configuration - model_group_info: Final[ModelGroupInfo | None] = self.llm_router.registry().group(model) - if model_group_info is None: + requested_group_info: Final[ModelGroupInfo | None] = self.llm_router.registry().group(model) + if requested_group_info is None: verbose_proxy_logger.debug("No model group info for %s, allowing request", model) return None + steered: Final = await self._steer_to_alternate(model, requested_group_info, user_api_key_dict, data, call_type) try: - # STEP 1: Check current saturation level - saturation: Final = await self._check_model_saturation(model, model_group_info) - - saturation_threshold: Final = _get_priority_settings().saturation_threshold - - verbose_proxy_logger.debug( - f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " - f"Threshold={saturation_threshold:.1%}, Priority={priority}" + admitted: Final = await self._admit_steered_or_preferred( + model, requested_group_info, steered, user_api_key_dict ) - - # STEP 2: Check rate limits in THREE phases - # Phase 1: Read-only check of ALL limits (no increments) - # Phase 2: Decide which limits to enforce (based on saturation) - # Phase 3: Increment ALL counters only if request will be allowed - # This prevents partial increments and ensures accurate tracking - await self._check_rate_limits( - model=model, - model_group_info=model_group_info, - user_api_key_dict=user_api_key_dict, - priority=priority, - saturation=saturation, - ) - except HTTPException: raise except Exception as e: @@ -862,11 +1113,155 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # Fail open on unexpected errors return None + if admitted is None: + return None + get_or_create_request_stash().steered_from_model = model + groups: Final = list(admitted.access_groups) + user_api_key_dict.matched_model_access_groups = groups # rebind-ok: spend charges the served model's groups + return _with_access_groups(dict(data, model=admitted.model), admitted.access_groups) + + async def _admit_steered_or_preferred( + self, + model: str, + requested_group_info: ModelGroupInfo, + steered: "SteeredModel | None", + user_api_key_dict: UserAPIKeyAuth, + ) -> "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) + 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) return None + async def _admit(self, model: str, model_group_info: ModelGroupInfo, user_api_key_dict: UserAPIKeyAuth) -> 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) + saturation_threshold: Final = _get_priority_settings().saturation_threshold + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " + f"Threshold={saturation_threshold:.1%}, Priority={priority}" + ) + await self._check_rate_limits( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user_api_key_dict, + priority=priority, + saturation=saturation, + ) + + async def _steer_to_alternate( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: CallTypesLiteral, + ) -> "SteeredModel | None": + """First team alternate below ``saturation_threshold`` once ``model`` is at or above it, else None. A + broken alternate check is logged and leaves the request on ``model``, as the limiter fails open elsewhere.""" + try: + return await self._pick_alternate(model, model_group_info, user_api_key_dict, data, call_type) + except Exception: # noqa: BLE001 # a failing alternate check must not fail a request that can stay put + verbose_proxy_logger.warning( + "Alternate check failed for team %s, staying on %s", user_api_key_dict.team_id, model, exc_info=True + ) + return None + + async def _pick_alternate( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: CallTypesLiteral, + ) -> "SteeredModel | None": + alternates: Final = alternates_for(user_api_key_dict, model) + if not alternates: + return None + threshold: Final = _get_priority_settings().saturation_threshold + if await self._check_model_saturation(model, model_group_info) < threshold: + return None + for alternate in alternates: + if alternate == model or not self._steerable(model, alternate, data, call_type, user_api_key_dict.team_id): + continue + if ( + candidate := await self._idle_permitted_alternate( + alternate, user_api_key_dict, threshold, model_group_info + ) + ) is not None: + verbose_proxy_logger.info("Steering team %s from %s to %s", user_api_key_dict.team_id, model, alternate) + return candidate + return None + + def _steerable( + self, model: str, alternate: str, data: Mapping[str, object], call_type: CallTypesLiteral, team_id: str | None + ) -> bool: + """Whether ``alternate`` may serve a request that already passed auth and the guardrail walk as ``model``: a + zero-cost ``model`` waived the aggregate budget checks, so only another zero-cost alternate may take it, an + unpriced alternate is blocked when unpriced models are, the alternate has to serve this endpoint, and an + alternate with a model-level guardrail the request has not run would reach it unchecked.""" + from litellm.proxy.auth.auth_checks import ( + _unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract + is_model_cost_zero, + ) + from litellm.proxy.route_llm_request import endpoint_unsupported_by_model + + if is_model_cost_zero(model=model, llm_router=self.llm_router) and not is_model_cost_zero( + model=alternate, llm_router=self.llm_router + ): + verbose_proxy_logger.info("Not steering zero-cost %s to priced %s", model, alternate) + return False + if litellm.block_requests_for_models_without_pricing and _unpriced_models_in_request( + model=alternate, llm_router=self.llm_router + ): + verbose_proxy_logger.info("Not steering %s to unpriced %s", model, alternate) + return False + if endpoint_unsupported_by_model(call_type, dict(data, model=alternate), self.llm_router, team_id): + verbose_proxy_logger.info( + "Not steering %s to %s, the alternate does not serve %s", model, alternate, call_type + ) + return False + probe: Final = REQUEST_BUCKET.validate_python( + check_and_merge_model_level_guardrails( + data=dict(data, model=alternate, metadata=_metadata_bucket(data, "metadata") or {}), + llm_router=self.llm_router, + trust_client_model_info=False, + ) + ) + if _guardrail_entries(probe) != _guardrail_entries(data): + verbose_proxy_logger.info( + "Not steering %s to %s, the alternate adds model-level guardrails", model, alternate + ) + return False + return True + + async def _idle_permitted_alternate( + self, alternate: str, user_api_key_dict: UserAPIKeyAuth, threshold: float, preferred: ModelGroupInfo + ) -> "SteeredModel | None": + group_info: Final = self.llm_router.registry().group(alternate) + if group_info is None: + return None + grant: Final = await self._is_model_authorized(user_api_key_dict, alternate) + if grant is None or not _reservation_covers(user_api_key_dict, grant.access_groups, preferred, group_info): + return None + if await self._check_model_saturation(alternate, group_info) >= threshold: + return None + return SteeredModel(model=alternate, group_info=group_info, access_groups=grant.access_groups) + @with_service_target(RATE_LIMITS_TARGET) async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response + self, + data: dict[str, object], # mutable-ok: the CustomLogger contract hands the v3 limiter a mutable dict + user_api_key_dict: UserAPIKeyAuth, + response, ) -> LLMResponseTypes: """ Post-call hook to add rate limit headers to response. @@ -896,6 +1291,21 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): verbose_proxy_logger.exception("Error in dynamic rate limiter v3 post-call hook: %s", e) return response + async def async_post_call_response_headers_hook( + self, + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + response: object, + request_headers: Mapping[str, str] | None = None, + litellm_call_info: Mapping[str, object] | None = None, + ) -> dict[str, str] | None: # mutable-ok: the CustomLogger contract returns a dict + """Names the model the caller asked for when steering served another, on streamed and buffered replies.""" + call_id: Final = data.get("litellm_call_id") + stash: Final = get_request_stash_for_call(call_id if isinstance(call_id, str) else None) + if stash is None or stash.steered_from_model is None or not _is_latin1_encodable(stash.steered_from_model): + return None + 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: """ diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6b53e952f98..4cbe08d0042 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -602,6 +602,7 @@ class RequestRateLimiterStash: reserved_model: RateLimitedModel | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) fallback_models: frozenset[str] = field(default_factory=frozenset) + steered_from_model: str | None = None ptu_ceiling: PTUTeamCeiling | None = None ptu_reserved_tokens: int = 0 ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset) @@ -706,7 +707,7 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash -def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: +def claim_request_stash_for_data(data: Mapping[str, object]) -> RequestRateLimiterStash: stash: Final = get_or_create_request_stash() owner_call_id: Final = data.get("litellm_call_id") if isinstance(owner_call_id, str): diff --git a/tests/integration/routing/test_team_model_preference_steering.py b/tests/integration/routing/test_team_model_preference_steering.py new file mode 100644 index 00000000000..c99cfcf660b --- /dev/null +++ b/tests/integration/routing/test_team_model_preference_steering.py @@ -0,0 +1,144 @@ +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, Scenario, eventually, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +OPENAI_MODEL: Final = "gpt-5.6" +PREFERRED_RPM: Final = 2 +ALTERNATE_RPM: Final = 1000 +SATURATION_THRESHOLD: Final = 0.5 +STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from" + + +def _reply() -> bytes: + return json.dumps( + { + "id": f"chatcmpl_{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "steered"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30}, + } + ).encode() + + +@pytest.fixture(scope="module") +def steering_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp("model-preference-steering") + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration: Final = { + **base, + "litellm_settings": { + **base["litellm_settings"], + "callbacks": ["dynamic_rate_limiter_v3"], + "priority_reservation_settings": { + "saturation_threshold": SATURATION_THRESHOLD, + "saturation_check_cache_ttl": 1, + }, + }, + } + path: Final = directory / "model_preference_steering.yaml" + path.write_text(yaml.safe_dump(configuration)) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions", request.target + return Reply(body=_reply()) + + with ( + gateway_from_environment() as rig, + wire_server(respond) as wire, + owned_proxy(rig, directory, {}, config=path) as owned, + ): + yield owned, wire.url + + +def _deployment(owned: Gateway, scenario: Scenario, upstream: str, rpm: int) -> tuple[str, str]: + name: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-key", rpm=rpm + ) + rows: Final = owned.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_name = %s', (name,)) + return name, str(rows[0]["model_id"]) + + +def _chat(owned: Gateway, key: str, model: str) -> httpx.Response: + return owned.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"model preference steering {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _spend_row(owned: Gateway, response: httpx.Response) -> list[dict[str, object]]: + return owned.rows( + 'SELECT model_group FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (str(response.json()["id"]),) + ) + + +def test_flexible_team_is_steered_to_its_permitted_alternate_once_the_preferred_model_saturates( + steering_proxy: tuple[Gateway, str], +) -> None: + owned, upstream = steering_proxy + with owned.scenario() as scenario: + preferred, preferred_id = _deployment(owned, scenario, upstream, PREFERRED_RPM) + alternate, alternate_id = _deployment(owned, scenario, upstream, ALTERNATE_RPM) + rigid_key: Final = scenario.key(team_id=scenario.team(models=[preferred, alternate])) + flexible_key: Final = scenario.key( + team_id=scenario.team( + models=[preferred, alternate], + metadata={"model_preferences": [{"model": preferred, "alternates": [alternate]}]}, + ) + ) + + warmup: Final = _chat(owned, rigid_key, preferred) + assert warmup.status_code == 200, warmup.text + assert warmup.headers["x-litellm-model-id"] == preferred_id + assert STEERED_FROM_HEADER not in warmup.headers + + steered: Final = eventually( + lambda: _chat(owned, flexible_key, preferred), + lambda response: STEERED_FROM_HEADER in response.headers, + seconds=15, + return_last_on_timeout=True, + ) + assert steered.status_code == 200, steered.text + assert steered.headers[STEERED_FROM_HEADER] == preferred + assert steered.headers["x-litellm-model-id"] == alternate_id + + rigid: Final = _chat(owned, rigid_key, preferred) + assert rigid.status_code == 200, rigid.text + assert rigid.headers["x-litellm-model-id"] == preferred_id + assert STEERED_FROM_HEADER not in rigid.headers + + charged: Final = eventually(lambda: _spend_row(owned, steered), lambda rows: len(rows) == 1, seconds=30) + assert charged[0]["model_group"] == alternate + + +def test_alternate_the_team_may_not_call_is_never_used(steering_proxy: tuple[Gateway, str]) -> None: + owned, upstream = steering_proxy + with owned.scenario() as scenario: + preferred, _ = _deployment(owned, scenario, upstream, PREFERRED_RPM) + forbidden, _ = _deployment(owned, scenario, upstream, ALTERNATE_RPM) + key: Final = scenario.key( + team_id=scenario.team( + models=[preferred], + metadata={"model_preferences": [{"model": preferred, "alternates": [forbidden]}]}, + ) + ) + + first: Final = _chat(owned, key, preferred) + assert first.status_code == 200, first.text + second: Final = eventually( + lambda: _chat(owned, key, preferred), lambda response: response.status_code != 200, seconds=15 + ) + assert second.status_code == 429, second.text + assert STEERED_FROM_HEADER not in second.headers 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 a8265680856..0bf0ea4d6ad 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -6,7 +6,9 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of import asyncio import time +from collections.abc import Mapping from datetime import datetime, timedelta +from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, patch @@ -14,9 +16,11 @@ import pytest import litellm from litellm import DualCache, Router +from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY 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 ( + AlternateGrant, PROXY_DynamicRateLimitHandlerV3, ) from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( @@ -1824,6 +1828,287 @@ async def test_post_call_success_hook_priority_header_is_always_http_encodable(t assert http_response.headers["x-litellm-rate-limiter-version"] == "v3" +def _steering_router( + preferred: str, + alternate: str, + preferred_rpm: int = 2, + alternate_rpm: int = 1000, + extra_params: Mapping[str, Mapping[str, object]] = MappingProxyType({}), + model_info: Mapping[str, Mapping[str, object]] = MappingProxyType({}), +) -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": rpm, + **extra_params.get(name, {}), + }, + "model_info": dict(model_info.get(name, {})), + } + for name, rpm in ((preferred, preferred_rpm), (alternate, alternate_rpm)) + ] + ) + + +def _team_user(team_models: list[str], preferences: object) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + team_id="flexible-team", team_models=team_models, team_metadata={"model_preferences": preferences} + ) + + +async def _team_allowlist(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None: + return AlternateGrant() if model in user_api_key_dict.team_models else None + + +async def _pooled_team_allowlist(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None: + return AlternateGrant(access_groups=(f"{model}-pool",)) if model in user_api_key_dict.team_models else None + + +def _steering_handler() -> DynamicRateLimitHandler: + return DynamicRateLimitHandler(internal_usage_cache=DualCache(), is_model_authorized=_team_allowlist) + + +async def _saturate(handler: DynamicRateLimitHandler, model: str, requests: int = 1) -> None: + for _ in range(requests): + assert await handler.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), {"model": model}, "completion") is None + + +@pytest.mark.asyncio +async def test_flexible_team_is_steered_to_its_alternate_once_the_preferred_model_saturates( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + handler.update_variables(llm_router=_steering_router("luna", "opus")) + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + + await _saturate(handler, "luna") + steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna", "n": 1}, "completion") + assert steered == {"model": "opus", "n": 1} + + headers: Final = await handler.async_post_call_response_headers_hook( + data={"model": "opus", "stream": True}, user_api_key_dict=flexible, response=None + ) + assert headers == {"x-litellm-model-steered-from": "luna"} + + from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash + + get_or_create_request_stash().steered_from_model = "luna\u2728" + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None + ) + + +@pytest.mark.asyncio +async def test_teams_without_a_permitted_idle_alternate_stay_on_the_preferred_model(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + handler.update_variables(llm_router=_steering_router("luna", "opus", preferred_rpm=10, alternate_rpm=1)) + await _saturate(handler, "luna", requests=5) + await _saturate(handler, "opus") + + forbidden: Final = _team_user(["luna"], [{"model": "luna", "alternates": ["opus"]}]) + unknown: Final = _team_user([], [{"model": "luna", "alternates": ["missing-model"]}]) + malformed: Final = _team_user([], {"luna": "opus"}) + both_saturated: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + for team in (forbidden, unknown, malformed, both_saturated): + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "luna"}, "completion") is None + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=team, response=None) is None + ) + + +@pytest.mark.asyncio +async def test_a_zero_cost_preferred_model_is_never_steered_to_a_priced_alternate(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + free_luna: Final = {"luna": {"input_cost_per_token": 0, "output_cost_per_token": 0}} + handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=free_luna)) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None + ) + + +@pytest.mark.asyncio +async def test_an_alternate_that_adds_model_level_guardrails_is_never_used(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + guarded_opus: Final = {"opus": {"guardrails": ["opus-only-input-scan"]}} + handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=guarded_opus)) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None + ) + + already_scanned: Final = {"model": "luna", "metadata": {"guardrails": ["opus-only-input-scan"]}} + steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), already_scanned, "completion") + assert isinstance(steered, dict) + assert steered["model"] == "opus" + assert steered["metadata"]["guardrails"] == ["opus-only-input-scan"] + + configured: Final = _steering_handler() + configured_opus: Final = {"opus": {"guardrails": [{"opus-only-input-scan": {"extra_body": {"strict": True}}}]}} + configured.update_variables(llm_router=_steering_router("luna", "opus", extra_params=configured_opus)) + await _saturate(configured, "luna") + plain_scan: Final = {"model": "luna", "metadata": {"guardrails": ["opus-only-input-scan"]}} + assert await configured.async_pre_call_hook(flexible, DualCache(), plain_scan, "completion") is None + + +@pytest.mark.asyncio +async def test_a_steered_request_is_charged_to_the_groups_that_granted_the_alternate(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = DynamicRateLimitHandler( + internal_usage_cache=DualCache(), is_model_authorized=_pooled_team_allowlist + ) + handler.update_variables(llm_router=_steering_router("luna", "opus")) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + auth_copy: Final = flexible.model_copy() + data: Final = {"model": "luna", "metadata": {"user_api_key_auth": auth_copy, "tags": ["t"]}} + steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), data, "completion") + assert isinstance(steered, dict) + assert steered["model"] == "opus" + assert steered["metadata"]["tags"] == ["t"] + assert steered["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] == ["opus-pool"] + assert steered["metadata"]["user_api_key_auth"].matched_model_access_groups == ["opus-pool"] + assert flexible.matched_model_access_groups == ["opus-pool"] + assert auth_copy.matched_model_access_groups is None + + +@pytest.mark.asyncio +async def test_an_alternate_granted_by_a_group_the_request_did_not_reserve_against_is_never_used( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = DynamicRateLimitHandler( + internal_usage_cache=DualCache(), is_model_authorized=_pooled_team_allowlist + ) + handler.update_variables(llm_router=_steering_router("luna", "opus")) + await _saturate(handler, "luna") + + def reserved_against(group: str) -> UserAPIKeyAuth: + user: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + user.budget_reservation = { + "entries": [ + { + "counter_key": f"spend:model_access_group:{group}", + "entity_type": "Model access group", + "entity_id": group, + "reserved_cost": 0.01, + "applied_adjustment": 0.0, + } + ] + } + return user + + assert ( + await handler.async_pre_call_hook(reserved_against("luna-pool"), DualCache(), {"model": "luna"}, "completion") + is None + ) + steered: Final = await handler.async_pre_call_hook( + reserved_against("opus-pool"), DualCache(), {"model": "luna"}, "completion" + ) + assert isinstance(steered, dict) + assert steered["model"] == "opus" + + +@pytest.mark.asyncio +async def test_a_pricier_alternate_is_never_used_while_a_budget_reservation_is_held(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + pricier_opus: Final = { + "luna": {"input_cost_per_token": 0.000001, "output_cost_per_token": 0.000002}, + "opus": {"input_cost_per_token": 0.000005, "output_cost_per_token": 0.000002}, + } + handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=pricier_opus)) + await _saturate(handler, "luna") + + reserved: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + reserved.budget_reservation = {"reserved_cost": 0.0001, "entries": []} + assert await handler.async_pre_call_hook(reserved, DualCache(), {"model": "luna"}, "completion") is None + + unreserved: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + steered: Final = await handler.async_pre_call_hook(unreserved, DualCache(), {"model": "luna"}, "completion") + assert isinstance(steered, dict) + assert steered["model"] == "opus" + + unpriced: Final = _steering_handler() + unpriced_opus: Final = {"opus": {"model": "openai/no-such-model-xyz"}} + unpriced.update_variables(llm_router=_steering_router("luna", "opus", extra_params=unpriced_opus)) + await _saturate(unpriced, "luna") + assert await unpriced.async_pre_call_hook(reserved, DualCache(), {"model": "luna"}, "completion") is None + + +@pytest.mark.asyncio +async def test_an_alternate_that_does_not_serve_the_endpoint_is_never_used(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + handler: Final = _steering_handler() + embeddings_only: Final = {"opus": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}} + handler.update_variables(llm_router=_steering_router("luna", "opus", model_info=embeddings_only)) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "acompletion") is None + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None + ) + + +@pytest.mark.asyncio +async def test_an_unpriced_alternate_is_never_used_while_unpriced_models_are_blocked(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + handler: Final = _steering_handler() + unpriced_opus: Final = {"opus": {"model": "openai/no-such-model-xyz"}} + handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=unpriced_opus)) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + + +@pytest.mark.asyncio +async def test_a_broken_alternate_check_keeps_the_request_on_the_preferred_model(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + + async def broken(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None: + raise RuntimeError("membership lookup is down") + + handler: Final = DynamicRateLimitHandler(internal_usage_cache=DualCache(), is_model_authorized=broken) + handler.update_variables(llm_router=_steering_router("luna", "opus")) + await _saturate(handler, "luna") + + flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}]) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + assert ( + await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None + ) + + def _borrowing_handler( monkeypatch: pytest.MonkeyPatch, borrows: bool,