Merge pull request #40930 from BerriAI/litellm_flat_retry_records

fix(router): record flat retry attempts and cap retries from attempted_retries
This commit is contained in:
Yassin Kortam 2026-09-14 13:37:10 -07:00 committed by GitHub
commit 97ddb9494e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 204 additions and 107 deletions

View file

@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
####### SECRET MANAGERS #####################
secret_manager_client: Optional[Any] = (
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.

View file

@ -303,6 +303,16 @@ def get_metadata_variable_name_from_kwargs(
return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_request: int | None) -> bool:
if num_retries_per_request is None:
return False
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
if not isinstance(metadata, Mapping):
return False
attempted_retries: Final = metadata.get("attempted_retries")
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
def get_or_create_metadata_bucket(
request_data: dict,
) -> tuple[Literal["metadata", "litellm_metadata"], dict]:

View file

@ -96,7 +96,6 @@ from litellm.litellm_core_utils.request_timeout_resolver import (
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.sensitive_data_masker import (
SensitiveDataMasker,
mask_credentials_in_payload,
mask_sensitive_structure,
)
from litellm.litellm_core_utils.token_counter import offload_token_count
@ -623,20 +622,6 @@ def _replay_live_router_model_cost() -> None:
set_live_deployment_replay(_replay_live_router_model_cost)
# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a
# breadcrumb entirely: the request payload, the proxy's snapshot of the inbound request (its body
# aliases the live request metadata, earlier breadcrumbs included, so copying it would nest every
# breadcrumb inside the next one), and the router-internal walk state. Credentials are handled
# separately by mask_credentials_in_payload, which scrubs credential-named values from whatever
# kwargs remain rather than trying to enumerate every credential-bearing key here.
RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
(
"messages",
"original_function",
"attempted_targets",
"proxy_server_request",
)
)
RETRY_BREADCRUMB_LIMIT: Final = 4
@ -8374,31 +8359,30 @@ class Router:
def log_retry(self, kwargs: dict, e: Exception) -> dict:
"""
When a retry or fallback happens, log the details of the just failed model call - similar to Sentry breadcrumbing
When a retry or fallback happens, record which model group, deployment and attempt just failed and why
"""
from litellm.types.router import RetryAttemptRecord
_metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var]
attempt_kwargs: Final = MappingProxyType(
{k: v for k, v in kwargs.items() if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS}
)
attempt_metadata: Final = MappingProxyType(
{k: v for k, v in request_metadata.items() if k != "previous_models"}
)
previous_model: Final = MappingProxyType(
{
"exception_type": type(e).__name__,
"exception_string": str(e),
**attempt_kwargs,
_metadata_var: attempt_metadata,
}
)
model_group: Final = kwargs.get("model")
model_info: Final = request_metadata.get("model_info")
deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None
attempted_retries: Final = request_metadata.get("attempted_retries")
attempt_record: Final[RetryAttemptRecord] = {
"model_group": model_group if isinstance(model_group, str) else None,
"deployment_id": deployment_id if isinstance(deployment_id, str) else None,
"exception_type": type(e).__name__,
"exception_string": str(e),
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
}
earlier_breadcrumbs: Final = request_metadata.get("previous_models")
kept_breadcrumbs: Final[tuple[object, ...]] = (
tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :]
if isinstance(earlier_breadcrumbs, (list, tuple))
else ()
)
breadcrumbs: Final = (*kept_breadcrumbs, mask_credentials_in_payload(previous_model))
breadcrumbs: Final = (*kept_breadcrumbs, attempt_record)
kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict
return kwargs

View file

