feat(proxy): router_settings.fallback_on_saturation walks key and team fallbacks before a model saturates, replacing team model_preferences (#51439)

* feat(proxy): fallback_on_saturation walks key and team fallbacks before a model saturates, replacing team_metadata.model_preferences

Team router_settings.fallbacks now fire on proxy-side rate limits: the pre-call
fallback walk reads router settings Key > Team > Global instead of the key's
column alone, and skips a fallback the caller may not call.

A key or team with router_settings.fallback_on_saturation: true walks that same
fallback list as soon as the requested model reaches saturation_threshold in
dynamic_rate_limiter_v3, and goes back to the requested model when no fallback
admits the request. Responses carry x-litellm-fallback-reason:
rate_limit:model_saturation_fallback.

Removes the team_metadata.model_preferences steering and the
x-litellm-model-steered-from header from #51350, which did the same job through
a second knob.

* fix(proxy): strict bool for fallback_on_saturation, drop a stale import and settle the pre-call tests

fallback_on_saturation only turns on for a real boolean true, so a string like "yes" in router_settings no longer enables it. The saturation unit test reads saturation through the limiter's cache, so it asserts the threshold was reached instead of an exact count. The response polling test now hands the pre-call path an awaitable router_settings lookup, as the pre-call fallback walk reads Key > Team settings before the first call.

* fix(proxy): pre-call fallbacks pass the budget and access-group checks a direct request does

A fallback the pre-call walk picks now has to clear every per-model budget layer and the budgets of the access
groups that grant it, and its spend is charged to those groups instead of the requested model's. At saturation a
fallback that fails for any reason counts as a miss and the request goes back to the requested model. The client
body can no longer opt itself into the saturation walk through router_settings_override, only the saved key or team
router_settings can. The integration test for a forbidden fallback waits for the preferred model's own 429 so it
cannot pass before saturation.

* refactor(proxy): one helper decides which refused fallbacks end the pre-call walk

_pre_call_with_fallbacks had grown past the strict C901 ceiling, so the two except clauses fold into one and _fallback_refusal_ends_the_walk carries the rule: at saturation every refusal is a miss, otherwise only a non-rate-limit failure or a hard per-model 429 reaches the caller.

* fix(proxy): keep the caller's request object through the pre-call fallback walk

Copying the request to drop a client fallback_on_saturation opt-in handed the hooks a different dict than the one the caller holds, so header metadata written into one never reached the other (test_add_litellm_metadata_from_request_headers). drop_client_saturation_opt_in now edits the request in place.

* fix(proxy): pre-call fallbacks keep the pricing and budget-reservation rules a direct request passes

A fallback is skipped when unpriced models are blocked and it has no pricing, when a zero-cost request would land on a priced model, and when the budget auth reserved for the requested model does not fit it (higher per-token prices, or reserved against an access group that does not grant it).

* refactor(proxy): read fallback prices through the router registry entry point

router_surface_check only lets the router surface shrink, so the per-token price lookup goes through registry().group() instead of the deprecated get_model_group_info.

* fix(proxy): a fallback's group budget check skips the groups the request already holds

The auth-time reservation sits in the same counter the plain budget read uses, so a fallback in a group the request already reserved against was refused by its own hold once the group was within one estimate of max_budget. Groups without a hold keep the plain read.

* test(dynamic_rate_limiter_v3): the throttle metric test from main uses fallback_on_saturation instead of the removed steering

A saturation handoff carries no rate_limited_by reason, so it is never counted as a throttle.

* fix(proxy): the saturation handoff only fires under the fallback walk

_pre_call_with_fallbacks arms the request before the hooks run, so routes that call the hooks directly (batches, fine-tuning, polling, images) keep admitting or refusing an opted-in key on a saturated model as before instead of surfacing the handoff 429.

* style(proxy): keep the cast-ok reason within the line length

* fix(proxy): never invent a metadata bucket when arming the saturation fallback walk, take main's limiter metric tests

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

* fix(proxy): arm the saturation fallback walk in the request context, refuse fallbacks for agent-owned keys

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

* fix(proxy): keep the saturation walk state out of the body, check the end user's allowlist on fallbacks, release a refused fallback's rate limit hold

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

* fix(proxy): only load the end user for fallback checks when a database is connected

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

* fix(proxy): grant pre-call fallbacks on the model an alias serves and release the rate limited model's hold before the walk

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

* fix(proxy): check a pre-call fallback's per-model budgets on its listed name and its served group, as auth does

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

* test(proxy): assert the listed fallback name's budget blocks an aliased fallback, not the order of budget reads

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

* fix(proxy): release the dynamic limiter's reservation when the fallback walk leaves a model, and replay an exhausted walk with the requested model

The dynamic rate limiter admits through its own inner v3 limiter, so the walk's hold release never reached its team-share
reservation. It now exposes async_post_call_failure_hook, which refunds that reservation when the proxy refuses a request
after admission, and _release_pre_call_hold calls it next to the registered limiter.

An exhausted saturation walk replayed with the already aliased model name, so chained model_alias_map entries moved the
request on to a further alias. It now replays with the model argument the first attempt saw.

The hold-release test drives the real v3 limiter with one slot per key instead of recording hook calls.

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

* fix(proxy): type the dynamic limiter's failure hook request dict for the type-discipline gate

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

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-12 21:57:18 -08:00 • committed by GitHub
parent efe349de9f
commit 172e16acde
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1848 additions and 881 deletions

View file

@ -5,7 +5,17 @@ import json
import logging
import math
import time
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
from collections.abc import (
AsyncGenerator,
AsyncIterator,
Awaitable,
Callable,
Coroutine,
Iterator,
Mapping,
MutableMapping,
Sequence,
)
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
@ -153,7 +163,11 @@ from litellm.proxy.common_utils.sse_keepalive import (
)
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import PROXY_DynamicRateLimitHandlerV3
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
get_or_create_request_stash,
)
from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost
from litellm.proxy.native_compaction import with_proxy_compaction_executor
from litellm.proxy.route_llm_request import (
@ -1022,6 +1036,7 @@ async def hierarchical_router_settings_updates(
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
_FALLBACK_LIST: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object], config=ConfigDict(strict=True))
_DECODED_JSON: Final = TypeAdapter(object)
_ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -2266,6 +2281,265 @@ def _timing_values(
return getattr(logging_obj, "response_timing_metrics", None) or {}
def _fallback_refusal_ends_the_walk(
exc: Exception, at_saturation: bool, general_settings: Mapping[str, object]
) -> bool:
"""At saturation the requested model still has room, so a fallback refused for any reason is a miss. Otherwise a
fallback's own non-rate-limit failure, or a hard per-model 429, is the answer the caller gets."""
from litellm.proxy.common_utils.proxy_rate_limit_error import (
PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS,
ProxyRateLimitError,
per_model_rate_limits_disable_fallbacks,
)
if at_saturation:
return False
if not isinstance(exc, ProxyRateLimitError):
return True
return exc.descriptor_key in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS and per_model_rate_limits_disable_fallbacks(
general_settings
)
async def _fallback_access_groups(
fallback_model: str, user_api_key_dict: UserAPIKeyAuth, llm_router: Router, proxy_logging_obj: ProxyLogging
) -> tuple[str, ...]:
from litellm.proxy.auth.auth_checks import (
collect_matched_model_access_groups,
get_project_object,
get_team_membership,
get_team_object,
)
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
if prisma_client is None:
return ()
team_id: Final = user_api_key_dict.team_id
user_id: Final = user_api_key_dict.user_id
project_id: Final = user_api_key_dict.project_id
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
)
return await collect_matched_model_access_groups(
model=fallback_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,
)
_PER_TOKEN_PRICES: Final[TypeAdapter[tuple[float, float]]] = TypeAdapter(tuple[float, float])
_RESERVATION_ENTRIES: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
_MODEL_ACCESS_GROUP_ENTITY: Final = "Model access group"
def _per_token_prices(llm_router: Router, model: str) -> tuple[float, float] | None:
info: Final = llm_router.registry().group(model)
if info is None:
return None
try:
return _PER_TOKEN_PRICES.validate_python((info.input_cost_per_token, info.output_cost_per_token))
except ValidationError:
return None
def _reserved_access_groups(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str] | None:
"""The model access groups the request's ``budget_reservation`` already holds, None when its entries are
unreadable."""
reservation: Final = user_api_key_dict.budget_reservation
if reservation is None:
return frozenset()
try:
entries: Final = _RESERVATION_ENTRIES.validate_python(reservation.get("entries") or ())
except ValidationError:
return None
return frozenset(
str(entry.get("entity_id")) for entry in entries if entry.get("entity_type") == _MODEL_ACCESS_GROUP_ENTITY
)
def _reservation_covers(
user_api_key_dict: UserAPIKeyAuth,
access_groups: Sequence[str],
held: frozenset[str] | None,
requested_model: str,
fallback_model: str,
llm_router: Router,
) -> bool:
"""Whether the budget auth reserved for ``requested_model`` still fits ``fallback_model``: the fallback costs no
more per token than the model the estimate was made for, and every access group ``held`` by the reservation
grants the fallback, so settlement never charges a group that did not serve."""
if user_api_key_dict.budget_reservation is None:
return True
if held is None:
return False
requested: Final = _per_token_prices(llm_router, requested_model)
fallback: Final = _per_token_prices(llm_router, fallback_model)
if requested is None or fallback is None or fallback[0] > requested[0] or fallback[1] > requested[1]:
return False
return held <= frozenset(access_groups)
def _fallback_pricing_admits(requested_model: str, fallback_model: str, llm_router: Router) -> bool:
"""A zero-cost ``requested_model`` waived the aggregate budget checks at auth, so only another zero-cost model may
take its place, and an unpriced fallback is refused while ``block_requests_for_models_without_pricing`` is on."""
from litellm.proxy.auth.auth_checks import (
_unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract
)
if is_model_cost_zero(model=requested_model, llm_router=llm_router) and not is_model_cost_zero(
model=fallback_model, llm_router=llm_router
):
return False
return not (
litellm.block_requests_for_models_without_pricing
and _unpriced_models_in_request(model=fallback_model, llm_router=llm_router)
)
async def _release_pre_call_hold(
proxy_logging_obj: ProxyLogging,
request_data: Mapping[str, object],
refusal: Exception,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Give back the slot and tokens the rate limiter holds for a model the walk is leaving: the requested model a
later hook rate limited after the limiter admitted it, or a fallback the proxy refused after admitting it. The
walk's next attempt overwrites the request's stash, so a hold not released here stays charged until the window
resets. The request has not failed, so only the limiters' releases run, not the failure logging and alerting of
``post_call_failure_hook``. The dynamic rate limiter admits through its own inner limiter, so its team-share
reservation is released alongside the registered limiter's slot and tokens."""
limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter")
dynamic_limiters: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(
PROXY_DynamicRateLimitHandlerV3
)
for holder in (limiter, *dynamic_limiters):
if not isinstance(holder, (PROXY_MaxParallelRequestsHandler_v3, PROXY_DynamicRateLimitHandlerV3)):
continue
await holder.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the limiter hooks take the proxy's bare request dict
request_data=dict(request_data), original_exception=refusal, user_api_key_dict=user_api_key_dict
)
async def _fallback_grant(
fallback_model: str,
user_api_key_dict: UserAPIKeyAuth,
llm_router: Router | None,
proxy_logging_obj: ProxyLogging,
requested_model: str,
listed_as: str | None = None,
) -> list[str] | None:
"""The model access groups a pre-call fallback to ``fallback_model`` charges, or None when the caller may not use
it. The caller passes the group the fallback is served by, through every alias rewrite, and the name the fallback
list gave it as ``listed_as``, so an alias never skips the budgets of the model that answers nor those set on the
listed name: an agent-owned key never leaves the model it asked for, as at auth, then the key, team and project allowlists
while ``enforce_fallback_model_access`` is on, the end user's allowlist, the pricing rules, every per-model budget layer, the budgets of the access groups that grant the model, and the budget auth reserved for
``requested_model``, the same checks a direct request passes at auth."""
from litellm.proxy.auth.auth_checks import (
_model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract
can_customer_access_model,
get_end_user_object,
)
from litellm.proxy.auth.fallback_model_access import is_model_authorized_for_token, router_fallback_access_check
from litellm.proxy.auth.user_api_key_auth import (
_is_within_every_model_budget, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract
)
from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client, user_api_key_cache
if llm_router is None:
return []
if user_api_key_dict.agent_id is not None:
return None
if router_fallback_access_check.is_enforced() and not await is_model_authorized_for_token(
model=fallback_model, valid_token=user_api_key_dict, llm_router=llm_router
):
return None
if not _fallback_pricing_admits(requested_model, fallback_model, llm_router):
return None
try:
end_user: Final = (
await get_end_user_object(
end_user_id=user_api_key_dict.end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if user_api_key_dict.end_user_id is not None and prisma_client is not None
else None
)
if end_user is not None and end_user.models:
can_customer_access_model(
model=fallback_model, end_user_object=end_user, llm_router=llm_router, valid_token=user_api_key_dict
)
if not is_model_cost_zero(model=fallback_model, llm_router=llm_router):
budget_names: Final = tuple(
dict.fromkeys(
(listed_as or fallback_model, fallback_model, llm_router.registry().resolve(fallback_model))
)
)
for budget_name in budget_names:
if not await _is_within_every_model_budget(
valid_token=user_api_key_dict, model_max_budget_limiter=model_max_budget_limiter, model=budget_name
):
return None
access_groups: Final = await _fallback_access_groups(
fallback_model, user_api_key_dict, llm_router, proxy_logging_obj
)
held: Final = _reserved_access_groups(user_api_key_dict)
await _model_access_group_max_budget_check(
matched_model_access_groups=tuple(group for group in access_groups if group not in (held or ())),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
if not _reservation_covers(user_api_key_dict, access_groups, held, requested_model, fallback_model, llm_router):
return None
except Exception as e: # noqa: BLE001 # fail closed: a grant the walk cannot settle must not admit the fallback
verbose_proxy_logger.warning(
"Skipping fallback to model=%s: %s",
fallback_model.replace("\r", "").replace("\n", ""),
e,
)
return None
return list(access_groups)
async def _team_member_within_fallback_model_budget(
user_api_key_dict: UserAPIKeyAuth,
model: str,
@ -2918,6 +3192,12 @@ class ProxyBaseLLMRequestProcessing:
PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS,
ProxyRateLimitError,
RateLimitFallback,
arm_saturation_fallback_walk,
disarm_saturation_fallback_walk,
drop_client_saturation_opt_in,
fallback_on_saturation_enabled,
is_saturation_fallback,
mark_saturation_fallbacks_exhausted,
per_model_rate_limits_disable_fallbacks,
with_rate_limit_fallback,
)
@ -2932,12 +3212,26 @@ class ProxyBaseLLMRequestProcessing:
and per_model_rate_limits_disable_fallbacks(general_settings_view)
)
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
router_settings: Final = (
await self._hierarchical_router_settings(user_api_key_dict, proxy_config, proxy_logging_obj)
if llm_router is not None
else None
)
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, router_settings=router_settings)
if llm_router is not None
else None
)
scrubbed: Final[dict] = self.data # mutable-ok: caller-defined request-payload shape
payload: Final = cast(MutableMapping[str, object], scrubbed) # cast-ok: untyped caller payload
drop_client_saturation_opt_in(payload)
walk: Final = arm_saturation_fallback_walk()
self.data = scrubbed
pristine: Final = (
independent_snapshot(scrubbed)
if configured_fallbacks or fallback_on_saturation_enabled(router_settings)
else None
)
try:
return await self.common_processing_pre_call_logic(
@ -2959,20 +3253,17 @@ class ProxyBaseLLMRequestProcessing:
except ProxyRateLimitError as original_exc:
rate_limited_data: Final = self.data
original_model: Final = rate_limited_data.get("model")
if (
pristine is None
or not configured_fallbacks
or rate_limited_data.get("disable_fallbacks")
or not isinstance(original_model, str)
or is_hard_per_model_limit(original_exc)
):
# The limiter raises this one while the requested model still has room, so a walk that admits nothing
# ends on that model instead of a 429.
at_saturation: Final = is_saturation_fallback(original_exc)
if pristine is None or not isinstance(original_model, str) or is_hard_per_model_limit(original_exc):
raise
fallback_models: Final = self._resolve_fallback_models(
model=original_model,
fallbacks=configured_fallbacks,
fallback_models: Final = (
self._resolve_fallback_models(model=original_model, fallbacks=configured_fallbacks)
if configured_fallbacks and not rate_limited_data.get("disable_fallbacks")
else None
)
if not fallback_models:
if not fallback_models and not at_saturation:
raise
verbose_proxy_logger.info(
@ -2988,11 +3279,26 @@ class ProxyBaseLLMRequestProcessing:
if isinstance(raw_router_settings, Mapping)
else None
)
await _release_pre_call_hold(
proxy_logging_obj, _REQUEST_FIELDS.validate_python(rate_limited_data), original_exc, user_api_key_dict
)
attempt_numbers: Final[Iterator[int]] = itertools.count(1)
requested_groups: Final = user_api_key_dict.matched_model_access_groups
try:
for fallback_model in fallback_models:
for fallback_model in fallback_models or ():
if fallback_model == original_model:
continue
if (
grant := await _fallback_grant(
pre_call_served_model(fallback_model, user_api_key_dict, group_alias),
user_api_key_dict,
llm_router,
proxy_logging_obj,
original_model,
listed_as=fallback_model,
)
) is None:
continue
if not await _team_member_within_fallback_model_budget(
user_api_key_dict=user_api_key_dict,
model=fallback_model,
@ -3012,6 +3318,7 @@ class ProxyBaseLLMRequestProcessing:
):
continue
attempt = next(attempt_numbers)
user_api_key_dict.matched_model_access_groups = grant # rebind-ok: bill the fallback's groups
admitted: tuple[
Mapping[str, object], LiteLLMLoggingObj
] = await self.common_processing_pre_call_logic(
@ -3031,9 +3338,16 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
rate_limited_model=original_model,
)
except ProxyRateLimitError as fallback_exc:
if is_hard_per_model_limit(fallback_exc):
except Exception as fallback_exc: # noqa: BLE001 # the helper decides which refusals end the walk
user_api_key_dict.matched_model_access_groups = requested_groups # rebind-ok: undo the fallback
if _fallback_refusal_ends_the_walk(fallback_exc, at_saturation, general_settings_view):
raise
await _release_pre_call_hold(
proxy_logging_obj,
self.data, # pyright: ignore[reportUnknownMemberType] # caller-defined request-payload shape
fallback_exc,
user_api_key_dict,
)
continue
else:
admitted_data, admitted_logging_obj = admitted
@ -3050,19 +3364,65 @@ class ProxyBaseLLMRequestProcessing:
)
self.data = stamped
return stamped, admitted_logging_obj
if not at_saturation:
self.data = rate_limited_data
raise original_exc
verbose_proxy_logger.info(
"No fallback of model=%s is below the saturation threshold, serving the request on it",
original_model.replace("\r", "").replace("\n", ""),
)
mark_saturation_fallbacks_exhausted()
self.data = independent_snapshot(dict(_REQUEST_FIELDS.validate_python(pristine)))
return await self.common_processing_pre_call_logic(
request=request,
general_settings=general_settings,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_config=proxy_config,
user_model=user_model,
user_temperature=user_temperature,
user_request_timeout=user_request_timeout,
user_max_tokens=user_max_tokens,
user_api_base=user_api_base,
model=model, # the requested name as the first attempt saw it, so its aliases apply once, not again
route_type=route_type, # pyright: ignore[reportArgumentType] # the proxy hands its route literal over as str, as the two calls above take it
llm_router=llm_router,
)
except BaseException:
user_api_key_dict.matched_model_access_groups = requested_groups # rebind-ok: undo the fallback
self.data = rate_limited_data
raise
self.data = rate_limited_data
raise original_exc
finally:
disarm_saturation_fallback_walk(walk)
@staticmethod
def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
key_router_settings: Final = user_api_key_dict.router_settings
key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.settings().fallbacks
return fallbacks if isinstance(fallbacks, list) and fallbacks else None
async def _hierarchical_router_settings(
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig, proxy_logging_obj: ProxyLogging
) -> Mapping[str, object] | None:
"""The key's ``router_settings``, else its team's, read through the same cached lookup the request body
gets them from, so the pre-call fallback walk and the router agree on Key > Team > Global."""
from litellm.proxy.proxy_server import prisma_client
return _ROUTER_SETTINGS_ADAPTER.validate_python(
await proxy_config.get_hierarchical_router_settings(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
)
@staticmethod
def _configured_fallbacks(llm_router: Router, router_settings: Mapping[str, object] | None) -> list[object] | None:
settings_fallbacks: Final = router_settings.get("fallbacks") if router_settings is not None else None
fallbacks: Final[object] = (
settings_fallbacks if settings_fallbacks is not None else llm_router.settings().fallbacks
)
try:
configured: Final = _FALLBACK_LIST.validate_python(fallbacks)
except ValidationError:
return None
return configured or None
@staticmethod
def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None:

View file

@ -36,10 +36,11 @@ This module provides a single proxy-side error class that:
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, MutableMapping
from contextvars import ContextVar, Token
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from typing import Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
@ -78,6 +79,87 @@ def per_model_rate_limits_disable_fallbacks(general_settings: Mapping[str, objec
return False
FALLBACK_ON_SATURATION_SETTING: Final = "fallback_on_saturation"
SATURATION_FALLBACK_DESCRIPTOR_KEY: Final = "model_saturation_fallback"
_FALLBACK_ON_SATURATION_FLAG: Final[TypeAdapter[bool | None]] = TypeAdapter(bool | None)
def fallback_on_saturation_enabled(router_settings: object) -> bool:
"""Whether the key's or team's ``router_settings`` ask for fallbacks once the requested model reaches
``saturation_threshold``, before the dynamic rate limiter admits or refuses the request on that model."""
raw_value: Final = _metadata_bucket(router_settings).get(FALLBACK_ON_SATURATION_SETTING)
try:
return _FALLBACK_ON_SATURATION_FLAG.validate_python(raw_value, strict=True) is True
except ValidationError:
verbose_proxy_logger.warning(
"router_settings.%s=%r is not a boolean, treating it as disabled", FALLBACK_ON_SATURATION_SETTING, raw_value
)
return False
def drop_client_saturation_opt_in(
request_data: MutableMapping[str, object], # mutable-ok: in-place edit keeps the caller's request object
) -> None:
"""Drop ``fallback_on_saturation`` from the client's ``router_settings_override``. Only the key's or team's saved
``router_settings`` turn the saturation walk on, and the proxy writes those over this field before the limiter
reads it."""
override: Final = _metadata_bucket(request_data.get("router_settings_override"))
if FALLBACK_ON_SATURATION_SETTING not in override:
return
request_data["router_settings_override"] = { # rebind-ok: the request object must stay the caller's
key: value for key, value in override.items() if key != FALLBACK_ON_SATURATION_SETTING
}
_SaturationWalkState: TypeAlias = Literal["off", "armed", "exhausted"]
_SATURATION_FALLBACK_WALK: Final[ContextVar[_SaturationWalkState]] = ContextVar(
"litellm_saturation_fallback_walk", default="off"
)
def saturation_fallback_walk_armed() -> bool:
"""Whether a caller that catches the handoff (``_pre_call_with_fallbacks``) is running this request's hooks and
has not walked the fallbacks yet."""
return _SATURATION_FALLBACK_WALK.get() == "armed"
def saturation_fallbacks_exhausted() -> bool:
"""Whether this request already tried every fallback at saturation and is back on the model it asked for."""
return _SATURATION_FALLBACK_WALK.get() == "exhausted"
def arm_saturation_fallback_walk() -> Token[_SaturationWalkState]:
"""Mark the current request as running under the fallback walk, so the limiter may hand a saturated model off to
it. The state lives in the request's context, never in its body, so nothing reaches the provider and no client
field can arm, skip or exhaust the walk. Routes that call the hooks directly never arm it and keep admitting or
refusing as before."""
return _SATURATION_FALLBACK_WALK.set("armed")
def mark_saturation_fallbacks_exhausted() -> None:
_SATURATION_FALLBACK_WALK.set("exhausted")
def disarm_saturation_fallback_walk(token: Token[_SaturationWalkState]) -> None:
_SATURATION_FALLBACK_WALK.reset(token)
def fallback_on_saturation_requested(request_data: Mapping[str, object]) -> bool:
"""Whether the dynamic rate limiter should hand a saturated model's request to the fallback walk instead of
admitting or refusing it: the key or team opted in through ``router_settings_override`` and the fallback walk
is the caller and has not already run for this request."""
return (
fallback_on_saturation_enabled(request_data.get("router_settings_override"))
and saturation_fallback_walk_armed()
)
def is_saturation_fallback(exc: "ProxyRateLimitError") -> bool:
"""Whether ``exc`` is the limiter asking for the fallback walk rather than refusing the request. The model it
names still has room, so when no fallback admits the request it goes back to that model instead of a 429."""
return exc.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY
RATE_LIMIT_FALLBACK_METADATA_KEY: Final = "rate_limit_fallback"
RATE_LIMIT_FALLBACK_REASON_HEADER: Final = "x-litellm-fallback-reason"
ATTEMPTED_FALLBACKS_HEADER: Final = "x-litellm-attempted-fallbacks"
@ -204,7 +286,7 @@ def _coerce_message(detail: object) -> str:
try:
return json.dumps(detail)
except (TypeError, ValueError):
return str(detail)
return f"{detail}"
return str(detail)

View file

@ -2,16 +2,15 @@
Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
"""
import json
import os
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from typing import TYPE_CHECKING, Final, Literal
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, ValidationError
from redis.exceptions import RedisError
import litellm
@ -27,11 +26,12 @@ 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 (
SATURATION_FALLBACK_DESCRIPTOR_KEY,
ProxyRateLimitError,
fallback_on_saturation_requested,
map_v3_rate_limit_type,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # legacy module exports
@ -56,7 +56,6 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l
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 (
@ -73,7 +72,7 @@ from litellm.proxy.hooks.rate_limiter_utils import (
resolve_llm_provider_for_rate_limit,
socket_session_usage,
)
from litellm.proxy.utils import InternalUsageCache, check_and_merge_model_level_guardrails
from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET, CounterIncrement, RateLimitCounter, RateLimitScope
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, Usage
@ -83,273 +82,6 @@ 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.customer_hierarchy_auth import check_customer_hierarchy_model_access
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
)
if user_api_key_dict.customer_hierarchy is not None:
check_customer_hierarchy_model_access(user_api_key_dict.customer_hierarchy, model, llm_router)
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."""
@ -572,10 +304,8 @@ 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)
self.active_teams = ActiveTeamRegistry(internal_usage_cache)
@ -1639,11 +1369,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
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:
admitted: Final = await self._admit_steered_or_preferred(
model, requested_group_info, steered, user_api_key_dict, data, call_type
)
await self._admit(model, requested_group_info, user_api_key_dict, data, call_type)
except ProxyRateLimitError as refusal:
raise self._throttled(refusal, model, user_api_key_dict) from None
except HTTPException:
@ -1652,34 +1379,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
verbose_proxy_logger.error("Error in dynamic rate limiter: %s, allowing request", e)
# 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,
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, 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, data, call_type)
return None
async def _admit(
@ -1690,7 +1389,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
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."""
"""Read saturation, then the three-phase rate limit check, raising the 429 when ``model`` is out of capacity.
A key or team with ``router_settings.fallback_on_saturation`` is sent to its fallbacks first once ``model`` is
at or above ``saturation_threshold``, and comes back here only when none of them admitted the request."""
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
@ -1698,6 +1399,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={priority}"
)
if saturation >= saturation_threshold and fallback_on_saturation_requested(data):
raise self._saturation_fallback_error(model, saturation, priority)
await self._check_rate_limits(
model=model,
model_group_info=model_group_info,
@ -1708,104 +1411,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
call_type=call_type,
)
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,
def _saturation_fallback_error(self, model: str, saturation: float, priority: str | None) -> ProxyRateLimitError:
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model)
return ProxyRateLimitError(
detail={
"error": f"Model {model} is at {saturation:.0%} of its capacity, at or above the saturation "
f"threshold of {_get_priority_settings().saturation_threshold:.0%}, trying fallbacks first. "
f"Priority: {priority}"
},
headers={"retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default"},
model=resolved_model,
llm_provider=llm_provider,
descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY,
)
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(
@ -1842,21 +1460,6 @@ 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: Mapping[str, object], response_obj: object, start_time: object, end_time: object
@ -1986,6 +1589,30 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
verbose_proxy_logger.exception("Error in dynamic rate limiter success event: %s", e)
@with_service_target(RATE_LIMITS_TARGET)
async def async_post_call_failure_hook(
self,
request_data: dict[str, object], # mutable-ok: the base hook declares this parameter as a dict
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
) -> None:
"""Give back the team-share tokens held at admission when the proxy refuses the request before the call runs,
as a later hook or the pre-call fallback walk does. ``async_log_failure_event`` is a completion callback and
never fires for those refusals, so without this the hold stays charged until the window resets. Idempotent:
the first release zeroes the stash's reservation, so a settlement that still follows charges only usage."""
stash: Final = get_request_stash()
if stash is None:
return
refund: Final = _team_model_settlement(stash, None, 0, self.v3_limiter.window_size)
stash.borrowing_team_model_reserved_tokens = 0
stash.borrowing_team_model_cap_reserved_tokens = 0
stash.borrowing_team_model_reserved_window = None
stash.borrowing_team_model_cap_reserved_window = None
if refund:
await self.v3_limiter.async_increment_reservation_aware_tokens(
pipeline_operations=refund, parent_otel_span=user_api_key_dict.parent_otel_span
)
async def async_log_failure_event(
self, kwargs: object, response_obj: object, start_time: object, end_time: object
) -> None:

View file

@ -620,7 +620,6 @@ 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)

View file

@ -237,6 +237,7 @@ class UpdateRouterConfig(LiteLLMBaseModel):
tag_routing_prefix: str | None = None
optional_pre_call_checks: OptionalPreCallChecks | None = None
fallback_on_unsupported_params: bool | None = None
fallback_on_saturation: bool | None = None
model_config = ConfigDict(protected_namespaces=())

View file

@ -51867,6 +51867,8 @@ export interface components {
cooldown_time?: number | null;
/** Enable Tag Filtering */
enable_tag_filtering?: boolean | null;
/** Fallback On Saturation */
fallback_on_saturation?: boolean | null;
/** Fallback On Unsupported Params */
fallback_on_unsupported_params?: boolean | null;
/** Fallbacks */

View file

@ -13,9 +13,10 @@ from integration._support.wire import Reply, Request, wire_server
OPENAI_MODEL: Final = "gpt-5.6"
PREFERRED_RPM: Final = 2
ALTERNATE_RPM: Final = 1000
FALLBACK_RPM: Final = 1000
SATURATION_THRESHOLD: Final = 0.5
STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from"
FALLBACK_REASON_HEADER: Final = "x-litellm-fallback-reason"
SATURATION_REASON: Final = "rate_limit:model_saturation_fallback"
def _reply() -> bytes:
@ -25,15 +26,15 @@ def _reply() -> bytes:
"object": "chat.completion",
"created": 1700000000,
"model": OPENAI_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "steered"}, "finish_reason": "stop"}],
"choices": [{"index": 0, "message": {"role": "assistant", "content": "served"}, "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")
def saturation_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]:
directory: Final = tmp_path_factory.mktemp("fallback-on-saturation")
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration: Final = {
**base,
@ -46,7 +47,7 @@ def steering_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[G
},
},
}
path: Final = directory / "model_preference_steering.yaml"
path: Final = directory / "fallback_on_saturation.yaml"
path.write_text(yaml.safe_dump(configuration))
def respond(request: Request) -> Reply:
@ -73,7 +74,7 @@ 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}"}]},
{"model": model, "messages": [{"role": "user", "content": f"fallback on saturation {uuid.uuid4().hex}"}]},
key=key,
)
@ -84,61 +85,69 @@ def _spend_row(owned: Gateway, response: httpx.Response) -> list[dict[str, objec
)
def test_flexible_team_is_steered_to_its_permitted_alternate_once_the_preferred_model_saturates(
steering_proxy: tuple[Gateway, str],
def test_a_team_with_fallback_on_saturation_is_served_by_its_fallback_once_the_model_saturates(
saturation_proxy: tuple[Gateway, str],
) -> None:
owned, upstream = steering_proxy
owned, upstream = saturation_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]))
fallback, fallback_id = _deployment(owned, scenario, upstream, FALLBACK_RPM)
rigid_key: Final = scenario.key(team_id=scenario.team(models=[preferred, fallback]))
flexible_key: Final = scenario.key(
team_id=scenario.team(
models=[preferred, alternate],
metadata={"model_preferences": [{"model": preferred, "alternates": [alternate]}]},
models=[preferred, fallback],
router_settings={"fallbacks": [{preferred: [fallback]}], "fallback_on_saturation": True},
)
)
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
assert FALLBACK_REASON_HEADER not in warmup.headers
steered: Final = eventually(
fell_back: Final = eventually(
lambda: _chat(owned, flexible_key, preferred),
lambda response: STEERED_FROM_HEADER in response.headers,
lambda response: response.headers.get(FALLBACK_REASON_HEADER) == SATURATION_REASON,
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
assert fell_back.status_code == 200, fell_back.text
assert fell_back.headers[FALLBACK_REASON_HEADER] == SATURATION_REASON
assert fell_back.headers["x-litellm-attempted-fallbacks"] == "1"
assert fell_back.headers["x-litellm-model-id"] == fallback_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
assert FALLBACK_REASON_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
charged: Final = eventually(lambda: _spend_row(owned, fell_back), lambda rows: len(rows) == 1, seconds=30)
assert charged[0]["model_group"] == fallback
def test_alternate_the_team_may_not_call_is_never_used(steering_proxy: tuple[Gateway, str]) -> None:
owned, upstream = steering_proxy
def test_a_fallback_the_team_may_not_call_is_never_used(saturation_proxy: tuple[Gateway, str]) -> None:
owned, upstream = saturation_proxy
with owned.scenario() as scenario:
preferred, _ = _deployment(owned, scenario, upstream, PREFERRED_RPM)
forbidden, _ = _deployment(owned, scenario, upstream, ALTERNATE_RPM)
preferred, preferred_id = _deployment(owned, scenario, upstream, PREFERRED_RPM)
forbidden, _ = _deployment(owned, scenario, upstream, FALLBACK_RPM)
key: Final = scenario.key(
team_id=scenario.team(
models=[preferred],
metadata={"model_preferences": [{"model": preferred, "alternates": [forbidden]}]},
router_settings={"fallbacks": [{preferred: [forbidden]}], "fallback_on_saturation": True},
)
)
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
def served_by_preferred_or_refused() -> httpx.Response:
response: Final = _chat(owned, key, preferred)
assert FALLBACK_REASON_HEADER not in response.headers, response.headers
if response.status_code == 200:
assert response.headers["x-litellm-model-id"] == preferred_id
return response
refused: Final = eventually(
served_by_preferred_or_refused,
lambda response: response.status_code != 200,
seconds=15,
return_last_on_timeout=True,
)
assert second.status_code == 429, second.text
assert STEERED_FROM_HEADER not in second.headers
assert refused.status_code == 429, refused.text

View file

@ -3,9 +3,18 @@ from typing import Final
import pytest
from litellm.proxy.common_utils.proxy_rate_limit_error import (
SATURATION_FALLBACK_DESCRIPTOR_KEY,
ProxyRateLimitError,
RateLimitFallback,
arm_saturation_fallback_walk,
disarm_saturation_fallback_walk,
fallback_on_saturation_enabled,
fallback_on_saturation_requested,
is_saturation_fallback,
mark_saturation_fallbacks_exhausted,
rate_limit_fallback_from_request_data,
rate_limit_fallback_headers,
saturation_fallbacks_exhausted,
)
STAMP: Final = {"original_model_group": "opus", "descriptor_key": "model_per_end_user"}
@ -56,3 +65,64 @@ def test_rate_limit_fallback_headers_count_every_candidate_a_limit_rejected(
)
assert headers["x-litellm-attempted-fallbacks"] == expected
@pytest.mark.parametrize(
("router_settings", "expected"),
(
({"fallback_on_saturation": True}, True),
({"fallback_on_saturation": False}, False),
({"fallback_on_saturation": "yes"}, False),
({"fallbacks": [{"luna": ["opus"]}]}, False),
(None, False),
("fallback_on_saturation", False),
),
)
def test_fallback_on_saturation_is_enabled_only_by_a_true_router_setting(
router_settings: object, expected: bool
) -> None:
assert fallback_on_saturation_enabled(router_settings) is expected
def test_a_request_asks_for_saturation_fallbacks_only_while_the_walk_is_armed_and_until_it_has_run() -> None:
request: Final = {"model": "luna", "router_settings_override": {"fallback_on_saturation": True}}
assert not fallback_on_saturation_requested(request)
walk: Final = arm_saturation_fallback_walk()
try:
assert fallback_on_saturation_requested(request)
assert not saturation_fallbacks_exhausted()
mark_saturation_fallbacks_exhausted()
assert saturation_fallbacks_exhausted()
assert not fallback_on_saturation_requested(request)
finally:
disarm_saturation_fallback_walk(walk)
assert not saturation_fallbacks_exhausted()
assert not fallback_on_saturation_requested(request)
def test_client_metadata_never_arms_skips_or_exhausts_the_walk() -> None:
forged: Final = {
"model": "luna",
"router_settings_override": {"fallback_on_saturation": True},
"metadata": {"saturation_fallback_walk": True, "saturation_fallbacks_exhausted": True},
"litellm_metadata": {"saturation_fallback_walk": True, "saturation_fallbacks_exhausted": True},
}
assert not fallback_on_saturation_requested(forged)
assert not saturation_fallbacks_exhausted()
walk: Final = arm_saturation_fallback_walk()
try:
assert fallback_on_saturation_requested(forged)
finally:
disarm_saturation_fallback_walk(walk)
def test_a_request_without_the_setting_never_asks_for_saturation_fallbacks() -> None:
assert not fallback_on_saturation_requested({"model": "luna"})
assert not fallback_on_saturation_requested({"model": "luna", "router_settings_override": {"fallbacks": []}})
def test_only_the_saturation_descriptor_marks_a_fallback_walk_the_model_can_still_serve() -> None:
assert is_saturation_fallback(ProxyRateLimitError(detail="x", descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY))
assert not is_saturation_fallback(ProxyRateLimitError(detail="x", descriptor_key="model_per_team"))
assert not is_saturation_fallback(ProxyRateLimitError(detail="x"))

View file

@ -6,9 +6,8 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of
import asyncio
import time
from collections.abc import Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, patch
@ -17,11 +16,9 @@ from prometheus_client import CollectorRegistry
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,
_PROXY_DynamicRateLimitHandlerV3, # pyright: ignore[reportPrivateUsage] # the public name is a Final alias, not a type
)
@ -1837,317 +1834,92 @@ 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 _saturation_fallback_handler(monkeypatch: pytest.MonkeyPatch) -> _PROXY_DynamicRateLimitHandlerV3:
return _borrowing_handler(monkeypatch, borrows=False, saturation_threshold=0.5, rpm=10)
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}
)
def _interactive_key(api_key: str, team_id: str) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key=api_key, team_id=team_id, metadata={"priority": "interactive"})
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
@pytest.mark.parametrize("team_allowed_models, granted", [(("luna", "opus"), True), (("luna",), False)])
async def test_an_alternate_is_granted_only_when_every_customer_level_permits_it(
monkeypatch: pytest.MonkeyPatch, team_allowed_models: tuple[str, ...], granted: bool
):
from litellm.proxy import proxy_server
from litellm.proxy.auth import auth_checks, user_api_key_auth
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import alternate_grant_for
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
monkeypatch.setattr(proxy_server, "llm_router", _steering_router("luna", "opus"))
monkeypatch.setattr(proxy_server, "prisma_client", None)
monkeypatch.setattr(user_api_key_auth, "_is_fallback_model_authorized", AsyncMock(return_value=True))
monkeypatch.setattr(user_api_key_auth, "_is_within_every_model_budget", AsyncMock(return_value=True))
monkeypatch.setattr(auth_checks, "collect_matched_model_access_groups", AsyncMock(return_value=()))
monkeypatch.setattr(auth_checks, "_model_access_group_max_budget_check", AsyncMock(return_value=None))
customer: Final = UserAPIKeyAuth(
customer_hierarchy=CustomerHierarchyAuth(
customer_member_id="member-1",
customer_team_id="team-1",
customer_org_id="org-1",
scopes=(
CustomerScopeAuth(key="customer_org", entity_id="org-1"),
CustomerScopeAuth(key="customer_team", entity_id="team-1", allowed_models=team_allowed_models),
CustomerScopeAuth(key="customer_member", entity_id="member-1", allowed_models=("luna", "opus")),
),
)
)
assert (await alternate_grant_for(customer, "opus") is not None) is granted
@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},
def _opted_in(model: str = "shared-model", **extra: object) -> dict[str, object]:
return {
"model": model,
"router_settings_override": {"fallback_on_saturation": True},
**extra,
}
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.fixture
def armed_walk() -> Iterator[None]:
from litellm.proxy.common_utils.proxy_rate_limit_error import (
arm_saturation_fallback_walk,
disarm_saturation_fallback_walk,
)
@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
walk: Final = arm_saturation_fallback_walk()
yield
disarm_saturation_fallback_walk(walk)
@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 test_an_opted_in_key_on_a_route_without_the_fallback_walk_is_admitted_on_the_saturated_model(
monkeypatch: pytest.MonkeyPatch,
):
handler = _saturation_fallback_handler(monkeypatch)
key: Final = _interactive_key("sk-flexible", "team-a")
await _saturate_to_threshold(handler, key)
direct: Final = {"model": "shared-model", "router_settings_override": {"fallback_on_saturation": True}}
async def broken(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None:
raise RuntimeError("membership lookup is down")
assert await handler.async_pre_call_hook(key, DualCache(), direct, "acompletion") is None
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
)
async def _saturate_to_threshold(handler: _PROXY_DynamicRateLimitHandlerV3, key: UserAPIKeyAuth) -> None:
for _ in range(5):
assert await handler.async_pre_call_hook(key, DualCache(), {"model": "shared-model"}, "acompletion") is None
group: Final = handler.llm_router.registry().group("shared-model")
assert group is not None
assert await handler._check_model_saturation("shared-model", group) >= 0.5
@pytest.mark.asyncio
async def test_a_key_with_fallback_on_saturation_is_sent_to_its_fallbacks_once_the_model_saturates(
monkeypatch: pytest.MonkeyPatch, armed_walk: None
):
from litellm.proxy.common_utils.proxy_rate_limit_error import SATURATION_FALLBACK_DESCRIPTOR_KEY
handler = _saturation_fallback_handler(monkeypatch)
key: Final = _interactive_key("sk-flexible", "team-a")
assert await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") is None
await _saturate_to_threshold(handler, key)
with pytest.raises(ProxyRateLimitError) as raised:
await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion")
assert raised.value.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY
assert "trying fallbacks first" in str(raised.value.detail)
@pytest.mark.asyncio
async def test_a_key_without_fallback_on_saturation_is_admitted_on_the_saturated_model(monkeypatch: pytest.MonkeyPatch):
handler = _saturation_fallback_handler(monkeypatch)
key: Final = _interactive_key("sk-rigid", "team-b")
await _saturate_to_threshold(handler, key)
assert await handler.async_pre_call_hook(key, DualCache(), {"model": "shared-model"}, "acompletion") is None
@pytest.mark.asyncio
async def test_a_request_back_from_an_exhausted_fallback_walk_is_admitted_on_the_saturated_model(
monkeypatch: pytest.MonkeyPatch, armed_walk: None
):
from litellm.proxy.common_utils.proxy_rate_limit_error import mark_saturation_fallbacks_exhausted
handler = _saturation_fallback_handler(monkeypatch)
key: Final = _interactive_key("sk-flexible", "team-a")
await _saturate_to_threshold(handler, key)
mark_saturation_fallbacks_exhausted()
assert await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") is None
def _borrowing_handler(
@ -2586,6 +2358,23 @@ async def test_a_reservation_at_saturation_is_refunded_when_the_call_fails(monke
assert await _team_tokens_remaining(handler) == 100
@pytest.mark.asyncio
async def test_a_reservation_at_saturation_is_refunded_when_the_proxy_refuses_the_request_before_the_call(
monkeypatch: pytest.MonkeyPatch,
):
"""A guardrail or the fallback walk turning the request away after admission gives the team its tokens back, and
a second release changes nothing."""
handler = _borrowing_handler(monkeypatch, borrows=True)
key: Final = _team_key()
await _check(handler, key, saturation=0.5, data=_prompt(words=30))
assert await _team_tokens_remaining(handler) < 70
for _ in range(2):
await handler.async_post_call_failure_hook(
request_data=_prompt(words=30), original_exception=Exception("guardrail refused"), user_api_key_dict=key
)
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,
@ -2806,7 +2595,7 @@ async def test_a_served_fallbacks_hold_is_not_refunded_again_under_the_name_the_
@pytest.mark.asyncio
async def test_a_steered_requests_images_are_counted_for_the_model_that_serves_it(monkeypatch: pytest.MonkeyPatch):
async def test_a_request_served_by_another_model_counts_its_images_for_that_model(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)
@ -3710,40 +3499,20 @@ async def test_a_class_share_refusal_is_counted_under_the_teams_priority(monkeyp
@pytest.mark.asyncio
async def test_an_alternates_refusal_is_not_counted_while_the_preferred_model_admits(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
monkeypatch.setattr(
litellm.priority_reservation_settings, "team_model_limit_borrows_idle_capacity_from_other_teams", True
)
registry = _prometheus_sink(monkeypatch)
handler: Final[_PROXY_DynamicRateLimitHandlerV3] = _steering_handler()
router: Final = _steering_router("luna", "opus", preferred_rpm=4)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
handler.update_variables(llm_router=router)
flexible: Final = UserAPIKeyAuth(
team_id="flexible-team",
team_models=["luna", "opus"],
team_metadata={
"model_preferences": [{"model": "luna", "alternates": ["opus"]}],
"model_rpm_limit": {"opus": 1, "luna": 1},
"team_model_limit_max_borrow_multiple": 1,
},
)
await _saturate(handler, "luna", requests=2)
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") == {
"model": "opus"
}
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team") is None
async def test_a_saturation_fallback_handoff_is_not_counted_as_a_throttle(
monkeypatch: pytest.MonkeyPatch, armed_walk: None
):
from litellm.proxy.common_utils.proxy_rate_limit_error import SATURATION_FALLBACK_DESCRIPTOR_KEY
refusal: Final = await _refused_at_the_hook(handler, flexible, model="luna")
assert refusal.headers is not None
assert refusal.headers["rate_limited_by"] == "model_per_team_saturated"
assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="luna") == 1
assert (
_sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="opus") is None
)
handler = _saturation_fallback_handler(monkeypatch)
registry = _prometheus_sink(monkeypatch)
key: Final = _interactive_key("sk-flexible", "team-a")
await _saturate_to_threshold(handler, key)
with pytest.raises(ProxyRateLimitError) as handed_off:
await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion")
assert handed_off.value.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY
assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="team-a") is None
@pytest.mark.asyncio

File diff suppressed because it is too large Load diff

View file

@ -92,7 +92,7 @@ class TestSkipPreCallLogic:
proxy_logging_obj=mock_proxy_logging,
llm_router=MagicMock(),
general_settings={},
proxy_config=MagicMock(),
proxy_config=MagicMock(get_hierarchical_router_settings=AsyncMock(return_value=None)),
)
except Exception:
pass