mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy): count a rate-limit fallback in x-litellm-attempted-fallbacks and name it in x-litellm-fallback-reason (#51138)
* feat(proxy): count a rate-limit fallback in x-litellm-attempted-fallbacks and name it in x-litellm-fallback-reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): count every fallback candidate a local rate limit rejected in x-litellm-attempted-fallbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): count fallback attempts without a mutable list and keep self.data on the served request Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
13089b191b
commit
fdafffb842
7 changed files with 304 additions and 8 deletions
|
|
@ -1,10 +1,11 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import itertools
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from functools import lru_cache
|
||||
|
|
@ -129,6 +130,11 @@ from litellm.proxy.common_utils.openai_error_payload import (
|
|||
openai_error_type,
|
||||
relayed_provider_specific_fields,
|
||||
)
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import (
|
||||
ATTEMPTED_FALLBACKS_HEADER,
|
||||
rate_limit_fallback_from_request_data,
|
||||
rate_limit_fallback_headers,
|
||||
)
|
||||
from litellm.proxy.common_utils.responses_stream_errors import (
|
||||
SSE_LINE_BREAK,
|
||||
ResponsesStreamErrorState,
|
||||
|
|
@ -2384,12 +2390,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
response_cost: float | str | None = None,
|
||||
hidden_params: Mapping[str, object] | None = None,
|
||||
fastest_response_batch_completion: bool | None = None,
|
||||
request_data: dict | None = {},
|
||||
request_data: Mapping[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None = None,
|
||||
read_timing_from_logging_obj: bool = True,
|
||||
**kwargs,
|
||||
) -> dict:
|
||||
**kwargs: object,
|
||||
) -> dict[str, str]:
|
||||
exclude_values: Final = {"", None, "None"}
|
||||
hidden_params = hidden_params or {}
|
||||
resolved_call_id: Final = (
|
||||
|
|
@ -2453,10 +2459,14 @@ class ProxyBaseLLMRequestProcessing:
|
|||
**{k: str(v) for k, v in kwargs.items()},
|
||||
}
|
||||
if request_data:
|
||||
remaining_tokens_header: Final = get_remaining_tokens_and_requests_from_request_data(request_data)
|
||||
remaining_tokens_header: Final = get_remaining_tokens_and_requests_from_request_data(dict(request_data))
|
||||
headers.update(remaining_tokens_header)
|
||||
|
||||
logging_caching_headers: Final = get_logging_caching_headers(request_data)
|
||||
quota_fallback: Final = rate_limit_fallback_from_request_data(request_data)
|
||||
if quota_fallback is not None:
|
||||
headers.update(rate_limit_fallback_headers(quota_fallback, headers.get(ATTEMPTED_FALLBACKS_HEADER)))
|
||||
|
||||
logging_caching_headers: Final = get_logging_caching_headers(dict(request_data))
|
||||
if logging_caching_headers:
|
||||
headers.update(logging_caching_headers)
|
||||
|
||||
|
|
@ -2864,7 +2874,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
from litellm.proxy.common_utils.proxy_rate_limit_error import (
|
||||
PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS,
|
||||
ProxyRateLimitError,
|
||||
RateLimitFallback,
|
||||
per_model_rate_limits_disable_fallbacks,
|
||||
with_rate_limit_fallback,
|
||||
)
|
||||
|
||||
general_settings_view: Final = cast( # cast-ok: this method keeps its legacy bare-dict settings parameter
|
||||
|
|
@ -2933,6 +2945,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if isinstance(raw_router_settings, Mapping)
|
||||
else None
|
||||
)
|
||||
attempt_numbers: Final[Iterator[int]] = itertools.count(1)
|
||||
try:
|
||||
for fallback_model in fallback_models:
|
||||
if fallback_model == original_model:
|
||||
|
|
@ -2955,7 +2968,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict.team_id,
|
||||
):
|
||||
continue
|
||||
return await self.common_processing_pre_call_logic(
|
||||
attempt = next(attempt_numbers)
|
||||
admitted: tuple[
|
||||
Mapping[str, object], LiteLLMLoggingObj
|
||||
] = await self.common_processing_pre_call_logic(
|
||||
request=request,
|
||||
general_settings=general_settings,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
|
|
@ -2976,6 +2992,21 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if is_hard_per_model_limit(fallback_exc):
|
||||
raise
|
||||
continue
|
||||
else:
|
||||
admitted_data, admitted_logging_obj = admitted
|
||||
stamped: dict = dict( # mutable-ok: caller-defined request-payload shape
|
||||
with_rate_limit_fallback(
|
||||
admitted_data,
|
||||
get_metadata_variable_name(request, admitted_data),
|
||||
RateLimitFallback(
|
||||
original_model_group=original_model,
|
||||
descriptor_key=original_exc.descriptor_key or "unknown",
|
||||
attempts=attempt,
|
||||
),
|
||||
)
|
||||
)
|
||||
self.data = stamped
|
||||
return stamped, admitted_logging_obj
|
||||
except BaseException:
|
||||
self.data = rate_limited_data
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -37,10 +37,12 @@ This module provides a single proxy-side error class that:
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType
|
||||
|
|
@ -73,6 +75,94 @@ def per_model_rate_limits_disable_fallbacks(general_settings: Mapping[str, objec
|
|||
return False
|
||||
|
||||
|
||||
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"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RateLimitFallback:
|
||||
original_model_group: str
|
||||
descriptor_key: str
|
||||
attempts: int = 1
|
||||
|
||||
def stamp(self) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
"original_model_group": self.original_model_group,
|
||||
"descriptor_key": self.descriptor_key,
|
||||
"attempts": self.attempts,
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def reason(self) -> str:
|
||||
return f"rate_limit:{self.descriptor_key}"
|
||||
|
||||
|
||||
class _StampedBucket(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
rate_limit_fallback: RateLimitFallback | None = None
|
||||
|
||||
|
||||
class _StampedRequest(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore", frozen=True)
|
||||
metadata: _StampedBucket | None = None
|
||||
litellm_metadata: _StampedBucket | None = None
|
||||
|
||||
|
||||
_STAMPED_REQUEST: Final[TypeAdapter[_StampedRequest]] = TypeAdapter(_StampedRequest)
|
||||
_ROUTER_HOPS: Final[TypeAdapter[int]] = TypeAdapter(int)
|
||||
|
||||
|
||||
def rate_limit_fallback_from_request_data(request_data: Mapping[str, object]) -> RateLimitFallback | None:
|
||||
try:
|
||||
stamped: Final = _STAMPED_REQUEST.validate_python(request_data)
|
||||
except ValidationError:
|
||||
return None
|
||||
buckets: Final = (stamped.metadata, stamped.litellm_metadata)
|
||||
return next(
|
||||
(bucket.rate_limit_fallback for bucket in buckets if bucket is not None and bucket.rate_limit_fallback), None
|
||||
)
|
||||
|
||||
|
||||
_METADATA_BUCKET: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _metadata_bucket(value: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _METADATA_BUCKET.validate_python(value)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def with_rate_limit_fallback(
|
||||
request_data: Mapping[str, object], bucket_name: str, fallback: RateLimitFallback
|
||||
) -> Mapping[str, object]:
|
||||
bucket: Final = _metadata_bucket(request_data.get(bucket_name))
|
||||
return MappingProxyType(
|
||||
{**request_data, bucket_name: {**bucket, RATE_LIMIT_FALLBACK_METADATA_KEY: dict(fallback.stamp())}}
|
||||
)
|
||||
|
||||
|
||||
def rate_limit_fallback_headers(fallback: RateLimitFallback, attempted_fallbacks: object) -> Mapping[str, str]:
|
||||
try:
|
||||
router_hops: Final = _ROUTER_HOPS.validate_python(attempted_fallbacks)
|
||||
except ValidationError:
|
||||
return MappingProxyType(
|
||||
{
|
||||
ATTEMPTED_FALLBACKS_HEADER: str(max(fallback.attempts, 1)),
|
||||
RATE_LIMIT_FALLBACK_REASON_HEADER: fallback.reason,
|
||||
}
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
ATTEMPTED_FALLBACKS_HEADER: str(max(router_hops, 0) + max(fallback.attempts, 1)),
|
||||
RATE_LIMIT_FALLBACK_REASON_HEADER: fallback.reason,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def map_v3_rate_limit_type(
|
||||
v3_value: str | None,
|
||||
) -> RateLimitType | None:
|
||||
|
|
|
|||
|
|
@ -536,6 +536,7 @@ ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
|
|||
"attempted_fallbacks",
|
||||
"original_model_group",
|
||||
"previous_models",
|
||||
"rate_limit_fallback",
|
||||
"request_retry_count",
|
||||
CLIENT_OUTPUT_CEILING_METADATA_KEY,
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,79 @@
|
|||
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 OwnedState, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import wire_server
|
||||
from integration.routing.test_per_model_rate_limit_fallbacks import (
|
||||
assert_served,
|
||||
chat,
|
||||
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
|
||||
second_fallback: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
suffix: Final = uuid.uuid4().hex[:10]
|
||||
primary: Final = f"quota-primary-{suffix}"
|
||||
fallback: Final = f"quota-fallback-{suffix}"
|
||||
second_fallback: Final = f"quota-second-fallback-{suffix}"
|
||||
directory: Final = tmp_path_factory.mktemp("rate-limit-fallback-headers")
|
||||
with ExitStack() as stack:
|
||||
rig_gateway: Final = stack.enter_context(gateway_from_environment())
|
||||
primary_wire: Final = stack.enter_context(wire_server(healthy_reply))
|
||||
fallback_wire: Final = stack.enter_context(wire_server(healthy_reply))
|
||||
second_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),
|
||||
model_entry(second_fallback, second_fallback_wire),
|
||||
),
|
||||
fallbacks=({primary: [fallback, second_fallback]},),
|
||||
general_settings={},
|
||||
)
|
||||
owned: Final = stack.enter_context(
|
||||
owned_proxy_process(
|
||||
rig_gateway, directory, {}, config=config, workers=2, state=OwnedState.SHARED_WITH_WORKER
|
||||
)
|
||||
)
|
||||
yield Rig(owned.gateway, primary, fallback, second_fallback)
|
||||
|
||||
|
||||
def test_quota_fallback_response_names_the_hop_and_its_reason(rig: Rig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(model_tpm_limit={rig.primary: 1})
|
||||
|
||||
response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}")
|
||||
assert_served(response, rig.fallback, "served pmrl-")
|
||||
assert response.headers.get("x-litellm-attempted-fallbacks") == "1", dict(response.headers)
|
||||
assert response.headers.get("x-litellm-fallback-reason") == "rate_limit:model_per_key", dict(response.headers)
|
||||
|
||||
|
||||
def test_quota_fallback_response_counts_every_candidate_a_limit_rejected(rig: Rig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
key: Final = scenario.key(model_tpm_limit={rig.primary: 1, rig.fallback: 1})
|
||||
|
||||
response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}")
|
||||
assert_served(response, rig.second_fallback, "served pmrl-")
|
||||
assert response.headers.get("x-litellm-attempted-fallbacks") == "2", dict(response.headers)
|
||||
assert response.headers.get("x-litellm-fallback-reason") == "rate_limit:model_per_key", dict(response.headers)
|
||||
58
tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py
Normal file
58
tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import (
|
||||
RateLimitFallback,
|
||||
rate_limit_fallback_from_request_data,
|
||||
rate_limit_fallback_headers,
|
||||
)
|
||||
|
||||
STAMP: Final = {"original_model_group": "opus", "descriptor_key": "model_per_end_user"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bucket", ("metadata", "litellm_metadata"))
|
||||
def test_rate_limit_fallback_is_read_from_either_metadata_bucket(bucket: str) -> None:
|
||||
fallback: Final = rate_limit_fallback_from_request_data({"model": "luna", bucket: {"rate_limit_fallback": STAMP}})
|
||||
|
||||
assert fallback is not None
|
||||
assert fallback == RateLimitFallback(original_model_group="opus", descriptor_key="model_per_end_user")
|
||||
assert fallback.reason == "rate_limit:model_per_end_user"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"request_data",
|
||||
(
|
||||
{"model": "luna"},
|
||||
{"model": "luna", "metadata": {}},
|
||||
{"model": "luna", "metadata": {"rate_limit_fallback": "opus"}},
|
||||
{"model": "luna", "metadata": {"rate_limit_fallback": {"original_model_group": "opus"}}},
|
||||
),
|
||||
)
|
||||
def test_a_request_without_a_complete_stamp_has_no_rate_limit_fallback(request_data: dict[str, object]) -> None:
|
||||
assert rate_limit_fallback_from_request_data(request_data) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("router_hops", "expected"), ((None, "1"), (0, "1"), ("0", "1"), (2, "3"), ("junk", "1")))
|
||||
def test_rate_limit_fallback_headers_count_the_hop_on_top_of_the_router(router_hops: object, expected: str) -> None:
|
||||
headers: Final = rate_limit_fallback_headers(
|
||||
RateLimitFallback(original_model_group="opus", descriptor_key="model_per_key"), router_hops
|
||||
)
|
||||
|
||||
assert headers == {
|
||||
"x-litellm-attempted-fallbacks": expected,
|
||||
"x-litellm-fallback-reason": "rate_limit:model_per_key",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("attempts", "router_hops", "expected"), ((2, None, "2"), (2, 0, "2"), (3, 1, "4"), (0, 0, "1"))
|
||||
)
|
||||
def test_rate_limit_fallback_headers_count_every_candidate_a_limit_rejected(
|
||||
attempts: int, router_hops: object, expected: str
|
||||
) -> None:
|
||||
headers: Final = rate_limit_fallback_headers(
|
||||
RateLimitFallback(original_model_group="opus", descriptor_key="model_per_key", attempts=attempts), router_hops
|
||||
)
|
||||
|
||||
assert headers["x-litellm-attempted-fallbacks"] == expected
|
||||
|
|
@ -13742,3 +13742,38 @@ class TestTotalTimeout:
|
|||
assert not prepared
|
||||
assert not called
|
||||
assert elapsed < 120
|
||||
|
||||
|
||||
@pytest.mark.parametrize("router_attempted_fallbacks", (0, 1))
|
||||
def test_get_custom_headers_count_a_quota_fallback_and_name_its_reason(router_attempted_fallbacks: int) -> None:
|
||||
auth: Final = ProxyUserAPIKeyAuth(api_key="sk-quota-fallback", tpm_limit=None, rpm_limit=None, max_budget=None)
|
||||
request_data: Final = {
|
||||
"model": "luna",
|
||||
"metadata": {"rate_limit_fallback": {"original_model_group": "opus", "descriptor_key": "model_per_key"}},
|
||||
}
|
||||
|
||||
custom_headers: Callable[..., dict[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers
|
||||
headers: Final = custom_headers(
|
||||
user_api_key_dict=auth,
|
||||
call_id="quota-fallback",
|
||||
request_data=request_data,
|
||||
**{"x-litellm-attempted-fallbacks": router_attempted_fallbacks},
|
||||
)
|
||||
|
||||
assert headers["x-litellm-attempted-fallbacks"] == str(router_attempted_fallbacks + 1)
|
||||
assert headers["x-litellm-fallback-reason"] == "rate_limit:model_per_key"
|
||||
|
||||
|
||||
def test_get_custom_headers_leave_a_plain_request_without_a_fallback_reason() -> None:
|
||||
auth: Final = ProxyUserAPIKeyAuth(api_key="sk-no-fallback", tpm_limit=None, rpm_limit=None, max_budget=None)
|
||||
|
||||
custom_headers: Callable[..., dict[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers
|
||||
headers: Final = custom_headers(
|
||||
user_api_key_dict=auth,
|
||||
call_id="no-fallback",
|
||||
request_data={"model": "opus", "metadata": {}},
|
||||
**{"x-litellm-attempted-fallbacks": 0},
|
||||
)
|
||||
|
||||
assert headers["x-litellm-attempted-fallbacks"] == "0"
|
||||
assert "x-litellm-fallback-reason" not in headers
|
||||
|
|
|
|||
|
|
@ -9369,6 +9369,7 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth:
|
|||
_PLANTED_STAMPS = {
|
||||
"attempted_fallbacks": 99,
|
||||
"original_model_group": "spoofed-group",
|
||||
"rate_limit_fallback": {"original_model_group": "spoofed-group", "descriptor_key": "model_per_key"},
|
||||
"request_retry_count": -100,
|
||||
"previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}],
|
||||
"_client_output_ceiling": {"api_base": "https://attacker.example"},
|
||||
|
|
@ -9404,6 +9405,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
|
|||
assert "litellm_metadata" not in updated
|
||||
assert "attempted_fallbacks" not in updated["metadata"]
|
||||
assert "original_model_group" not in updated["metadata"]
|
||||
assert "rate_limit_fallback" not in updated["metadata"]
|
||||
assert "_client_output_ceiling" not in updated["metadata"]
|
||||
assert "request_retry_count" not in updated["metadata"]
|
||||
assert "previous_models" not in updated["metadata"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue