mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(health): keep team public names to the owning team and let model_id win over model
A public name a team publishes its own deployment copy under now targets that copy only for a caller from that team, so an admin or another team probing the shared name gets the global deployment alone model_id wins when paired with model: a foreign id still gets the 403, and an id no deployment carries gets the 404 of the lone-id path, before any probe runs or a result is stored under it cache_health_check_results accepts the Mapping sequences perform_health_check returns
This commit is contained in:
parent
787f2dee0c
commit
6209f0694b
6 changed files with 231 additions and 46 deletions
|
|
@ -258,15 +258,18 @@ 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) -> bool:
|
||||
"""True when `model_name` is the deployment's model_name or the public name a team key reaches it by."""
|
||||
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."""
|
||||
info: Final = deployment.get("model_info")
|
||||
public_name: Final = info.get("team_public_model_name") if isinstance(info, Mapping) else None
|
||||
return model_name in (deployment.get("model_name"), public_name)
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
def _narrow_to_target(
|
||||
model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None
|
||||
model_list: Sequence[Mapping[str, object]], model: str | None, model_id: str | None, team_id: str | None
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""Narrow to the requested deployment. An id matching nothing keeps the whole list."""
|
||||
if model_id is not None:
|
||||
|
|
@ -275,7 +278,7 @@ def _narrow_to_target(
|
|||
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))
|
||||
return by_param or tuple(x for x in model_list if deployment_answers_to(x, model, team_id))
|
||||
|
||||
|
||||
def _is_strategy_router_deployment(litellm_params: Mapping[str, object]) -> bool:
|
||||
|
|
@ -820,13 +823,15 @@ async def perform_health_check(
|
|||
instrumentation_context: dict | None = None,
|
||||
health_check_skip_disabled_background_models: bool = False,
|
||||
router: "Router | None" = None,
|
||||
team_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
Perform a health check on the system.
|
||||
|
||||
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.
|
||||
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 ``health_check_skip_disabled_background_models`` is True (via
|
||||
``general_settings.health_check_skip_disabled_background_models``), deployments
|
||||
|
|
@ -857,7 +862,7 @@ async def perform_health_check(
|
|||
cycle_start_time: Final = time.monotonic()
|
||||
requested_model_count: Final = len(model_list)
|
||||
skip_disabled: Final = health_check_skip_disabled_background_models
|
||||
narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id), skip_disabled)
|
||||
narrowed: Final = _health_check_eligible(_narrow_to_target(model_list, model, model_id, team_id), skip_disabled)
|
||||
if not narrowed:
|
||||
if instrumentation_enabled:
|
||||
logger.debug(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -143,8 +144,8 @@ class SharedHealthCheckManager:
|
|||
|
||||
async def cache_health_check_results(
|
||||
self,
|
||||
healthy_endpoints: list[dict[str, Any]],
|
||||
unhealthy_endpoints: list[dict[str, Any]],
|
||||
healthy_endpoints: Sequence[Mapping[str, object]],
|
||||
unhealthy_endpoints: Sequence[Mapping[str, object]],
|
||||
) -> None:
|
||||
"""
|
||||
Cache health check results in Redis.
|
||||
|
|
|
|||
|
|
@ -925,41 +925,31 @@ def _caller_may_probe_deployment(
|
|||
)
|
||||
|
||||
|
||||
def _resolve_targeted_model_ids(model_list: list, model: str | None, model_id: str | None) -> set | None:
|
||||
def _resolve_targeted_model_ids(
|
||||
model_list: list, model: str | None, model_id: str | None, team_id: str | None
|
||||
) -> set | None:
|
||||
"""
|
||||
Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of
|
||||
deployment IDs the response should be scoped to.
|
||||
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.
|
||||
|
||||
Mirrors the live-path semantics in ``perform_health_check()``: ``model``
|
||||
matches the deployment's ``model_name`` alias, its ``litellm_params.model``
|
||||
provider string, or the ``model_info.team_public_model_name`` a team key
|
||||
reaches it by. ``model_id`` matches ``model_info.id``.
|
||||
|
||||
Both query params are validated against the supplied ``model_list``.
|
||||
Callers pass an already-scoped list (filtered to the caller's allowed
|
||||
models for non-admins, full list for admins), so a ``model_id`` that
|
||||
isn't present resolves to an empty set rather than a single-element
|
||||
set — preventing a non-admin from reading another deployment's cached
|
||||
health entry by guessing its ID.
|
||||
|
||||
Returns ``None`` when no targeting is requested — callers should treat
|
||||
that as "no filter."
|
||||
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.
|
||||
Returns ``None`` when no targeting is requested.
|
||||
"""
|
||||
if not model and not model_id:
|
||||
if model_id:
|
||||
return {i for m in model_list if (i := (m.get("model_info") or {}).get("id")) == model_id}
|
||||
if not model:
|
||||
return None
|
||||
target_ids: Final[set] = set()
|
||||
for m in model_list:
|
||||
deployment_id = (m.get("model_info") or {}).get("id")
|
||||
if not deployment_id:
|
||||
continue
|
||||
if model_id and deployment_id == model_id:
|
||||
target_ids.add(deployment_id)
|
||||
continue
|
||||
if model:
|
||||
litellm_model = (m.get("litellm_params") or {}).get("model")
|
||||
if litellm_model == model or deployment_answers_to(m, model):
|
||||
target_ids.add(deployment_id)
|
||||
return target_ids
|
||||
return {
|
||||
i
|
||||
for m in model_list
|
||||
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))
|
||||
}
|
||||
|
||||
|
||||
def _filter_health_check_results_by_model_ids(results: dict, allowed_model_ids: set) -> dict:
|
||||
|
|
@ -1040,8 +1030,12 @@ def _health_endpoint_resolve_target_model_name(
|
|||
model_id: str | None,
|
||||
llm_router,
|
||||
) -> str | None:
|
||||
"""Map ``model_id`` (without ``model``) to ``model_name`` for live health checks."""
|
||||
if not model_id or model:
|
||||
"""Map ``model_id`` to its deployment's ``model_name`` for live health checks.
|
||||
|
||||
``model_id`` wins over ``model``, so an id no deployment carries is a 404 even
|
||||
when it is paired with a known name.
|
||||
"""
|
||||
if not model_id:
|
||||
return model
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -1161,7 +1155,7 @@ async def health_endpoint(
|
|||
if not restrict_to_allowed_models
|
||||
or _caller_may_probe_deployment(m, allowed_models, llm_router, user_api_key_dict.team_id, is_admin)
|
||||
]
|
||||
targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id)
|
||||
targeted_ids: Final = _resolve_targeted_model_ids(_llm_model_list, model, model_id, user_api_key_dict.team_id)
|
||||
if restrict_to_allowed_models and targeted_ids is not None and not targeted_ids:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
|
|
@ -1228,6 +1222,7 @@ async def health_endpoint(
|
|||
model_id=model_id,
|
||||
max_concurrency=health_check_concurrency,
|
||||
router=llm_router,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
**_hc_filter,
|
||||
)
|
||||
return _post_process(router_result)
|
||||
|
|
|
|||
|
|
@ -2357,6 +2357,7 @@ async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_e
|
|||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
model="model-a",
|
||||
model_id=None,
|
||||
)
|
||||
|
||||
assert response.status_code == 503
|
||||
|
|
@ -2417,6 +2418,7 @@ async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endp
|
|||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
model="model-a",
|
||||
model_id=None,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
|
@ -3088,6 +3090,161 @@ async def test_health_endpoint_targets_both_deployments_behind_a_shared_public_n
|
|||
assert [ep["model_id"] for ep in result["healthy_endpoints"]] == ["id-bedrock", "id-team-b"]
|
||||
|
||||
|
||||
async def _live_narrowed_model_ids(
|
||||
model_list: Sequence[Mapping[str, object]],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
model: str | None = None,
|
||||
model_id: str | None = None,
|
||||
) -> set[str]:
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
async def fake_probe(model_list, details=True, max_concurrency=None, instrumentation_context=None):
|
||||
probed = [{"model": m["litellm_params"]["model"], "model_id": m["model_info"]["id"]} for m in model_list]
|
||||
return probed, [], {}
|
||||
|
||||
with (
|
||||
_proxy_health_globals(model_list, _router_for(model_list)),
|
||||
patch( # test-quality-ok: the probe is the provider edge; which deployments reach it is the assertion
|
||||
"litellm.proxy.health_check._perform_health_check", side_effect=fake_probe
|
||||
),
|
||||
):
|
||||
result = await health_endpoint(
|
||||
response=Response(), user_api_key_dict=user_api_key_dict, model=model, model_id=model_id
|
||||
)
|
||||
|
||||
return {ep["model_id"] for ep in result["healthy_endpoints"]}
|
||||
|
||||
|
||||
_ADMIN_OUTSIDE_TEAM_B = UserAPIKeyAuth(api_key="hashed-test-key", models=[], user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies():
|
||||
"""
|
||||
An admin outside team-b asks for ``bedrock-nova``. Team-b's copy answers to
|
||||
that name only for team-b (routing keys public names by team), so probing
|
||||
it too would spend team-b's credentials and let a healthy team copy mask a
|
||||
down global deployment as 200.
|
||||
"""
|
||||
probed = await _live_narrowed_model_ids(_TEAM_MODEL_LIST, _ADMIN_OUTSIDE_TEAM_B, model="bedrock-nova")
|
||||
|
||||
assert probed == {"id-bedrock"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_probes_both_deployments_behind_a_shared_public_name_for_the_owning_team():
|
||||
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"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_keeps_an_admin_probe_by_name_off_other_teams_public_copies_on_background_cache_path():
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
with _proxy_health_globals(
|
||||
_TEAM_MODEL_LIST,
|
||||
_router_for(_TEAM_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-bedrock"]
|
||||
|
||||
|
||||
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"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("use_background_health_checks", [False, True])
|
||||
async def test_health_endpoint_rejects_an_in_scope_model_paired_with_a_foreign_model_id(use_background_health_checks):
|
||||
"""
|
||||
A key scoped to ``bedrock-nova`` pairs that name with another team's
|
||||
deployment id. The in-scope name must not carry the foreign id past the
|
||||
403: the live path narrows by id first, so the caller's own deployment
|
||||
would be probed and its result stored under the foreign id.
|
||||
"""
|
||||
from fastapi import HTTPException, Response
|
||||
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
with (
|
||||
_proxy_health_globals(
|
||||
_TEAM_MODEL_LIST,
|
||||
_router_for(_TEAM_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="id-team-b",
|
||||
)
|
||||
|
||||
assert refused.value.status_code == 403
|
||||
assert "id-team-b" in str(refused.value.detail)
|
||||
probe.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_background_health_checks", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_endpoint_returns_404_for_a_model_paired_with_an_unknown_model_id(use_background_health_checks):
|
||||
"""
|
||||
``model_id`` wins over ``model``: pairing a known name with an id no
|
||||
deployment carries gets the same 404 as the lone unknown id, before any
|
||||
probe runs or a result is stored under the unknown id.
|
||||
"""
|
||||
from fastapi import HTTPException, Response
|
||||
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
|
||||
|
||||
with (
|
||||
_proxy_health_globals(
|
||||
_TEAM_MODEL_LIST,
|
||||
_router_for(_TEAM_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=_ADMIN_OUTSIDE_TEAM_B,
|
||||
model="bedrock-nova",
|
||||
model_id="id-nobody-has",
|
||||
)
|
||||
|
||||
assert refused.value.status_code == 404
|
||||
assert "id-nobody-has" in str(refused.value.detail)
|
||||
probe.assert_not_awaited()
|
||||
|
||||
|
||||
def test_health_test_connection_keeps_error_and_raw_request_through_the_allowlist(monkeypatch):
|
||||
"""
|
||||
The dashboard's Test Connect button reads ``result.error`` and
|
||||
|
|
@ -3198,6 +3355,8 @@ async def test_health_endpoint_result_survives_non_json_safe_deployment_params()
|
|||
result = await health_endpoint(
|
||||
response=Response(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin-key", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
model=None,
|
||||
model_id=None,
|
||||
)
|
||||
|
||||
encoded = jsonable_encoder(result)
|
||||
|
|
|
|||
|
|
@ -644,7 +644,7 @@ async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_n
|
|||
"litellm.proxy.health_check._perform_health_check", probe
|
||||
):
|
||||
healthy, unhealthy, _ = await perform_health_check(
|
||||
model_list=[team_deployment, other_deployment], model="bedrock-nova"
|
||||
model_list=[team_deployment, other_deployment], model="bedrock-nova", team_id="team-b"
|
||||
)
|
||||
|
||||
assert [m["model_info"]["id"] for m in probe.call_args.args[0]] == ["id-team-b"]
|
||||
|
|
@ -652,6 +652,31 @@ async def test_perform_health_check_narrows_to_a_team_deployment_by_its_public_n
|
|||
assert unhealthy == []
|
||||
|
||||
|
||||
@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."""
|
||||
from litellm.proxy.health_check import perform_health_check
|
||||
|
||||
team_deployment = {
|
||||
"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"},
|
||||
}
|
||||
probe = AsyncMock(return_value=([{"model": "bedrock/us.amazon.nova-2-lite-v1:0", "model_id": "id-team-b"}], [], {}))
|
||||
|
||||
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=[team_deployment], model="bedrock-nova", team_id=team_id
|
||||
)
|
||||
|
||||
probe.assert_not_awaited()
|
||||
assert healthy == []
|
||||
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
|
||||
|
||||
|
|
|
|||
|
|
@ -872,9 +872,9 @@ def test_narrowing_by_an_id_that_matches_nothing_keeps_the_whole_list():
|
|||
"""Pinned because the disabled-dependency fix moved this filter into its own helper."""
|
||||
deployments = [{"model_name": "a", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a-1"}}]
|
||||
|
||||
assert hc_module._narrow_to_target(deployments, None, "no-such-id") == tuple(deployments)
|
||||
assert hc_module._narrow_to_target(deployments, None, "a-1") == tuple(deployments)
|
||||
assert hc_module._narrow_to_target(deployments, "a", None) == tuple(deployments)
|
||||
assert hc_module._narrow_to_target(deployments, None, "no-such-id", None) == tuple(deployments)
|
||||
assert hc_module._narrow_to_target(deployments, None, "a-1", None) == tuple(deployments)
|
||||
assert hc_module._narrow_to_target(deployments, "a", None, None) == tuple(deployments)
|
||||
|
||||
|
||||
def _nested_router_fixture(parent_tier: str):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue