This commit is contained in:
Aidan Singh 2026-10-22 02:10:17 +00:00 • committed by GitHub
commit 9bd707e8e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 341 additions and 17 deletions

View file

@ -357,6 +357,8 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"fallback_depth",
"_target_order",
"attempted_targets",
"_excluded_deployment_ids",
"_retry_skipped_deployment_ids",
"proxy_server_request",
"standard_logging_object",
"secret_fields",
@ -438,6 +440,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"_pipeline_managed_guardrails",
"client_disconnected",
"error_information",
"_failover_excluded_ids",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)

View file

@ -210,6 +210,7 @@ from litellm.router_utils.clientside_credential_handler import (
)
from litellm.router_utils.common_utils import (
MID_STREAM_CONTINUATION_KWARG,
as_retry_skipped_deployment_ids,
deployments_the_caller_may_use,
filter_team_based_models,
filter_web_search_deployments,
@ -673,10 +674,6 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj
return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {})
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
if silent_model is None:
return ()
@ -8858,7 +8855,7 @@ class Router:
return ()
if litellm.should_retry(status_code):
return ()
already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped)
already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped)
skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id))))
verbose_router_logger.debug(
"Retry skips deployments that already answered %s to this request: %s", status_code, skipped
@ -14935,7 +14932,7 @@ class Router:
## RETRY SKIP ## -> drop deployments that already refused this request with a
## non-retryable status, unless that leaves nothing, so the caller still gets
## the provider's own error instead of a no-deployments error.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
_retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
@ -16169,7 +16166,7 @@ class Router:
)
## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
_retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (

View file

@ -4,7 +4,7 @@ from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeVar
from pydantic import TypeAdapter
from pydantic import TypeAdapter, ValidationError
if TYPE_CHECKING:
from litellm.types.llms.openai import OpenAIFileObject
@ -38,6 +38,7 @@ _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapt
_DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(
tuple[Mapping[str, object], ...]
)
_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...])
def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment:
@ -57,6 +58,31 @@ def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
return {key: value for key, value in kwargs.items() if key not in ROUTER_ONLY_CALL_KWARGS}
def as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
if not isinstance(value, tuple):
return ()
items: Final[tuple[object, ...]] = value
return tuple(item for item in items if isinstance(item, str))
def as_excluded_deployment_ids(value: object) -> tuple[str, ...]:
if value is None:
return ()
try:
return _DEPLOYMENT_IDS_ADAPTER.validate_python(value)
except ValidationError:
return ()
def deployment_ids_failed_in_request(request_kwargs: Mapping[str, object]) -> frozenset[str]:
return frozenset(
(
*as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")),
*as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")),
)
)
def request_user_api_key_auth(request_kwargs: Mapping[str, object]) -> object:
"""The caller's auth object, from whichever metadata bucket this surface writes to."""
metadata_value: Final = request_kwargs.get("metadata")

View file

