fix(proxy): price earlier timed-out retry attempts at their own deployment's model rate (#47648)

* fix(proxy): bill every earlier timed-out retry attempt at its own deployment's model rate

* fix(router): read the retry record's base_model through the typed metadata helper

* fix(router): keep the four-record retry breadcrumb cap

* fix(router): carry a deployment's litellm_params base_model on its retry record

* chore(router): place the cast-ok reason on the cast line and list the base_model helper as covered through log_retry

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-20 15:26:09 -07:00 • committed by GitHub
parent 0c6731eb5c
commit 27efc387f7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 174 additions and 17 deletions

View file

@ -1161,6 +1161,43 @@ def _timed_out_earlier_records(
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
def _record_pricing(
model_call_details: Mapping[str, object], record: Mapping[str, object], base_model: str | None
) -> tuple[Mapping[str, object], str | None]:
"""The ``(model_call_details, base_model)`` to price one earlier attempt with: the record's own
deployment model and ``base_model`` when it served a different model than the final attempt
(the provider is then inferred from that model), else the final attempt's."""
record_model: Final = record.get("model")
if not isinstance(record_model, str) or record_model == model_call_details.get("model"):
return model_call_details, base_model
record_base_model: Final = record.get("base_model")
return (
{**model_call_details, "model": record_model, "custom_llm_provider": None},
record_base_model if isinstance(record_base_model, str) else None,
)
def _billed_earlier_attempt(
model_call_details: Mapping[str, object],
record: Mapping[str, object],
index: int,
estimated_usage: Usage,
base_model: str | None,
) -> BilledFailedAttempt:
pricing_details, pricing_base_model = _record_pricing(model_call_details, record, base_model)
return BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
pricing_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=pricing_base_model,
),
)
def _billed_earlier_attempts(
model_call_details: Mapping[str, object],
litellm_params: Mapping[str, object] | None,
@ -1174,23 +1211,14 @@ def _billed_earlier_attempts(
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the
model). Only the exception class survives on the record, so an earlier connect timeout is
billed too; any other failure class is not."""
attempt, at that deployment's input rate (else the cost map's rate for the record's own
``base_model`` or model, which a fallback hop can make a different model than the final one).
Only the exception class survives on the record, so an earlier connect timeout is billed too;
any other failure class is not."""
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
return tuple(
BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
model_call_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=base_model,
),
)
_billed_earlier_attempt(model_call_details, record, index, estimated_usage, base_model)
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
)

View file

