mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #39539 from BerriAI/litellm_fix_health_check_db_storm
fix(proxy): dedup latest health checks in SQL and gate the DB save per window
This commit is contained in:
ryan-crabbe-berri
GitHub
commit
33fa195949
9 changed files with 767 additions and 284 deletions
|
|
@ -1654,6 +1654,7 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data"
|
|||
MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data"
|
||||
CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000))
|
||||
SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
|
||||
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME: Final = "background_health_check_db_save"
|
||||
KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
|
||||
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
|
||||
|
|
|
|||
96
litellm/proxy/db/health_check_latest.py
Normal file
96
litellm/proxy/db/health_check_latest.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""
|
||||
Latest health-check row per model, deduplicated by Postgres.
|
||||
|
||||
prisma-client-py's ``find_many(distinct=...)`` dedups client-side: the emitted
|
||||
SQL carries no DISTINCT, so the whole append-only history table streams to the
|
||||
worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per
|
||||
(model_id, model_name) and is served by the matching descending index.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, field_validator
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
LATEST_HEALTH_CHECKS_SQL: Final = """
|
||||
SELECT DISTINCT ON ("model_id", "model_name")
|
||||
"health_check_id", "model_name", "model_id", "status",
|
||||
"healthy_count", "unhealthy_count", "error_message",
|
||||
"response_time_ms", "details", "checked_by",
|
||||
"checked_at", "created_at", "updated_at"
|
||||
FROM "LiteLLM_HealthCheckTable"
|
||||
ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC
|
||||
"""
|
||||
|
||||
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL: Final = """
|
||||
SELECT DISTINCT ON ("model_id", "model_name")
|
||||
"health_check_id", "model_name", "model_id", "status",
|
||||
"healthy_count", "unhealthy_count", "error_message",
|
||||
"response_time_ms", "details", "checked_by",
|
||||
"checked_at", "created_at", "updated_at"
|
||||
FROM "LiteLLM_HealthCheckTable"
|
||||
WHERE "model_name" = ANY($1)
|
||||
ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC
|
||||
"""
|
||||
|
||||
|
||||
class LatestHealthCheckRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, protected_namespaces=())
|
||||
|
||||
health_check_id: str
|
||||
model_name: str
|
||||
model_id: str | None = None
|
||||
status: str
|
||||
healthy_count: int = 0
|
||||
unhealthy_count: int = 0
|
||||
error_message: str | None = None
|
||||
response_time_ms: float | None = None
|
||||
details: JsonValue | None = None
|
||||
checked_by: str | None = None
|
||||
checked_at: datetime
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
|
||||
@field_validator("details", mode="before")
|
||||
@classmethod
|
||||
def _decode_json_text(cls, value: object) -> object:
|
||||
return json.loads(value) if isinstance(value, str) else value
|
||||
|
||||
@field_validator("checked_at", "created_at", "updated_at")
|
||||
@classmethod
|
||||
def _assume_utc(cls, value: datetime) -> datetime:
|
||||
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value
|
||||
|
||||
|
||||
_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...])
|
||||
|
||||
|
||||
async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
|
||||
try:
|
||||
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
|
||||
return _ROWS_ADAPTER.validate_python(rows)
|
||||
except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them
|
||||
verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err)
|
||||
return ()
|
||||
|
||||
|
||||
async def fetch_latest_health_checks_for_models(
|
||||
prisma_client: PrismaClient, model_names: Sequence[str]
|
||||
) -> tuple[LatestHealthCheckRow, ...]:
|
||||
if not model_names:
|
||||
return ()
|
||||
try:
|
||||
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names))
|
||||
return _ROWS_ADAPTER.validate_python(rows)
|
||||
except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration
|
||||
verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err)
|
||||
return ()
|
||||
|
|
@ -7,7 +7,7 @@ import secrets
|
|||
import time
|
||||
import traceback
|
||||
from collections.abc import Iterable, Mapping
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Final, Literal, TypedDict, cast
|
||||
|
||||
import fastapi
|
||||
|
|
@ -41,6 +41,7 @@ from litellm.proxy.auth.auth_utils import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
|
||||
from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
|
||||
from litellm.proxy.health_check import (
|
||||
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
|
||||
|
|
@ -747,13 +748,42 @@ def _aggregate_health_check_results(
|
|||
return model_results
|
||||
|
||||
|
||||
class _AggregatedHealthResult(TypedDict):
|
||||
"""One entry of ``_aggregate_health_check_results``: a model's counts for this cycle."""
|
||||
|
||||
model_name: ReadOnly[str]
|
||||
model_id: ReadOnly[str | None]
|
||||
healthy_count: ReadOnly[int]
|
||||
unhealthy_count: ReadOnly[int]
|
||||
error_message: ReadOnly[str | None]
|
||||
|
||||
|
||||
def _new_health_status(result: _AggregatedHealthResult) -> str:
|
||||
return "healthy" if result["healthy_count"] > 0 else "unhealthy"
|
||||
|
||||
|
||||
def _should_persist_health_check_result(
|
||||
result: _AggregatedHealthResult, latest_checks_map: Mapping[str, LatestHealthCheckRow]
|
||||
) -> bool:
|
||||
"""
|
||||
True when this result has to be written: no previous row, the status changed, or the
|
||||
previous row is older than one hour (periodic refresh while the status is stable).
|
||||
"""
|
||||
lookup_key: Final = result["model_id"] if result["model_id"] else result["model_name"]
|
||||
last_check: Final = latest_checks_map.get(lookup_key)
|
||||
if last_check is None or last_check.status != _new_health_status(result):
|
||||
return True
|
||||
time_since_last_check: Final = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds()
|
||||
return time_since_last_check >= 3600 # 1 hour threshold
|
||||
|
||||
|
||||
async def _save_health_check_results_if_changed(
|
||||
prisma_client,
|
||||
model_results: dict,
|
||||
latest_checks_map: dict,
|
||||
start_time: float,
|
||||
checked_by: str | None = None,
|
||||
):
|
||||
) -> bool:
|
||||
"""
|
||||
Save health check results to database, but only if status changed or >1 hour since last save.
|
||||
|
||||
|
|
@ -764,47 +794,39 @@ async def _save_health_check_results_if_changed(
|
|||
- Status changes: Immediate write (no delay)
|
||||
- Result: ~92% reduction in DB writes for stable systems, while maintaining real-time updates on changes
|
||||
|
||||
The writes are awaited rather than detached so the caller learns whether this cycle's
|
||||
persistence completed.
|
||||
|
||||
Args:
|
||||
prisma_client: Database client
|
||||
model_results: Dictionary of aggregated health check results per model
|
||||
latest_checks_map: Dictionary mapping model_id/model_name to latest health check
|
||||
start_time: Start time of health check for calculating response time
|
||||
checked_by: Identifier for who/what performed the check
|
||||
|
||||
Returns:
|
||||
True when every row that needed writing was written (including when nothing needed
|
||||
writing); False when any write failed.
|
||||
"""
|
||||
for result in model_results.values():
|
||||
new_status = "healthy" if result["healthy_count"] > 0 else "unhealthy"
|
||||
|
||||
# Check if we should save this result
|
||||
should_save = True
|
||||
lookup_key = result["model_id"] if result["model_id"] else result["model_name"]
|
||||
if lookup_key in latest_checks_map:
|
||||
last_check = latest_checks_map[lookup_key]
|
||||
# Only save if status changed or if it's been a while since last check
|
||||
if last_check.status == new_status:
|
||||
# Check if last check was recent (within 1 hour)
|
||||
if last_check.checked_at:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
time_since_last_check = (datetime.now(timezone.utc) - last_check.checked_at).total_seconds()
|
||||
# Only skip if status unchanged AND checked recently (within 1 hour)
|
||||
# This ensures we still get periodic updates even if status is stable
|
||||
if time_since_last_check < 3600: # 1 hour threshold
|
||||
should_save = False
|
||||
|
||||
if should_save:
|
||||
asyncio.create_task(
|
||||
prisma_client.save_health_check_result(
|
||||
model_name=result["model_name"],
|
||||
model_id=result["model_id"],
|
||||
status=new_status,
|
||||
healthy_count=result["healthy_count"],
|
||||
unhealthy_count=result["unhealthy_count"],
|
||||
error_message=result["error_message"],
|
||||
response_time_ms=(time.time() - start_time) * 1000,
|
||||
details=None,
|
||||
checked_by=checked_by,
|
||||
)
|
||||
)
|
||||
to_write: Final = tuple(
|
||||
result for result in model_results.values() if _should_persist_health_check_result(result, latest_checks_map)
|
||||
)
|
||||
writes: Final = tuple(
|
||||
prisma_client.save_health_check_result(
|
||||
model_name=result["model_name"],
|
||||
model_id=result["model_id"],
|
||||
status=_new_health_status(result),
|
||||
healthy_count=result["healthy_count"],
|
||||
unhealthy_count=result["unhealthy_count"],
|
||||
error_message=result["error_message"],
|
||||
response_time_ms=(time.time() - start_time) * 1000,
|
||||
details=None,
|
||||
checked_by=checked_by,
|
||||
)
|
||||
for result in to_write
|
||||
)
|
||||
rows: Final = await asyncio.gather(*writes)
|
||||
return all(row is not None for row in rows)
|
||||
|
||||
|
||||
async def _save_background_health_checks_to_db(
|
||||
|
|
@ -814,7 +836,7 @@ async def _save_background_health_checks_to_db(
|
|||
unhealthy_endpoints: list,
|
||||
start_time: float,
|
||||
checked_by: str | None = None,
|
||||
):
|
||||
) -> bool:
|
||||
"""
|
||||
Save background health check results to database for each model.
|
||||
|
||||
|
|
@ -823,9 +845,13 @@ async def _save_background_health_checks_to_db(
|
|||
|
||||
OPTIMIZATION: Only saves to database if the status has changed from the last saved check.
|
||||
This dramatically reduces database writes when health status remains stable.
|
||||
|
||||
Returns:
|
||||
True when this cycle's persistence completed; False when it was skipped or any step
|
||||
failed. Never raises: a database failure must not break the health check loop.
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return
|
||||
return False
|
||||
|
||||
try:
|
||||
# Step 1: Build mapping from model parameter to model info
|
||||
|
|
@ -848,7 +874,7 @@ async def _save_background_health_checks_to_db(
|
|||
latest_checks_map[key] = check
|
||||
|
||||
# Step 4: Save aggregated results, but only if status changed
|
||||
await _save_health_check_results_if_changed(
|
||||
return await _save_health_check_results_if_changed(
|
||||
prisma_client,
|
||||
model_results,
|
||||
latest_checks_map,
|
||||
|
|
@ -858,6 +884,7 @@ async def _save_background_health_checks_to_db(
|
|||
except Exception as db_error:
|
||||
verbose_proxy_logger.warning("Failed to save background health checks to database: %s", db_error)
|
||||
# Continue execution - don't let database save failure break health checks
|
||||
return False
|
||||
|
||||
|
||||
_PROXY_ADMIN_ROLES: Final = frozenset(
|
||||
|
|
|
|||
|
|
@ -16,7 +16,16 @@ import threading
|
|||
import time
|
||||
import traceback
|
||||
import warnings
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence
|
||||
from collections.abc import (
|
||||
AsyncGenerator,
|
||||
AsyncIterator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Collection,
|
||||
Mapping,
|
||||
MutableMapping,
|
||||
Sequence,
|
||||
)
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType, UnionType
|
||||
|
|
@ -51,6 +60,7 @@ from litellm.constants import (
|
|||
AIOHTTP_NEEDS_CLEANUP_CLOSED,
|
||||
AIOHTTP_TTL_DNS_CACHE,
|
||||
AUDIO_SPEECH_CHUNK_SIZE,
|
||||
BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
|
||||
BASE_MCP_ROUTE,
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
|
||||
DEFAULT_MAX_RECURSE_DEPTH,
|
||||
|
|
@ -151,6 +161,7 @@ if TYPE_CHECKING:
|
|||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
from litellm.proxy.health_check_utils.shared_health_check_manager import SharedHealthCheckManager
|
||||
|
||||
Span = _Span | Any
|
||||
else:
|
||||
|
|
@ -233,7 +244,7 @@ def generate_feedback_box():
|
|||
import contextlib
|
||||
from collections import defaultdict
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from functools import lru_cache, partial
|
||||
|
||||
import litellm
|
||||
import litellm._redis
|
||||
|
|
@ -415,6 +426,7 @@ from litellm.proxy.config_resolvers.alerting import (
|
|||
)
|
||||
from litellm.proxy.container_endpoints.endpoints import router as container_router
|
||||
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
|
||||
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
|
||||
SpendLogCleanup,
|
||||
|
|
@ -3747,13 +3759,50 @@ async def _run_direct_health_check_with_instrumentation(
|
|||
raise AssertionError("perform_health_check rejected every optional argument")
|
||||
|
||||
|
||||
async def _window_gated_health_check_db_save(
|
||||
save: Callable[[], Awaitable[bool]],
|
||||
pod_lock_manager: PodLockManager | None,
|
||||
lock_ttl: int | None,
|
||||
) -> None:
|
||||
"""
|
||||
Persist at most once per window fleet-wide. A completed save keeps the lock as the
|
||||
"this window's save is done" marker, so it is deliberately never released and expires
|
||||
with the interval. A save that reports failure or is cancelled releases the lock so
|
||||
another pod's cycle in the same window can retry, instead of the fleet going a whole
|
||||
window without a write.
|
||||
"""
|
||||
if pod_lock_manager is None or pod_lock_manager.redis_cache is None:
|
||||
await save()
|
||||
return
|
||||
acquired: Final = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
|
||||
ttl=lock_ttl,
|
||||
allow_reentrant=False,
|
||||
)
|
||||
if not acquired:
|
||||
verbose_proxy_logger.debug("background_health_check_db_save_skipped another pod persisted this window")
|
||||
return
|
||||
try:
|
||||
persisted: Final = await save()
|
||||
except BaseException:
|
||||
await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME)
|
||||
raise
|
||||
if not persisted:
|
||||
verbose_proxy_logger.warning(
|
||||
"background_health_check_db_save_incomplete released the window lock so another pod can retry"
|
||||
)
|
||||
await pod_lock_manager.release_lock(cronjob_id=BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME)
|
||||
|
||||
|
||||
def _schedule_background_health_check_db_save(
|
||||
prisma_client,
|
||||
shared_health_manager,
|
||||
prisma_client: PrismaClient | None,
|
||||
shared_health_manager: "SharedHealthCheckManager | None",
|
||||
model_list: list,
|
||||
healthy_endpoints: list,
|
||||
unhealthy_endpoints: list,
|
||||
):
|
||||
pod_lock_manager: PodLockManager | None = None,
|
||||
lock_ttl: int | None = None,
|
||||
) -> None:
|
||||
"""Fire-and-forget: persist health check results to DB if prisma is available."""
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
|
@ -3765,16 +3814,16 @@ def _schedule_background_health_check_db_save(
|
|||
|
||||
checked_by: Final = shared_health_manager.pod_id if shared_health_manager is not None else "background_health_check"
|
||||
start_time: Final = time_module.time()
|
||||
asyncio.create_task(
|
||||
_save_background_health_checks_to_db(
|
||||
prisma_client,
|
||||
model_list,
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
start_time,
|
||||
checked_by=checked_by,
|
||||
)
|
||||
save: Final = partial(
|
||||
_save_background_health_checks_to_db,
|
||||
prisma_client,
|
||||
model_list,
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
start_time,
|
||||
checked_by=checked_by,
|
||||
)
|
||||
asyncio.create_task(_window_gated_health_check_db_save(save, pod_lock_manager, lock_ttl))
|
||||
|
||||
|
||||
def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int:
|
||||
|
|
@ -4081,6 +4130,8 @@ async def _run_background_health_check():
|
|||
_llm_model_list,
|
||||
healthy_endpoints,
|
||||
unhealthy_endpoints,
|
||||
pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager,
|
||||
lock_ttl=health_check_interval,
|
||||
)
|
||||
|
||||
# Write health state to router cache for health-check-driven routing
|
||||
|
|
|
|||
|
|
@ -125,6 +125,11 @@ from litellm.proxy.db.exception_handler import (
|
|||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
)
|
||||
from litellm.proxy.db.health_check_latest import (
|
||||
LatestHealthCheckRow,
|
||||
fetch_latest_health_checks,
|
||||
fetch_latest_health_checks_for_models,
|
||||
)
|
||||
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
||||
from litellm.proxy.db.prisma_client import (
|
||||
PrismaWrapper,
|
||||
|
|
@ -6462,48 +6467,13 @@ class PrismaClient:
|
|||
verbose_proxy_logger.error("Error getting health check history: %s", e)
|
||||
return []
|
||||
|
||||
async def get_all_latest_health_checks(self) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
|
||||
"""
|
||||
Get the latest health check for each model.
|
||||
async def get_all_latest_health_checks(self) -> tuple[LatestHealthCheckRow, ...]:
|
||||
"""Latest health check per (model_id, model_name), deduplicated in Postgres."""
|
||||
return await fetch_latest_health_checks(self)
|
||||
|
||||
Uses DB-level DISTINCT ON (model_id, model_name) with ORDER BY checked_at DESC
|
||||
(via Prisma ``distinct`` + ``order``) so we never load the full history into memory.
|
||||
"""
|
||||
try:
|
||||
return await HealthCheckRepository(self).table.find_many(
|
||||
distinct=["model_id", "model_name"],
|
||||
order=[
|
||||
{"model_id": "asc"},
|
||||
{"model_name": "asc"},
|
||||
{"checked_at": "desc"},
|
||||
],
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error getting all latest health checks: %s", e)
|
||||
return []
|
||||
|
||||
async def get_latest_health_checks_for_models(
|
||||
self, model_names: "Sequence[str]"
|
||||
) -> "Sequence[prisma_models.LiteLLM_HealthCheckTable]":
|
||||
"""
|
||||
Get the latest health check for each of the named models.
|
||||
|
||||
Same DISTINCT ON as ``get_all_latest_health_checks``, bounded to the models asked
|
||||
about, so a paged caller reads health for its page instead of for the whole table.
|
||||
"""
|
||||
if not model_names:
|
||||
return ()
|
||||
latest_first: Final = (("model_id", "asc"), ("model_name", "asc"), ("checked_at", "desc"))
|
||||
order: Final = [{field: direction} for field, direction in latest_first] # mutable-ok: prisma order is a list
|
||||
try:
|
||||
return await HealthCheckRepository(self).table.find_many(
|
||||
where={"model_name": {"in": list(model_names)}}, # mutable-ok: prisma filters are dicts and lists
|
||||
distinct=["model_id", "model_name"], # mutable-ok: prisma distinct takes a list
|
||||
order=order,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # health decorates a list; a driver error must not fail the page
|
||||
verbose_proxy_logger.error("Error getting latest health checks for models: %s", e)
|
||||
return ()
|
||||
async def get_latest_health_checks_for_models(self, model_names: Sequence[str]) -> tuple[LatestHealthCheckRow, ...]:
|
||||
"""Same as ``get_all_latest_health_checks``, bounded to the named models."""
|
||||
return await fetch_latest_health_checks_for_models(self, model_names)
|
||||
|
||||
|
||||
### HELPER FUNCTIONS ###
|
||||
|
|
|
|||
111
tests/test_litellm/proxy/db/test_health_check_latest.py
Normal file
111
tests/test_litellm/proxy/db/test_health_check_latest.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.health_check_latest import (
|
||||
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL,
|
||||
LATEST_HEALTH_CHECKS_SQL,
|
||||
fetch_latest_health_checks,
|
||||
fetch_latest_health_checks_for_models,
|
||||
)
|
||||
|
||||
|
||||
def _prisma(rows):
|
||||
prisma = MagicMock()
|
||||
prisma.db.query_raw = AsyncMock(return_value=rows)
|
||||
return prisma
|
||||
|
||||
|
||||
def _raw_row(**overrides):
|
||||
row = {
|
||||
"health_check_id": "hc-1",
|
||||
"model_name": "gpt-4",
|
||||
"model_id": "deployment-abc",
|
||||
"status": "healthy",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
"response_time_ms": 12.5,
|
||||
"details": None,
|
||||
"checked_by": "pod-1",
|
||||
"checked_at": "2026-08-25T00:00:00+00:00",
|
||||
"created_at": "2026-08-25T00:00:00+00:00",
|
||||
"updated_at": "2026-08-25T00:00:00+00:00",
|
||||
}
|
||||
return {**row, **overrides}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_all_runs_one_distinct_on_query_with_no_parameters():
|
||||
"""The dedup must be in the SQL: prisma find_many(distinct=...) streams the whole history table."""
|
||||
prisma = _prisma([])
|
||||
assert await fetch_latest_health_checks(prisma) == ()
|
||||
assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_SQL,)
|
||||
assert 'DISTINCT ON ("model_id", "model_name")' in LATEST_HEALTH_CHECKS_SQL
|
||||
assert '"checked_at" DESC' in LATEST_HEALTH_CHECKS_SQL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_datetimes_come_back_tz_aware_with_or_without_an_offset():
|
||||
"""The save path subtracts checked_at from datetime.now(timezone.utc); a naive value would TypeError."""
|
||||
naive = _raw_row(health_check_id="hc-naive", model_id=None, checked_at="2026-08-25T00:00:00")
|
||||
aware = _raw_row(health_check_id="hc-aware", checked_at="2026-08-25T01:00:00+02:00")
|
||||
rows = await fetch_latest_health_checks(_prisma([naive, aware]))
|
||||
assert {row.health_check_id: (row.model_id, row.checked_at) for row in rows} == {
|
||||
"hc-naive": (None, datetime(2026, 8, 25, 0, 0, tzinfo=timezone.utc)),
|
||||
"hc-aware": ("deployment-abc", datetime(2026, 8, 24, 23, 0, tzinfo=timezone.utc)),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_details_decode_from_text_and_pass_through_as_dict():
|
||||
rows = await fetch_latest_health_checks(
|
||||
_prisma(
|
||||
[
|
||||
_raw_row(health_check_id="text", details='{"region": "eu"}'),
|
||||
_raw_row(health_check_id="dict", details={"region": "us"}),
|
||||
_raw_row(health_check_id="none", details=None),
|
||||
]
|
||||
)
|
||||
)
|
||||
assert {row.health_check_id: row.details for row in rows} == {
|
||||
"text": {"region": "eu"},
|
||||
"dict": {"region": "us"},
|
||||
"none": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_all_degrades_to_no_rows_when_the_query_fails():
|
||||
prisma = _prisma([])
|
||||
prisma.db.query_raw.side_effect = RuntimeError("db down")
|
||||
assert await fetch_latest_health_checks(prisma) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_all_degrades_to_no_rows_for_a_malformed_row():
|
||||
assert await fetch_latest_health_checks(_prisma([{"unexpected": "shape"}])) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_for_models_binds_the_page_as_the_only_parameter():
|
||||
prisma = _prisma([_raw_row()])
|
||||
rows = await fetch_latest_health_checks_for_models(prisma, ("gpt-4", "claude-opus"))
|
||||
assert prisma.db.query_raw.await_args.args == (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-4", "claude-opus"])
|
||||
assert [row.model_name for row in rows] == ["gpt-4"]
|
||||
assert 'WHERE "model_name" = ANY($1)' in LATEST_HEALTH_CHECKS_FOR_MODELS_SQL
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_for_models_skips_the_database_for_an_empty_page():
|
||||
prisma = _prisma([])
|
||||
assert await fetch_latest_health_checks_for_models(prisma, ()) == ()
|
||||
prisma.db.query_raw.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_for_models_degrades_to_no_rows_when_the_query_fails():
|
||||
prisma = _prisma([])
|
||||
prisma.db.query_raw.side_effect = RuntimeError("db down")
|
||||
assert await fetch_latest_health_checks_for_models(prisma, ("gpt-4",)) == ()
|
||||
|
|
@ -20,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.constants import BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME
|
||||
from litellm.proxy.proxy_server import (
|
||||
_adaptive_router_flusher_loop,
|
||||
_get_endpoint_exception_status,
|
||||
|
|
@ -111,13 +112,11 @@ async def test_run_direct_health_check_with_instrumentation_returns_results(
|
|||
lambda _gs: {},
|
||||
)
|
||||
|
||||
healthy, unhealthy, exceptions = (
|
||||
await _run_direct_health_check_with_instrumentation(
|
||||
model_list=[{"model_name": "gpt-4"}],
|
||||
details=False,
|
||||
max_concurrency=1,
|
||||
instrumentation_context={"source": "test"},
|
||||
)
|
||||
healthy, unhealthy, exceptions = await _run_direct_health_check_with_instrumentation(
|
||||
model_list=[{"model_name": "gpt-4"}],
|
||||
details=False,
|
||||
max_concurrency=1,
|
||||
instrumentation_context={"source": "test"},
|
||||
)
|
||||
|
||||
assert normalize(
|
||||
|
|
@ -245,6 +244,137 @@ async def test_schedule_background_health_check_db_save_invalid_no_event_loop_ra
|
|||
)
|
||||
|
||||
|
||||
def _lock_manager(redis_cache, acquired):
|
||||
manager = MagicMock()
|
||||
manager.redis_cache = redis_cache
|
||||
manager.acquire_lock = AsyncMock(return_value=acquired)
|
||||
manager.release_lock = AsyncMock()
|
||||
return manager
|
||||
|
||||
|
||||
def _capture_saves(monkeypatch, persisted=True):
|
||||
saves = []
|
||||
|
||||
async def _fake_save(*_args, **kwargs):
|
||||
saves.append(kwargs)
|
||||
return persisted
|
||||
|
||||
import litellm.proxy.health_endpoints._health_endpoints as he
|
||||
|
||||
monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
|
||||
return saves
|
||||
|
||||
|
||||
def _cancel_during_save(monkeypatch):
|
||||
async def _fake_save(*_args, **_kwargs):
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
import litellm.proxy.health_endpoints._health_endpoints as he
|
||||
|
||||
monkeypatch.setattr(he, "_save_background_health_checks_to_db", _fake_save)
|
||||
|
||||
|
||||
def _schedule_with(lock_manager):
|
||||
_schedule_background_health_check_db_save(
|
||||
prisma_client=MagicMock(),
|
||||
shared_health_manager=None,
|
||||
model_list=[],
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=[],
|
||||
pod_lock_manager=lock_manager,
|
||||
lock_ttl=300,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_background_health_check_db_save_skips_a_window_another_pod_persisted(monkeypatch):
|
||||
saves = _capture_saves(monkeypatch)
|
||||
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=False)
|
||||
|
||||
_schedule_with(lock_manager)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert saves == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_background_health_check_db_save_holds_the_window_lock_for_the_whole_interval(monkeypatch):
|
||||
"""The lock is the "saved this window" marker: never reentrant, TTL = interval, and never released."""
|
||||
saves = _capture_saves(monkeypatch)
|
||||
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
||||
|
||||
_schedule_with(lock_manager)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert normalize(
|
||||
{
|
||||
"saves": len(saves),
|
||||
"lock_request": lock_manager.acquire_lock.await_args.kwargs,
|
||||
"released": lock_manager.release_lock.await_count,
|
||||
}
|
||||
) == {
|
||||
"saves": 1,
|
||||
"lock_request": {
|
||||
"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME,
|
||||
"ttl": 300,
|
||||
"allow_reentrant": False,
|
||||
},
|
||||
"released": 0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_reports_failure(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A failed save must not burn the window: release the lock so another pod's cycle can retry."""
|
||||
saves = _capture_saves(monkeypatch, persisted=False)
|
||||
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
||||
|
||||
_schedule_with(lock_manager)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert normalize(
|
||||
{
|
||||
"saves": len(saves),
|
||||
"release_request": lock_manager.release_lock.await_args.kwargs,
|
||||
"release_count": lock_manager.release_lock.await_count,
|
||||
}
|
||||
) == {
|
||||
"saves": 1,
|
||||
"release_request": {"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME},
|
||||
"release_count": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_background_health_check_db_save_releases_the_window_lock_when_the_save_is_cancelled(
|
||||
monkeypatch,
|
||||
):
|
||||
"""A pod shutting down mid-save releases the lock instead of holding it until the TTL."""
|
||||
_cancel_during_save(monkeypatch)
|
||||
lock_manager = _lock_manager(redis_cache=MagicMock(), acquired=True)
|
||||
|
||||
_schedule_with(lock_manager)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert (
|
||||
lock_manager.release_lock.await_args.kwargs,
|
||||
lock_manager.release_lock.await_count,
|
||||
) == ({"cronjob_id": BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME}, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schedule_background_health_check_db_save_runs_ungated_without_redis(monkeypatch):
|
||||
saves = _capture_saves(monkeypatch, persisted=False)
|
||||
lock_manager = _lock_manager(redis_cache=None, acquired=True)
|
||||
|
||||
_schedule_with(lock_manager)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert (len(saves), lock_manager.acquire_lock.await_count, lock_manager.release_lock.await_count) == (1, 0, 0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _get_endpoint_exception_status
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -319,13 +449,9 @@ def test_write_health_state_to_router_cache_sets_states(monkeypatch):
|
|||
|
||||
_write_health_state_to_router_cache(healthy, unhealthy, exceptions)
|
||||
|
||||
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(
|
||||
fake_states
|
||||
)
|
||||
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states)
|
||||
|
||||
call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[
|
||||
0
|
||||
][0]
|
||||
call_args = fake_router.health_state_cache.set_deployment_health_states.call_args[0][0]
|
||||
assert normalize(
|
||||
{
|
||||
"states_keys": sorted(call_args.keys()),
|
||||
|
|
@ -367,9 +493,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp
|
|||
fake_router.cooldown_time = 30
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "general_settings", {"model_list_healthy_only": True}
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"model_list_healthy_only": True})
|
||||
|
||||
fake_states = {"m1": {"is_healthy": True}, "m2": {"is_healthy": False}}
|
||||
|
||||
|
|
@ -403,9 +527,7 @@ def test_write_health_state_to_router_cache_populates_for_listing_filter(monkeyp
|
|||
{"m2": SimpleNamespace(status_code=500)},
|
||||
)
|
||||
|
||||
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(
|
||||
fake_states
|
||||
)
|
||||
fake_router.health_state_cache.set_deployment_health_states.assert_called_once_with(fake_states)
|
||||
assert cooldowns == []
|
||||
assert failures == []
|
||||
|
||||
|
|
@ -415,9 +537,7 @@ def test_write_health_state_to_router_cache_swallows_internal_failures(monkeypat
|
|||
fake_router = MagicMock()
|
||||
fake_router.enable_health_check_routing = True
|
||||
fake_router.health_check_ignore_transient_errors = False
|
||||
fake_router.health_state_cache.set_deployment_health_states.side_effect = (
|
||||
RuntimeError("cache exploded")
|
||||
)
|
||||
fake_router.health_state_cache.set_deployment_health_states.side_effect = RuntimeError("cache exploded")
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
|
||||
|
|
@ -447,9 +567,7 @@ async def test_adaptive_router_flusher_loop_flushes_each_router(monkeypatch):
|
|||
from litellm.types.router import TaggedPreRoutingStrategy
|
||||
|
||||
fake_router = MagicMock()
|
||||
fake_router.adaptive_routers = {
|
||||
"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]
|
||||
}
|
||||
fake_router.adaptive_routers = {"alpha": [TaggedPreRoutingStrategy(tags=(), strategy=fake_ar)]}
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
|
||||
|
|
@ -547,12 +665,8 @@ async def test_run_background_health_check_runs_one_cycle_then_cancels(monkeypat
|
|||
"_run_direct_health_check_with_instrumentation",
|
||||
_fake_direct,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"health_check_filter_kwargs_from_general_settings",
|
||||
|
|
@ -630,12 +744,8 @@ async def test_run_background_health_check_probes_only_listed_model_groups(monke
|
|||
"_run_direct_health_check_with_instrumentation",
|
||||
_fake_direct,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "_schedule_background_health_check_db_save", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(proxy_server, "_write_health_state_to_router_cache", lambda *a, **kw: None)
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"health_check_filter_kwargs_from_general_settings",
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
|
||||
from litellm.proxy.health_endpoints._health_endpoints import (
|
||||
_aggregate_health_check_results,
|
||||
_build_model_param_to_info_mapping,
|
||||
|
|
@ -13,6 +15,7 @@ from litellm.proxy.health_endpoints._health_endpoints import (
|
|||
_save_background_health_checks_to_db,
|
||||
_save_health_check_results_if_changed,
|
||||
_save_health_check_to_db,
|
||||
latest_health_checks_endpoint,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
|
@ -21,12 +24,8 @@ from litellm.proxy.utils import PrismaClient
|
|||
def mock_prisma():
|
||||
"""Simplified mock PrismaClient with bound methods"""
|
||||
client = MagicMock()
|
||||
client.db.litellm_healthchecktable.create = AsyncMock(
|
||||
return_value={"id": "test-id"}
|
||||
)
|
||||
client.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[{"id": "1", "model_name": "test"}]
|
||||
)
|
||||
client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"})
|
||||
client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}])
|
||||
|
||||
# Bind actual methods
|
||||
import types
|
||||
|
|
@ -52,14 +51,10 @@ def mock_prisma():
|
|||
("healthy", 1, 0, False), # Database error case
|
||||
],
|
||||
)
|
||||
async def test_save_health_check_result(
|
||||
mock_prisma, status, healthy, unhealthy, should_succeed
|
||||
):
|
||||
async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed):
|
||||
"""Test health check result saving with various scenarios"""
|
||||
if not should_succeed:
|
||||
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception(
|
||||
"DB Error"
|
||||
)
|
||||
mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error")
|
||||
|
||||
result = await mock_prisma.save_health_check_result(
|
||||
model_name="test-model",
|
||||
|
|
@ -187,9 +182,7 @@ def test_aggregate_health_check_results():
|
|||
{"model": "gpt-4", "error": "Rate limit exceeded"},
|
||||
]
|
||||
|
||||
result = _aggregate_health_check_results(
|
||||
model_param_to_info, healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints)
|
||||
|
||||
# Check gpt-3.5-turbo is healthy
|
||||
gpt35_key = ("model-123", "gpt-3.5-turbo")
|
||||
|
|
@ -220,9 +213,7 @@ def test_aggregate_health_check_results_multiple_endpoints():
|
|||
]
|
||||
unhealthy_endpoints = []
|
||||
|
||||
result = _aggregate_health_check_results(
|
||||
model_param_to_info, healthy_endpoints, unhealthy_endpoints
|
||||
)
|
||||
result = _aggregate_health_check_results(model_param_to_info, healthy_endpoints, unhealthy_endpoints)
|
||||
|
||||
key = ("model-123", "gpt-3.5-turbo")
|
||||
assert result[key]["healthy_count"] == 2
|
||||
|
|
@ -398,7 +389,7 @@ async def test_save_background_health_checks_to_db():
|
|||
|
||||
start_time = 1234567890.0
|
||||
|
||||
await _save_background_health_checks_to_db(
|
||||
persisted = await _save_background_health_checks_to_db(
|
||||
mock_prisma,
|
||||
model_list,
|
||||
healthy_endpoints,
|
||||
|
|
@ -407,7 +398,8 @@ async def test_save_background_health_checks_to_db():
|
|||
"background_health_check",
|
||||
)
|
||||
|
||||
# Should call get_all_latest_health_checks and save_health_check_result
|
||||
# Should call get_all_latest_health_checks and save_health_check_result, and report completion
|
||||
assert persisted is True
|
||||
mock_prisma.get_all_latest_health_checks.assert_called_once()
|
||||
mock_prisma.save_health_check_result.assert_called_once()
|
||||
|
||||
|
|
@ -418,22 +410,112 @@ async def test_save_background_health_checks_to_db():
|
|||
assert call_kwargs["checked_by"] == "background_health_check"
|
||||
|
||||
|
||||
def _two_model_results():
|
||||
return {
|
||||
("model-1", "gpt-4"): {
|
||||
"model_name": "gpt-4",
|
||||
"model_id": "model-1",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
},
|
||||
("model-2", "gpt-4o"): {
|
||||
"model_name": "gpt-4o",
|
||||
"model_id": "model-2",
|
||||
"healthy_count": 0,
|
||||
"unhealthy_count": 1,
|
||||
"error_message": "boom",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_health_check_results_if_changed_awaits_every_write_and_reports_success():
|
||||
"""Writes are awaited, not detached, so the caller can tell the cycle's persistence completed."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"})
|
||||
|
||||
persisted = await _save_health_check_results_if_changed(
|
||||
mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check"
|
||||
)
|
||||
|
||||
assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_health_check_results_if_changed_reports_failure_when_a_write_returns_none():
|
||||
"""save_health_check_result swallows DB errors and returns None; that must surface as False."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock(side_effect=[{"id": "row"}, None])
|
||||
|
||||
persisted = await _save_health_check_results_if_changed(
|
||||
mock_prisma, _two_model_results(), {}, 1234567890.0, "background_health_check"
|
||||
)
|
||||
|
||||
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_health_check_results_if_changed_reports_success_when_nothing_needed_writing():
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
model_results = {
|
||||
("model-1", "gpt-4"): {
|
||||
"model_name": "gpt-4",
|
||||
"model_id": "model-1",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
},
|
||||
}
|
||||
latest_checks_map = {
|
||||
"model-1": MagicMock(status="healthy", checked_at=datetime.now(timezone.utc) - timedelta(minutes=5)),
|
||||
}
|
||||
|
||||
persisted = await _save_health_check_results_if_changed(
|
||||
mock_prisma, model_results, latest_checks_map, 1234567890.0, "background_health_check"
|
||||
)
|
||||
|
||||
assert (persisted, mock_prisma.save_health_check_result.await_count) == (True, 0)
|
||||
|
||||
|
||||
def _one_model_setup():
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"model_info": {"id": "model-123"},
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
},
|
||||
]
|
||||
return model_list, [{"model": "gpt-3.5-turbo"}], []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_background_health_checks_to_db_returns_false_when_a_write_fails():
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
|
||||
mock_prisma.save_health_check_result = AsyncMock(return_value=None)
|
||||
model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup()
|
||||
|
||||
persisted = await _save_background_health_checks_to_db(
|
||||
mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, 1234567890.0, "background_health_check"
|
||||
)
|
||||
|
||||
assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_background_health_checks_to_db_no_prisma():
|
||||
"""Test graceful handling when no prisma client"""
|
||||
result = await _save_background_health_checks_to_db(
|
||||
None, [], [], [], 0.0, "background_health_check"
|
||||
)
|
||||
assert result is None
|
||||
result = await _save_background_health_checks_to_db(None, [], [], [], 0.0, "background_health_check")
|
||||
assert result is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_background_health_checks_to_db_exception_handling():
|
||||
"""Test exception handling in background health check save"""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(
|
||||
side_effect=Exception("DB Error")
|
||||
)
|
||||
mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error"))
|
||||
|
||||
model_list = [
|
||||
{
|
||||
|
|
@ -443,104 +525,134 @@ async def test_save_background_health_checks_to_db_exception_handling():
|
|||
},
|
||||
]
|
||||
|
||||
# Should not raise exception, should handle gracefully
|
||||
await _save_background_health_checks_to_db(
|
||||
# Must not raise (the health check loop has to survive a DB outage) but must report
|
||||
# the failure, so the window lock can be released for another pod to retry
|
||||
persisted = await _save_background_health_checks_to_db(
|
||||
mock_prisma, model_list, [], [], 0.0, "background_health_check"
|
||||
)
|
||||
|
||||
# Function should complete without raising
|
||||
assert persisted is False
|
||||
|
||||
|
||||
def _raw_latest_row(model_name: str, model_id, checked_at: datetime) -> dict:
|
||||
return {
|
||||
"health_check_id": f"hc-{model_id or 'no-id'}-{model_name}",
|
||||
"model_name": model_name,
|
||||
"model_id": model_id,
|
||||
"status": "healthy",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
"response_time_ms": 10.0,
|
||||
"details": None,
|
||||
"checked_by": "pod-1",
|
||||
"checked_at": checked_at.isoformat(),
|
||||
"created_at": checked_at.isoformat(),
|
||||
"updated_at": checked_at.isoformat(),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_with_model_id(mock_prisma):
|
||||
"""Test get_all_latest_health_checks properly groups by model_id"""
|
||||
mock_check2 = MagicMock()
|
||||
mock_check2.model_id = "model-456"
|
||||
mock_check2.model_name = "gpt-3.5-turbo"
|
||||
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5)
|
||||
|
||||
mock_check3 = MagicMock()
|
||||
mock_check3.model_id = "model-123"
|
||||
mock_check3.model_name = "gpt-3.5-turbo"
|
||||
mock_check3.checked_at = datetime.now(timezone.utc) - timedelta(
|
||||
minutes=1
|
||||
) # Latest for model-123
|
||||
|
||||
# Order by checked_at desc
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[mock_check3, mock_check2]
|
||||
)
|
||||
|
||||
result = await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
# Should return 2 unique models (by model_id)
|
||||
assert len(result) == 2
|
||||
|
||||
# Should have latest check for each model_id
|
||||
model_ids = {check.model_id for check in result}
|
||||
assert "model-123" in model_ids
|
||||
assert "model-456" in model_ids
|
||||
|
||||
# model-123 should have the latest check (1 minute ago)
|
||||
model123_check = next(c for c in result if c.model_id == "model-123")
|
||||
assert model123_check.checked_at == mock_check3.checked_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_without_model_id(mock_prisma):
|
||||
"""Test get_all_latest_health_checks groups by model_name when model_id is None"""
|
||||
mock_check2 = MagicMock()
|
||||
mock_check2.model_id = None
|
||||
mock_check2.model_name = "gpt-3.5-turbo"
|
||||
mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest
|
||||
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[mock_check2]
|
||||
)
|
||||
|
||||
result = await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
# Should return 1 unique model (by model_name)
|
||||
assert len(result) == 1
|
||||
assert result[0].model_name == "gpt-3.5-turbo"
|
||||
assert result[0].checked_at == mock_check2.checked_at # Latest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_same_name_with_and_without_model_id(
|
||||
mock_prisma,
|
||||
):
|
||||
async def test_get_all_latest_health_checks_keeps_every_distinct_group_with_its_own_checked_at(mock_prisma):
|
||||
"""
|
||||
Same model_name can appear twice after DISTINCT ON: once keyed by (model_id, name)
|
||||
and once by (NULL, name) — different Postgres groups than a single row with id.
|
||||
Postgres owns the dedup. (id, name), (other id, name) and (NULL, name) are distinct groups and each row
|
||||
must arrive typed, with its own checked_at, for the 1h re-save compare and the id-or-name lookup key.
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
with_id = MagicMock()
|
||||
with_id.model_id = "deployment-abc"
|
||||
with_id.model_name = "gpt-4"
|
||||
with_id.checked_at = now - timedelta(minutes=2)
|
||||
|
||||
without_id = MagicMock()
|
||||
without_id.model_id = None
|
||||
without_id.model_name = "gpt-4"
|
||||
without_id.checked_at = now - timedelta(minutes=1)
|
||||
|
||||
mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
return_value=[without_id, with_id]
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
_raw_latest_row("gpt-3.5-turbo", "model-123", now - timedelta(minutes=1)),
|
||||
_raw_latest_row("gpt-3.5-turbo", "model-456", now - timedelta(minutes=5)),
|
||||
_raw_latest_row("gpt-4", "deployment-abc", now - timedelta(minutes=2)),
|
||||
_raw_latest_row("gpt-4", None, now - timedelta(minutes=3)),
|
||||
]
|
||||
)
|
||||
|
||||
result = await mock_prisma.get_all_latest_health_checks()
|
||||
|
||||
assert len(result) == 2
|
||||
names = {r.model_name for r in result}
|
||||
assert names == {"gpt-4"}
|
||||
ids = {r.model_id for r in result}
|
||||
assert "deployment-abc" in ids
|
||||
assert None in ids
|
||||
assert {(check.model_id, check.model_name): check.checked_at for check in result} == {
|
||||
("model-123", "gpt-3.5-turbo"): now - timedelta(minutes=1),
|
||||
("model-456", "gpt-3.5-turbo"): now - timedelta(minutes=5),
|
||||
("deployment-abc", "gpt-4"): now - timedelta(minutes=2),
|
||||
(None, "gpt-4"): now - timedelta(minutes=3),
|
||||
}
|
||||
|
||||
by_key = {(r.model_id, r.model_name): r for r in result}
|
||||
assert by_key[("deployment-abc", "gpt-4")].checked_at == with_id.checked_at
|
||||
assert by_key[(None, "gpt-4")].checked_at == without_id.checked_at
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_background_health_checks_compares_raw_checked_at_against_utc_now(mock_prisma):
|
||||
"""
|
||||
Raw rows carry ISO strings and the engine may omit the offset. A naive checked_at would TypeError
|
||||
inside the 1h compare, be swallowed, and silently stop every save; a stale row must still re-save.
|
||||
"""
|
||||
stale = (datetime.now(timezone.utc) - timedelta(hours=2)).replace(tzinfo=None)
|
||||
fresh = datetime.now(timezone.utc) - timedelta(minutes=5)
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
_raw_latest_row("stale-model", "stale-id", stale),
|
||||
_raw_latest_row("fresh-model", "fresh-id", fresh),
|
||||
]
|
||||
)
|
||||
mock_prisma.save_health_check_result = AsyncMock()
|
||||
model_list = [
|
||||
{"model_name": "stale-model", "model_info": {"id": "stale-id"}, "litellm_params": {"model": "openai/stale"}},
|
||||
{"model_name": "fresh-model", "model_info": {"id": "fresh-id"}, "litellm_params": {"model": "openai/fresh"}},
|
||||
]
|
||||
|
||||
await _save_background_health_checks_to_db(
|
||||
mock_prisma,
|
||||
model_list,
|
||||
[{"model": "openai/stale"}, {"model": "openai/fresh"}],
|
||||
[],
|
||||
time.time(),
|
||||
"pod-1",
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert [call.kwargs["model_id"] for call in mock_prisma.save_health_check_result.await_args_list] == ["stale-id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_latest_health_checks_endpoint_serialises_raw_rows(monkeypatch):
|
||||
row = LatestHealthCheckRow(
|
||||
health_check_id="hc-1",
|
||||
model_name="gpt-4",
|
||||
model_id="deployment-abc",
|
||||
status="healthy",
|
||||
healthy_count=1,
|
||||
unhealthy_count=0,
|
||||
error_message=None,
|
||||
response_time_ms=12.5,
|
||||
details='{"region": "eu"}',
|
||||
checked_by="pod-1",
|
||||
checked_at=datetime(2026, 8, 25),
|
||||
created_at=datetime(2026, 8, 25, tzinfo=timezone.utc),
|
||||
updated_at=datetime(2026, 8, 25, tzinfo=timezone.utc),
|
||||
)
|
||||
prisma = MagicMock()
|
||||
prisma.get_all_latest_health_checks = AsyncMock(return_value=(row,))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
|
||||
|
||||
response = await latest_health_checks_endpoint(user_api_key_dict=UserAPIKeyAuth())
|
||||
|
||||
assert response == {
|
||||
"latest_health_checks": {
|
||||
"deployment-abc": {
|
||||
"health_check_id": "hc-1",
|
||||
"model_name": "gpt-4",
|
||||
"model_id": "deployment-abc",
|
||||
"status": "healthy",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
"response_time_ms": 12.5,
|
||||
"details": {"region": "eu"},
|
||||
"checked_by": "pod-1",
|
||||
"checked_at": "2026-08-25T00:00:00+00:00",
|
||||
"created_at": "2026-08-25T00:00:00+00:00",
|
||||
}
|
||||
},
|
||||
"total_models": 1,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -628,12 +740,7 @@ def test_parse_background_health_check_model_groups_unset_returns_none():
|
|||
|
||||
assert parse_background_health_check_model_groups(None) is None
|
||||
assert parse_background_health_check_model_groups({}) is None
|
||||
assert (
|
||||
parse_background_health_check_model_groups(
|
||||
{"background_health_check_model_groups": None}
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert parse_background_health_check_model_groups({"background_health_check_model_groups": None}) is None
|
||||
|
||||
|
||||
def test_parse_background_health_check_model_groups_list_returns_frozenset():
|
||||
|
|
@ -650,9 +757,7 @@ def test_parse_background_health_check_model_groups_malformed_raises(bad_value):
|
|||
from litellm.proxy.health_check import parse_background_health_check_model_groups
|
||||
|
||||
with pytest.raises(ValueError, match="must be a list of model group names"):
|
||||
parse_background_health_check_model_groups(
|
||||
{"background_health_check_model_groups": bad_value}
|
||||
)
|
||||
parse_background_health_check_model_groups({"background_health_check_model_groups": bad_value})
|
||||
|
||||
|
||||
def test_filter_deployments_to_model_groups():
|
||||
|
|
@ -665,9 +770,7 @@ def test_filter_deployments_to_model_groups():
|
|||
]
|
||||
|
||||
assert filter_deployments_to_model_groups(model_list, None) == tuple(model_list)
|
||||
assert filter_deployments_to_model_groups(
|
||||
model_list, frozenset({"prod-openai"})
|
||||
) == (model_list[0], model_list[2])
|
||||
assert filter_deployments_to_model_groups(model_list, frozenset({"prod-openai"})) == (model_list[0], model_list[2])
|
||||
assert filter_deployments_to_model_groups(model_list, frozenset()) == ()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,10 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.health_check_latest import (
|
||||
LATEST_HEALTH_CHECKS_FOR_MODELS_SQL,
|
||||
LATEST_HEALTH_CHECKS_SQL,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
|
|
@ -261,36 +265,49 @@ async def test_get_health_check_history_db_error_returns_empty_list(
|
|||
assert await prisma_client.get_health_check_history() == []
|
||||
|
||||
|
||||
def _raw_health_check_row(model_name: str = "gpt-4", model_id: str | None = "deployment-abc") -> dict[str, Any]:
|
||||
return {
|
||||
"health_check_id": f"hc-{model_name}",
|
||||
"model_name": model_name,
|
||||
"model_id": model_id,
|
||||
"status": "healthy",
|
||||
"healthy_count": 1,
|
||||
"unhealthy_count": 0,
|
||||
"error_message": None,
|
||||
"response_time_ms": 12.5,
|
||||
"details": None,
|
||||
"checked_by": "pod-1",
|
||||
"checked_at": "2026-08-25T00:00:00+00:00",
|
||||
"created_at": "2026-08-25T00:00:00+00:00",
|
||||
"updated_at": "2026-08-25T00:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_uses_distinct(
|
||||
async def test_get_all_latest_health_checks_dedups_in_postgres(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
rows = [MagicMock(name=f"row-{i}") for i in range(3)]
|
||||
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=rows)
|
||||
"""A revert to prisma find_many(distinct=...) streams the whole history table; the SQL must own the DISTINCT."""
|
||||
prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row()])
|
||||
result = await prisma_client.get_all_latest_health_checks()
|
||||
kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs
|
||||
actual = {
|
||||
"len": len(result),
|
||||
"distinct": kwargs["distinct"],
|
||||
"order_len": len(kwargs["order"]),
|
||||
"first_order": kwargs["order"][0],
|
||||
"query": prisma_client.db.query_raw.await_args.args,
|
||||
"rows": [(row.model_id, row.model_name, row.status) for row in result],
|
||||
"row_type": type(result[0]).__name__,
|
||||
}
|
||||
assert actual == {
|
||||
"len": 3,
|
||||
"distinct": ["model_id", "model_name"],
|
||||
"order_len": 3,
|
||||
"first_order": {"model_id": "asc"},
|
||||
"query": (LATEST_HEALTH_CHECKS_SQL,),
|
||||
"rows": [("deployment-abc", "gpt-4", "healthy")],
|
||||
"row_type": "LatestHealthCheckRow",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_latest_health_checks_db_error_returns_empty_list(
|
||||
async def test_get_all_latest_health_checks_db_error_returns_no_rows(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(
|
||||
side_effect=RuntimeError("oops")
|
||||
)
|
||||
assert await prisma_client.get_all_latest_health_checks() == []
|
||||
prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops"))
|
||||
assert await prisma_client.get_all_latest_health_checks() == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -298,18 +315,15 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod
|
|||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
"""A paged caller reads health for its page; an unbounded read is the bug this exists to avoid."""
|
||||
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
|
||||
await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"])
|
||||
kwargs = prisma_client.db.litellm_healthchecktable.find_many.await_args.kwargs
|
||||
prisma_client.db.query_raw = AsyncMock(return_value=[_raw_health_check_row(model_name="gpt-5")])
|
||||
result = await prisma_client.get_latest_health_checks_for_models(["gpt-5", "claude-opus"])
|
||||
actual = {
|
||||
"where": kwargs["where"],
|
||||
"distinct": kwargs["distinct"],
|
||||
"order": kwargs["order"],
|
||||
"query": prisma_client.db.query_raw.await_args.args,
|
||||
"rows": [row.model_name for row in result],
|
||||
}
|
||||
assert actual == {
|
||||
"where": {"model_name": {"in": ["gpt-5", "claude-opus"]}},
|
||||
"distinct": ["model_id", "model_name"],
|
||||
"order": [{"model_id": "asc"}, {"model_name": "asc"}, {"checked_at": "desc"}],
|
||||
"query": (LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, ["gpt-5", "claude-opus"]),
|
||||
"rows": ["gpt-5"],
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -317,14 +331,14 @@ async def test_get_latest_health_checks_for_models_bounds_the_query_to_those_mod
|
|||
async def test_get_latest_health_checks_for_models_does_not_query_for_an_empty_page(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.query_raw = AsyncMock(return_value=[])
|
||||
assert await prisma_client.get_latest_health_checks_for_models([]) == ()
|
||||
assert prisma_client.db.litellm_healthchecktable.find_many.await_count == 0
|
||||
assert prisma_client.db.query_raw.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_latest_health_checks_for_models_db_error_returns_empty_list(
|
||||
async def test_get_latest_health_checks_for_models_db_error_returns_no_rows(
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
prisma_client.db.litellm_healthchecktable.find_many = AsyncMock(side_effect=RuntimeError("oops"))
|
||||
prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("oops"))
|
||||
assert await prisma_client.get_latest_health_checks_for_models(["gpt-5"]) == ()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue