scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): stop health probes from rebuilding Prisma clients

Browse source
This commit is contained in:
ZXYxc 2026-08-26 16:34:55 +08:00
parent 137311ffd6
commit 08ae3eec71
6 changed files with 88 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/db/prisma_client.py
View file

@ -51,6 +51,8 @@ class _PrismaEngine(Protocol):
async def rollback_transaction(self, tx_id: str) -> None: ...
async def request(self, method: str, path: str) -> object: ...
class _PrismaClient(Protocol):
_Prisma__engine: _PrismaEngine
@ -308,6 +310,10 @@ class PrismaWrapper:
await self._original_prisma.connect(timeout)
self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma)
async def query_engine_status(self) -> object:
engine: Final = self._read_engine(self._original_prisma)
return await engine.request("GET", "/status")
@staticmethod
async def _kill_engine_process(pid: int) -> None:
"""Force-kill the engine subprocess to prevent DB connection pool leaks.

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

@ -33,7 +33,6 @@ from litellm.proxy.auth.auth_utils import (
_BANNED_REQUEST_BODY_PARAMS, # pyright: ignore[reportPrivateUsage] # one canonical list, shared with the request-body check
)
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.proxy_worker_heartbeat import count_live_proxy_workers
from litellm.proxy.health_check import (
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
@ -1339,21 +1338,11 @@ async def _db_health_readiness_check():
await prisma_client.health_check()
db_health_cache = {"status": "connected", "last_updated": datetime.now()}
return db_health_cache
except Exception as e:
except Exception:
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
if PrismaDBExceptionHandler.is_database_transport_error(e):
try:
verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect")
await prisma_client.attempt_db_reconnect(reason="health_readiness_check")
await prisma_client.health_check()
verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded")
db_health_cache = {
"status": "connected",
"last_updated": datetime.now(),
}
return db_health_cache
except Exception:
verbose_proxy_logger.error("_db_health_readiness_check: reconnect failed")
verbose_proxy_logger.warning(
"_db_health_readiness_check: health_check failed; returning disconnected without reconnect"
)
return db_health_cache

4
litellm/proxy/utils.py
View file

@ -5468,7 +5468,7 @@ class PrismaClient:
async def start_db_health_watchdog_task(self) -> None:
"""Start background tasks that monitor DB health:
- A periodic SELECT 1 probe that triggers reconnect on network/connection failure.
- A periodic query-engine status probe that triggers reconnect on local engine failure.
- A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.
"""
if self._db_health_watchdog_enabled is not True:
@ -5509,7 +5509,7 @@ class PrismaClient:
try:
await asyncio.sleep(self._db_health_watchdog_interval_seconds)
await asyncio.wait_for(
self.db.query_raw("SELECT 1"),
self.writer_db.query_engine_status(),
timeout=self._db_health_watchdog_probe_timeout_seconds,
)
if isinstance(self.db, RoutingPrismaWrapper) and self.db.writer_unavailable:

14
tests/test_litellm/proxy/db/test_prisma_self_heal.py
View file

@ -269,13 +269,15 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
@pytest.mark.asyncio
async def test_db_health_watchdog_should_trigger_reconnect_on_db_error(
async def test_db_health_watchdog_should_trigger_reconnect_on_engine_error(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.query_raw = AsyncMock(side_effect=Exception("db connection dropped"))
client.writer_db.query_engine_status = AsyncMock(
side_effect=Exception("query engine unavailable")
)
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_watchdog_reconnect_timeout_seconds = 7.0
@ -306,7 +308,9 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout(
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.query_raw = AsyncMock(side_effect=asyncio.TimeoutError())
client.writer_db.query_engine_status = AsyncMock(
side_effect=asyncio.TimeoutError()
)
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_watchdog_reconnect_timeout_seconds = 9.0
@ -563,7 +567,7 @@ async def test_db_health_watchdog_should_reconnect_degraded_writer(
)
writer = MagicMock()
reader = MagicMock()
reader.query_raw = AsyncMock(return_value=[{"result": 1}])
writer.query_engine_status = AsyncMock(return_value={"status": "ok"})
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
routing._writer_unavailable = True
client.db = routing
@ -596,7 +600,7 @@ async def test_db_health_watchdog_should_not_reconnect_healthy_writer(
)
writer = MagicMock()
reader = MagicMock()
reader.query_raw = AsyncMock(return_value=[{"result": 1}])
writer.query_engine_status = AsyncMock(return_value={"status": "ok"})
routing = RoutingPrismaWrapper(writer=writer, reader=reader)
client.db = routing
client.attempt_db_reconnect = AsyncMock(return_value=True)

79
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -141,9 +141,7 @@ async def test_db_health_transport_error_never_raises(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check"
)
mock_prisma.attempt_db_reconnect.assert_not_called()
@pytest.mark.asyncio
@ -155,49 +153,14 @@ async def test_db_health_transport_error_never_raises(transport_error):
HTTPClientClosedError(),
],
)
async def test_db_health_transport_error_reconnect_succeeds(transport_error):
async def test_db_health_transport_error_does_not_reconnect(transport_error):
"""
When health_check raises a transport error and attempt_db_reconnect
succeeds, the second health_check passes and we return 'connected'.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=[transport_error, None])
mock_prisma.attempt_db_reconnect = AsyncMock(return_value=True)
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert result["status"] == "connected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check"
)
assert mock_prisma.health_check.call_count == 2
@pytest.mark.asyncio
@pytest.mark.parametrize(
"transport_error",
[
httpx.ConnectError("All connection attempts failed"),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_transport_error_reconnect_fails(transport_error):
"""
When health_check raises a transport error and attempt_db_reconnect also
fails, return 'disconnected' without raising.
A readiness probe reports a transport error as disconnected without
rebuilding the Prisma query engine or retrying the probe inline.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=transport_error)
mock_prisma.attempt_db_reconnect = AsyncMock(
side_effect=RuntimeError("reconnect failed")
)
mock_prisma.attempt_db_reconnect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
@ -208,6 +171,38 @@ async def test_db_health_transport_error_reconnect_fails(transport_error):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"transport_error",
[
httpx.ConnectError("All connection attempts failed"),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_transport_error_keeps_disconnected_cache(transport_error):
"""
A failed readiness probe refreshes the disconnected cache timestamp so
callers receive the current state without launching reconnect work.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=transport_error)
mock_prisma.attempt_db_reconnect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
assert _health_endpoints_module.db_health_cache["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_not_called()
@pytest.mark.asyncio

58
tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py
View file

@ -447,58 +447,60 @@ async def test_stop_db_health_watchdog_task_noop_when_no_task(
async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The watchdog loop reconnects when ``wait_for`` raises TimeoutError
or a recognized DB connection error.
"""
"""A local query-engine status timeout triggers reconnect."""
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock(return_value=True)
prisma_client.writer_db.query_engine_status = AsyncMock(
side_effect=asyncio.TimeoutError()
)
call_count = {"n": 0}
async def _timeout_then_cancel(*args: Any, **kwargs: Any) -> None:
call_count["n"] += 1
if call_count["n"] >= 2:
raise asyncio.CancelledError()
raise asyncio.TimeoutError()
monkeypatch.setattr("asyncio.wait_for", _timeout_then_cancel)
await prisma_client._db_health_watchdog_loop()
with monkeypatch.context() as patch_context:
patch_context.setattr(
asyncio,
"sleep",
AsyncMock(side_effect=[None, asyncio.CancelledError()]),
)
await prisma_client._db_health_watchdog_loop()
pinned = {
"reconnect_called": prisma_client.attempt_db_reconnect.await_count,
"reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[
"reason"
],
"wait_for_calls": call_count["n"],
"status_calls": prisma_client.writer_db.query_engine_status.await_count,
"loop_exited_clean": True,
}
assert pinned == {
"reconnect_called": 1,
"reconnect_reason": "db_health_watchdog_connection_error",
"wait_for_calls": 2,
"status_calls": 1,
"loop_exited_clean": True,
}
@pytest.mark.asyncio
async def test_db_health_watchdog_loop_swallows_non_db_errors(
async def test_db_health_watchdog_loop_probes_engine_status_without_querying_db(
prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A non-DB error during the probe should NOT trigger reconnect; the
loop logs and continues until cancellation.
"""
"""A healthy local engine does not issue a PostgreSQL query."""
prisma_client._db_health_watchdog_interval_seconds = 0
prisma_client.attempt_db_reconnect = AsyncMock()
prisma_client.writer_db.query_engine_status = AsyncMock(
return_value={"status": "ok"}
)
prisma_client.db.query_raw = AsyncMock(
side_effect=AssertionError("watchdog must not query PostgreSQL")
)
call_count = {"n": 0}
with monkeypatch.context() as patch_context:
patch_context.setattr(
asyncio,
"sleep",
AsyncMock(side_effect=[None, asyncio.CancelledError()]),
)
await prisma_client._db_health_watchdog_loop()
async def _raise_then_cancel(*args: Any, **kwargs: Any) -> None:
call_count["n"] += 1
if call_count["n"] >= 2:
raise asyncio.CancelledError()
raise ValueError("not a db error")
monkeypatch.setattr("asyncio.wait_for", _raise_then_cancel)
await prisma_client._db_health_watchdog_loop()
prisma_client.writer_db.query_engine_status.assert_awaited_once()
prisma_client.db.query_raw.assert_not_awaited()
assert prisma_client.attempt_db_reconnect.await_count == 0