scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

feat(proxy): enhance PrismaClient health watchdog with reconnect cooldown and tests

Browse source 
- Introduced a reconnect cooldown for the health watchdog to prevent rapid reconnection attempts during database errors.
- Updated health watchdog parameters to allow for configurable intervals and timeouts via environment variables.
- Enhanced logging to provide clearer information on the watchdog's operation.
- Added tests to verify the reconnect behavior during cooldown periods and ensure proper functionality of the health watchdog.
This commit is contained in:
harish-berri 2026-05-19 00:32:08 +00:00
parent 72d7814695
commit f6d743fb18
2 changed files with 51 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/proxy/utils.py
View file

@ -2851,21 +2851,26 @@ class PrismaClient:
self._db_reconnect_lock = asyncio.Lock()
self._db_health_watchdog_task: Optional[asyncio.Task] = None
self._db_last_reconnect_attempt_ts: float = 0.0
self._db_last_health_watchdog_reconnect_attempt_ts: float = 0.0
self._db_reconnect_cooldown_seconds: int = max(
1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))
)
self._db_health_watchdog_interval_seconds: int = max(
5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "30"))
5, int(os.getenv("PRISMA_HEALTH_WATCHDOG_INTERVAL_SECONDS", "120"))
)
self._db_health_watchdog_enabled: bool = (
str_to_bool(os.getenv("PRISMA_HEALTH_WATCHDOG_ENABLED", "true")) is True
)
self._db_health_watchdog_probe_timeout_seconds: float = max(
0.5,
float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "5.0")),
float(os.getenv("PRISMA_HEALTH_WATCHDOG_PROBE_TIMEOUT_SECONDS", "15")),
)
self._db_health_watchdog_reconnect_cooldown_seconds: int = max(
1,
int(os.getenv("PRISMA_HEALTH_WATCHDOG_RECONNECT_COOLDOWN_SECONDS", "120")),
)
self._db_watchdog_reconnect_timeout_seconds: float = max(
1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "30.0"))
1.0, float(os.getenv("PRISMA_WATCHDOG_RECONNECT_TIMEOUT_SECONDS", "60.0"))
)
self._db_auth_reconnect_timeout_seconds: float = max(
0.5, float(os.getenv("PRISMA_AUTH_RECONNECT_TIMEOUT_SECONDS", "2.0"))
@ -4772,9 +4777,9 @@ class PrismaClient:
self._db_health_watchdog_loop()
)
verbose_proxy_logger.info(
"Started Prisma DB health watchdog (interval=%ss, reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)",
"Started Prisma DB health watchdog (interval=%ss, watchdog_reconnect_cooldown=%ss, probe_timeout=%ss, reconnect_timeout=%ss)",
self._db_health_watchdog_interval_seconds,
self._db_reconnect_cooldown_seconds,
self._db_health_watchdog_reconnect_cooldown_seconds,
self._db_health_watchdog_probe_timeout_seconds,
self._db_watchdog_reconnect_timeout_seconds,
)
@ -4807,8 +4812,19 @@ class PrismaClient:
if isinstance(
e, asyncio.TimeoutError
) or PrismaDBExceptionHandler.is_database_connection_error(e):
now = time.time()
if (
now - self._db_last_health_watchdog_reconnect_attempt_ts
< self._db_health_watchdog_reconnect_cooldown_seconds
):
verbose_proxy_logger.debug(
"Skipping DB health watchdog reconnect due to watchdog cooldown."
)
continue
self._db_last_health_watchdog_reconnect_attempt_ts = now
await self.attempt_db_reconnect(
reason="db_health_watchdog_connection_error",
force=True,
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
)
else:

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

@ -271,6 +271,7 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_db_error(
client.attempt_db_reconnect.assert_awaited_once_with(
reason="db_health_watchdog_connection_error",
force=True,
timeout_seconds=7.0,
)
@ -302,10 +303,39 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout(
client.attempt_db_reconnect.assert_awaited_once_with(
reason="db_health_watchdog_connection_error",
force=True,
timeout_seconds=9.0,
)
@pytest.mark.asyncio
async def test_db_health_watchdog_should_skip_reconnect_during_watchdog_cooldown(
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.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_health_watchdog_reconnect_cooldown_seconds = 3600
client._db_last_health_watchdog_reconnect_attempt_ts = time.time()
with (
patch(
"litellm.proxy.utils.asyncio.sleep",
AsyncMock(side_effect=[None, asyncio.CancelledError()]),
),
patch(
"litellm.proxy.db.exception_handler.PrismaDBExceptionHandler.is_database_connection_error",
return_value=True,
),
):
await client._db_health_watchdog_loop()
client.attempt_db_reconnect.assert_not_called()
@pytest.mark.asyncio
async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
client = PrismaClient(