diff --git a/litellm/proxy/db/health_check_latest.py b/litellm/proxy/db/health_check_latest.py index bc9120e3b4e..1c3188210f0 100644 --- a/litellm/proxy/db/health_check_latest.py +++ b/litellm/proxy/db/health_check_latest.py @@ -10,7 +10,7 @@ worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per from __future__ import annotations import json -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Collection, Iterable, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -109,11 +109,20 @@ async def fetch_latest_health_checks_for_models( return () -def latest_by_model_name(rows: Iterable[LatestHealthCheckRow]) -> Mapping[str, LatestHealthCheckRow]: +def latest_by_model_name( + rows: Iterable[LatestHealthCheckRow], + active_model_ids: Mapping[str, Collection[str]] | None = None, +) -> Mapping[str, LatestHealthCheckRow]: """The newest row per model name. The queries above return one row per (model_id, model_name), so a model group with several deployments comes back as several rows, plus one without a model_id if it was ever checked by name. They arrive in model_id order with NULL last, which says nothing about which check is newest. + Active IDs exclude retired deployments while keeping legacy name-only probes eligible. """ - return MappingProxyType({row.model_name: row for row in sorted(rows, key=lambda row: row.checked_at)}) + eligible_rows: Final = ( + row + for row in rows + if active_model_ids is None or row.model_id is None or row.model_id in active_model_ids.get(row.model_name, ()) + ) + return MappingProxyType({row.model_name: row for row in sorted(eligible_rows, key=lambda row: row.checked_at)}) diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index cb1e51d0deb..ec719422d4b 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -5,6 +5,7 @@ import re from collections.abc import Awaitable, Callable, Mapping, Sequence from functools import cache from importlib.resources import files +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, HTTPException, Request @@ -302,7 +303,13 @@ async def public_model_hub(): if prisma_client is not None: try: latest_checks: Final = await prisma_client.get_all_latest_health_checks() - for model_name, check in latest_by_model_name(latest_checks).items(): + active_model_ids: Final = MappingProxyType( + { + group.model_group: frozenset(llm_router.registry().names(group.model_group, kind="ids")) + for group in model_groups + } + ) + for model_name, check in latest_by_model_name(latest_checks, active_model_ids).items(): health_checks_map[model_name] = convert_health_check_to_dict(check) except Exception: pass diff --git a/litellm/proxy/public_endpoints/public_v1/model_hub.py b/litellm/proxy/public_endpoints/public_v1/model_hub.py index e58149fba0b..a9f76f969f0 100644 --- a/litellm/proxy/public_endpoints/public_v1/model_hub.py +++ b/litellm/proxy/public_endpoints/public_v1/model_hub.py @@ -25,6 +25,7 @@ from litellm.proxy.list_api.list_framework import ( handle_list, ) from litellm.proxy.utils import PrismaClient +from litellm.router_api import RegistryReader from litellm.types.proxy.management_endpoints.management_v1 import ( FacetListResponse, ListResponse, @@ -55,9 +56,13 @@ class HealthSnapshotLookup(Protocol): @dataclass(frozen=True, slots=True) class PrismaHealthSnapshotLookup: prisma_client: PrismaClient + router: RegistryReader async def latest_for(self, model_groups: Sequence[str]) -> Mapping[str, HealthSnapshot]: checks: Final = await self.prisma_client.get_latest_health_checks_for_models(model_groups) + active_model_ids: Final = MappingProxyType( + {name: frozenset(self.router.registry().names(name, kind="ids")) for name in model_groups} + ) return MappingProxyType( { model_group: HealthSnapshot( @@ -65,7 +70,7 @@ class PrismaHealthSnapshotLookup: response_time_ms=check.response_time_ms, checked_at=check.checked_at.isoformat() if check.checked_at else None, ) - for model_group, check in latest_by_model_name(checks).items() + for model_group, check in latest_by_model_name(checks, active_model_ids).items() } ) @@ -218,13 +223,14 @@ def _published_rows() -> Sequence[ModelGroupInfoProxy]: def _executor( rows: Sequence[ModelGroupInfoProxy], prisma_client: PrismaClient | None, + router: RegistryReader | None, ) -> InMemoryListExecutor[ModelGroupInfoProxy]: - if prisma_client is None: + if prisma_client is None or router is None: return InMemoryListExecutor(rows=rows, cells=_cells) return InMemoryListExecutor( rows=rows, cells=_cells, - enrich_page=HealthEnricher(lookup=PrismaHealthSnapshotLookup(prisma_client=prisma_client)), + enrich_page=HealthEnricher(lookup=PrismaHealthSnapshotLookup(prisma_client=prisma_client, router=router)), ) @@ -253,11 +259,11 @@ async def public_model_hub_list( ``` """ try: - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import llm_router, prisma_client return await handle_list( spec=MODEL_HUB_LIST_SPEC, - executor=_executor(_published_rows(), prisma_client), + executor=_executor(_published_rows(), prisma_client, llm_router), request=request, caller=user_api_key_dict, ) diff --git a/tests/unit/proxy/db/test_health_check_latest.py b/tests/unit/proxy/db/test_health_check_latest.py index 2f27be4fe61..284b458ca39 100644 --- a/tests/unit/proxy/db/test_health_check_latest.py +++ b/tests/unit/proxy/db/test_health_check_latest.py @@ -1,6 +1,7 @@ from collections.abc import Awaitable, Callable from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock +from typing import Final +from unittest.mock import MagicMock import pytest @@ -164,3 +165,36 @@ def test_latest_by_model_name_lets_a_newer_name_probe_speak_for_the_group(): ) assert latest_by_model_name((per_deployment, name_probe))["gpt-4"] is name_probe + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("retired_status,active_status", [("unhealthy", "healthy"), ("healthy", "unhealthy")]) +def test_active_deployment_selection_excludes_retired_rows( + reverse: bool, retired_status: str, active_status: str +) -> None: + active: Final = _latest_row(model_id="active", status=active_status, checked_at="2026-10-16T11:10:00+00:00") + retired: Final = _latest_row(model_id="retired", status=retired_status, checked_at="2026-10-16T12:00:00+00:00") + rows: Final = (active, retired) + + latest: Final = latest_by_model_name(rows[::-1] if reverse else rows, {"gpt-4": frozenset({"active"})}) + + assert latest == {"gpt-4": active} + + +def test_active_deployment_selection_does_not_reuse_an_id_from_another_group() -> None: + retired: Final = _latest_row(model_id="reused", status="unhealthy") + + assert latest_by_model_name((retired,), {"gpt-4": frozenset(), "claude": frozenset({"reused"})}) == {} + + +def test_active_deployment_selection_keeps_legacy_name_probes() -> None: + name_probe: Final = _latest_row(model_id=None) + + assert latest_by_model_name((name_probe,), {"gpt-4": frozenset({"active"})}) == {"gpt-4": name_probe} + + +def test_active_deployment_selection_keeps_the_newest_active_provider() -> None: + older: Final = _latest_row(model_id="provider-a", status="unhealthy", checked_at="2026-10-16T11:10:00+00:00") + newer: Final = _latest_row(model_id="provider-b", checked_at="2026-10-16T12:00:00+00:00") + + assert latest_by_model_name((newer, older), {"gpt-4": frozenset({"provider-a", "provider-b"})}) == {"gpt-4": newer} diff --git a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py index 143f5bfb83b..e55e5b97a67 100644 --- a/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py +++ b/tests/unit/proxy/public_endpoints/public_v1/test_model_hub.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -33,6 +34,9 @@ class _FakeRouter: def get_model_list_of_routed_group(self, model_group: str, team_id: str | None = None) -> list[DeploymentTypedDict]: return [] + def get_model_ids(self, model_name: str) -> tuple[str, ...]: + return (f"{model_name}-deployment",) if model_name in self.infos else () + def _info( name: str, @@ -168,7 +172,7 @@ def test_the_latest_health_check_lands_on_its_row(monkeypatch): def test_a_stale_row_sorting_after_the_fresh_one_does_not_speak_for_the_group(monkeypatch): """The read orders by model_id with NULL last; that order used to pick the group's health.""" fresh = _health_check( - "model-001", model_id="deployment-a", checked_at=datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc) + "model-001", model_id="model-001-deployment", checked_at=datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc) ) stale = _health_check("model-001", status="unhealthy", checked_at=datetime(2026, 9, 4, 11, 20, tzinfo=timezone.utc)) prisma_client, _ = _recording_prisma([fresh, stale]) @@ -180,6 +184,32 @@ def test_a_stale_row_sorting_after_the_fresh_one_does_not_speak_for_the_group(mo assert rows["model-001"]["health_checked_at"] == "2026-10-16T11:10:00+00:00" +@pytest.mark.parametrize("retired_status,active_status", [("unhealthy", "healthy"), ("healthy", "unhealthy")]) +@pytest.mark.parametrize("retired_hour", [10, 12], ids=["older-retired-check", "newer-retired-check"]) +def test_retired_deployment_cannot_override_active_health( + monkeypatch: pytest.MonkeyPatch, retired_status: str, active_status: str, retired_hour: int +) -> None: + active: Final = _health_check( + "model-001", + status=active_status, + model_id="model-001-deployment", + checked_at=datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc), + ) + retired: Final = _health_check( + "model-001", + status=retired_status, + model_id="retired-deployment", + checked_at=datetime(2026, 10, 16, retired_hour, 0, tzinfo=timezone.utc), + ) + prisma_client, _ = _recording_prisma((active, retired)) + _publish(monkeypatch, _named(3), prisma_client=prisma_client) + + rows: Final = {row["model_group"]: row for row in _get().json()["data"]} + + assert rows["model-001"]["health_status"] == active_status + assert rows["model-001"]["health_checked_at"] == "2026-10-16T11:10:00+00:00" + + def test_a_health_read_that_returns_nothing_still_serves_the_page(monkeypatch): prisma_client, read = _recording_prisma() read.return_value = [] @@ -424,7 +454,9 @@ def test_features_filter_matches_a_model_with_any_of_the_named_features(monkeypa def test_a_single_feature_filter_selects_only_models_with_it(monkeypatch): - _publish(monkeypatch, (_info("sees", supports_vision=True), _info("plain"), _info("reasons", supports_reasoning=True))) + _publish( + monkeypatch, (_info("sees", supports_vision=True), _info("plain"), _info("reasons", supports_reasoning=True)) + ) assert _groups(_get("filter[features][in]=vision")) == ["sees"] assert _groups(_get("filter[features][in]=reasoning")) == ["reasons"] @@ -459,11 +491,14 @@ def test_a_facet_serves_the_distinct_values_of_its_column(monkeypatch, facet): response = _facet(facet) assert response.status_code == 200, response.text - assert response.json()["data"] == { - "providers": ["anthropic", "mistral", "openai"], - "modes": ["chat", "embedding"], - "features": ["vision"], - }[facet] + assert ( + response.json()["data"] + == { + "providers": ["anthropic", "mistral", "openai"], + "modes": ["chat", "embedding"], + "features": ["vision"], + }[facet] + ) def test_a_facet_offers_only_values_the_table_can_show(monkeypatch): diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 27ba8968200..3cb23d8910f 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -959,6 +959,41 @@ def _stored_health_check(model_id: str | None, status: str, checked_at: datetime ) +@pytest.mark.parametrize("retired_status,active_status", [("unhealthy", "healthy"), ("healthy", "unhealthy")]) +@pytest.mark.parametrize("retired_hour", [10, 12], ids=["older-retired-check", "newer-retired-check"]) +def test_public_model_hub_ignores_checks_for_retired_deployments( + retired_status: str, active_status: str, retired_hour: int +) -> None: + app: Final = FastAPI() + app.include_router(router) + client: Final = TestClient(app) + active: Final = _stored_health_check( + "deployment-a", active_status, datetime(2026, 10, 16, 11, 10, tzinfo=timezone.utc) + ) + retired: Final = _stored_health_check( + "deployment-z", retired_status, datetime(2026, 10, 16, retired_hour, 0, tzinfo=timezone.utc) + ) + mock_prisma: Final = MagicMock() + mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=(active, retired)) + mock_router: Final = MagicMock() + mock_router.registry.return_value.names.return_value = ("deployment-a",) + + with ( + patch("litellm.public_model_groups", ["gpt-4"]), + patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + mock_get_info.return_value = [ + ModelGroupInfoProxy(model_group="gpt-4", providers=["openai"], is_public_model_group=True) + ] + response: Final = client.get("/public/model_hub") + + assert response.status_code == 200 + assert response.json()[0]["health_status"] == active_status + assert response.json()[0]["health_checked_at"] == "2026-10-16T11:10:00+00:00" + + def test_public_model_hub_reports_the_newest_check_of_a_group_not_the_last_row_read(): """Rows come back in model_id order with NULL last, so a stale name-probe row used to overwrite the fresh one.""" app = FastAPI() @@ -969,11 +1004,13 @@ def test_public_model_hub_reports_the_newest_check_of_a_group_not_the_last_row_r stale = _stored_health_check(None, "unhealthy", datetime(2026, 9, 4, 11, 20, tzinfo=timezone.utc)) mock_prisma = MagicMock() mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=(fresh, stale)) + mock_router: Final = MagicMock() + mock_router.registry.return_value.names.return_value = ("deployment-a",) with ( patch("litellm.public_model_groups", ["gpt-4"]), patch("litellm.proxy.proxy_server.get_model_group_info") as mock_get_info, - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", mock_router), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), ): mock_get_info.return_value = [