diff --git a/litellm/router.py b/litellm/router.py index 3340e06627..fac48b45fb 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -321,6 +321,7 @@ class Router: enable_health_check_routing: bool = False, health_check_staleness_threshold: Optional[int] = None, health_check_ignore_transient_errors: bool = False, + enable_weighted_failover: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -356,6 +357,7 @@ class Router: provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. + enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. Returns: Router: An instance of the litellm.Router class. @@ -524,6 +526,7 @@ class Router: ) self.disable_cooldowns = disable_cooldowns self.enable_health_check_routing = enable_health_check_routing + self.enable_weighted_failover = enable_weighted_failover self.health_check_ignore_transient_errors = health_check_ignore_transient_errors _staleness = health_check_staleness_threshold or ( DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER @@ -1860,6 +1863,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) + self._set_failed_deployment_id_on_exception(e, deployment) raise e def _get_silent_experiment_kwargs(self, **kwargs) -> dict: @@ -2525,6 +2529,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) + self._set_failed_deployment_id_on_exception(e, deployment) raise e except Exception as e: verbose_router_logger.info( @@ -2535,6 +2540,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) + self._set_failed_deployment_id_on_exception(e, deployment) raise e def _update_kwargs_before_fallbacks( @@ -2579,6 +2585,27 @@ class Router: except (ValueError, TypeError): pass # Skip if value can't be converted to int + def _set_failed_deployment_id_on_exception( + self, exception: Exception, deployment: dict + ) -> None: + """ + Stamp the failed deployment's `model_info.id` on the exception so the + fallback layer can exclude it from subsequent re-picks within the same + request (used by weighted-routing failover). + + Idempotent: never overwrites an existing value, so the id of the + deployment that *first* failed in a chain is preserved if multiple + layers re-raise. + """ + if getattr(exception, "failed_deployment_id", None): + return + deployment_id = (deployment.get("model_info") or {}).get("id") + if deployment_id: + try: + exception.failed_deployment_id = deployment_id # type: ignore[attr-defined] + except Exception: + pass + def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: Optional[str] = "metadata" ) -> None: @@ -5641,6 +5668,85 @@ class Router: #### [END] ASSISTANTS API #### + async def _maybe_run_weighted_failover( + self, + exception: Exception, + original_model_group: str, + all_deployments: List[DeploymentTypedDict], + args: tuple, + kwargs: dict, + input_kwargs: dict, + ) -> Optional[Any]: + """Same-model-group retry after a failed deployment; returns None if not applicable.""" + strategy, _ = self._get_routing_context(original_model_group) + if strategy != "simple-shuffle": + return None + + failed_id: Optional[str] = getattr(exception, "failed_deployment_id", None) + if not failed_id: + return None + + metadata_variable_name = self._get_metadata_variable_name_from_kwargs(kwargs) + meta = kwargs.get(metadata_variable_name) + if meta is None: + meta = {} + kwargs[metadata_variable_name] = meta + if not isinstance(meta, dict): + return None + prev_excluded = set(meta.get("_failover_excluded_ids") or []) + excluded = prev_excluded | {failed_id} + + all_ids = { + (d.get("model_info") or {}).get("id") + for d in all_deployments + if (d.get("model_info") or {}).get("id") is not None + } + # Only consider deployments that are currently healthy (not in cooldown). + # Using all_ids here would cause a wasteful run_async_fallback invocation + # that fails with RouterRateLimitError whenever the "remaining" entries + # are all in cooldown — the inner async_get_healthy_deployments call + # would find an empty list and raise immediately. + cooldown_ids = set( + await _async_get_cooldown_deployments( + litellm_router_instance=self, parent_otel_span=None + ) + ) + remaining = (all_ids - cooldown_ids) - excluded + if not remaining: + return None + + verbose_router_logger.debug( + f"Weighted failover: exclude={excluded!r}, remaining={len(remaining)} " + f"for model_group={original_model_group!r}" + ) + + meta["_failover_excluded_ids"] = list(excluded) + + entry = { + "model": original_model_group, + "_excluded_deployment_ids": list(excluded), + } + # Build a local copy so the weighted-failover keys do not leak back to + # the caller's shared kwargs dict (any downstream fallback path reads + # the same dict and must not inherit our `_excluded_deployment_ids` + # entry). + failover_kwargs = { + **input_kwargs, + "fallback_model_group": [entry], + "original_model_group": original_model_group, + } + try: + return await run_async_fallback(*args, **failover_kwargs) + except (openai.APIError, RouterRateLimitError, RouterRateLimitErrorBasic): + # Expected model-level failure on the retried deployment. All + # litellm provider errors derive from openai.APIError; if every + # remaining deployment in the group is in cooldown the router + # raises RouterRateLimitError (a ValueError, not an APIError). + # In either case defer to the regular fallback path. Programming + # errors (AttributeError, KeyError, TypeError, etc.) intentionally + # propagate so they remain visible. + return None + async def async_function_with_fallbacks_common_utils( # noqa: PLR0915 self, e: Exception, @@ -5743,6 +5849,23 @@ class Router: ) return response + # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. + if ( + self.enable_weighted_failover + and not _skip_order_fallback + and original_model_group is not None + ): + response = await self._maybe_run_weighted_failover( + exception=e, + original_model_group=original_model_group, + all_deployments=all_deployments, + args=args, + kwargs=kwargs, + input_kwargs=input_kwargs, + ) + if response is not None: + return response + try: verbose_router_logger.info("Trying to fallback b/w models") @@ -9344,6 +9467,7 @@ class Router: "model_group_retry_policy", "retry_policy", "model_group_alias", + "enable_weighted_failover", ] for var in vars_to_include: @@ -9380,6 +9504,7 @@ class Router: "context_window_fallbacks", "model_group_retry_policy", "model_group_alias", + "enable_weighted_failover", ] _int_settings = [ @@ -10073,6 +10198,17 @@ class Router: cast(List[Dict], healthy_deployments), target_order=_target_order ) + ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in + ## this request via weighted-failover. Always honored, regardless of the + ## router-level flag, so a stale exclusion key on kwargs cannot escape. + _excluded_deployment_ids = (request_kwargs or {}).pop( + "_excluded_deployment_ids", None + ) + healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + cast(List[Dict], healthy_deployments), + excluded_deployment_ids=_excluded_deployment_ids, + ) + if len(healthy_deployments) == 0: exception = await async_raise_no_deployment_exception( litellm_router_instance=self, @@ -10464,6 +10600,17 @@ class Router: healthy_deployments, target_order=_target_order ) + ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in + ## this request via weighted-failover. See async counterpart in + ## async_get_healthy_deployments for details. + _excluded_deployment_ids = (request_kwargs or {}).pop( + "_excluded_deployment_ids", None + ) + healthy_deployments = litellm.utils._get_excluded_filtered_deployments( + healthy_deployments, + excluded_deployment_ids=_excluded_deployment_ids, + ) + if len(healthy_deployments) == 0: model_ids = self.get_model_ids(model_name=model) _cooldown_time = self.cooldown_cache.get_min_cooldown( diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index 9827522747..f78acbfbd0 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -48,6 +48,13 @@ def simple_shuffle( ] verbose_router_logger.debug(f"\nweight {weights}") total_weight = sum(weights) + if total_weight <= 0: + # All remaining candidates have weight 0 for this metric (e.g. + # after a weighted-failover exclusion left only zero-weight + # backups). Skip to the next metric (rpm/tpm) which may still + # provide a meaningful weighted pick; if none do, we fall + # through to the uniform random pick at the end. + continue weights = [weight / total_weight for weight in weights] verbose_router_logger.debug(f"\n weights {weights} by {weight_by}") # Perform weighted random pick diff --git a/litellm/utils.py b/litellm/utils.py index da80e4ae16..796547cb04 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4983,6 +4983,33 @@ def _get_order_filtered_deployments( return healthy_deployments +def _get_excluded_filtered_deployments( + healthy_deployments: List[Dict], + excluded_deployment_ids: Optional[Iterable[str]] = None, +) -> List: + """ + Filter out deployments whose `model_info.id` appears in `excluded_deployment_ids`. + + Used by weighted-routing failover so a single logical request can re-pick + across the remaining deployments in the same model group after one of them + has failed. + + If the filter would leave no deployments, an empty list is returned so the + caller raises its usual no-deployments error and the weighted-failover + helper falls through to the cross-group fallback path. Returning the + original unfiltered list here would re-include the just-failed deployment. + """ + if not excluded_deployment_ids: + return healthy_deployments + + excluded_set = set(excluded_deployment_ids) + return [ + d + for d in healthy_deployments + if (d.get("model_info") or {}).get("id") not in excluded_set + ] + + def _get_model_region( custom_llm_provider: str, litellm_params: LiteLLM_Params ) -> Optional[str]: diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py new file mode 100644 index 0000000000..8faf6bcd9c --- /dev/null +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -0,0 +1,771 @@ +""" +Tests for weighted-routing failover (router_settings.enable_weighted_failover). + +When enabled and the routing strategy is "simple-shuffle", a retryable failure +on one deployment causes the request to re-pick a different deployment in the +SAME model group (weighted across the remaining deployments) before any +cross-group fallback runs. +""" + +from collections import Counter +from typing import Optional +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm import Router +from litellm.utils import _get_excluded_filtered_deployments + + +# --------------------------------------------------------------------------- +# Unit tests for _get_excluded_filtered_deployments +# --------------------------------------------------------------------------- + + +def _make_dep(dep_id: str, weight: Optional[int] = None) -> dict: + params: dict = {"model": "gpt-4o", "api_key": "key"} + if weight is not None: + params["weight"] = weight + return { + "model_name": "test-model", + "litellm_params": params, + "model_info": {"id": dep_id}, + } + + +class TestGetExcludedFilteredDeployments: + def test_no_excluded_returns_all(self): + deps = [_make_dep("a"), _make_dep("b")] + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=None) + assert len(result) == 2 + + def test_empty_excluded_returns_all(self): + deps = [_make_dep("a"), _make_dep("b")] + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=[]) + assert len(result) == 2 + + def test_drops_excluded(self): + deps = [_make_dep("a"), _make_dep("b"), _make_dep("c")] + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + ids = sorted(d["model_info"]["id"] for d in result) + assert ids == ["a", "c"] + + def test_all_excluded_returns_empty(self): + # When every healthy deployment has been excluded, the helper must + # return an empty list so the caller raises its usual no-deployments + # error. Returning the original list here would re-include the + # just-failed deployment and let weighted failover re-pick it. + deps = [_make_dep("a"), _make_dep("b")] + result = _get_excluded_filtered_deployments( + deps, excluded_deployment_ids=["a", "b"] + ) + assert result == [] + + def test_excluded_set_with_unknown_ids(self): + deps = [_make_dep("a"), _make_dep("b")] + result = _get_excluded_filtered_deployments( + deps, excluded_deployment_ids=["zzz"] + ) + assert len(result) == 2 + + def test_handles_missing_model_info(self): + deps = [ + {"model_name": "x", "litellm_params": {"model": "gpt-4o"}}, # no model_info + _make_dep("b"), + ] + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["b"]) + assert len(result) == 1 + + +# --------------------------------------------------------------------------- +# Router helpers (router_code_coverage.py requires these names in a *router* test file) +# --------------------------------------------------------------------------- + + +def test_set_failed_deployment_id_on_exception(): + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + exc = Exception("fail") + dep = _make_dep("dep-a") + router._set_failed_deployment_id_on_exception(exc, dep) + assert getattr(exc, "failed_deployment_id", None) == "dep-a" + router._set_failed_deployment_id_on_exception(exc, _make_dep("dep-b")) + assert exc.failed_deployment_id == "dep-a" + + +@pytest.mark.asyncio +async def test_maybe_run_weighted_failover_returns_none_without_failed_id(): + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + enable_weighted_failover=True, + ) + result = await router._maybe_run_weighted_failover( + exception=Exception("fail"), + original_model_group="test-model", + all_deployments=[_make_dep("A"), _make_dep("B")], + args=(), + kwargs={"metadata": {}}, + input_kwargs={}, + ) + assert result is None + + +@pytest.mark.asyncio +async def test_maybe_run_weighted_failover_persists_excluded_ids_to_kwargs(monkeypatch): + """Regression: writing to the metadata dict returned by `setdefault` must + update the dict in `kwargs` itself so the next hop sees prior exclusions. + Previously `setdefault(..., {}) or {}` returned a disconnected dict on the + first hop, dropping `_failover_excluded_ids` writes. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + enable_weighted_failover=True, + ) + + async def _stub_run_async_fallback(*args, **kwargs): + return "ok" + + monkeypatch.setattr("litellm.router.run_async_fallback", _stub_run_async_fallback) + + exc = Exception("fail") + exc.failed_deployment_id = "A" + kwargs: dict = {"metadata": {}} + await router._maybe_run_weighted_failover( + exception=exc, + original_model_group="test-model", + all_deployments=[_make_dep("A"), _make_dep("B")], + args=(), + kwargs=kwargs, + input_kwargs={}, + ) + # The dict inside kwargs must reflect the write — proves `meta` was the + # same object as kwargs["metadata"] (no disconnected copy). + assert kwargs["metadata"].get("_failover_excluded_ids") == ["A"] + + +# --------------------------------------------------------------------------- +# Integration tests for weighted-failover end-to-end via Router +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_no_failover_when_flag_off(): + """Default behavior: a failure on the picked deployment surfaces to caller.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("region-A failed"), + "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, # weight=0 so A is always picked + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + # enable_weighted_failover defaults to False + ) + + with pytest.raises(Exception): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.asyncio +async def test_failover_lands_on_other_deployment_when_flag_on(): + """Flag on: when A fails, request must succeed via B in the same call.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("region-A down"), + "weight": 1, # always picked first (B has weight 0) + }, + "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=0, + enable_weighted_failover=True, + ) + + response = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert response._hidden_params["model_id"] == "B" + + +@pytest.mark.asyncio +async def test_failover_chain_three_deployments(): + """A and B fail, request succeeds on C.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("A down"), + "weight": 1_000_000, # A always picked first + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("B down"), + "weight": 1, # picked when A is excluded + }, + "model_info": {"id": "B"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from C", + "weight": 0, + }, + "model_info": {"id": "C"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + response = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert response._hidden_params["model_id"] == "C" + + +@pytest.mark.asyncio +async def test_failover_exhausted_raises_original_error_class(): + """When ALL deployments fail, the request raises (does not hang).""" + router = 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": "bad", + "mock_response": Exception("B down"), + "weight": 1, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + with pytest.raises(Exception): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.asyncio +async def test_failover_falls_through_to_external_fallback(): + """When all deployments in the group fail, external fallback still runs.""" + router = 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": "bad", + "mock_response": Exception("B down"), + "weight": 1, + }, + "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, + enable_weighted_failover=True, + fallbacks=[{"test-model": ["fallback-model"]}], + ) + + response = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert response._hidden_params["model_id"] == "fallback" + + +@pytest.mark.asyncio +async def test_weights_respected_when_all_healthy(): + """With both regions healthy, the picker should still honor configured + weights — failover must not change the steady-state load shape.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "from A", + "weight": 80, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "from B", + "weight": 20, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + counts: Counter = Counter() + for _ in range(1000): + resp = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + counts[resp._hidden_params["model_id"]] += 1 + + # Expect ~80/20 split. Loose bounds to keep the test stable under CI load. + assert counts["A"] > counts["B"] * 2 # A should heavily dominate + assert counts["B"] > 50 # but B should still get a meaningful share + + +@pytest.mark.asyncio +async def test_failover_skipped_for_non_simple_shuffle(): + """Weighted failover is only wired up for `simple-shuffle`. With another + strategy, a failure on the picked deployment must NOT silently retry the + other deployment in the same group. Both deployments fail here to keep the + test deterministic regardless of which one the strategy picks first. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("A down"), + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("B down"), + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="latency-based-routing", + num_retries=0, + enable_weighted_failover=True, + ) + + with pytest.raises(Exception): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + +@pytest.mark.asyncio +async def test_failover_skipped_for_context_window_error(): + """ContextWindowExceededError must NOT trigger weighted failover — + it has its own dedicated fallback path. Uses the router's built-in + `mock_testing_context_fallbacks` to deterministically raise the right + exception class. + """ + import litellm + + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from A", + "weight": 1, + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "ok from B", + "weight": 1, + }, + "model_info": {"id": "B"}, + }, + ], + routing_strategy="simple-shuffle", + num_retries=0, + enable_weighted_failover=True, + ) + + with pytest.raises(litellm.ContextWindowExceededError): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + mock_testing_context_fallbacks=True, + ) + + +@pytest.mark.asyncio +async def test_user_config_two_region_failover(): + """Mirrors the user's actual proxy_server_config.yaml shape: two Azure + regions weighted 50/50, num_retries=0. With the flag on, a failure in + one region is recovered by the other in the same request.""" + router = Router( + model_list=[ + { + "model_name": "gpt-5.4-mini", + "litellm_params": { + "model": "azure/deployment-eastus2", + "api_key": "bad", + "api_base": "https://eastus2.example", + "mock_response": Exception("eastus2 5xx"), + "weight": 50, + }, + "model_info": {"id": "eastus2"}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": { + "model": "azure/deployment-northcentralus", + "api_key": "good", + "api_base": "https://northcentralus.example", + "mock_response": "ok from northcentralus", + "weight": 50, + }, + "model_info": {"id": "northcentralus"}, + }, + ], + routing_strategy="simple-shuffle", + cooldown_time=120, + num_retries=0, + enable_pre_call_checks=True, + disable_cooldowns=False, + allowed_fails=5, + enable_weighted_failover=True, + ) + + # Force eastus2 to be picked first by leaving its weight intact and + # asserting we always end up on northcentralus when eastus2 errors. + # Run several requests and ensure we never see an unhandled failure. + successes = Counter() + for _ in range(20): + resp = await router.acompletion( + model="gpt-5.4-mini", + messages=[{"role": "user", "content": "hi"}], + ) + successes[resp._hidden_params["model_id"]] += 1 + + # With one region permanently failing, every request must land on the + # other region (either directly because it was picked first, or via + # failover because eastus2 was picked first). + assert successes["northcentralus"] == 20 + assert successes["eastus2"] == 0 + + +# --------------------------------------------------------------------------- +# Tests for healthy-deployment-only check in _maybe_run_weighted_failover +# (Issue: weighted failover checked all deployments, not just healthy ones) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_maybe_run_weighted_failover_skips_when_remaining_all_in_cooldown( + monkeypatch, +): + """When every non-excluded deployment is in cooldown, _maybe_run_weighted_failover + must return None immediately without invoking run_async_fallback. + + Previously the check was against all_deployments (including cooldown ones), so + run_async_fallback would be called unnecessarily and would raise RouterRateLimitError. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "B"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "C"}, + }, + ], + routing_strategy="simple-shuffle", + enable_weighted_failover=True, + ) + + # A just failed; B and C are both in cooldown. + exc = Exception("A down") + exc.failed_deployment_id = "A" + + run_async_fallback_called = False + + async def _should_not_be_called(*args, **kwargs): + nonlocal run_async_fallback_called + run_async_fallback_called = True + return "should not reach here" + + monkeypatch.setattr("litellm.router.run_async_fallback", _should_not_be_called) + + # Patch cooldown so B and C appear in cooldown. + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["B", "C"]), + ): + result = await router._maybe_run_weighted_failover( + exception=exc, + original_model_group="test-model", + all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], + args=(), + kwargs={"metadata": {}}, + input_kwargs={}, + ) + + assert ( + result is None + ), "Should return None when all remaining deployments are in cooldown" + assert ( + not run_async_fallback_called + ), "run_async_fallback must NOT be called when no healthy deployments remain" + + +@pytest.mark.asyncio +async def test_maybe_run_weighted_failover_proceeds_when_one_healthy_remains( + monkeypatch, +): + """When at least one non-excluded deployment is healthy (not in cooldown), + _maybe_run_weighted_failover should still invoke run_async_fallback normally. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "B"}, + }, + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, + "model_info": {"id": "C"}, + }, + ], + routing_strategy="simple-shuffle", + enable_weighted_failover=True, + ) + + # A just failed; B is in cooldown; C is healthy. + exc = Exception("A down") + exc.failed_deployment_id = "A" + + run_async_fallback_called = False + + async def _stub_run_async_fallback(*args, **kwargs): + nonlocal run_async_fallback_called + run_async_fallback_called = True + return "ok from C" + + monkeypatch.setattr("litellm.router.run_async_fallback", _stub_run_async_fallback) + + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["B"]), + ): + result = await router._maybe_run_weighted_failover( + exception=exc, + original_model_group="test-model", + all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], + args=(), + kwargs={"metadata": {}}, + input_kwargs={}, + ) + + assert result == "ok from C" + assert ( + run_async_fallback_called + ), "run_async_fallback must be called when a healthy deployment remains" + + +@pytest.mark.asyncio +async def test_failover_falls_through_to_external_fallback_when_remaining_in_cooldown(): + """End-to-end: when the only non-failed deployments are in cooldown, + weighted failover must fall through to the configured cross-group fallback. + + Without the fix the _maybe_run_weighted_failover would invoke run_async_fallback + unnecessarily (because it counted cooldown deployments as "remaining"), get back + RouterRateLimitError, return None, and reach the same fallback path — but only + incidentally. With the fix the early-exit path is taken directly. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("A down"), + "weight": 1_000_000, # always picked first + }, + "model_info": {"id": "A"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("B down"), + "weight": 1, + }, + "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, + enable_weighted_failover=True, + fallbacks=[{"test-model": ["fallback-model"]}], + ) + + # Put B in cooldown so weighted failover can't use it after A fails. + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["B"]), + ): + response = await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + + assert response._hidden_params["model_id"] == "fallback"