diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index 820db52d5a1..9f9550c1943 100644 --- a/litellm/router_rust/surface/registry.py +++ b/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 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 8a0d5b9a1d6..f93338d3b0d 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/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", + ) diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py index 19d58ff10d4..af5ab46e65e 100644 --- a/tests/unit/test_router/core/test_registry_characterization_lookups.py +++ b/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) == []