diff --git a/litellm/router.py b/litellm/router.py index 8865543badd..4929b17f7fc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1553,6 +1553,16 @@ class Router: return False return sum(len(self.model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 + def team_model_has_alternatives(self, deployment_id: str) -> bool: + deployment: Final = self.get_deployment(model_id=deployment_id) + if deployment is None: + return False + team_id: Final = deployment.model_info.team_id + public_model_name: Final = deployment.model_info.team_public_model_name + if team_id is None or public_model_name is None: + return False + return len(self.team_model_to_deployment_indices.get((team_id, public_model_name)) or ()) > 1 + _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index f722b6fd20c..027f0a9ca05 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -343,8 +343,9 @@ def _should_cooldown_deployment( model_group: Final = litellm_router_instance.get_model_group(id=deployment) is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: - is_single_deployment_model_group = not litellm_router_instance.routing_group_has_alternatives( - requested_model_group + is_single_deployment_model_group = not ( + litellm_router_instance.routing_group_has_alternatives(requested_model_group) + or litellm_router_instance.team_model_has_alternatives(deployment) ) ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 7ee0ed3701b..6f66bb863cc 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -437,3 +437,53 @@ class TestRoutingGroupCooldownAlternatives: ) is False ) + + +class TestTeamModelCooldownAlternatives: + def _router(self, team_deployments: int): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": f"model_name_team-1_{i}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": { + "id": f"team-deploy-{i}", + "team_id": "team-1", + "team_public_model_name": "team-gpt-4o-mini", + }, + } + for i in range(team_deployments) + ] + ) + + def test_429_on_team_deployment_with_sibling_cools_down(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=2) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is True + ) + + def test_429_on_only_team_deployment_keeps_single_deployment_exemption(self): + from litellm.router_utils.cooldown_handlers import _should_cooldown_deployment + + router = self._router(team_deployments=1) + assert ( + _should_cooldown_deployment( + litellm_router_instance=router, + deployment="team-deploy-0", + exception_status=429, + original_exception=Exception("rate limited"), + requested_model_group="team-gpt-4o-mini", + ) + is False + )