@ -99,23 +99,13 @@ def setup(
def check_limits(kwargs: Mapping[str, object]) -> None:
import litellm
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor
if litellm.max_budget and current_cost > litellm.max_budget:
raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget)
metadata: Final = kwargs.get("metadata")
if isinstance(metadata, Mapping):
typed_metadata: Final = cast( # cast-ok: runtime Mapping check establishes read-only metadata
Mapping[str, object], metadata
)
previous: Final = typed_metadata.get("previous_models")
if (
isinstance(previous, list)
and litellm.num_retries_per_request is not None
and len(cast(list[object], previous)) # cast-ok: runtime list check establishes the retry history
>= litellm.num_retries_per_request
):
raise RuntimeError("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise RuntimeError("Max retries per request hit!")
def finalize(

View file

@ -908,6 +908,14 @@ class RouterModelGroupAliasItem(TypedDict):
hidden: bool # if 'True', don't return on `.get_model_list`
class RetryAttemptRecord(TypedDict):
model_group: ReadOnly[str | None]
deployment_id: ReadOnly[str | None]
exception_type: ReadOnly[str]
exception_string: ReadOnly[str]
attempted_retries: ReadOnly[int | None]
VALID_LITELLM_ENVIRONMENTS = [
"development",
"staging",

View file

@ -1491,6 +1491,8 @@ def post_call_processing(
def client(original_function):
from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit
Rules: Final = litellm_utils.Rules
rules_obj: Final = Rules()
@ -1501,12 +1503,8 @@ def client(original_function):
call_type = original_function.__name__
if _is_async_request(kwargs):
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
if litellm.num_retries_per_request is not None:
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
if previous_models is not None:
if litellm.num_retries_per_request <= len(previous_models):
raise Exception("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise Exception("Max retries per request hit!")
# MODEL CALL
result = original_function(*args, **kwargs)
@ -1565,12 +1563,8 @@ def client(original_function):
)
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
if litellm.num_retries_per_request is not None:
# check if previous_models passed in as ['litellm_params']['metadata]['previous_models']
previous_models = (kwargs.get("metadata") or {}).get("previous_models", None)
if previous_models is not None:
if litellm.num_retries_per_request <= len(previous_models):
raise Exception("Max retries per request hit!")
if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request):
raise Exception("Max retries per request hit!")
# [OPTIONAL] CHECK CACHE
print_verbose(

View file

@ -1,3 +1,4 @@
import json
import os
import traceback
from dotenv import load_dotenv
@ -628,17 +629,29 @@ def test_deployment_callback_respects_cooldown_time(model_list):
assert mock_set.call_args.kwargs["time_to_cooldown"] == 0
def test_log_retry(model_list):
"""Test if the '_log_retry' function is working correctly"""
import time
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
def test_log_retry(model_list, metadata_key):
"""log_retry appends one flat record per failed attempt and copies neither the request kwargs nor
the request metadata into it"""
router = Router(model_list=model_list)
new_kwargs = router.log_retry(
kwargs={"metadata": {}},
e=Exception(),
kwargs={
"model": "gpt-3.5-turbo",
"api_key": "sk-must-not-be-recorded",
"messages": [{"role": "user", "content": "hi"}],
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
},
e=litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-3.5-turbo"),
)
assert "metadata" in new_kwargs
assert "previous_models" in new_kwargs["metadata"]
assert json.loads(json.dumps(new_kwargs[metadata_key]["previous_models"])) == [
{
"model_group": "gpt-3.5-turbo",
"deployment_id": "deployment-1",
"exception_type": "RateLimitError",
"exception_string": "litellm.RateLimitError: slow down",
"attempted_retries": 2,
}
]
def test_update_usage(model_list):

View file

@ -0,0 +1,30 @@
from typing import Final
import pytest
import litellm
from litellm.rust_bridge.lifecycle import check_limits
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize(
"cap, attempted_retries, refused",
[(5, 5, True), (5, 4, False), (0, 0, False), (0, 1, True)],
ids=[
"cap-above-four-reached",
"cap-above-four-not-reached",
"first-attempt-passes-cap-of-zero",
"cap-of-zero-refuses-first-retry",
],
)
def test_check_limits_reads_attempted_retries(
monkeypatch: pytest.MonkeyPatch, metadata_key: str, cap: int, attempted_retries: int, refused: bool
) -> None:
monkeypatch.setattr(litellm, "num_retries_per_request", cap)
monkeypatch.setattr(litellm, "max_budget", None)
kwargs: Final = {"model": "mistral/mistral-ocr-latest", metadata_key: {"attempted_retries": attempted_retries}}
if refused:
with pytest.raises(RuntimeError, match="Max retries per request hit!"):
check_limits(kwargs)
else:
check_limits(kwargs)

View file

@ -10706,6 +10706,7 @@ def _cyclic_fallback_router(num_retries=0):
"api_key": "sk-fake",
"mock_response": "litellm.InternalServerError",
},
"model_info": {"id": f"{group}-deployment"},
}
for group in groups
],
@ -10755,28 +10756,37 @@ async def test_cyclic_fallback_graph_does_not_amplify_one_request():
assert sum(len(message) for message in capture.messages) < 5_000
_FLAT_ATTEMPT_RECORD_KEYS = frozenset(
{"model_group", "deployment_id", "exception_type", "exception_string", "attempted_retries"}
)
_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip"
@pytest.mark.asyncio
async def test_retry_breadcrumbs_do_not_carry_the_walk_state():
"""log_retry copies every kwarg into previous_models, which reaches spend logs and
logging callbacks. The set of already-attempted groups is router-internal walk state
with no diagnostic value there, and it is the one entry that is not a plain scalar.
A retry has to be configured for the walk state to reach log_retry at all."""
async def test_retry_records_are_flat_and_name_the_failed_group_on_fallback_hops():
"""Each failed attempt leaves a flat record in previous_models, which reaches spend logs and
logging callbacks. Nothing downstream reads the failed attempt's kwargs or metadata, and copying
them is what carried client credentials and multiplied the payload on every retry. A fallback hop
calls log_retry too, so the record has to name the group that failed, not the one taken next."""
router = _cyclic_fallback_router(num_retries=1)
capture = _LogCapture(logging.ERROR)
recorder = _FallbackAttemptRecorder()
await _drive_cyclic_fallback(router, capture, recorder)
breadcrumbs = [breadcrumb for hop in recorder.breadcrumbs_per_target for breadcrumb in hop]
assert breadcrumbs, "no retry breadcrumbs were recorded"
assert any(
"fallback_depth" in breadcrumb for breadcrumb in breadcrumbs
), "no breadcrumb carried router walk state, so this test cannot see the leak"
for breadcrumb in breadcrumbs:
assert "attempted_targets" not in breadcrumb
_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip"
records = [record for hop in recorder.breadcrumbs_per_target for record in hop]
assert records, "no retry records were recorded"
for record in records:
assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS
assert record["exception_type"] == "InternalServerError"
assert record["deployment_id"] == f"{record['model_group']}-deployment"
group_failed_before_hop = {"group-b": "group-a", "group-c": "group-b", "group-d": "group-c"}
for failed_target, hop_records in zip(recorder.failed_targets, recorder.breadcrumbs_per_target):
groups = [record["model_group"] for record in hop_records]
first_own_attempt = groups.index(failed_target)
assert groups[first_own_attempt - 1] == group_failed_before_hop[failed_target]
assert set(groups[first_own_attempt:]) == {failed_target}
assert [record["attempted_retries"] for record in hop_records[first_own_attempt:]][:2] == [0, 1]
@pytest.mark.parametrize(
@ -10802,22 +10812,20 @@ _BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doN
],
)
@pytest.mark.asyncio
async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_key, request_kwargs):
"""log_retry copies kwargs into previous_models, which reaches spend logs and logging callbacks.
Any of these kwargs can carry a client's forwarded Authorization token or a provider key, and a
breadcrumb has no diagnostic use for the raw secret. A denylist of key names is always one new
credential kwarg behind, so log_retry scrubs credential-named values by pattern instead: the
container still reaches the breadcrumb, but the raw secret never does, whatever key holds it."""
async def test_retry_records_never_carry_a_forwarded_credential(container_key, request_kwargs):
"""previous_models reaches spend logs and logging callbacks. Any request kwarg can carry a client's
forwarded Authorization token or a provider key, so the record must not carry request kwargs at
all: neither the credential-bearing container nor the raw secret, whatever key holds it."""
router = _cyclic_fallback_router(num_retries=1)
capture = _LogCapture(logging.ERROR)
metadata = {}
await _drive_cyclic_fallback(router, capture, metadata=metadata, **request_kwargs)
breadcrumbs = metadata["previous_models"]
assert breadcrumbs, "no retry breadcrumbs were recorded"
dumped = json.dumps(breadcrumbs, default=str)
assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak"
records = metadata["previous_models"]
assert records, "no retry records were recorded"
dumped = json.dumps(records)
assert container_key not in dumped
assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped
@ -10837,12 +10845,12 @@ def _always_failing_router(num_retries):
)
async def _fail_one_proxy_shaped_request(router, request_marker):
async def _fail_one_proxy_shaped_request(router, request_marker, expected_error=litellm.InternalServerError):
"""The proxy hands the router a metadata dict and a proxy_server_request whose body is a
shallow copy of the request, so body["metadata"] is the very same dict the router later
stamps previous_models onto."""
metadata = {"request_marker": request_marker}
with pytest.raises(litellm.InternalServerError):
with pytest.raises(expected_error):
await router.acompletion(
model="broken-group",
messages=[{"role": "user", "content": "hi"}],
@ -10868,34 +10876,49 @@ def _nested_breadcrumb_lists(node):
@pytest.mark.asyncio
async def test_retry_breadcrumbs_stay_per_request_and_flat_across_failing_requests():
"""Every failed attempt appends a breadcrumb to metadata["previous_models"], and the proxy's
async def test_retry_records_stay_per_request_and_flat_across_failing_requests():
"""Every failed attempt appends a record to metadata["previous_models"], and the proxy's
request snapshot aliases that same metadata dict. Kept on the Router and copied wholesale,
each breadcrumb embedded every earlier one from every earlier request, so the breadcrumb
each breadcrumb once embedded every earlier one from every earlier request, so the breadcrumb
tree, and with it the debug repr of the kwargs, roughly doubled on each failed attempt until
a single-worker proxy spent minutes in the redaction regex and stopped answering."""
router = _always_failing_router(num_retries=2)
breadcrumbs_per_request = [
records_per_request = [
await _fail_one_proxy_shaped_request(router, f"request-{request_number}") for request_number in range(1, 7)
]
for request_number, breadcrumbs in enumerate(breadcrumbs_per_request, start=1):
assert len(breadcrumbs) == 3, "one initial attempt plus two retries failed, each leaving one breadcrumb"
assert {breadcrumb["metadata"]["request_marker"] for breadcrumb in breadcrumbs} == {f"request-{request_number}"}
for breadcrumb in breadcrumbs:
assert _nested_breadcrumb_lists(breadcrumb) == []
assert len({len(repr(breadcrumbs)) for breadcrumbs in breadcrumbs_per_request}) == 1
for records in records_per_request:
assert [record["attempted_retries"] for record in records] == [0, 1, 2]
for record in records:
assert set(record) == _FLAT_ATTEMPT_RECORD_KEYS
assert _nested_breadcrumb_lists(record) == []
assert len({len(repr(records)) for records in records_per_request}) == 1
@pytest.mark.asyncio
async def test_retry_breadcrumbs_keep_only_the_last_four_attempts():
async def test_retry_records_keep_only_the_last_four_attempts():
router = _always_failing_router(num_retries=6)
breadcrumbs = await _fail_one_proxy_shaped_request(router, "request-1")
records = await _fail_one_proxy_shaped_request(router, "request-1")
assert len(breadcrumbs) == 4
assert [breadcrumb["metadata"]["attempted_retries"] for breadcrumb in breadcrumbs] == [3, 4, 5, 6]
assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6]
@pytest.mark.asyncio
async def test_num_retries_per_request_stops_retries_at_caps_above_four(monkeypatch):
monkeypatch.setattr(litellm, "num_retries_per_request", 5)
router = _always_failing_router(num_retries=6)
records = await _fail_one_proxy_shaped_request(router, "request-1", expected_error=litellm.APIConnectionError)
assert [record["attempted_retries"] for record in records] == [3, 4, 5, 6]
assert ["Max retries per request hit!" in record["exception_string"] for record in records] == [
False,
False,
True,
True,
]
@pytest.mark.asyncio

View file

@ -4063,6 +4063,51 @@ class TestMetadataNoneHandling:
assert metadata == {}
_RETRY_CAP_CASES: Final = (
pytest.param(5, {"attempted_retries": 5}, True, id="cap-above-four-reached"),
pytest.param(5, {"attempted_retries": 4}, False, id="cap-above-four-not-reached"),
pytest.param(0, {"attempted_retries": 0}, False, id="first-attempt-passes-cap-of-zero"),
pytest.param(0, {"attempted_retries": 1}, True, id="cap-of-zero-refuses-first-retry"),
pytest.param(5, {"previous_models": ("a", "b", "c", "d", "e")}, False, id="breadcrumb-count-is-not-the-cap"),
pytest.param(5, None, False, id="metadata-none"),
)
def _capped_completion_kwargs(metadata_key: str, metadata: object) -> dict[str, object]:
return {
"model": "openai/gpt-4o-mini",
"messages": [{"role": "user", "content": "hi"}],
"api_key": "sk-fake",
"mock_response": "ok",
metadata_key: metadata,
}
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES)
def test_num_retries_per_request_reads_attempted_retries_sync(monkeypatch, metadata_key, cap, metadata, refused):
monkeypatch.setattr(litellm, "num_retries_per_request", cap)
kwargs: Final = _capped_completion_kwargs(metadata_key, metadata)
if refused:
with pytest.raises(Exception, match="Max retries per request hit!"):
litellm.completion(**kwargs)
else:
assert litellm.completion(**kwargs).choices[0].message.content == "ok"
@pytest.mark.asyncio
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
@pytest.mark.parametrize("cap, metadata, refused", _RETRY_CAP_CASES)
async def test_num_retries_per_request_reads_attempted_retries_async(monkeypatch, metadata_key, cap, metadata, refused):
monkeypatch.setattr(litellm, "num_retries_per_request", cap)
kwargs: Final = _capped_completion_kwargs(metadata_key, metadata)
if refused:
with pytest.raises(Exception, match="Max retries per request hit!"):
await litellm.acompletion(**kwargs)
else:
assert (await litellm.acompletion(**kwargs)).choices[0].message.content == "ok"
class TestValidateAndFixThinkingParam:
"""Tests for validate_and_fix_thinking_param."""

View file

@ -806,7 +806,7 @@ async def test_shared_call_limits_still_reject_before_reading_ocr_file(
monkeypatch.setattr(litellm, "_current_cost", 2)
monkeypatch.setattr(litellm, "num_retries_per_request", 1 if limit == "retries" else None)
expected: Final = litellm.BudgetExceededError if limit == "budget" else RuntimeError
arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"previous_models": ["earlier"]}}
arguments: Final = {"document": {"type": "file", "file": File()}, "metadata": {"attempted_retries": 1}}
with pytest.raises(expected, match=r"Budget has been exceeded|Max retries per request hit"):
await call_aocr(ocr_server, **arguments) if asynchronous else call_ocr(ocr_server, **arguments)
assert reads == []