From 00cc4b67cdec9890243b0aaeea6c46937e17f673 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:46:35 -0700 Subject: [PATCH] fix(router): release affinity pins on deployments that failed this request Delete a session or API key pin whose deployment already failed the request, so the deployment that serves it claims the pin. Prompt caching pins skip failed deployments too, and the proxy strips the internal failed-id fields from client bodies. --- litellm/proxy/litellm_pre_call_utils.py | 2 ++ litellm/router.py | 11 +++---- litellm/router_utils/common_utils.py | 28 +++++++++++++++- .../deployment_affinity_check.py | 30 ++++------------- .../prompt_caching_deployment_check.py | 5 ++- .../unit/proxy/test_litellm_pre_call_utils.py | 2 ++ .../test_deployment_affinity_check.py | 33 +++++++++++++++++++ .../test_prompt_caching_deployment_check.py | 21 ++++++++++++ tests/unit/test_router/test_router.py | 4 +-- tests/unit/test_router_weighted_failover.py | 4 +++ 10 files changed, 106 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index f8e6498502b..96bef31512b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -355,6 +355,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", diff --git a/litellm/router.py b/litellm/router.py index 74facf200d2..c02b9a0b15d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -209,6 +209,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, @@ -668,10 +669,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 () @@ -8730,7 +8727,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 @@ -14794,7 +14791,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 = ( @@ -16028,7 +16025,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 = ( diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index a00d84babf3..a84ea67a2ec 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -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,8 @@ _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, ...]) +_OBJECT_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -57,6 +59,30 @@ 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 () + return tuple(item for item in _OBJECT_TUPLE_ADAPTER.validate_python(value) 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") diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 342a2720f12..927f1aa7e16 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,7 +16,6 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast -from pydantic import TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -31,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 @@ -39,8 +39,6 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] -_DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) - VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -330,24 +328,6 @@ class DeploymentAffinityCheck(CustomLogger): return stored return None - @staticmethod - def _deployment_ids(value: object) -> tuple[str, ...]: - if value is None: - return () - try: - return _DEPLOYMENT_IDS_ADAPTER.validate_python(value) - except ValidationError: - return () - - @classmethod - def _deployment_ids_failed_in_request(cls, request_kwargs: Mapping[str, object]) -> frozenset[str]: - return frozenset( - ( - *cls._deployment_ids(request_kwargs.get("_excluded_deployment_ids")), - *cls._deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), - ) - ) - def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) @@ -426,7 +406,7 @@ 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 = self._deployment_ids_failed_in_request( + 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 = ( @@ -450,7 +430,9 @@ class DeploymentAffinityCheck(CustomLogger): elif isinstance(session_cache_result, str): session_model_id = session_cache_result - if session_model_id and session_model_id not in failed_deployment_ids: + if session_model_id in failed_deployment_ids: + await self.cache.async_delete_cache(key=session_cache_key) + elif session_model_id: session_deployment: Final = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, model_id=session_model_id, @@ -485,6 +467,8 @@ class DeploymentAffinityCheck(CustomLogger): # Backwards / safety: allow raw string values. model_id = cache_result + if model_id in failed_deployment_ids: + await self.cache.async_delete_cache(key=cache_key) if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index bb3475cdcbe..c04bab4d360 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -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: diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c2a5b01fe4e..28d62fc3084 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1109,6 +1109,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( diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index f1393a17cbe..a960f9ef5be 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -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, ) @@ -1266,3 +1267,35 @@ async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_ ] 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, + 0, + ] + + +@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]] diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 8d893160780..9487e4d6321 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -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(): """ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d645ab4566d..e3a12b4eb63 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -20980,9 +20980,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( diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index adc4dd587d6..1808541666f 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -1008,3 +1008,7 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy ] assert [response._hidden_params["model_id"] for response in responses] == ["B", "B"] + assert [response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] for response in responses] == [ + 1, + 0, + ]