@ -30,6 +30,7 @@ from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import deployment_ids_failed_in_request
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes
@ -405,6 +406,9 @@ class DeploymentAffinityCheck(CustomLogger):
session_affinity_active: Final = (
enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None
)
failed_deployment_ids: Final = deployment_ids_failed_in_request(
cast(Mapping[str, object], request_kwargs) # cast-ok: request kwargs keys are param names
)
user_key: Final = (
self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if (session_affinity_active or enable_user_key)
@ -426,7 +430,7 @@ class DeploymentAffinityCheck(CustomLogger):
elif isinstance(session_cache_result, str):
session_model_id = session_cache_result
if session_model_id:
if session_model_id and session_model_id not in failed_deployment_ids:
session_deployment: Final = self._find_deployment_by_model_id(
healthy_deployments=typed_healthy_deployments,
model_id=session_model_id,
@ -461,7 +465,7 @@ class DeploymentAffinityCheck(CustomLogger):
# Backwards / safety: allow raw string values.
model_id = cache_result
if not model_id:
if not model_id or model_id in failed_deployment_ids:
return typed_healthy_deployments
deployment = self._find_deployment_by_model_id(

View file

@ -51,6 +51,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
strip_encrypted_reasoning_from_messages,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import as_excluded_deployment_ids, as_retry_skipped_deployment_ids
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import Deployment
from litellm.utils import get_order_filtered_deployments
@ -182,6 +183,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
None,
)
@staticmethod
def _deployment_model_id(deployment: Mapping[str, object]) -> str | None:
model_info: Final = deployment.get("model_info")
if not isinstance(model_info, Mapping):
return None
deployment_id: Final = cast(Mapping[str, object], model_info).get("id") # cast-ok: narrowed by isinstance
return None if deployment_id is None else str(deployment_id)
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
for deployment in healthy_deployments:
@ -345,7 +354,8 @@ class EncryptedContentAffinityCheck(CustomLogger):
typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments)
if not self._is_enabled_for_model_group(model):
return typed_healthy_deployments
target_order: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs).get("_target_order")
typed_request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(request_kwargs)
target_order: Final = typed_request_kwargs.get("_target_order")
candidates: Final = (
get_order_filtered_deployments(typed_healthy_deployments, target_order=target_order)
if isinstance(target_order, int)
@ -375,8 +385,22 @@ class EncryptedContentAffinityCheck(CustomLogger):
model_id,
)
excluded_ids: Final = frozenset(
as_excluded_deployment_ids(typed_request_kwargs.get("_excluded_deployment_ids"))
)
retry_skipped_ids: Final = frozenset(
as_retry_skipped_deployment_ids(typed_request_kwargs.get("_retry_skipped_deployment_ids"))
)
not_excluded_candidates: Final = [
candidate for candidate in candidates if self._deployment_model_id(candidate) not in excluded_ids
]
untried_candidates: Final = [
candidate
for candidate in not_excluded_candidates
if self._deployment_model_id(candidate) not in retry_skipped_ids
] or not_excluded_candidates
deployment: Final = self._find_deployment_by_model_id(
healthy_deployments=candidates,
healthy_deployments=untried_candidates,
model_id=model_id,
)
if deployment is not None:
@ -392,7 +416,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
# Follow-up switched model_name (LIT-2531): pin by Azure resource instead.
boundary_matches, _originating = self._find_deployments_on_same_encryption_boundary(
healthy_deployments=candidates,
healthy_deployments=untried_candidates,
model_id=model_id,
)
if boundary_matches:

View file

@ -19,6 +19,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model
from litellm.router_utils.common_utils import deployment_ids_failed_in_request
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import CallTypes, StandardLoggingPayload
from litellm.utils import (
@ -152,7 +153,9 @@ class PromptCachingDeploymentCheck(CustomLogger):
request_tools,
),
)
if model_id_dict is not None:
if model_id_dict is not None and model_id_dict["model_id"] not in deployment_ids_failed_in_request(
request_data or {}
):
model_id: Final = model_id_dict["model_id"]
for deployment in healthy_deployments:
if deployment["model_info"]["id"] == model_id:

View file

@ -1129,6 +1129,8 @@ async def test_add_litellm_data_to_request_strips_string_encoded_admin_injection
("fallback_depth", True),
("_target_order", 2),
("attempted_targets", ["forged"]),
("_excluded_deployment_ids", ["pinned-deployment"]),
("_retry_skipped_deployment_ids", ["pinned-deployment"]),
],
)
async def test_add_litellm_data_to_request_strips_forged_fallback_hop_state(
@ -1213,6 +1215,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_managed_background_response": True,
"_guardrail_pipelines": [{"name": "spoofed"}],
"_pipeline_managed_guardrails": ["evaded"],
"_failover_excluded_ids": ["pinned-deployment"],
"safe_user_metadata": "kept",
}
data = {
@ -1273,6 +1276,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"_managed_background_response",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
"_failover_excluded_ids",
}
assert "litellm_metadata" not in updated
for stripped_key in stripped_keys:

View file

