diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1e7cdc44c70..0f8ac58861b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -363,6 +363,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", @@ -444,6 +446,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( "_pipeline_managed_guardrails", "client_disconnected", "error_information", + "_failover_excluded_ids", PRE_CALL_EXECUTED_GUARDRAILS_KEY, ) diff --git a/litellm/router.py b/litellm/router.py index 18c5c67abce..e53bab4659e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -212,6 +212,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, @@ -685,10 +686,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 () @@ -8942,7 +8939,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 @@ -15186,7 +15183,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 = ( @@ -16418,7 +16415,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..91b9e5cb114 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,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") 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 8693eaa5151..daf7d6bb345 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -31,6 +31,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 @@ -429,6 +430,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) @@ -458,7 +462,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, @@ -493,7 +497,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( diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 5d47bf63a38..806c3db7601 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -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,31 @@ class EncryptedContentAffinityCheck(CustomLogger): model_id, ) + previous_response_id: Final = typed_request_kwargs.get("previous_response_id") + previous_response_origin: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id( + None if previous_response_id is None else str(previous_response_id) + ) + previous_response_binds_origin: Final = previous_response_origin == model_id + excluded_ids: Final[frozenset[str]] = ( + frozenset() + if previous_response_binds_origin + else frozenset(as_excluded_deployment_ids(typed_request_kwargs.get("_excluded_deployment_ids"))) + ) + retry_skipped_ids: Final[frozenset[str]] = ( + frozenset() + if previous_response_binds_origin + else 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 +425,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: 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 a44373b87c8..f743a14c7a2 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1393,6 +1393,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( @@ -1477,6 +1479,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 = { @@ -1537,6 +1540,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: 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 d8bc4c45ab8..fa75f980cc0 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, ) @@ -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]] diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index 87fbf4433a1..a3745c7b23f 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -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,111 @@ 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?"}, + ] + + +def _previous_response_id_from(origin: str) -> str: + return ResponsesAPIRequestUtils._build_responses_api_response_id("openai", origin, "resp_test") + + +_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 +@pytest.mark.parametrize("continuing", [{}, {"previous_response_id": _previous_response_id_from("region-b")}]) +async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_request( + failed_in_request: dict[str, object], continuing: 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, **continuing}, + ) + + assert result is group + assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} + + +@pytest.mark.asyncio +@_ORIGIN_FAILED_THIS_REQUEST +async def test_affinity_keeps_origin_and_reasoning_when_previous_response_id_binds_the_failed_origin( + failed_in_request: dict[str, object], +): + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + pool: Final = (_region_deployment("region-a"), _region_deployment("region-b")) + reasoning_input: Final = _reasoning_input_minted_by("region-a") + original_reasoning: Final = dict(reasoning_input[0]) + + assert await EncryptedContentAffinityCheck(router=litellm.Router(model_list=list(pool))).async_filter_deployments( + model="gpt-5.4", + healthy_deployments=list(pool), + messages=None, + request_kwargs={ + "litellm_metadata": {}, + "input": reasoning_input, + "previous_response_id": _previous_response_id_from("region-a"), + **failed_in_request, + }, + ) == [pool[0]] + assert reasoning_input[0] == original_reasoning + + +@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 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 d9b655a0aff..30d675e130d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -21239,9 +21239,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 9cb38b00e1e..8b7fc81dbde 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -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"