mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
06943b6468
* feat(router): make routing groups callable as virtual models and list them in /v1/models * fix(router): traffic-scoped cooldown exemption, live model_names on delete, group-info cache invalidation * fix(router): share one recognized-model predicate across proxy gates, resolve aliases in group cooldown, read metadata via the dual-bucket owner * fix(router): close the gate and cache families for callable groups, strip member access_groups from group rows, prove cooldown wiring end to end * refactor(router): cache materialized group rows under the model-group cache owner and drop the redundant wiring test * fix(router): warn-and-shadow on group name collisions, name-level test coverage for group helpers, faithful router doubles in a2a and cursor tests * test(router): pin group cooldown metadata across the retry path
379 lines
17 KiB
Python
379 lines
17 KiB
Python
from unittest.mock import MagicMock, patch
|
|
|
|
import litellm
|
|
from litellm.router_utils.cooldown_handlers import (
|
|
_get_deployment_cooldown_policy,
|
|
_resolve_allowed_fails_from_policy,
|
|
_should_cooldown_based_on_deployment_policy,
|
|
should_cooldown_based_on_allowed_fails_policy,
|
|
)
|
|
|
|
|
|
class TestGetDeploymentCooldownPolicy:
|
|
def _make_router(self, deployment_id: str, model_info: dict | None = None):
|
|
router = MagicMock()
|
|
if model_info is None:
|
|
router.get_model_info.return_value = None
|
|
else:
|
|
router.get_model_info.return_value = {"model_info": model_info}
|
|
return router
|
|
|
|
def test_deployment_not_found_returns_none_none(self):
|
|
router = self._make_router("dep-1")
|
|
policy, allowed = _get_deployment_cooldown_policy(router, "dep-1")
|
|
assert policy is None
|
|
assert allowed is None
|
|
|
|
def test_no_model_info_returns_none_none(self):
|
|
router = MagicMock()
|
|
router.get_model_info.return_value = {"model_info": {}}
|
|
policy, allowed = _get_deployment_cooldown_policy(router, "dep-1")
|
|
assert policy is None
|
|
assert allowed is None
|
|
|
|
def test_returns_policy_dict_and_allowed_fails(self):
|
|
router = self._make_router(
|
|
"dep-1",
|
|
{"allowed_fails_policy": {"RateLimitErrorAllowedFails": 2}, "allowed_fails": 3},
|
|
)
|
|
policy, allowed = _get_deployment_cooldown_policy(router, "dep-1")
|
|
assert policy == {"RateLimitErrorAllowedFails": 2}
|
|
assert allowed == 3
|
|
|
|
def test_non_dict_policy_treated_as_none(self):
|
|
router = self._make_router("dep-1", {"allowed_fails_policy": "invalid", "allowed_fails": 5})
|
|
policy, allowed = _get_deployment_cooldown_policy(router, "dep-1")
|
|
assert policy is None
|
|
assert allowed == 5
|
|
|
|
def test_allowed_fails_only(self):
|
|
router = self._make_router("dep-1", {"allowed_fails": 1})
|
|
policy, allowed = _get_deployment_cooldown_policy(router, "dep-1")
|
|
assert policy is None
|
|
assert allowed == 1
|
|
|
|
|
|
class TestResolveAllowedFailsFromPolicy:
|
|
def test_none_policy_returns_none(self):
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(None, exc) is None
|
|
|
|
def test_matching_rate_limit_error(self):
|
|
policy = {"RateLimitErrorAllowedFails": 3}
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 3
|
|
|
|
def test_matching_internal_server_error(self):
|
|
policy = {"InternalServerErrorAllowedFails": 5}
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 5
|
|
|
|
def test_matching_service_unavailable_error(self):
|
|
policy = {"ServiceUnavailableErrorAllowedFails": 4}
|
|
exc = litellm.ServiceUnavailableError("503", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 4
|
|
|
|
def test_matching_bad_gateway_error(self):
|
|
policy = {"BadGatewayErrorAllowedFails": 2}
|
|
exc = litellm.BadGatewayError("502", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 2
|
|
|
|
def test_matching_not_found_error(self):
|
|
policy = {"NotFoundErrorAllowedFails": 1}
|
|
exc = litellm.NotFoundError("404", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 1
|
|
|
|
def test_unmatched_exception_returns_none(self):
|
|
policy = {"RateLimitErrorAllowedFails": 3}
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) is None
|
|
|
|
def test_field_absent_from_policy_returns_none(self):
|
|
policy: dict[str, int] = {}
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) is None
|
|
|
|
def test_content_policy_violation_not_shadowed_by_bad_request_error(self):
|
|
"""ContentPolicyViolationError subclasses BadRequestError, so if
|
|
BadRequestError were checked first, this would incorrectly resolve to
|
|
BadRequestErrorAllowedFails (10) instead of
|
|
ContentPolicyViolationErrorAllowedFails (2)."""
|
|
policy = {"BadRequestErrorAllowedFails": 10, "ContentPolicyViolationErrorAllowedFails": 2}
|
|
exc = litellm.ContentPolicyViolationError("flagged content", "openai", "gpt-4")
|
|
assert _resolve_allowed_fails_from_policy(policy, exc) == 2
|
|
|
|
|
|
class TestShouldCooldownBasedOnDeploymentPolicy:
|
|
def _make_router(self, model_info: dict | None = None):
|
|
router = MagicMock()
|
|
if model_info is None:
|
|
router.get_model_info.return_value = None
|
|
else:
|
|
router.get_model_info.return_value = model_info
|
|
return router
|
|
|
|
def test_policy_match_uses_exception_type_as_cache_key_suffix(self):
|
|
policy = {"RateLimitErrorAllowedFails": 0}
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
result = _should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, policy, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
assert result is True
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["allowed_fails_override"] == 0
|
|
assert call_kwargs["cache_key_suffix"] == "RateLimitError"
|
|
|
|
def test_no_policy_match_uses_dep_allowed_fails_and_generic_suffix(self):
|
|
policy: dict[str, int] = {}
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = False
|
|
result = _should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, policy, dep_allowed_fails=3, is_single_deployment_model_group=False
|
|
)
|
|
|
|
assert result is False
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["allowed_fails_override"] == 3
|
|
assert call_kwargs["cache_key_suffix"] == "generic"
|
|
|
|
def test_dep_allowed_fails_on_single_deployment_group_does_not_cooldown(self):
|
|
"""A generic, deployment-wide allowed_fails predates the per-exception-type
|
|
policy and is a less deliberate opt-in, so on a single-deployment model group
|
|
it must not silently disable the "avoid cooldowns on single deployment model
|
|
groups" safety net."""
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
result = _should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, dep_allowed_fails=3, is_single_deployment_model_group=True
|
|
)
|
|
|
|
assert result is False
|
|
mock_sc.assert_not_called()
|
|
|
|
def test_named_policy_on_single_deployment_group_still_cools_down(self):
|
|
"""Unlike a generic allowed_fails, an explicit per-exception-type policy entry
|
|
is a deliberate opt-in and must still apply on a single-deployment group."""
|
|
policy = {"RateLimitErrorAllowedFails": 0}
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
result = _should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, policy, None, is_single_deployment_model_group=True
|
|
)
|
|
|
|
assert result is True
|
|
mock_sc.assert_called_once()
|
|
|
|
def test_no_policy_and_no_dep_allowed_fails_defers_to_router_level(self):
|
|
"""When neither a deployment policy nor a deployment-wide allowed_fails covers
|
|
this exception, defer to router-level behavior instead of forcing an
|
|
immediate cooldown (allowed_fails_override=0 would trip on the first failure
|
|
of any exception type the deployment's config doesn't mention)."""
|
|
exc = litellm.InternalServerError("500", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["allowed_fails_override"] is None
|
|
assert call_kwargs["cache_key_suffix"] is None
|
|
|
|
def test_partial_policy_without_dep_allowed_fails_defers_for_uncovered_exception(self):
|
|
"""A deployment that only sets RateLimitErrorAllowedFails must not force a
|
|
zero-fail threshold on an unrelated TimeoutError; it should defer to
|
|
router-level behavior for exception types its policy doesn't mention."""
|
|
policy = {"RateLimitErrorAllowedFails": 0}
|
|
exc = litellm.Timeout("timed out", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = False
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, policy, dep_allowed_fails=None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["allowed_fails_override"] is None
|
|
assert call_kwargs["cache_key_suffix"] is None
|
|
|
|
def test_cooldown_time_from_model_info_passed_through(self):
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {}, "model_info": {"cooldown_time": 120.0}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["cooldown_time_override"] == 120.0
|
|
|
|
def test_cooldown_time_from_litellm_params_used_as_fallback(self):
|
|
"""cooldown_time has pre-existing litellm_params support on the primary
|
|
failure path, so it must still be honored here when model_info doesn't
|
|
set it."""
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {"cooldown_time": 120.0}, "model_info": {}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["cooldown_time_override"] == 120.0
|
|
|
|
def test_cooldown_time_from_model_info_takes_priority_over_litellm_params(self):
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router({"litellm_params": {"cooldown_time": 120.0}, "model_info": {"cooldown_time": 15.0}})
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = True
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["cooldown_time_override"] == 15.0
|
|
|
|
def test_model_info_none_passes_none_cooldown_time(self):
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
router = self._make_router(None)
|
|
|
|
with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc:
|
|
mock_sc.return_value = False
|
|
_should_cooldown_based_on_deployment_policy(
|
|
router, "dep-1", exc, None, None, is_single_deployment_model_group=False
|
|
)
|
|
|
|
call_kwargs = mock_sc.call_args[1]
|
|
assert call_kwargs["cooldown_time_override"] is None
|
|
|
|
|
|
class TestShouldCooldownBasedOnAllowedFailsPolicy:
|
|
def _make_router(self, cooldown_time: float = 60.0) -> MagicMock:
|
|
router = MagicMock()
|
|
router.cooldown_time = cooldown_time
|
|
router.allowed_fails = 0
|
|
router.allowed_fails_policy = None
|
|
router.get_allowed_fails_from_policy.return_value = None
|
|
router.failed_calls.get_cache.return_value = None
|
|
return router
|
|
|
|
def test_cooldown_time_override_zero_is_not_falsy(self):
|
|
"""cooldown_time_override=0 must be honored; it must not fall through to the router-level value."""
|
|
router = self._make_router(cooldown_time=60.0)
|
|
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
|
|
|
should_cooldown_based_on_allowed_fails_policy(
|
|
litellm_router_instance=router,
|
|
deployment="dep-1",
|
|
original_exception=exc,
|
|
allowed_fails_override=5,
|
|
cooldown_time_override=0.0,
|
|
)
|
|
|
|
set_cache_call = router.failed_calls.set_cache.call_args
|
|
assert set_cache_call is not None
|
|
assert set_cache_call[1]["ttl"] == 0.0, (
|
|
"cooldown_time_override=0 should be used as TTL, not the router-level 60.0"
|
|
)
|
|
|
|
|
|
class TestRoutingGroupCooldownAlternatives:
|
|
def _router(self, routing_groups=None):
|
|
from litellm import Router
|
|
|
|
return Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "solo-member",
|
|
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"},
|
|
"model_info": {"id": "cg-deploy-1"},
|
|
},
|
|
{
|
|
"model_name": "other-member",
|
|
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
|
|
"model_info": {"id": "cg-deploy-2"},
|
|
},
|
|
],
|
|
routing_groups=routing_groups,
|
|
)
|
|
|
|
def test_group_call_429_cools_down_member_with_alternatives(self):
|
|
from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
|
|
|
|
router = self._router(
|
|
routing_groups=[
|
|
{
|
|
"group_name": "grouped",
|
|
"models": ["solo-member", "other-member"],
|
|
"routing_strategy": "simple-shuffle",
|
|
}
|
|
]
|
|
)
|
|
assert (
|
|
_should_cooldown_deployment(
|
|
litellm_router_instance=router,
|
|
deployment="cg-deploy-1",
|
|
exception_status=429,
|
|
original_exception=Exception("rate limited"),
|
|
requested_model_group="grouped",
|
|
)
|
|
is True
|
|
)
|
|
|
|
def test_direct_member_429_keeps_single_deployment_exemption(self):
|
|
from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
|
|
|
|
router = self._router(
|
|
routing_groups=[
|
|
{
|
|
"group_name": "grouped",
|
|
"models": ["solo-member", "other-member"],
|
|
"routing_strategy": "simple-shuffle",
|
|
}
|
|
]
|
|
)
|
|
assert (
|
|
_should_cooldown_deployment(
|
|
litellm_router_instance=router,
|
|
deployment="cg-deploy-1",
|
|
exception_status=429,
|
|
original_exception=Exception("rate limited"),
|
|
requested_model_group="solo-member",
|
|
)
|
|
is False
|
|
)
|
|
|
|
def test_429_without_request_context_keeps_exemption(self):
|
|
from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment
|
|
|
|
router = self._router(routing_groups=None)
|
|
assert (
|
|
_should_cooldown_deployment(
|
|
litellm_router_instance=router,
|
|
deployment="cg-deploy-1",
|
|
exception_status=429,
|
|
original_exception=Exception("rate limited"),
|
|
)
|
|
is False
|
|
)
|