@ -9,6 +9,7 @@ import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -1209,3 +1210,92 @@ async def test_virtual_key_hash_wins_over_user_id_for_affinity():
parent_otel_span=None,
)
assert other_key_same_user == healthy_deployments
class _ForbiddenError(Exception):
status_code: Final = 403
@pytest.mark.asyncio
@pytest.mark.parametrize(
("pre_call_check", "metadata"),
[
("session_affinity", {"session_id": "sess-retry"}),
("deployment_affinity", {"user_api_key_hash": "hashed-key-retry"}),
],
)
async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_refused(
pre_call_check: str, metadata: dict[str, str]
):
router: Final = litellm.Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "bad",
"mock_response": _ForbiddenError("A refuses"),
"weight": 1,
},
"model_info": {"id": "A"},
},
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "good",
"mock_response": "ok from B",
"weight": 0,
},
"model_info": {"id": "B"},
},
],
routing_strategy="simple-shuffle",
num_retries=1,
retry_after=0,
disable_cooldowns=True,
optional_pre_call_checks=[pre_call_check],
)
responses: Final = [
await router.acompletion(
model="test-model",
messages=[{"role": "user", "content": "hi"}],
metadata=dict(metadata),
)
for _ in range(2)
]
assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"]
assert [response._hidden_params["additional_headers"]["x-litellm-attempted-retries"] for response in responses] == [
1,
1,
]
@pytest.mark.asyncio
async def test_previous_response_id_pin_holds_even_after_its_deployment_failed():
callback: Final = DeploymentAffinityCheck(
cache=DualCache(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=True,
)
healthy_deployments: Final = [
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "A"}},
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}, "model_info": {"id": "B"}},
]
assert await callback.async_filter_deployments(
model="gpt-4o",
healthy_deployments=healthy_deployments,
messages=None,
request_kwargs={
"previous_response_id": ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="A", response_id="resp_upstream"
),
"_excluded_deployment_ids": ["A"],
"metadata": {"user_api_key_hash": "hashed-key"},
},
parent_otel_span=None,
) == [healthy_deployments[0]]

View file

@ -16,6 +16,7 @@ The mechanism works without any cache and supports two encoding strategies:
"""
import time
from typing import Final
from unittest.mock import AsyncMock, patch
import pytest
@ -2606,3 +2607,79 @@ async def test_affinity_strips_for_team_and_pattern_routes():
assert result is sibling_pool
assert not request_kwargs["input"][1].get("encrypted_content")
def _region_deployment(region: str) -> dict[str, object]:
return {
"model_name": "gpt-5.4",
"litellm_params": {
"model": "openai/gpt-5.4",
"api_base": f"https://{region}.example.com/v1",
"api_key": f"key-{region}",
},
"model_info": {"id": region},
}
def _reasoning_input_minted_by(origin: str) -> list[dict[str, object]]:
return [
{
"type": "reasoning",
"id": ResponsesAPIRequestUtils._build_encrypted_item_id(origin, "rs_test"),
"encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("gAAAAA-blob", origin),
"summary": [{"type": "summary_text", "text": "scattering"}],
},
{"role": "user", "content": "and sunsets?"},
]
_ORIGIN_FAILED_THIS_REQUEST: Final = pytest.mark.parametrize(
"failed_in_request",
[{"_excluded_deployment_ids": ["region-a"]}, {"_retry_skipped_deployment_ids": ("region-a",)}],
)
@pytest.mark.asyncio
@_ORIGIN_FAILED_THIS_REQUEST
async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request(
failed_in_request: dict[str, object],
):
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
router: Final = litellm.Router(model_list=[_region_deployment("region-a"), _region_deployment("region-b")])
group: Final = router.model_list
reasoning_input: Final = _reasoning_input_minted_by("region-a")
result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments(
model="gpt-5.4",
healthy_deployments=group,
messages=None,
request_kwargs={"litellm_metadata": {}, "input": reasoning_input, **failed_in_request},
)
assert result is group
assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]}
@pytest.mark.asyncio
@pytest.mark.parametrize("regions", [("region-a",), ("region-a", "region-b")])
async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_the_request(regions: tuple[str, ...]):
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
)
router: Final = litellm.Router(model_list=[_region_deployment(region) for region in regions])
reasoning_input: Final = _reasoning_input_minted_by("region-a")
original_reasoning: Final = dict(reasoning_input[0])
result: Final = await EncryptedContentAffinityCheck(router=router).async_filter_deployments(
model="gpt-5.4",
healthy_deployments=router.model_list,
messages=None,
request_kwargs={"litellm_metadata": {}, "input": reasoning_input, "_retry_skipped_deployment_ids": regions},
)
assert [deployment["model_info"]["id"] for deployment in result] == ["region-a"]
assert reasoning_input[0] == original_reasoning

View file

@ -313,6 +313,27 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
assert filtered == deployments
@pytest.mark.asyncio
@pytest.mark.parametrize(
"failed_in_request",
[{"_excluded_deployment_ids": ["dep-2"]}, {"_retry_skipped_deployment_ids": ("dep-2",)}],
)
async def test_async_filter_deployments_does_not_pin_to_a_deployment_that_failed_this_request(
failed_in_request: dict[str, object],
):
cache: Final = DualCache()
check: Final = PromptCachingDeploymentCheck(cache=cache)
messages: Final = _messages(word_count=5000)
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
assert await check.async_filter_deployments(
model=MODEL_GROUP_ALIAS,
healthy_deployments=_deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6"),
messages=messages,
request_kwargs=failed_in_request,
) == _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
@pytest.mark.asyncio
async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower():
"""

