Merge 674b66a6b9 into 5103d8a0ee

litellm/router_rust/surface/registry.py

@@ -1822,10 +1822,16 @@ class HostRegistry(HostModelInfo):
     def _core_get_model_ids(
         self, model_name: str | None = None, exclude_team_models: bool = False
     ) -> list[str]:  # mutable-ok: a fresh list
+        """The ids of every deployment, or of the deployments ``model_name`` routes to: it may be any
+        name the router serves (see ``_resolve_to_deployment_model_names``)."""
         if model_name is not None and _index_key(model_name) is None:
             return []
-        ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models)
-        self._raise_at(raw)
+        names: Final = (None,) if model_name is None else self._resolve_to_deployment_model_names(model=model_name)
+        ids: Final[list[str]] = []  # mutable-ok: a fresh list
+        for name in names:
+            found, raw = self._lookup().lookup_model_ids(name, exclude_team_models)
+            self._raise_at(raw)
+            ids.extend(found)
         return ids
 
     # -- resolving a name -------------------------------------------------------------------------
@@ -1833,6 +1839,22 @@ class HostRegistry(HostModelInfo):
     def _core_get_model_from_alias(self, model: str) -> str | None:
         return resolve_model_group_alias(self._core_model_group_alias, model)
 
+    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_routable_model_group(self, model_name: str) -> str:
         target: Final = self._core_get_model_from_alias(model_name)
         return target if target is not None else model_name

tests/unit/router_strategy/test_router_routing_groups.py

@@ -10,7 +10,6 @@ import datetime
 import uuid
 from collections.abc import Mapping
 from typing import Final, Literal
-from unittest.mock import patch
 
 import pytest
 from pydantic import ValidationError
@@ -21,7 +20,7 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.router_api import RoutingContext
 from litellm.router_strategy.least_busy import picks_reserve
-from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup
+from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RouterRateLimitError, RoutingGroup
 from litellm.utils import Rules, function_setup
 
 
@@ -1150,3 +1149,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._registry_state._resolve_to_deployment_model_names(model="quality") == (
+        "filtered-model",
+        "other-model",
+    )
+    # an alias onto a group composes: alias -> group -> members
+    assert router._registry_state._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._registry_state._resolve_to_deployment_model_names(model="plain-alias") == (
+        "other-model",
+    )
+    # a plain `model_name` resolves to itself
+    assert router._registry_state._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._registry_state._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._registry_state._resolve_to_deployment_model_names(model="quality") == ("shared-model",)
+    assert router._registry_state._resolve_to_deployment_model_names(model="shared-model_team-a_uuid") == (
+        "shared-model_team-a_uuid",
+    )

tests/unit/test_router/core/test_registry_characterization_lookups.py

@@ -190,10 +190,10 @@ def test_ids_can_leave_out_team_deployments(make_core: MakeCore) -> None:
     assert registry.names(kind="ids", exclude_team_models=True) == ["a", "b", "e", "w"]
 
 
-def test_ids_of_a_group_read_only_the_group_index(make_core: MakeCore) -> None:
+def test_ids_of_a_group_follow_its_alias_and_read_only_the_group_index(make_core: MakeCore) -> None:
     registry = _core(make_core).registry()
 
-    assert registry.names("fast", kind="ids") == []
+    assert registry.names("fast", kind="ids") == ["a", "b"]
     assert registry.names("pub", kind="ids") == []
     assert registry.names("int-a", kind="ids") == ["t1"]
     assert registry.names("int-a", kind="ids", exclude_team_models=True) == []