Merge 09b084fe9f into 83be155b65

litellm/router_python_core.py

@@ -2537,6 +2537,23 @@ class PythonCore(DeprecatedRouterMembers):
             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._core_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 _core_routing_group_has_alternatives(self, model_group: str | None) -> bool:
         """
         True when `model_group` names a callable routing group whose member
@@ -13789,6 +13806,11 @@ class PythonCore(DeprecatedRouterMembers):
 
         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.
         """
@@ -13796,11 +13818,12 @@ class PythonCore(DeprecatedRouterMembers):
 
         if model_name is not None:
             # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name
-            for model in self._core_deployments_named(model_name) or ():
-                if "model_info" in model and "id" in model["model_info"]:
-                    if exclude_team_models and model["model_info"].get("team_id"):
-                        continue
-                    ids.append(model["model_info"]["id"])
+            for deployment_model_name in self._resolve_to_deployment_model_names(model=model_name):
+                for model in self._core_deployments_named(deployment_model_name) or ():
+                    if "model_info" in model and "id" in model["model_info"]:
+                        if exclude_team_models and model["model_info"].get("team_id"):
+                            continue
+                        ids.append(model["model_info"]["id"])
         else:
             # When model_name is None, return all model IDs
             # Use the index map keys for O(n) where n = total deployments

tests/unit/router_strategy/test_router_routing_groups.py

@@ -21,7 +21,13 @@ from litellm import Router
 from litellm.caching.redis_cache import RedisPipelineIncrementOperation
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.router_api import RoutingContext
-from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy
+from litellm.types.router import (
+    DeploymentTypedDict,
+    FallbackAccessCheck,
+    RouterRateLimitError,
+    RoutingGroup,
+    RoutingStrategy,
+)
 from litellm.utils import Rules, function_setup
 
 
@@ -2255,3 +2261,175 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se
                 await t
             except asyncio.CancelledError:
                 pass
+
+
+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"]
+
+
+def test_resolve_to_deployment_model_names_covers_alias_group_and_plain_names():
+    """
+    `_resolve_to_deployment_model_names` is what lookups keyed by a served name
+    go through, so each of its three documented branches is pinned here: an
+    alias resolves to its target, a routing group resolves to its member
+    `model_name`s, and anything else resolves to itself.
+    """
+    router = _build_router(
+        routing_groups=[
+            {
+                "group_name": "quality",
+                "models": ["filtered-model", "other-model"],
+                "routing_strategy": "simple-shuffle",
+            }
+        ],
+    )
+    router.model_group_alias = {"quality-alias": "quality", "plain-alias": "other-model"}
+
+    # a group resolves to its members, in the order the group declares them
+    assert router._resolve_to_deployment_model_names(model="quality") == (
+        "filtered-model",
+        "other-model",
+    )
+    # an alias onto a group composes: alias -> group -> members
+    assert router._resolve_to_deployment_model_names(model="quality-alias") == (
+        "filtered-model",
+        "other-model",
+    )
+    # an alias onto a plain deployment resolves to that deployment's name
+    assert router._resolve_to_deployment_model_names(model="plain-alias") == (
+        "other-model",
+    )
+    # a plain `model_name` resolves to itself
+    assert router._resolve_to_deployment_model_names(model="filtered-model") == (
+        "filtered-model",
+    )
+    # an unknown name is passed through rather than dropped, so the caller's
+    # own lookup decides it misses
+    assert router._resolve_to_deployment_model_names(model="no-such-model") == (
+        "no-such-model",
+    )
+
+
+def test_resolve_to_deployment_model_names_keeps_team_deployments_on_their_own_name():
+    """
+    The team-scoped counterpart of the above. A team deployment is indexed under
+    its internal `model_name`, so resolving a group never yields it, and it is
+    reached only when that internal name is requested.
+    """
+    router = Router(
+        model_list=[
+            {
+                "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"},
+            },
+            {
+                "model_name": "shared-model_team-a_uuid",
+                "litellm_params": {
+                    "model": "openai/gpt-4o",
+                    "api_key": "sk-test-2",
+                    "api_base": "https://example.invalid",
+                },
+                "model_info": {
+                    "id": "team-a-1",
+                    "team_id": "team-a",
+                    "team_public_model_name": "shared-model",
+                },
+            },
+        ],
+        routing_groups=[
+            {
+                "group_name": "quality",
+                "models": ["shared-model"],
+                "routing_strategy": "simple-shuffle",
+            }
+        ],
+    )
+
+    assert router._resolve_to_deployment_model_names(model="quality") == ("shared-model",)
+    assert router._resolve_to_deployment_model_names(model="shared-model_team-a_uuid") == (
+        "shared-model_team-a_uuid",
+    )