scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #26225 from BerriAI/litellm_dbReconnectNonBlocking

Browse source 
[Fix] Proxy: reconnect Prisma DB without blocking the event loop
This commit is contained in:
yuneng-jiang 2026-04-29 16:09:22 -07:00 • committed by GitHub
commit fc0cc9c581
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 150 additions and 110 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -52,18 +52,25 @@ class PrismaWrapper:
engine = self._original_prisma._engine
process = getattr(engine, "process", None) if engine is not None else None
if process is not None:
return process.pid
pid = process.pid
if isinstance(pid, int):
return pid
except (AttributeError, TypeError):
pass
return 0
@staticmethod
async def _kill_engine_process(pid: int) -> None:
"""Force-kill an orphaned engine subprocess to prevent DB connection pool leaks.
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
Called when disconnect() fails and the old engine process may still be
holding open connections. Sends SIGTERM for graceful shutdown, waits
briefly, then SIGKILL as a backstop.
Called on every reconnect (in `recreate_prisma_client`) to retire the
old query-engine subprocess without invoking prisma-client-py's
synchronous `disconnect()` — which blocks the asyncio event loop on
`subprocess.Popen.wait()` for 30-120+ seconds when the engine is
stuck on TCP close.
Sends SIGTERM for graceful shutdown, waits briefly, then SIGKILL as
a backstop.
"""
if pid <= 0:
return
@ -72,7 +79,7 @@ class PrismaWrapper:
except (ProcessLookupError, PermissionError, OSError):
return # Already dead or inaccessible
verbose_proxy_logger.warning(
"Sent SIGTERM to orphaned prisma-query-engine PID %s after failed disconnect.",
"Sent SIGTERM to prisma-query-engine PID %s during reconnect.",
pid,
)
# Brief wait for graceful shutdown, then force-kill
@ -217,15 +224,18 @@ class PrismaWrapper:
async def recreate_prisma_client(
self, new_db_url: str, http_client: Optional[Any] = None
):
"""Disconnect and reconnect the Prisma client with a new database URL."""
"""Disconnect and reconnect the Prisma client with a new database URL.
Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than
calling `disconnect()`. prisma-client-py's `disconnect()` calls a
synchronous `subprocess.Popen.wait()` that can freeze the asyncio event
loop for 30-120+ seconds when the engine is stuck on TCP close,
breaking `/health/liveliness` and causing Kubernetes pod restarts.
"""
from prisma import Prisma # type: ignore
old_engine_pid = self._get_engine_pid()
try:
await self._original_prisma.disconnect()
except Exception as e:
verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}")
if old_engine_pid > 0:
await self._kill_engine_process(old_engine_pid)
if http_client is not None:

30
litellm/proxy/utils.py
View file

