diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 9ca31c8b626..d774fd6e191 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -258,14 +258,47 @@ def _deployment_model(deployment: Mapping[str, object]) -> str | None: return params.get("model") if isinstance(params, Mapping) else None -def deployment_answers_to(deployment: Mapping[str, object], model_name: str, team_id: str | None) -> bool: - """True when `model_name` is the deployment's model_name or the public name the caller's own team reaches it by.""" +def _owner_team_id(deployment: Mapping[str, object]) -> str | None: info: Final = deployment.get("model_info") - public_name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None - owner_team_id: Final = info.get("team_id") if isinstance(info, Mapping) else None - return model_name == deployment.get("model_name") or ( - model_name == public_name and team_id is not None and owner_team_id == team_id + owner: Final = info.get("team_id") if isinstance(info, Mapping) else None + return owner if isinstance(owner, str) else None + + +def _team_public_model_name(deployment: Mapping[str, object]) -> str | None: + info: Final = deployment.get("model_info") + name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None + return name if isinstance(name, str) else None + + +def _deployments_routed_by_name( + model_list: Sequence[Mapping[str, object]], model_name: str, team_id: str | None +) -> tuple[Mapping[str, object], ...]: + """The deployments a request for ``model_name`` from this caller routes to. + + A team's own copies published under that name win, then deployments carrying it as + ``model_name``. A caller with no team reaches a public name only when nothing carries + it as ``model_name``, and only an admin still has another team's deployment in a + scoped ``model_list`` by then. + """ + own_copies: Final = tuple( + x + for x in model_list + if team_id is not None and _owner_team_id(x) == team_id and _team_public_model_name(x) == model_name ) + if own_copies: + return own_copies + by_name: Final = tuple(x for x in model_list if x.get("model_name") == model_name) + if by_name or team_id is not None: + return by_name + return tuple(x for x in model_list if _team_public_model_name(x) == model_name) + + +def deployments_targeted_by_name( + model_list: Sequence[Mapping[str, object]], model: str, team_id: str | None +) -> tuple[Mapping[str, object], ...]: + """``model`` targets deployments by ``litellm_params.model`` first, then the way a request for it routes.""" + by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model) + return by_param or _deployments_routed_by_name(model_list, model, team_id) def _narrow_to_target( @@ -277,8 +310,7 @@ def _narrow_to_target( return by_id or tuple(model_list) if model is None: return tuple(model_list) - by_param: Final = tuple(x for x in model_list if _deployment_model(x) == model) - return by_param or tuple(x for x in model_list if deployment_answers_to(x, model, team_id)) + return deployments_targeted_by_name(model_list, model, team_id) def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool: @@ -830,8 +862,10 @@ async def perform_health_check( When model_id is provided, only the deployment with that id is checked (so models that share the same name but have different ids are checked separately). - When model (name) is provided, all deployments matching that name are checked; - a team's public model name only matches for a caller from that team (``team_id``). + When model (name) is provided, the deployments a request for that name from the + caller (``team_id``) would route to are checked: the caller's team copies published + under that name, else the deployments named that way, else a public name that only + another team's deployment carries. When ``health_check_skip_disabled_background_models`` is True (via ``general_settings.health_check_skip_disabled_background_models``), deployments diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index dcef0fcd58a..002c4debb46 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -50,7 +50,7 @@ from litellm.proxy.health_check import ( ADMIN_ONLY_HEALTH_DISPLAY_PARAMS, _clean_endpoint_data, _update_litellm_params_for_health_check, - deployment_answers_to, + deployments_targeted_by_name, health_check_filter_kwargs_from_general_settings, perform_health_check, run_with_timeout, @@ -932,9 +932,9 @@ def _resolve_targeted_model_ids( Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of deployment IDs the response should be scoped to, mirroring the live-path narrowing in ``perform_health_check()``: ``model_id`` wins when given and - matches ``model_info.id`` only; ``model`` matches the deployment's - ``model_name`` alias, its ``litellm_params.model`` provider string, or the - ``model_info.team_public_model_name`` the caller's own team reaches it by. + matches ``model_info.id`` only; ``model`` targets deployments by their + ``litellm_params.model`` provider string, else the deployments a request + for that name from the caller would route to (``deployments_targeted_by_name``). Callers pass an already-scoped list, so a ``model_id`` outside the caller's scope resolves to an empty set and never to the unvalidated id. @@ -946,9 +946,8 @@ def _resolve_targeted_model_ids( return None return { i - for m in model_list + for m in deployments_targeted_by_name(model_list, model, team_id) if (i := (m.get("model_info") or {}).get("id")) - and ((m.get("litellm_params") or {}).get("model") == model or deployment_answers_to(m, model, team_id)) } diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index c5950f6f30e..4c920dfb4d1 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -3069,7 +3069,7 @@ async def test_health_endpoint_returns_a_team_only_deployment_by_its_public_name @pytest.mark.asyncio -async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_name_on_background_cache_path(): +async def test_health_endpoint_returns_only_the_owning_teams_copy_behind_a_shared_public_name_on_cache_path(): from fastapi import Response from litellm.proxy.health_endpoints._health_endpoints import health_endpoint @@ -3087,7 +3087,7 @@ async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_n model_id=None, ) - assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock", "id-team-b"] + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] async def _live_narrowed_model_ids( @@ -3134,14 +3134,15 @@ async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_publ @pytest.mark.asyncio -async def test_health_endpoint_probes_both_deployments_behind_a_shared_public_name_for_the_owning_team(): +async def test_health_endpoint_probes_only_the_owning_teams_copy_behind_a_shared_public_name(): + """Team-b's requests for ``bedrock-nova`` route to its copy alone, so its health probe reaches only that copy.""" probed = await _live_narrowed_model_ids( _TEAM_MODEL_LIST, UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"], team_id="team-b"), model="bedrock-nova", ) - assert probed == {"id-bedrock", "id-team-b"} + assert probed == {"id-team-b"} @pytest.mark.asyncio @@ -3163,13 +3164,80 @@ async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_publ assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock"] +@pytest.mark.asyncio +async def test_health_endpoint_probes_a_team_only_public_name_for_an_admin_on_live_path(): + """ + An admin's request for a public name only team-b's deployment carries routes + to that deployment, so the health probe for that name must reach it too + instead of answering an empty 503. + """ + probed = await _live_narrowed_model_ids(_TEAM_ONLY_MODEL_LIST, _ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova") + + assert probed == {"id-team-b"} + + +@pytest.mark.asyncio +async def test_health_endpoint_returns_a_team_only_public_name_for_an_admin_on_background_cache_path(): + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=True, + health_check_results=_TEAM_CACHED_RESULTS, + ): + result = await health_endpoint( + response=Response(), user_api_key_dict=_ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova", model_id=None + ) + + assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-team-b"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_background_health_checks", [False, True]) +async def test_health_endpoint_keeps_a_team_only_public_name_off_a_team_less_key(use_background_health_checks): + """ + A key with no team holds the name ``bedrock-nova`` but never sees team-b's + deployment, so the public-name fallback an admin gets must not open it up. + """ + from fastapi import HTTPException, Response + + from litellm.proxy.health_endpoints._health_endpoints import health_endpoint + + with ( + _proxy_health_globals( + _TEAM_ONLY_MODEL_LIST, + _router_for(_TEAM_ONLY_MODEL_LIST), + use_background_health_checks=use_background_health_checks, + health_check_results=_TEAM_CACHED_RESULTS, + ), + patch( # test-quality-ok: the probe must never run; the endpoint has no injection seam for it + "litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save", new_callable=AsyncMock + ) as probe, + pytest.raises(HTTPException) as refused, + ): + await health_endpoint( + response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-test-key", models=["bedrock-nova"]), + model="bedrock-nova", + model_id=None, + ) + + assert refused.value.status_code == 403 + assert "bedrock-nova" in str(refused.value.detail) + probe.assert_not_awaited() + + def test_resolve_targeted_model_ids_lets_model_id_win_over_model(): resolve = _health_endpoints_module._resolve_targeted_model_ids assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", "id-team-b", None) == {"id-team-b"} assert resolve([_TEAM_MODEL_LIST[0]], "bedrock-nova", "id-team-b", None) == set() assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, None) == {"id-bedrock"} - assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, "team-b") == {"id-bedrock", "id-team-b"} + assert resolve(_TEAM_MODEL_LIST, "bedrock-nova", None, "team-b") == {"id-team-b"} + assert resolve(_TEAM_ONLY_MODEL_LIST, "bedrock-nova", None, None) == {"id-team-b"} @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index af20b80a718..e7abfcf126c 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -653,9 +653,8 @@ async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_n @pytest.mark.asyncio -@pytest.mark.parametrize("team_id", [None, "team-a"]) -async def test_perform_health_check_keeps_a_public_name_to_the_owning_team(team_id): - """A team's public model name is not a global alias: callers outside the team must not probe its deployment.""" +async def test_perform_health_check_keeps_a_public_name_off_another_team(): + """A team's public model name is not a global alias: a caller from another team must not probe its deployment.""" from litellm.proxy.health_check import perform_health_check team_deployment = { @@ -669,7 +668,7 @@ async def test_perform_health_check_keeps_a_public_name_to_the_owning_team(team_ "litellm.proxy.health_check._perform_health_check", probe ): healthy, unhealthy, _ = await perform_health_check( - model_list=[team_deployment], model="bedrock-nova", team_id=team_id + model_list=[team_deployment], model="bedrock-nova", team_id="team-a" ) probe.assert_not_awaited() @@ -677,6 +676,50 @@ async def test_perform_health_check_keeps_a_public_name_to_the_owning_team(team_ assert unhealthy == [] +_GLOBAL_DEPLOYMENT = { + "model_name": "bedrock-nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-bedrock"}, +} +_TEAM_B_COPY = { + "model_name": "bedrock-nova_team-b_9f2c", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}, + "model_info": {"id": "id-team-b", "team_id": "team-b", "team_public_model_name": "bedrock-nova"}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_id", "model_list", "expected_ids"), + [ + (None, [_TEAM_B_COPY], ["id-team-b"]), + (None, [_GLOBAL_DEPLOYMENT, _TEAM_B_COPY], ["id-bedrock"]), + ("team-b", [_GLOBAL_DEPLOYMENT, _TEAM_B_COPY], ["id-team-b"]), + ], + ids=[ + "a team-less caller reaches a public name nothing else carries", + "model_name wins over a public name for a team-less caller", + "a team's own copy wins over the global model_name", + ], +) +async def test_perform_health_check_targets_a_name_the_way_a_request_for_it_routes(team_id, model_list, expected_ids): + """``/health?model=`` probes the deployments a request for that name from the same caller would route to.""" + from litellm.proxy.health_check import perform_health_check + + probe = AsyncMock( + return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": i} for i in expected_ids], [], {}) + ) + + with patch( # test-quality-ok: the deployments handed to the probe are the assertion; no injection seam + "litellm.proxy.health_check._perform_health_check", probe + ): + healthy, unhealthy, _ = await perform_health_check(model_list=model_list, model="bedrock-nova", team_id=team_id) + + assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == expected_ids + assert [ep["model_id"] for ep in healthy] == expected_ids + assert unhealthy == [] + + def test_parse_background_health_check_model_groups_unset_returns_none(): from litellm.proxy.health_check import parse_background_health_check_model_groups