Merge pull request #41145 from BerriAI/litellm_health_check_skip_save_on_failed_read_main

litellm/proxy/db/health_check_latest.py

@@ -74,10 +74,14 @@ class LatestHealthCheckRow(BaseModel):
 _ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...])
 
 
+async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
+    rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
+    return _ROWS_ADAPTER.validate_python(rows)
+
+
 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)
+        return await query_latest_health_checks(prisma_client)
     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 ()

litellm/proxy/health_endpoints/_health_endpoints.py

@@ -45,7 +45,10 @@ from litellm.proxy.auth.auth_utils import (
 from litellm.proxy.auth.model_checks import get_key_models
 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.health_check_latest import (
+    LatestHealthCheckRow,
+    query_latest_health_checks,
+)
 from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
 from litellm.proxy.health_check import (
     ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
@@ -876,7 +879,7 @@ async def _save_background_health_checks_to_db(
         )
 
         # Step 3: Get latest health checks for all models in one query to compare status
-        latest_checks: Final = await prisma_client.get_all_latest_health_checks()
+        latest_checks: Final = await query_latest_health_checks(prisma_client)
         latest_checks_map: Final = {}
         for check in latest_checks:
             # Use model_id as primary key, fallback to model_name

tests/test_litellm/proxy/db/test_health_check_latest.py

@@ -8,6 +8,7 @@ from litellm.proxy.db.health_check_latest import (
     LATEST_HEALTH_CHECKS_SQL,
     fetch_latest_health_checks,
     fetch_latest_health_checks_for_models,
+    query_latest_health_checks,
 )
 
 
@@ -83,6 +84,15 @@ async def test_fetch_all_degrades_to_no_rows_when_the_query_fails():
     assert await fetch_latest_health_checks(prisma) == ()
 
 
+@pytest.mark.asyncio
+async def test_query_all_raises_when_the_query_fails_instead_of_reading_as_an_empty_table():
+    """The background save decides what to write from this read; a failure has to be told apart from no rows."""
+    prisma = _prisma([])
+    prisma.db.query_raw.side_effect = RuntimeError("db down")
+    with pytest.raises(RuntimeError, match="db down"):
+        await query_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"}])) == ()

tests/test_litellm/proxy/test_health_check_functions.py

@@ -374,7 +374,7 @@ async def test_save_background_health_checks_to_db():
     """Test the main background health check save function"""
     mock_prisma = MagicMock()
     mock_prisma.save_health_check_result = AsyncMock()
-    mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[])
+    mock_prisma.db.query_raw = AsyncMock(return_value=[])
 
     model_list = [
         {
@@ -398,9 +398,9 @@ async def test_save_background_health_checks_to_db():
         "background_health_check",
     )
 
-    # Should call get_all_latest_health_checks and save_health_check_result, and report completion
+    # Should read the latest rows and save_health_check_result, and report completion
     assert persisted is True
-    mock_prisma.get_all_latest_health_checks.assert_called_once()
+    mock_prisma.db.query_raw.assert_awaited_once()
     mock_prisma.save_health_check_result.assert_called_once()
 
     call_kwargs = mock_prisma.save_health_check_result.call_args[1]
@@ -493,7 +493,7 @@ def _one_model_setup():
 @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.db.query_raw = AsyncMock(return_value=[])
     mock_prisma.save_health_check_result = AsyncMock(return_value=None)
     model_list, healthy_endpoints, unhealthy_endpoints = _one_model_setup()
 
@@ -504,6 +504,23 @@ async def test_save_background_health_checks_to_db_returns_false_when_a_write_fa
     assert (persisted, mock_prisma.save_health_check_result.await_count) == (False, 1)
 
 
+@pytest.mark.asyncio
+async def test_save_background_health_checks_to_db_writes_nothing_when_the_latest_row_read_fails(mock_prisma):
+    """
+    A failed dedup read must not read as an empty table. Treated that way, every model was written on every
+    cycle by every pod while the read kept failing, which is what filled the table in production.
+    """
+    mock_prisma.db.query_raw = AsyncMock(side_effect=RuntimeError("db down"))
+    mock_prisma.save_health_check_result = AsyncMock(return_value={"id": "row"})
+    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, 0)
+
+
 @pytest.mark.asyncio
 async def test_save_background_health_checks_to_db_no_prisma():
     """Test graceful handling when no prisma client"""
@@ -515,7 +532,7 @@ async def test_save_background_health_checks_to_db_no_prisma():
 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.db.query_raw = AsyncMock(side_effect=Exception("DB Error"))
 
     model_list = [
         {