From 93046930bf25f9a66bf0e8d5564688b25d15fa4e Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Sat, 31 Oct 2026 11:28:27 -0700 Subject: [PATCH] refactor(health): reuse existing records for aliases (#49219) --- .../health_endpoints/_health_endpoints.py | 111 +++------------ .../health_endpoints/test_health_endpoints.py | 133 +++++------------- 2 files changed, 51 insertions(+), 193 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 6a409b6b284..c9000b36f9e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -9,7 +9,6 @@ import traceback from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone -from enum import Enum from itertools import chain, groupby from types import MappingProxyType from typing import Final, Literal, TypedDict @@ -76,6 +75,7 @@ from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, + latest_by_model_name, query_latest_health_checks, ) from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers @@ -1795,92 +1795,6 @@ async def health_check_history_endpoint( ) -class _ModelHealthStatus(str, Enum): - HEALTHY = "healthy" - UNHEALTHY = "unhealthy" - UNKNOWN = "unknown" - - -class _ModelHealth(TypedDict): - status: ReadOnly[_ModelHealthStatus] - healthy_count: ReadOnly[int] - unhealthy_count: ReadOnly[int] - checked_at: ReadOnly[str | None] - - -def _model_health_status(healthy_count: int, unhealthy_count: int, deployment_count: int) -> _ModelHealthStatus: - if healthy_count > 0: - return _ModelHealthStatus.HEALTHY - if unhealthy_count > 0 and unhealthy_count >= deployment_count: - return _ModelHealthStatus.UNHEALTHY - return _ModelHealthStatus.UNKNOWN - - -def _latest_model_checks( - checks: Sequence[LatestHealthCheckRow], deployments: Sequence[Mapping[str, object]], model_name: str -) -> tuple[LatestHealthCheckRow, ...]: - ids: Final = frozenset( - model_id for deployment in deployments if (model_id := _deployment_id(deployment)) is not None - ) - names: Final = frozenset(deployment.get("model_name") for deployment in deployments) | {model_name} - by_id: Final = { - check.model_id: check for check in sorted(checks, key=lambda row: row.checked_at) if check.model_id in ids - } - group_check: Final = max( - (check for check in checks if check.model_id is None and check.model_name in names), - key=lambda row: row.checked_at, - default=None, - ) - if group_check is not None and all(group_check.checked_at >= check.checked_at for check in by_id.values()): - return (group_check,) - return tuple(by_id.values()) - - -def _summarize_model_health(checks: Sequence[LatestHealthCheckRow], deployment_count: int) -> _ModelHealth: - healthy_count: Final = sum(check.healthy_count for check in checks) - unhealthy_count: Final = sum(check.unhealthy_count for check in checks) - checked_at: Final = min((check.checked_at for check in checks), default=None) - return { - "status": _model_health_status(healthy_count, unhealthy_count, deployment_count), - "healthy_count": healthy_count, - "unhealthy_count": unhealthy_count, - "checked_at": checked_at.isoformat() if checked_at is not None else None, - } - - -async def _latest_health_by_model( - checks: Sequence[LatestHealthCheckRow], caller: UserAPIKeyAuth -) -> Mapping[str, _ModelHealth]: - from litellm.proxy.proxy_server import llm_router - - aliases: Final = await _caller_aliases(caller) - configured_names: Final = ( - frozenset(llm_router.registry().names(team_id=caller.team_id)) - if llm_router is not None - else frozenset(check.model_name for check in checks) - ) - names: Final = await _names_key_may_request( - configured_names | requestable_alias_names(aliases), aliases, caller, await _key_owner(caller), llm_router - ) - summaries: Final = ((name, _served_group(name, aliases, llm_router)) for name in sorted(names)) - - def summarize(model_name: str) -> _ModelHealth: - deployments: Final = ( - _DEPLOYMENTS.validate_python(llm_router.registry().deployments(model_name, team_id=caller.team_id)) - if llm_router is not None - else () - ) - visible: Final = tuple( - deployment for deployment in deployments if deployment_usable_by_team(deployment, caller.team_id) - ) - if llm_router is not None and not visible: - return _summarize_model_health((), 0) - latest: Final = _latest_model_checks(checks, visible, model_name) - return _summarize_model_health(latest, len(visible)) - - return MappingProxyType({name: summarize(model_name) for name, model_name in summaries}) - - @router.get("/health/latest", tags=["health"], dependencies=[Depends(user_api_key_auth)]) async def latest_health_checks_endpoint( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1890,21 +1804,34 @@ async def latest_health_checks_endpoint( Returns the most recent health check result for each model. """ + from litellm.proxy.proxy_server import llm_router + prisma_client: Final = _check_prisma_client() try: latest_checks: Final = await prisma_client.get_all_latest_health_checks() - # Convert to dict format for JSON response using helper function + checks_by_id: Final = {(check.model_id or check.model_name): check for check in latest_checks} + checks_by_name: Final = latest_by_model_name(latest_checks) + aliases: Final = await _caller_aliases(user_api_key_dict) + alias_names: Final = await _names_key_may_request( + _alias_names_caller_may_request(aliases, llm_router), + aliases, + user_api_key_dict, + await _key_owner(user_api_key_dict), + llm_router, + ) + alias_targets: Final = { + alias: checks_by_name.get(_served_group(alias, aliases, llm_router)) for alias in alias_names + } + alias_checks: Final = {alias: check for alias, check in alias_targets.items() if check is not None} checks_data: Final = { - (check.model_id if check.model_id else check.model_name): convert_health_check_to_dict(check) - for check in latest_checks + name: convert_health_check_to_dict(check) for name, check in {**alias_checks, **checks_by_id}.items() } return { "latest_health_checks": checks_data, - "total_models": len(checks_data), - "model_health": await _latest_health_by_model(latest_checks, user_api_key_dict), + "total_models": len(checks_by_id), } except Exception as e: verbose_proxy_logger.error("Error getting latest health checks: %s", e) diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index a8b2af6953c..f714bfc15ee 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -6257,127 +6257,58 @@ async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( @pytest.mark.asyncio -@pytest.mark.parametrize( - ("deployment_statuses", "expected_status", "healthy_count", "unhealthy_count"), - [ - (("healthy", "healthy"), "healthy", 2, 0), - (("healthy", "unhealthy"), "healthy", 1, 1), - (("unhealthy", "unhealthy"), "unhealthy", 0, 2), - (("unhealthy", None), "unknown", 0, 1), - ((None, None), "unknown", 0, 0), - ], -) @pytest.mark.parametrize("alias_layer", ["router", "key", "team", "hierarchical-router"]) -@pytest.mark.parametrize("check_granularity", ["deployment", "group"]) -async def test_latest_health_resolves_aliases_and_aggregates_their_deployments( - monkeypatch, deployment_statuses, expected_status, healthy_count, unhealthy_count, alias_layer, check_granularity -): +async def test_latest_health_alias_mirrors_the_targets_latest_record(monkeypatch, alias_layer): from litellm.proxy import proxy_server - checked_at: Final = "2026-10-31T12:00:00+00:00" - model_list: Final = [ - { - "model_name": "chat", - "litellm_params": {"model": "openai/gpt-5.6"}, - "model_info": {"id": f"chat-{index}"}, - } - for index in range(2) - ] llm_router: Final = Router( - model_list=model_list, - model_group_alias={"default": {"model": "chat", "hidden": False}} if alias_layer == "router" else {}, + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-5.6"}, "model_info": {"id": name}} + for name in ("chat", "unchecked", "private") + ], + model_group_alias={ + **({"default": "chat"} if alias_layer == "router" else {}), + "uncached": "unchecked", + "restricted": "private", + "old-deployment": "chat", + }, ) - deployment_rows: Final = [ + rows: Final = [ { - "health_check_id": f"check-{index}", - "model_name": "chat", - "model_id": f"chat-{index}", - "status": health_status, - "healthy_count": int(health_status == "healthy"), - "unhealthy_count": int(health_status == "unhealthy"), + "health_check_id": model_id, + "model_name": model_name, + "model_id": model_id, + "status": status, + "healthy_count": int(status == "healthy"), + "unhealthy_count": int(status == "unhealthy"), "checked_at": checked_at, "created_at": checked_at, "updated_at": checked_at, } - for index, health_status in enumerate(deployment_statuses) - if health_status is not None + for model_id, model_name, status, checked_at in ( + ("new-deployment", "chat", "healthy", "2026-10-31T12:00:00+00:00"), + ("old-deployment", "chat", "unhealthy", "2026-10-30T12:00:00+00:00"), + ("private", "private", "healthy", "2026-10-31T12:00:00+00:00"), + ) ] - rows: Final = ( - [{**deployment_rows[0], "model_id": None, "healthy_count": healthy_count, "unhealthy_count": unhealthy_count}] - if check_granularity == "group" and deployment_rows - else deployment_rows - ) prisma: Final = PrismaClient( database_url="postgresql://localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj ) monkeypatch.setattr(prisma, "db", SimpleNamespace(query_raw=engine_call(rows))) caller: Final = UserAPIKeyAuth( - models=["chat"], + models=["chat", "unchecked"], aliases={"default": "chat"} if alias_layer == "key" else {}, team_model_aliases={"default": "chat"} if alias_layer == "team" else {}, router_settings={"model_group_alias": {"default": "chat"}} if alias_layer == "hierarchical-router" else {}, ) - with _proxy_health_globals(model_list, llm_router, prisma_client=prisma): + with _proxy_health_globals([], llm_router, prisma_client=prisma): result: Final = await _health_endpoints_module.latest_health_checks_endpoint(user_api_key_dict=caller) - expected_health: Final = { - "status": expected_status, - "healthy_count": healthy_count, - "unhealthy_count": unhealthy_count, - "checked_at": checked_at if rows else None, - } - assert result["model_health"] == {"chat": expected_health, "default": expected_health} - assert result["total_models"] == len(rows) - assert result["latest_health_checks"] == { - row["model_id"] or row["model_name"]: { - **{key: value for key, value in row.items() if key != "updated_at"}, - "error_message": None, - "response_time_ms": None, - "details": None, - "checked_by": None, - } - for row in rows - } - - -@pytest.mark.asyncio -async def test_latest_health_alias_uses_current_target_and_omits_inaccessible_aliases(monkeypatch): - from litellm.proxy import proxy_server - - checked_at: Final = "2026-10-31T12:00:00+00:00" - llm_router: Final = Router( - model_list=[ - {"model_name": name, "litellm_params": {"model": "openai/gpt-5.6"}, "model_info": {"id": name}} - for name in ("old", "current", "private") - ], - model_group_alias={"default": "current", "restricted": "private"}, - ) - rows: Final = [ - { - "health_check_id": name, - "model_name": "current" if name == "deleted" else name, - "model_id": name, - "status": "healthy" if name in ("old", "deleted") else "unhealthy", - "healthy_count": int(name in ("old", "deleted")), - "unhealthy_count": int(name not in ("old", "deleted")), - "checked_at": checked_at, - "created_at": checked_at, - "updated_at": checked_at, - } - for name in ("old", "current", "private", "deleted") - ] - prisma: Final = PrismaClient( - database_url="postgresql://localhost/test", proxy_logging_obj=proxy_server.proxy_logging_obj - ) - monkeypatch.setattr(prisma, "db", SimpleNamespace(query_raw=engine_call(rows))) - - with _proxy_health_globals([], llm_router, prisma_client=prisma): - result: Final = await _health_endpoints_module.latest_health_checks_endpoint( - user_api_key_dict=UserAPIKeyAuth(models=["current"]) - ) - - assert result["model_health"] == { - name: {"status": "unhealthy", "healthy_count": 0, "unhealthy_count": 1, "checked_at": checked_at} - for name in ("current", "default") - } + checks: Final = result["latest_health_checks"] + assert set(checks) == {"new-deployment", "old-deployment", "private", "default"} + assert checks["default"] == checks["new-deployment"] + assert checks["default"]["status"] == "healthy" + assert checks["default"]["checked_at"] == "2026-10-31T12:00:00+00:00" + assert checks["old-deployment"]["status"] == "unhealthy" + assert result["total_models"] == 3