@ -9498,17 +9498,24 @@ class Router:
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
and count it toward the request-wide num_retries_per_request cap
"""
from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata
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]
request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs
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
base_model: Final = get_base_model_from_litellm_call_metadata(
metadata=request_metadata
) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None)
deployment_model: Final = request_metadata.get("deployment")
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,
"model": deployment_model if isinstance(deployment_model, str) else None,
"base_model": base_model if isinstance(base_model, str) else None,
"exception_type": type(e).__name__,
"exception_string": str(e),
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
@ -9539,6 +9546,18 @@ class Router:
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
return kwargs
def _deployment_params_base_model(self, deployment_id: str | None) -> str | None:
"""The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment
shape), the same key the success path's ``get_base_model_from_metadata`` reads."""
deployment: Final = self.get_deployment(model_id=deployment_id) if deployment_id is not None else None
if deployment is None:
return None
extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below
"Mapping[str, object]", deployment.litellm_params.model_extra or {}
)
base_model: Final = extra_params.get("base_model")
return base_model if isinstance(base_model, str) else None
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
"""
Update deployment rpm for that minute

View file

@ -1184,6 +1184,8 @@ class RouterModelGroupAliasItem(TypedDict):
class RetryAttemptRecord(TypedDict):
model_group: ReadOnly[str | None]
deployment_id: ReadOnly[str | None]
model: ReadOnly[str | None]
base_model: ReadOnly[str | None]
exception_type: ReadOnly[str]
exception_string: ReadOnly[str]
attempted_retries: ReadOnly[int | None]

View file

@ -93,6 +93,7 @@ ignored_function_names = [
"_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name)
"_wait_for_scheduler_turn", # Tested through prioritized acompletion and atext_completion in test_router.py
"arm_routing_read_prefetch", # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name)
"_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py
"_configured_model_info", # Tested through get_configured_service_tiers in test_router.py
"_routable_deployments", # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py
"_async_get_available_deployment", # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py

View file

@ -2347,6 +2347,74 @@ def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(mo
assert all(a.usage is usage for a in attempts)
def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_the_final_models(monkeypatch):
"""After a fallback hop the final attempt runs on another model group, so the records of the group
that timed out first name their own deployment model (and its ``base_model`` when the deployment
name is not in the cost map). Each is billed at that model's input rate, not at the final model's."""
from litellm.proxy.utils import _billed_earlier_attempts
for pricing_model, provider, input_rate in (
("lit145-final-model", "openai", 0.5),
("lit145-first-model", "openai", 0.02),
("azure/lit145-first-model", "azure", 0.02),
):
monkeypatch.setitem(
litellm.model_cost,
pricing_model,
{"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0},
)
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
attempts = _billed_earlier_attempts(
model_call_details={"model": "openai/lit145-final-model", "custom_llm_provider": "openai"},
litellm_params={
"metadata": {
"attempted_retries": 1,
"model_info": {"id": "dep-lit145-final"},
"previous_models": (
{
"deployment_id": "dep-lit145-first",
"model": "openai/lit145-first-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 0,
},
{
"deployment_id": "dep-lit145-first-azure",
"model": "azure/lit145-first-deployment",
"base_model": "azure/lit145-first-model",
"exception_type": "Timeout",
"attempted_retries": 1,
},
{
"deployment_id": "dep-lit145-final",
"model": "openai/lit145-final-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 0,
},
{
"deployment_id": "dep-lit145-final",
"model": "openai/lit145-final-model",
"base_model": None,
"exception_type": "Timeout",
"attempted_retries": 1,
},
),
}
},
estimated_usage=usage,
base_model=None,
deployment_id="dep-lit145-final",
)
assert [(a.attempted_retries, a.deployment_id, a.response_cost) for a in attempts] == [
(0, "dep-lit145-first", pytest.approx(24 * 0.02)),
(1, "dep-lit145-first-azure", pytest.approx(24 * 0.02)),
(0, "dep-lit145-final", pytest.approx(24 * 0.5)),
]
def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage():
"""A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks
delivered); the earlier timed-out attempts never produced output, so they are priced from the

View file

@ -15195,7 +15195,7 @@ async def test_cyclic_fallback_graph_does_not_amplify_one_request():
_FLAT_ATTEMPT_RECORD_KEYS = frozenset(
{"model_group", "deployment_id", "exception_type", "exception_string", "attempted_retries"}
{"model_group", "deployment_id", "model", "base_model", "exception_type", "exception_string", "attempted_retries"}
)
_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip"

View file

@ -837,7 +837,12 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
"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"},
metadata_key: {
"model_info": {"id": "deployment-1", "base_model": "azure/gpt-3.5-turbo"},
"deployment": "azure/my-35-deployment",
"attempted_retries": 2,
"user_api_key": "sk-proxy",
},
},
e=rate_limit_error,
)
@ -845,6 +850,8 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
{
"model_group": "gpt-3.5-turbo",
"deployment_id": "deployment-1",
"model": "azure/my-35-deployment",
"base_model": "azure/gpt-3.5-turbo",
"exception_type": "RateLimitError",
"exception_string": "litellm.RateLimitError: slow down",
"attempted_retries": 2,
@ -856,6 +863,38 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> None:
"""An Azure named deployment declares ``base_model`` under ``litellm_params``, not ``model_info``, and
that is the model the failure row is priced at, so the retry record has to carry it too"""
router = Router(
model_list=[
{
"model_name": "azure-group",
"litellm_params": {
"model": "azure/my-named-deployment",
"api_key": "sk-fake",
"api_base": "https://example.openai.azure.com",
"base_model": "azure/gpt-4o-mini",
},
"model_info": {"id": "dep-azure-named"},
}
]
)
new_kwargs = router.log_retry(
kwargs={
"model": "azure-group",
"metadata": {
"model_info": {"id": "dep-azure-named"},
"deployment": "azure/my-named-deployment",
"attempted_retries": 0,
},
},
e=litellm.Timeout(message="slow", llm_provider="azure", model="azure/my-named-deployment"),
)
record = new_kwargs["metadata"]["previous_models"][0]
assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini")
@pytest.mark.usefixtures("router_minute_pinned")
def test_update_usage(model_list):
"""Test if the '_update_usage' function is working correctly"""