diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index c362c6899db..330500789a4 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -185,7 +185,14 @@ def _should_cooldown_based_on_deployment_policy( override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`. """ allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception) - if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group: + if ( + allowed_fails_from_policy is None + and is_single_deployment_model_group + and ( + dep_allowed_fails is not None + or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None + ) + ): return False allowed_fails_override: Final[int | None] = ( @@ -399,6 +406,34 @@ def _should_run_cooldown_logic( return True +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...]) + + +def _fallback_target_name(target: str | Mapping[str, object]) -> str: + return target if isinstance(target, str) else str(target.get("model", "")) + + +def _fallback_chain_names( + entry: Mapping[str, Sequence[str | Mapping[str, object]]], +) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()} + + +def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + + configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None) + if requested_model_group is None or not configured: + return False + chains: Final = tuple( + _fallback_chain_names(entry) + for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + if isinstance(entry, Mapping) + ) + fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) + return bool(fallback_group) + + def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, @@ -431,6 +466,9 @@ def _should_cooldown_deployment( litellm_router_instance.routing_group_has_alternatives(requested_model_group) or litellm_router_instance.team_model_has_alternatives(deployment) ) + is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( + litellm_router_instance, requested_model_group + ) exception_status_int: Final = cast_exception_status_to_int(exception_status) dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) @@ -450,7 +488,7 @@ def _should_cooldown_deployment( original_exception, dep_policy, dep_allowed_fails, - is_single_deployment_model_group, + is_lone_deployment, ) ## BASE CASE - single deployment @@ -495,6 +533,8 @@ def _should_cooldown_deployment( return False else: + if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None: + return False return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, deployment=deployment, diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index bc5cf1a5913..865edf31260 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import json import re import uuid @@ -390,6 +391,79 @@ def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_coun assert markers_received(refusing) == (first.marker, second.marker) +def failing_reply(request: Request) -> Reply: + return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down") + + +def assert_failed(outcome: Outcome) -> None: + assert outcome.status == 500, (outcome.status, outcome.detail) + assert "upstream down" in outcome.detail, outcome.detail + + +@pytest.mark.timeout(900) +def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback( + tmp_path: Path, +) -> None: + lone_group: Final = new_group() + pair_group: Final = new_group() + primary_group: Final = new_group() + backup_group: Final = new_group() + with ExitStack() as stack: + upstream: Final = stack.enter_context(gateway_from_environment()) + cache: Final = stack.enter_context(owned_redis(tmp_path)) + settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}]) + owned: Final = stack.enter_context( + owned_proxy_process( + upstream, + tmp_path, + {}, + config=cooldown_config(tmp_path, settings), + remove_environment=("PROMETHEUS_MULTIPROC_DIR",), + ) + ) + failing: Final = stack.enter_context(wire_server(failing_reply)) + healthy: Final = stack.enter_context(wire_server(healthy_reply)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing)) + failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4)) + for outcome in failed: + assert_failed(outcome) + assert not cooled(cache, lone_id) + assert markers_received(failing) == markers(failed) + + paired_id: Final = two_deployments(scenario, pair_group, failing, healthy) + history: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, paired_id), + (raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history] + assert sum(outcome.status == 500 for outcome in history) == 1 + served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10)) + assert_served(served, chat_identity) + assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served)) + + primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing)) + deployment(scenario, backup_group, openai_params(healthy)) + failing.drain() + healthy.drain() + fell_back: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, primary_id), + (raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back] + assert_served(fell_back, chat_identity) + backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4)) + assert_served(backed, chat_identity) + received: Final = markers_received(failing) + assert received[: len(fell_back)] == markers(fell_back) + assert len(received) <= len(fell_back) + 1, received + assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy)) + + def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None: group: Final = new_group() policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}} diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0c9d7ab9f85..bab6d257579 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -2931,10 +2931,65 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( assert tuple(cooldown_deployments) == () -@pytest.mark.parametrize("use_policy", [False, True], ids=["allowed_fails", "allowed_fails_policy"]) +TIMED_OUT_DEPLOYMENT: Final = { + "model_name": "test-model", + "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, +} +BACKUP_DEPLOYMENT: Final = { + "model_name": "backup-model", + "litellm_params": {"model": "openai/backup-model", "api_key": "test-key", "api_base": "https://backup.test/v1"}, + "model_info": {"id": "backup-deployment"}, +} +BACKUP_REPLY: Final = { + "id": "chatcmpl-backup", + "object": "chat.completion", + "created": 1700000000, + "model": "backup-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +async def _lone_call_outcome(router: Router) -> type[Exception]: + try: + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "timeout test"}]) + except (litellm.Timeout, RouterRateLimitError) as error: + return type(error) + pytest.fail("the timed-out deployment answered") + + +BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, RouterRateLimitError, RouterRateLimitError) + + +@pytest.mark.parametrize( + ("router_kwargs", "model_info", "outcomes"), + [ + ({"allowed_fails": 1}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails": 1}, {"allowed_fails_policy": {}}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=1)}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)}, {}, BENCHED), + ({"allowed_fails": 1, "fallbacks": [{"test-model": ["backup-model"]}]}, {}, BENCHED), + ( + {"allowed_fails": 1, "fallbacks": [{"test-model": [{"model": "backup-model", "temperature": 0}]}]}, + {}, + BENCHED, + ), + ], + ids=[ + "allowed_fails", + "allowed_fails_with_an_empty_deployment_policy", + "policy_names_another_exception", + "policy_names_timeout", + "allowed_fails_with_a_fallback_group", + "allowed_fails_with_a_dict_fallback_target", + ], +) @pytest.mark.asyncio -async def test_allowed_fails_cools_a_single_deployment( - use_policy: bool, +async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_policy_names_the_exception( + router_kwargs: Mapping[str, object], + model_info: Mapping[str, object], + outcomes: tuple[type[Exception], ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2942,38 +2997,52 @@ async def test_allowed_fails_cools_a_single_deployment( route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=httpx.ReadTimeout("upstream timeout") ) - policy_kwargs: Final = ( - {"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)} - if use_policy - else {"allowed_fails": 1} + deployment: Final = {**TIMED_OUT_DEPLOYMENT, "model_info": {**TIMED_OUT_DEPLOYMENT["model_info"], **model_info}} + router: Final = Router(model_list=[deployment], cooldown_time=300, num_retries=0, **router_kwargs) + + seen: Final = tuple([await _lone_call_outcome(router) for _ in range(5)]) + + assert seen == outcomes + assert len(route.calls) == outcomes.count(litellm.Timeout) + + +async def _group_call_outcome(router: Router) -> str: + try: + response: Final = await router.acompletion( + model="test-model", messages=[{"role": "user", "content": "timeout test"}] + ) + except litellm.Timeout: + return "timeout" + return response.id + + +@pytest.mark.asyncio +async def test_router_allowed_fails_still_benches_a_deployment_that_has_a_sibling_in_its_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + timed_out: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=httpx.ReadTimeout("upstream timeout") ) + respx_mock.post("https://backup.test/v1/chat/completions").mock(return_value=httpx.Response(200, json=BACKUP_REPLY)) router: Final = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - } - ], + model_list=[TIMED_OUT_DEPLOYMENT, {**BACKUP_DEPLOYMENT, "model_name": "test-model"}], cooldown_time=300, num_retries=0, - **policy_kwargs, + allowed_fails=1, ) - for _ in range(2): - with pytest.raises(litellm.Timeout): - await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "timeout test"}], - ) - + while len(timed_out.calls) < 2: + await _group_call_outcome(router) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, - parent_otel_span=None, + litellm_router_instance=router, parent_otel_span=None ) + served: Final = tuple([await _group_call_outcome(router) for _ in range(5)]) - assert len(route.calls) == 2 assert tuple(cooldown_deployments) == ("test-deployment",) + assert served == ("chatcmpl-backup",) * 5 + assert len(timed_out.calls) == 2 def test_per_user_session_upstream_errors_never_cool_down_the_shared_deployment(single_deployment_router): diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 53375ee0ab4..c5a476c450b 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -1255,7 +1255,12 @@ class TestTriggerCooldownForFailedDeployment: "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, @@ -1345,7 +1350,12 @@ async def test_a_fallback_deployment_whose_stream_stalled_stays_out_of_cooldown( "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e2ee80b7874..d9b655a0aff 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -12212,7 +12212,12 @@ class TestBackgroundResponseCostPollCooldown: "model_name": "gpt-4.1", "litellm_params": {"model": "openai/gpt-4.1"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, ) @@ -12340,7 +12345,12 @@ class TestCallerTimeoutCooldown: "model_name": "slow-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "slow-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, cooldown_time=120,