@ -4198,8 +4198,11 @@ class PrismaClient:
Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll
handlers) to choose between heavy reconnect (engine dead -- recreate
Prisma client, re-arm watcher) and lightweight reconnect (network
blip -- disconnect, connect, SELECT 1).
Prisma client, re-arm watcher) and direct reconnect (network blip --
recreate Prisma client, re-arm watcher, SELECT 1). Both paths recreate
the client via the non-blocking kill-then-construct flow rather than
calling disconnect(), which blocks the event loop on the synchronous
subprocess.Popen.wait() inside prisma-client-py (see issue #26191).
"""
effective_timeout = (
timeout_seconds
@ -4243,17 +4246,20 @@ class PrismaClient:
)
async def _do_direct_reconnect() -> None:
old_pid = self._get_engine_pid()
try:
await self.db.disconnect()
except Exception as disconnect_err:
verbose_proxy_logger.warning(
"Prisma DB disconnect before reconnect failed: %s",
disconnect_err,
db_url = os.getenv("DATABASE_URL", "")
if not db_url:
verbose_proxy_logger.error(
"DATABASE_URL not set; cannot reconnect Prisma client."
)
await PrismaWrapper._kill_engine_process(old_pid)
await self.db.connect()
raise RuntimeError("DATABASE_URL not set")
# Fresh Prisma client + new engine subprocess. The previous
# "lightweight" path called `disconnect()` which blocks the
# event loop on `subprocess.Popen.wait()`; since that call
# ends up killing the engine anyway, we do it non-blockingly
# via `_kill_engine_process` inside `recreate_prisma_client`.
self._cleanup_engine_watcher()
await self.db.recreate_prisma_client(db_url)
await self._start_engine_watcher()
await self.db.query_raw("SELECT 1")
await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout)

77
tests/litellm/proxy/test_prisma_engine_watchdog.py
View file

@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive(
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
engine_client,
) -> None:
"""_run_reconnect_cycle uses disconnect/connect when engine is alive."""
engine_client._engine_pid = 1234
"""Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1.
with patch.object(engine_client, "_is_engine_alive", return_value=True):
The old "lightweight" path called `disconnect()` + `connect()`, which
blocks the event loop on the sync `process.wait()` inside aclose().
The fix routes both engine-alive and engine-dead paths through
`recreate_prisma_client`, which non-blockingly kills the old engine.
"""
engine_client._engine_pid = 1234
engine_client._start_engine_watcher = AsyncMock()
with (
patch.object(engine_client, "_is_engine_alive", return_value=True),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown(
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
engine_client,
) -> None:
"""_run_reconnect_cycle uses lightweight path when engine PID is not tracked."""
"""When the engine PID is not tracked, direct reconnect still runs."""
engine_client._engine_pid = 0
engine_client._start_engine_watcher = AsyncMock()
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.connect.assert_awaited_once()
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test"
)
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
@pytest.mark.asyncio
async def test_escalation_after_consecutive_lightweight_failures(engine_client):
"""After N consecutive lightweight reconnect failures, _engine_confirmed_dead
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
engine_client._reconnect_escalation_threshold = 3
engine_client._consecutive_reconnect_failures = 0
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
engine_client._start_engine_watcher = AsyncMock(return_value=None)
# Make lightweight reconnect fail every time
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed"))
# Make direct reconnect fail every time
engine_client.db.recreate_prisma_client = AsyncMock(
side_effect=Exception("recreate failed")
)
# Run 3 failed reconnect attempts
for i in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
for _ in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
assert engine_client._consecutive_reconnect_failures == 3
# Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle
# Next attempt should escalate to the heavy path (recreate_prisma_client still
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client._start_engine_watcher = AsyncMock(return_value=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test_escalation", timeout_seconds=5.0
)
# Heavy reconnect should have been attempted (recreate_prisma_client called)
engine_client.db.recreate_prisma_client.assert_awaited_once()
@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client):
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
engine_client._consecutive_reconnect_failures = 2
engine_client._db_reconnect_cooldown_seconds = 0
engine_client._start_engine_watcher = AsyncMock()
# Make reconnect succeed
engine_client.db.disconnect = AsyncMock(return_value=None)
engine_client.db.connect = AsyncMock(return_value=None)
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is True
assert engine_client._consecutive_reconnect_failures == 0

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

