diff --git a/litellm/constants.py b/litellm/constants.py index b29c22810e7..25eb5d9b854 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -89,6 +89,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset( "ignore_invalid_deployments", "fallback_access_check", "fallback_budget_check", + "fallback_rate_limit_check", "auto_router_capability_limit", } ) diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index 89d6729766f..154c3942edc 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -299,9 +299,13 @@ def format_fallback_outcome_message( over_budget_targets: Final = tuple( dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.OVER_BUDGET) ) + rate_limited_targets: Final = tuple( + dict.fromkeys(target.target for target in skipped if target.reason is FallbackSkipReason.RATE_LIMITED) + ) attempted_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in attempted_targets) not_authorized_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in not_authorized_targets) over_budget_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in over_budget_targets) + rate_limited_text: Final = ", ".join(str(mask_sensitive_structure(target)) for target in rate_limited_targets) attempted_failure: Final = ( (f"Fallback to {attempted_text} also failed: {fallback_failure_detail}",) if fallback_failure_detail and attempted_targets @@ -323,7 +327,15 @@ def format_fallback_outcome_message( if over_budget_targets else () ) - sentences: Final = attempted_failure + not_authorized_sentence + over_budget_sentence + rate_limited_sentence: Final = ( + ( + f"Fallback to {rate_limited_text} was skipped because the caller is over its per-model rate limit for " + f"{'it' if len(rate_limited_targets) == 1 else 'them'}.", + ) + if rate_limited_targets + else () + ) + sentences: Final = attempted_failure + not_authorized_sentence + over_budget_sentence + rate_limited_sentence if sentences: return f"{lead} {' '.join(sentences)}" diff --git a/litellm/proxy/auth/fallback_rate_limit.py b/litellm/proxy/auth/fallback_rate_limit.py new file mode 100644 index 00000000000..1d4a33b2994 --- /dev/null +++ b/litellm/proxy/auth/fallback_rate_limit.py @@ -0,0 +1,48 @@ +""" +Re-check per-model rate limits before the router takes a fallback. + +The v3 limiter admits a request once, against the model group the caller asked for. When that +group fails at the provider and the router moves to a fallback group, nothing re-checked the +caller's per-model limits on the fallback, so a provider outage on one model removed the cap on +the other for every caller at once. `RouterFallbackRateLimitCheck` runs the limiter's +model-scoped descriptors for the fallback target and reserves capacity on them, so the router +skips a target the caller has exhausted. +""" + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.auth.fallback_budget import ( + _user_api_key_auth_from_request, # pyright: ignore[reportPrivateUsage] # shared auth lookup for fallback checks +) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # the proxy's limiter type +) +from litellm.router_api import RoutingContext + + +@dataclass(frozen=True, slots=True) +class RouterFallbackRateLimitCheck: + """ + `FallbackRateLimitCheck` for the proxy's router: a fallback target is attempted only when the + caller still has per-model rate limit capacity on it, and the attempt reserves that capacity. + Requests that carry no key (for example internal health checks) and proxies running without + the v3 limiter are not restricted. + """ + + limiter: Callable[[], object] + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RoutingContext) -> bool: + limiter: Final = self.limiter() + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + return True + valid_token: Final = _user_api_key_auth_from_request(request_kwargs) + if valid_token is None: + return True + try: + return await limiter.async_fallback_target_allowed(valid_token, model, request_kwargs) + except Exception as e: # noqa: BLE001 # fail closed: a limiter failure must not hand out unchecked capacity + verbose_proxy_logger.warning("Skipping fallback to model=%s: rate limit check failed: %s", model, e) + return False diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 02da6f706d1..28d28a5c9bd 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -49,6 +49,7 @@ PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS: Final = frozenset( { "model_per_key", "model_per_team", + "model_per_end_user", "model_per_organization", "model_per_project", "model_per_project_itpm", diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index ae9fa548cc9..a3227bca235 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -79,6 +79,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.callback_utils import get_model_group_from_litellm_kwargs from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.proxy.common_utils.proxy_rate_limit_error import ( + PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, ProxyRateLimitError, map_v3_rate_limit_type, ) @@ -3777,6 +3778,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if response["overall_code"] == "OK": await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span) + async def async_fallback_target_allowed( + self, + user_api_key_dict: UserAPIKeyAuth, + model: str, + request_data: Mapping[str, object], + ) -> bool: + """Reserve the caller's per-model limits on a router fallback target. False when any is exhausted.""" + data: Final = MappingProxyType({**request_data, "model": model}) + hidden_stash: Final = _request_stash.set(None) + try: + built: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None) + finally: + _request_stash.reset(hidden_stash) + descriptors: Final = [ + descriptor for descriptor in built if descriptor["key"] in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS + ] + if not descriptors: + return True + response: Final = await self.should_rate_limit(descriptors, parent_otel_span=user_api_key_dict.parent_otel_span) + return response["overall_code"] == "OK" + @asynccontextmanager async def request_capacity( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index be90b028cf4..888276f7c31 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -454,6 +454,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check from litellm.proxy.auth.fallback_model_access import router_fallback_access_check +from litellm.proxy.auth.fallback_rate_limit import RouterFallbackRateLimitCheck from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.legacy_empty_models_check import warn_about_legacy_empty_models from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck @@ -7656,6 +7657,9 @@ class ProxyConfig: is_enforced=router_fallback_budget_check.is_enforced, team_member_model_budget_limiter=model_max_budget_limiter, ), + fallback_rate_limit_check=RouterFallbackRateLimitCheck( + limiter=lambda: proxy_logging_obj.get_proxy_hook("parallel_request_limiter"), + ), auto_router_capability_limit=_license_check.auto_router_capability_limit, ) ) @@ -8162,6 +8166,9 @@ class ProxyConfig: is_enforced=router_fallback_budget_check.is_enforced, team_member_model_budget_limiter=model_max_budget_limiter, ), + fallback_rate_limit_check=RouterFallbackRateLimitCheck( + limiter=lambda: proxy_logging_obj.get_proxy_hook("parallel_request_limiter"), + ), auto_router_capability_limit=_license_check.auto_router_capability_limit, ) ) diff --git a/litellm/router.py b/litellm/router.py index 0d46326b904..98b016f590f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -34,6 +34,7 @@ from litellm.types.router import ( DeploymentTypedDict, FallbackAccessCheck, FallbackBudgetCheck, + FallbackRateLimitCheck, GuardrailTypedDict, LiteLLM_Params, ModelGroupInfo, @@ -120,6 +121,7 @@ def _router_init( enable_mid_stream_fallback_continuation: bool = False, fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, + fallback_rate_limit_check: FallbackRateLimitCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, session_affinity_repin_on_failover: bool = False, _host: PythonHost = python_host, @@ -166,6 +168,7 @@ def _router_init( enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). + fallback_rate_limit_check (Optional[FallbackRateLimitCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over its per-model rate limit is skipped. Defaults to None (rate limits are not re-checked on fallback). Returns: Router: An instance of the litellm.Router class. diff --git a/litellm/router_api.py b/litellm/router_api.py index c85a764151f..a75397ea43d 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -80,6 +80,7 @@ if TYPE_CHECKING: DeploymentTypedDict, FallbackAccessCheck, FallbackBudgetCheck, + FallbackRateLimitCheck, GuardrailTypedDict, ModelGroupInfo, PreRoutingHookResponse, @@ -273,6 +274,7 @@ ROUTER_INIT_ARGS: Final[tuple[str, ...]] = ( "enable_mid_stream_fallback_continuation", "fallback_access_check", "fallback_budget_check", + "fallback_rate_limit_check", "auto_router_capability_limit", "session_affinity_repin_on_failover", ) @@ -1204,6 +1206,7 @@ SETTING_FIELDS: Final[tuple[str, ...]] = ( "provider_default_deployment_ids", "fallback_access_check", "fallback_budget_check", + "fallback_rate_limit_check", "kubernetes_pod_discovery", "search_tools", "guardrail_list", @@ -1264,6 +1267,7 @@ class RouterSettings: provider_default_deployment_ids: list[str] = None # pyright: ignore[reportAssignmentType] # unset only when built directly (see above) # mutable-ok: the router's live list fallback_access_check: "FallbackAccessCheck | None" = None fallback_budget_check: "FallbackBudgetCheck | None" = None + fallback_rate_limit_check: "FallbackRateLimitCheck | None" = None kubernetes_pod_discovery: "KubernetesPodDiscovery" = None # pyright: ignore[reportAssignmentType] # unset only when built directly (see above) search_tools: "list[SearchToolTypedDict]" = None # pyright: ignore[reportAssignmentType] # unset only when built directly (see above) # mutable-ok: the router's live list guardrail_list: "list[GuardrailTypedDict]" = None # pyright: ignore[reportAssignmentType] # unset only when built directly (see above) # mutable-ok: the router's live list @@ -1522,7 +1526,7 @@ def _takes_legacy_llm_router(check: object) -> bool: async def _run_fallback_check( - check: "FallbackAccessCheck | FallbackBudgetCheck", + check: "FallbackAccessCheck | FallbackBudgetCheck | FallbackRateLimitCheck", ctx: "RoutingContext", *, model: str, @@ -1666,6 +1670,10 @@ class RoutingContext: check: Final = self._router.settings().fallback_budget_check return check is None or await _run_fallback_check(check, self, model=model, request_kwargs=request_kwargs) + async def fallback_rate_limit_allowed(self, *, model: str, request_kwargs: Mapping[str, object]) -> bool: + check: Final = self._router.settings().fallback_rate_limit_check + return check is None or await _run_fallback_check(check, self, model=model, request_kwargs=request_kwargs) + def pre_call_checks( self, model: str, diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 618ac4321d8..2cf111da2a5 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -345,6 +345,7 @@ class DeprecatedRouterMembers: ) fallback_access_check = _setting("fallback_access_check", lambda r: r.settings().fallback_access_check) fallback_budget_check = _setting("fallback_budget_check", lambda r: r.settings().fallback_budget_check) + fallback_rate_limit_check = _setting("fallback_rate_limit_check", lambda r: r.settings().fallback_rate_limit_check) # These three are declared with the attribute type Router gave them rather than the descriptor's: # protocols in router_utils and the proxy match Router on them, and protocol matching does not # see through a descriptor. Reads and writes still go through _setting at runtime. diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 682aa557007..aef4db05ca9 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -330,6 +330,7 @@ _SERVED_ARGUMENTS: Final = frozenset( "default_litellm_params", "fallback_access_check", "fallback_budget_check", + "fallback_rate_limit_check", "enable_mid_stream_fallback_continuation", "redis_url", "redis_host", @@ -418,6 +419,7 @@ _DYNAMIC_SETTINGS: Final = frozenset( "default_litellm_params", "fallback_access_check", "fallback_budget_check", + "fallback_rate_limit_check", "routing_strategy_args", "enable_tag_filtering", "tag_filtering_match_any", @@ -664,6 +666,7 @@ class Resolved: default_litellm_params: Mapping[str, object] fallback_access_check: object fallback_budget_check: object + fallback_rate_limit_check: object enable_mid_stream_fallback_continuation: bool = False routing_groups: "tuple[RoutingGroup, ...]" = () enable_tag_filtering: bool = False @@ -779,6 +782,7 @@ class Resolved: default_litellm_params=MappingProxyType(dict(as_str_mapping(get("default_litellm_params")) or {})), fallback_access_check=get("fallback_access_check"), fallback_budget_check=get("fallback_budget_check"), + fallback_rate_limit_check=get("fallback_rate_limit_check"), enable_mid_stream_fallback_continuation=continuation_flag(get("enable_mid_stream_fallback_continuation")), enable_tag_filtering=bool(get("enable_tag_filtering")), tag_filtering_match_any=_true_unless_given(get("tag_filtering_match_any")), @@ -1911,6 +1915,7 @@ class RustCore( tag_routing_prefix=r.read("tag_routing_prefix"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it fallback_access_check=r.fallback_access_check, # pyright: ignore[reportArgumentType] # the check as passed fallback_budget_check=r.fallback_budget_check, # pyright: ignore[reportArgumentType] # the check as passed + fallback_rate_limit_check=r.fallback_rate_limit_check, # pyright: ignore[reportArgumentType] # the check as passed cache_responses=self._endpoints.cache_responses, search_tools=self._endpoints.search_tools, # pyright: ignore[reportArgumentType] # the configured entries, as passed guardrail_list=self._endpoints.guardrail_list, # pyright: ignore[reportArgumentType] # same @@ -5414,6 +5419,11 @@ class _Driver: if not await ctx.fallback_budget_allowed(model=target, request_kwargs=request): self._walk.targets().record_skip(target, FallbackSkipReason.OVER_BUDGET) return False + if target not in self._walk.targets() and not await ctx.fallback_rate_limit_allowed( + model=target, request_kwargs=request + ): + self._walk.targets().record_skip(target, FallbackSkipReason.RATE_LIMITED) + return False except Exception as error: # noqa: BLE001 # handed to the Rust call loop, which ends the chain self._level_errors[self._level + 1] = error return {**attempt_failure(error), "raised": True} diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 068782f7a08..88687269d4b 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -753,6 +753,47 @@ async def _is_fallback_target_within_budget( return False +async def _is_fallback_target_admitted( + ctx: "RoutingContext", + fallback_entry: str | Mapping[str, object], + original_model_group: str, + kwargs: Mapping[str, object], + *, + attempted: AttemptedFallbackTargets, +) -> bool: + return ( + await _is_fallback_target_authorized(ctx, fallback_entry, original_model_group, kwargs, attempted=attempted) + and await _is_fallback_target_within_budget( + ctx, fallback_entry, original_model_group, kwargs, attempted=attempted + ) + and await _is_fallback_target_within_rate_limit( + ctx, fallback_entry, original_model_group, kwargs, attempted=attempted + ) + ) + + +async def _is_fallback_target_within_rate_limit( + ctx: "RoutingContext", + fallback_entry: str | Mapping[str, object], + original_model_group: str, + kwargs: Mapping[str, object], + *, + attempted: AttemptedFallbackTargets, +) -> bool: + rate_limit_check: Final = ctx.settings().fallback_rate_limit_check + target: Final = _get_fallback_target_model_group(fallback_entry) + if rate_limit_check is None or target is None or target == original_model_group: + return True + if await ctx.fallback_rate_limit_allowed(model=target, request_kwargs=kwargs): + return True + verbose_router_logger.info( + "Skipping fallback to model_group = %s: caller is over its per-model rate limit", + mask_sensitive_structure(fallback_entry), + ) + attempted.record_skip(target, FallbackSkipReason.RATE_LIMITED) + return False + + def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ True when a file, batch, or fine-tuning job operation names an id that only exists @@ -863,18 +904,16 @@ async def run_async_fallback( original_model_group, ) continue - if not await _is_fallback_target_authorized(ctx, mg, original_model_group, kwargs, attempted=attempted): - continue - if not await _is_fallback_target_within_budget(ctx, mg, original_model_group, kwargs, attempted=attempted): - continue attempt_key = fallback_attempt_key(mg) + if attempt_key is not None and attempt_key in attempted: + verbose_router_logger.info( + "Skipping fallback to model_group = %s, already attempted for this request", + mask_sensitive_structure(mg), + ) + continue + if not await _is_fallback_target_admitted(ctx, mg, original_model_group, kwargs, attempted=attempted): + continue if attempt_key is not None: - if attempt_key in attempted: - verbose_router_logger.info( - "Skipping fallback to model_group = %s, already attempted for this request", - mask_sensitive_structure(mg), - ) - continue attempted.record(attempt_key) try: # LOGGING diff --git a/litellm/types/router.py b/litellm/types/router.py index 6f5652db771..bc715759c45 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -89,6 +89,7 @@ class SilentExperimentFlag(TypedDict, total=False): class FallbackSkipReason(str, enum.Enum): NOT_AUTHORIZED = "not_authorized" OVER_BUDGET = "over_budget" + RATE_LIMITED = "rate_limited" @dataclass(frozen=True, slots=True) @@ -1352,6 +1353,19 @@ class FallbackBudgetCheck(Protocol): async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: "RouterContext") -> bool: ... +class FallbackRateLimitCheck(Protocol): + """ + Decides whether the caller behind `request_kwargs` still has per-model rate limit capacity on fallback `model`. + + The limiter admits a request once, against the *requested* model group, so a provider failure would hand + the caller a fallback group whose per-model limits were never checked. The router runs this before every + cross-model-group fallback attempt, skips rejected targets, and an implementation reserves capacity when + it admits so concurrent failovers cannot all pass. + """ + + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: "RouterContext") -> bool: ... + + class AutoRouterCapabilityLimit(Protocol): """ Resolves how many complexity routers may claim each licensed capability right now; None means unlimited. diff --git a/tests/integration/routing/test_fallback_rate_limit_check.py b/tests/integration/routing/test_fallback_rate_limit_check.py new file mode 100644 index 00000000000..e7db24f5464 --- /dev/null +++ b/tests/integration/routing/test_fallback_rate_limit_check.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.routing.test_payment_required_mapping import is_model_info_probe, model_list_reply +from integration.routing.test_per_model_rate_limit_fallbacks import ( + assert_served, + chat, + error_message, + healthy_reply, + model_entry, + proxy_config, +) + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + primary: str + fallback: str + fallback_wire: Wire + + +def provider_error_reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + return Reply(status=500, body=b'{"error": {"message": "scripted provider outage", "type": "server_error"}}') + + +def served_requests(wire: Wire) -> int: + return sum(1 for request in wire.drain() if not is_model_info_probe(request)) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:10] + primary: Final = f"fallback-check-primary-{suffix}" + fallback: Final = f"fallback-check-fallback-{suffix}" + directory: Final = tmp_path_factory.mktemp("fallback-rate-limit-check") + with ExitStack() as stack: + rig_gateway: Final = stack.enter_context(gateway_from_environment()) + primary_wire: Final = stack.enter_context(wire_server(provider_error_reply)) + fallback_wire: Final = stack.enter_context(wire_server(healthy_reply)) + config: Final = proxy_config( + directory, + model_list=(model_entry(primary, primary_wire), model_entry(fallback, fallback_wire)), + fallbacks=({primary: [fallback]},), + general_settings={}, + ) + owned: Final = stack.enter_context(owned_proxy_process(rig_gateway, directory, {}, config=config, workers=2)) + yield Rig(owned.gateway, primary, fallback, fallback_wire) + + +def test_provider_error_fallback_honors_the_keys_per_model_rpm_limit_on_the_fallback_model(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_rpm_limit={rig.fallback: 1}) + served_requests(rig.fallback_wire) + + first: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}") + assert_served(first, rig.fallback, "served pmrl-") + + second: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}") + assert second.status_code == 500, (second.status_code, second.text) + assert f"Fallback to {rig.fallback} was skipped because the caller is over its per-model rate limit" in ( + error_message(second.text) + ), second.text + assert served_requests(rig.fallback_wire) == 1 diff --git a/tests/unit/proxy/auth/test_fallback_rate_limit.py b/tests/unit/proxy/auth/test_fallback_rate_limit.py new file mode 100644 index 00000000000..b1173a758cd --- /dev/null +++ b/tests/unit/proxy/auth/test_fallback_rate_limit.py @@ -0,0 +1,77 @@ +from typing import Final + +import pytest + +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.fallback_rate_limit import RouterFallbackRateLimitCheck +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 +from litellm.proxy.utils import InternalUsageCache +from litellm.router_api import RoutingContext + +PRIMARY: Final = {"model_name": "primary", "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}} +FALLBACK: Final = {"model_name": "fallback", "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}} + + +def _router() -> Router: + return Router(model_list=[PRIMARY, FALLBACK], fallbacks=[{"primary": ["fallback"]}]) + + +def _limiter() -> PROXY_MaxParallelRequestsHandler_v3: + return PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache())) + + +def _request(token: UserAPIKeyAuth) -> dict[str, object]: + return { + "model": "primary", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"user_api_key_auth": token}, + } + + +@pytest.mark.asyncio +async def test_fallback_target_is_admitted_until_the_keys_per_model_rpm_limit_is_spent() -> None: + limiter: Final = _limiter() + check: Final = RouterFallbackRateLimitCheck(limiter=lambda: limiter) + token: Final = UserAPIKeyAuth(api_key="hashed", token="hashed", metadata={"model_rpm_limit": {"fallback": 2}}) + ctx: Final = RoutingContext(_router()) + verdicts: Final = tuple([await check(model="fallback", request_kwargs=_request(token), ctx=ctx) for _ in range(3)]) + assert verdicts == (True, True, False) + + +@pytest.mark.asyncio +async def test_fallback_target_without_a_per_model_limit_is_always_admitted() -> None: + limiter: Final = _limiter() + check: Final = RouterFallbackRateLimitCheck(limiter=lambda: limiter) + token: Final = UserAPIKeyAuth(api_key="hashed", token="hashed", metadata={"model_rpm_limit": {"other": 1}}) + ctx: Final = RoutingContext(_router()) + verdicts: Final = tuple([await check(model="fallback", request_kwargs=_request(token), ctx=ctx) for _ in range(3)]) + assert verdicts == (True, True, True) + + +@pytest.mark.asyncio +async def test_request_without_a_key_is_unrestricted() -> None: + check: Final = RouterFallbackRateLimitCheck(limiter=_limiter) + assert await check(model="fallback", request_kwargs={"model": "primary"}, ctx=RoutingContext(_router())) is True + + +@pytest.mark.asyncio +async def test_proxy_without_the_v3_limiter_is_unrestricted() -> None: + check: Final = RouterFallbackRateLimitCheck(limiter=lambda: None) + token: Final = UserAPIKeyAuth(api_key="hashed", token="hashed", metadata={"model_rpm_limit": {"fallback": 0}}) + assert await check(model="fallback", request_kwargs=_request(token), ctx=RoutingContext(_router())) is True + + +@pytest.mark.asyncio +async def test_a_failing_limiter_fails_closed() -> None: + class BrokenLimiter(PROXY_MaxParallelRequestsHandler_v3): + async def async_fallback_target_allowed( + self, user_api_key_dict: object, model: str, request_data: object + ) -> bool: + raise RuntimeError("redis down") + + limiter: Final = BrokenLimiter(internal_usage_cache=InternalUsageCache(DualCache())) + check: Final = RouterFallbackRateLimitCheck(limiter=lambda: limiter) + token: Final = UserAPIKeyAuth(api_key="hashed", token="hashed", metadata={"model_rpm_limit": {"fallback": 5}}) + assert await check(model="fallback", request_kwargs=_request(token), ctx=RoutingContext(_router())) is False diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 6263292fb40..9c034b77fe7 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -5,7 +5,7 @@ import time from collections.abc import AsyncIterator, Mapping from datetime import datetime, timedelta from types import MappingProxyType -from typing import Any, Final, Literal, NoReturn +from typing import Any, Final, NoReturn from unittest.mock import MagicMock, patch import httpx @@ -20,7 +20,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_api import RouterSettings, RoutingContext from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure -from litellm.types.utils import Choices, Message, ModelResponse, StandardLoggingPayload, Usage from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, PRE_ROUTING_SELECTED_MODEL_KEY, @@ -45,6 +44,7 @@ from litellm.router_utils.fallback_event_handlers import ( run_async_fallback, ) from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget +from litellm.types.utils import Choices, Message, ModelResponse, StandardLoggingPayload, Usage from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -2501,3 +2501,113 @@ def test_is_bad_request_without_fallback_only_stops_on_a_real_400_when_the_flag_ ) -> None: assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=False) is stops_fallbacks assert is_bad_request_without_fallback(exception, fallback_on_unsupported_params=True) is False + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_a_target_the_rate_limit_check_rejects() -> None: + checked: Final = [] # mutable-ok: records the targets the check saw + + async def refuse(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + checked.append(model) + return False + + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + { + "model_name": "fallback", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "mock_response": "served"}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + fallback_rate_limit_check=refuse, + ) + with pytest.raises( + Exception, match="Fallback to fallback was skipped because the caller is over its per-model rate limit" + ): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + assert checked == ["fallback"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_serves_a_target_the_rate_limit_check_admits() -> None: + async def admit(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + return True + + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + { + "model_name": "fallback", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "mock_response": "served"}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + num_retries=0, + fallback_rate_limit_check=admit, + ) + response: Final = await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + assert response.choices[0].message.content == "served" + + +@pytest.mark.asyncio +async def test_run_async_fallback_reserves_a_repeated_target_once() -> None: + checked: Final = [] # mutable-ok: records the targets the check saw + + async def admit(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + checked.append(model) + return True + + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + { + "model_name": "fallback", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + ], + fallbacks=[{"primary": ["fallback", "fallback"]}], + num_retries=0, + fallback_rate_limit_check=admit, + ) + with pytest.raises(Exception, match="model group 'primary' failed"): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + assert checked == ["fallback"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_never_re_reserves_the_requested_model_group() -> None: + checked: Final = [] # mutable-ok: records the targets the check saw + + async def admit(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + checked.append(model) + return True + + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + { + "model_name": "fallback", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9"}, + }, + ], + fallbacks=[{"primary": ["fallback"]}, {"fallback": ["primary"]}], + num_retries=0, + fallback_rate_limit_check=admit, + ) + with pytest.raises(Exception, match="model group 'primary' failed"): + await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) + assert checked == ["fallback"] diff --git a/tests/unit/test_router/core/test_deprecated_members.py b/tests/unit/test_router/core/test_deprecated_members.py index 01cc6b3f6d9..7d5f634cda5 100644 --- a/tests/unit/test_router/core/test_deprecated_members.py +++ b/tests/unit/test_router/core/test_deprecated_members.py @@ -103,6 +103,7 @@ _CONFIGURED_SETTINGS: Final[dict[str, object]] = { # mutable-ok: expected value "provider_default_deployment_ids": ["w"], "fallback_access_check": None, "fallback_budget_check": None, + "fallback_rate_limit_check": None, "search_tools": [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily"}}], "guardrail_list": [{"guardrail_name": "g", "litellm_params": {"guardrail": "x"}}], "cache_responses": True,