View file

@ -21051,9 +21051,9 @@ def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry(
],
)
def test_router_as_retry_skipped_deployment_ids_keeps_only_a_tuple_of_strings(value, expected):
from litellm.router import _as_retry_skipped_deployment_ids
from litellm.router_utils.common_utils import as_retry_skipped_deployment_ids
assert _as_retry_skipped_deployment_ids(value) == expected
assert as_retry_skipped_deployment_ids(value) == expected
@pytest.mark.parametrize(

View file

@ -8,13 +8,14 @@ cross-group fallback runs.
"""
from collections import Counter
from typing import Optional
from typing import Final, Optional
from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm import Router
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.utils import get_excluded_filtered_deployments
@ -944,3 +945,77 @@ async def test_failover_falls_through_to_external_fallback_when_remaining_in_coo
)
assert response._hidden_params["model_id"] == "fallback"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("pre_call_check", "metadata"),
[
("session_affinity", {"session_id": "sess-1"}),
("deployment_affinity", {"user_api_key_hash": "hashed-key-1"}),
],
)
async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deployment(
pre_call_check: str, metadata: dict[str, str]
):
router: Final = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "bad",
"mock_response": Exception("A down"),
"weight": 1,
},
"model_info": {"id": "A"},
},
{
"model_name": "test-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "good",
"mock_response": "ok from B",
"weight": 0,
},
"model_info": {"id": "B"},
},
{
"model_name": "fallback-model",
"litellm_params": {
"model": "gpt-4o",
"api_key": "good",
"mock_response": "ok from fallback",
},
"model_info": {"id": "fallback"},
},
],
routing_strategy="simple-shuffle",
num_retries=0,
disable_cooldowns=True,
enable_weighted_failover=True,
optional_pre_call_checks=[pre_call_check],
fallbacks=[{"test-model": ["fallback-model"]}],
)
while_a_is_down: Final = await router.acompletion(
model="test-model",
messages=[{"role": "user", "content": "hi"}],
metadata=dict(metadata),
)
router.upsert_deployment(
Deployment(
model_name="test-model",
litellm_params=LiteLLM_Params(model="gpt-4o", api_key="bad", mock_response="ok from A", weight=1),
model_info=ModelInfo(id="A"),
)
)
after_a_recovers: Final = await router.acompletion(
model="test-model",
messages=[{"role": "user", "content": "hi"}],
metadata=dict(metadata),
)
assert while_a_is_down._hidden_params["model_id"] == "B"
assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
assert after_a_recovers._hidden_params["model_id"] == "A"