@ -35,18 +35,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
result = await client.attempt_db_reconnect(
reason="unit_test_reconnect_success",
force=True,
)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await client.attempt_db_reconnect(
reason="unit_test_reconnect_success",
force=True,
)
assert result is True
client.db.disconnect.assert_awaited_once()
client.db.connect.assert_awaited_once()
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
client.db.query_raw.assert_awaited_once_with("SELECT 1")
@ -141,15 +141,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
)
client._db_last_reconnect_attempt_ts = 0.0
client._db_reconnect_cooldown_seconds = 10
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
# Use a counter-based mock to avoid StopIteration when time.time() is called
# more times than expected (varies by Python version / internal code paths).
fake_clock = iter(range(100, 10000))
with patch(
"litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock))
with (
patch(
"litellm.proxy.utils.time.time",
side_effect=lambda: float(next(fake_clock)),
),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
result = await client.attempt_db_reconnect(
reason="unit_test_cooldown_timestamp_after_attempt",
@ -163,23 +167,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
@pytest.mark.asyncio
async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops(
async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
mock_proxy_logging,
):
"""Direct reconnect goes through recreate_prisma_client (which non-blockingly
kills the old engine) instead of calling disconnect() — see issue #26191.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used"))
client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used"))
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.disconnect = AsyncMock(
side_effect=AssertionError("disconnect must not be called")
)
client.db.recreate_prisma_client = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
client._start_engine_watcher = AsyncMock()
await client._run_reconnect_cycle(timeout_seconds=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await client._run_reconnect_cycle(timeout_seconds=None)
client.db.disconnect.assert_awaited_once()
client.db.connect.assert_awaited_once()
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test")
client.db.query_raw.assert_awaited_once_with("SELECT 1")
client.db.disconnect.assert_not_awaited()
@pytest.mark.asyncio
@ -190,19 +199,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client._db_watchdog_reconnect_timeout_seconds = 0.1
client.db.disconnect = AsyncMock(return_value=None)
client._start_engine_watcher = AsyncMock()
async def _slow_connect():
async def _slow_recreate(_db_url):
await asyncio.sleep(0.08)
async def _slow_query(_query: str):
await asyncio.sleep(0.08)
return [{"result": 1}]
client.db.connect = AsyncMock(side_effect=_slow_connect)
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
client.db.query_raw = AsyncMock(side_effect=_slow_query)
with pytest.raises(asyncio.TimeoutError):
with (
pytest.raises(asyncio.TimeoutError),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await client._run_reconnect_cycle(timeout_seconds=None)
@ -213,19 +225,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.disconnect = AsyncMock(return_value=None)
client._start_engine_watcher = AsyncMock()
async def _slow_connect():
async def _slow_recreate(_db_url):
await asyncio.sleep(0.08)
async def _slow_query(_query: str):
await asyncio.sleep(0.08)
return [{"result": 1}]
client.db.connect = AsyncMock(side_effect=_slow_connect)
client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate)
client.db.query_raw = AsyncMock(side_effect=_slow_query)
with pytest.raises(asyncio.TimeoutError):
with (
pytest.raises(asyncio.TimeoutError),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await client._run_reconnect_cycle(timeout_seconds=0.1)
@ -320,45 +335,35 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
@pytest.mark.asyncio
async def test_lightweight_reconnect_kills_engine_on_disconnect_failure(
async def test_recreate_prisma_client_kills_old_engine_without_disconnect(
mock_proxy_logging,
):
"""Lightweight reconnect must kill the old engine PID when disconnect() fails."""
"""recreate_prisma_client SIGTERMs the old engine PID directly rather than
calling `disconnect()`, which blocks the asyncio event loop on the sync
`subprocess.Popen.wait()` inside prisma-client-py — see issue #26191.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed"))
client.db.connect = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
disconnect_mock = AsyncMock(
side_effect=AssertionError("disconnect must not be called on reconnect path")
)
client.db._original_prisma.disconnect = disconnect_mock
with (
patch.object(client, "_get_engine_pid", return_value=9999),
patch("os.kill") as mock_kill,
patch("asyncio.sleep", new_callable=AsyncMock),
patch.object(client.db, "_get_engine_pid", return_value=9999),
patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill,
patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock),
):
await client._run_reconnect_cycle(timeout_seconds=5.0)
# Return a Prisma instance whose connect() is awaitable.
fake_new_prisma = MagicMock()
fake_new_prisma.connect = AsyncMock(return_value=None)
with patch("prisma.Prisma", return_value=fake_new_prisma):
await client.db.recreate_prisma_client("postgresql://test")
mock_kill.assert_any_call(9999, signal.SIGTERM)
client.db.connect.assert_awaited_once()
client.db.query_raw.assert_awaited_once_with("SELECT 1")
@pytest.mark.asyncio
async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
mock_proxy_logging,
):
"""Lightweight reconnect must NOT kill when disconnect() succeeds."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
with patch("os.kill") as mock_kill:
await client._run_reconnect_cycle(timeout_seconds=5.0)
mock_kill.assert_not_called()
disconnect_mock.assert_not_awaited()
fake_new_prisma.connect.assert_awaited_once()
# ---------------------------------------------------------------------------