scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 128

refactor(health): reuse existing records for aliases (#49219)

Browse source
This commit is contained in:
ishaan-j-berri 2026-10-31 11:28:27 -07:00
parent d8b3b62071
commit 93046930bf
2 changed files with 51 additions and 193 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -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)

133
tests/unit/proxy/health_endpoints/test_health_endpoints.py
View file

@ -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