Merge 01210921f6 into e52f05566d

litellm/router.py

@@ -1322,6 +1322,23 @@ class Router:
             or self.get_routing_group(model) is not None
         )
 
+    def _resolve_to_deployment_model_names(self, model: str) -> tuple[str, ...]:
+        """
+        The deployment `model_name`s behind a requested name, so lookups keyed
+        by a served name reach the deployments the router would route to:
+
+        - a `model_group_alias` resolves to its target
+        - a callable routing group resolves to its member `model_name`s
+        - any other name resolves to itself
+
+        Team-scoped deployments keep their own internal `model_name`, so they
+        are reached only when that name is requested, as is already the case
+        for a plain `model_name` lookup.
+        """
+        resolved: Final = self._get_model_from_alias(model=model) or model
+        group: Final = self.get_routing_group(resolved)
+        return tuple(group.models) if group is not None else (resolved,)
+
     def routing_group_has_alternatives(self, model_group: str | None) -> bool:
         """
         True when `model_group` names a callable routing group whose member
@@ -10312,6 +10329,11 @@ class Router:
 
         Returns list of model id's.
 
+        `model_name` may be any name the router serves: a deployment
+        `model_name`, a `model_group_alias`, or a callable routing group, which
+        resolve to the deployments they route to (see
+        `_resolve_to_deployment_model_names`).
+
         Optimized with O(1) or O(k) index lookup when model_name provided,
         instead of O(n) linear scan.
         """
@@ -10319,9 +10341,8 @@ class Router:
 
         if model_name is not None:
             # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name
-            if model_name in self.model_name_to_deployment_indices:
-                indices: Final = self.model_name_to_deployment_indices[model_name]
-                for idx in indices:
+            for deployment_model_name in self._resolve_to_deployment_model_names(model=model_name):
+                for idx in self.model_name_to_deployment_indices.get(deployment_model_name) or ():
                     model = self.model_list[idx]
                     if "model_info" in model and "id" in model["model_info"]:
                         if exclude_team_models and model["model_info"].get("team_id"):

tests/test_litellm/router_strategy/test_router_routing_groups.py

@@ -12,7 +12,7 @@ import pytest
 
 import litellm
 from litellm import Router
-from litellm.types.router import RoutingGroup, RoutingStrategy
+from litellm.types.router import RouterRateLimitError, RoutingGroup, RoutingStrategy
 
 
 def _model_list():
@@ -1065,3 +1065,86 @@ async def test_group_call_429_cools_down_member_across_retries():
     )
     cooldown_ids = await _call_and_get_cooldowns(router, "quality")
     assert "deploy-3" in cooldown_ids
+
+
+def test_get_model_ids_resolves_group_and_alias_to_member_deployments():
+    router = Router(
+        model_list=_model_list(),
+        model_group_alias={"quality-alias": "quality"},
+        routing_groups=_quality_group("simple-shuffle"),
+    )
+    assert sorted(router.get_model_ids(model_name="quality")) == ["deploy-1", "deploy-2", "deploy-3"]
+    assert sorted(router.get_model_ids(model_name="quality-alias")) == ["deploy-1", "deploy-2", "deploy-3"]
+    assert sorted(router.get_model_ids(model_name="filtered-model")) == ["deploy-1", "deploy-2"]
+    assert router.get_model_ids(model_name="not-a-model") == []
+
+
+@pytest.mark.asyncio
+async def test_group_call_reports_member_cooldown_time_when_every_member_is_cooling_down():
+    """
+    Regression: the group name has to resolve to its members' deployment ids, else the
+    429 raised for a fully cooled-down group reports the router's default cooldown
+    instead of the members' own, and the proxy sends that as `retry-after`.
+    """
+    deployment_cooldown_time = 120
+    router = Router(
+        model_list=[
+            {**deployment, "model_info": {**deployment["model_info"], "cooldown_time": deployment_cooldown_time}}
+            for deployment in _model_list()
+        ],
+        routing_groups=_quality_group("simple-shuffle"),
+        num_retries=0,
+    )
+    assert router.cooldown_cache.default_cooldown_time != deployment_cooldown_time
+
+    for deployment_id in ("deploy-1", "deploy-2", "deploy-3"):
+        router.cooldown_cache.add_deployment_to_cooldown(
+            model_id=deployment_id,
+            original_exception=litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4o"),
+            exception_status=429,
+            cooldown_time=deployment_cooldown_time,
+        )
+
+    with pytest.raises(RouterRateLimitError) as exc_info:
+        await router.async_get_available_deployment(model="quality", request_kwargs={})
+
+    assert exc_info.value.cooldown_time == deployment_cooldown_time
+
+
+def test_get_model_ids_for_a_group_does_not_reach_team_scoped_deployments():
+    """
+    Group resolution goes through the `model_name` index, and a team deployment is
+    indexed under its own internal `model_name` (its `team_public_model_name` lives
+    in a separate index), so resolving a group cannot pull in another team's
+    deployment ids. Unchanged from a plain `model_name` lookup.
+    """
+    shared_deployment = {
+        "model_name": "shared-model",
+        "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test-1", "api_base": "https://example.invalid"},
+        "model_info": {"id": "shared-1"},
+    }
+    team_deployments = [
+        {
+            "model_name": f"shared-model_{team_id}_uuid",
+            "litellm_params": {
+                "model": "openai/gpt-4o",
+                "api_key": f"sk-test-{team_id}",
+                "api_base": "https://example.invalid",
+            },
+            "model_info": {
+                "id": f"{team_id}-1",
+                "team_id": team_id,
+                "team_public_model_name": "shared-model",
+            },
+        }
+        for team_id in ("team-a", "team-b")
+    ]
+    router = Router(
+        model_list=[shared_deployment, *team_deployments],
+        routing_groups=[{"group_name": "quality", "models": ["shared-model"], "routing_strategy": "simple-shuffle"}],
+    )
+
+    assert router.get_model_ids(model_name="quality") == ["shared-1"]
+    assert router.get_model_ids(model_name="shared-model") == ["shared-1"]
+    # the team deployments are still routable under their own names
+    assert router.get_model_ids(model_name="shared-model_team-a_uuid") == ["team-a-1"]