scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(proxy): use Popen.poll() instead of os.kill(pid, 0) for engine liveness on Windows

Browse source 
On POSIX, os.kill(pid, 0) is a no-op existence check. On Windows, Python's
os.kill routes any signal other than CTRL_C_EVENT/CTRL_BREAK_EVENT through
TerminateProcess, so os.kill(pid, 0) terminates the target process rather than
checking liveness.

_poll_engine_proc called os.kill(self._engine_pid, 0) on every poll iteration.
On Windows this killed the Prisma query-engine sidecar on the first event-loop
tick after start_db_health_watchdog_task() yielded, causing the proxy to exit
within a second of startup whenever a database was configured.

Fix: add _is_engine_alive_windows() that uses subprocess.Popen.poll() (None
while running, an integer exit code once exited). _is_engine_alive() delegates
to it on Windows. _poll_engine_proc now calls _is_engine_alive() instead of
os.kill directly, making the poll loop Windows-safe.
This commit is contained in:
NTZ_ARM 2026-06-19 17:50:05 +07:00
parent 7e5699c7ab
commit 026beed001
2 changed files with 149 additions and 87 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/utils.py
View file

@ -4394,6 +4394,8 @@ class PrismaClient:
def _is_engine_alive(self) -> bool:
if self._engine_pid <= 0:
return True
if sys.platform == "win32":
return self._is_engine_alive_windows()
try:
os.kill(self._engine_pid, 0)
return True
@ -4402,6 +4404,19 @@ class PrismaClient:
except (PermissionError, OSError):
return True
def _is_engine_alive_windows(self) -> bool:
# os.kill(pid, 0) on Windows calls TerminateProcess — it kills the
# process rather than checking liveness. Use Popen.poll() instead,
# which returns None while the process is still running.
try:
engine = self.writer_db._original_prisma._engine # type: ignore[attr-defined]
process = getattr(engine, "process", None) if engine is not None else None
if process is not None:
return process.poll() is None
except (AttributeError, TypeError):
pass
return True
@staticmethod
def _reap_all_zombies() -> set:
"""Reap ALL zombie child processes via waitpid(-1, WNOHANG).
@ -4615,14 +4630,12 @@ class PrismaClient:
)
async def _poll_engine_proc(self) -> None:
"""poll via os.kill(pid, 0) every 1s.
"""Poll engine process liveness every 1s via _is_engine_alive().
Only used when BOTH waitpid thread and pidfd are unavailable
(e.g., PID is not our child process and pidfd_open fails)
(e.g., on Windows, or when PID is not our child process).
"""
while self._watching_engine and self._engine_pid > 0:
try:
os.kill(self._engine_pid, 0)
except ProcessLookupError:
if not self._is_engine_alive():
dead_pid = self._engine_pid
if self._consume_expected_death(dead_pid):
verbose_proxy_logger.info(
@ -4644,13 +4657,6 @@ class PrismaClient:
force=True,
)
return
except (PermissionError, OSError):
verbose_proxy_logger.debug(
"Cannot signal PID %s; stopping engine poll.",
self._engine_pid,
)
self._cleanup_engine_watcher()
return
await asyncio.sleep(1)
def _cleanup_engine_watcher(self) -> None:

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

@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -32,16 +30,12 @@ def mock_proxy_logging():
@pytest.mark.asyncio
async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.recreate_prisma_client = AsyncMock(return_value=True)
# Probe fails (connection genuinely broken) so the direct path proceeds to
# recreate; the post-recreate smoke test then succeeds. A healthy probe
# would instead skip the recreate (covered in test_prisma_client_reconnect).
client.db.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"result": 1}]]
)
client.db.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"result": 1}]])
client._start_engine_watcher = AsyncMock()
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
@ -51,17 +45,13 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging):
)
assert result is True
client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test", expected_generation=0
)
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test", expected_generation=0)
assert client.db.query_raw.await_count == 2
@pytest.mark.asyncio
async def test_attempt_db_reconnect_should_skip_when_in_cooldown(mock_proxy_logging):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=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.query_raw = AsyncMock(return_value=[{"result": 1}])
@ -83,9 +73,7 @@ async def test_attempt_db_reconnect_should_skip_when_in_cooldown(mock_proxy_logg
async def test_attempt_db_reconnect_should_skip_when_lock_timeout_expires(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=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.query_raw = AsyncMock(return_value=[{"result": 1}])
@ -111,9 +99,7 @@ async def test_attempt_db_reconnect_should_skip_when_lock_timeout_expires(
async def test_attempt_db_reconnect_should_not_leak_lock_on_timeout_race(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=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.query_raw = AsyncMock(return_value=[{"result": 1}])
@ -143,9 +129,7 @@ async def test_attempt_db_reconnect_should_not_leak_lock_on_timeout_race(
async def test_attempt_db_reconnect_should_set_cooldown_after_attempt(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._db_last_reconnect_attempt_ts = 0.0
client._db_reconnect_cooldown_seconds = 10
client.db.recreate_prisma_client = AsyncMock(return_value=None)
@ -180,26 +164,18 @@ async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
"""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.db.disconnect = AsyncMock(
side_effect=AssertionError("disconnect must not be called")
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.disconnect = AsyncMock(side_effect=AssertionError("disconnect must not be called"))
client.db.recreate_prisma_client = AsyncMock(return_value=True)
# Probe fails so we proceed to recreate (and verify disconnect is never
# used — issue #26191); the post-recreate smoke test then succeeds.
client.db.query_raw = AsyncMock(
side_effect=[ConnectionError("probe failed"), [{"result": 1}]]
)
client.db.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"result": 1}]])
client._start_engine_watcher = AsyncMock()
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await client._run_reconnect_cycle(timeout_seconds=None)
client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test", expected_generation=0
)
client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test", expected_generation=0)
assert client.db.query_raw.await_count == 2
client.db.disconnect.assert_not_awaited()
@ -208,9 +184,7 @@ async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client(
async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._db_watchdog_reconnect_timeout_seconds = 0.1
client._start_engine_watcher = AsyncMock()
@ -242,9 +216,7 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget(
async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._start_engine_watcher = AsyncMock()
async def _slow_recreate(_db_url, **_kwargs):
@ -275,9 +247,7 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget(
async def test_db_health_watchdog_should_trigger_reconnect_on_db_error(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=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
@ -306,9 +276,7 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_db_error(
async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout(
mock_proxy_logging,
):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.query_raw = AsyncMock(side_effect=asyncio.TimeoutError())
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
@ -335,9 +303,7 @@ async def test_db_health_watchdog_should_trigger_reconnect_on_probe_timeout(
@pytest.mark.asyncio
async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._db_health_watchdog_enabled = True
client._db_health_watchdog_interval_seconds = 3600
@ -350,9 +316,7 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging):
coro.close()
return dummy_task
with patch(
"litellm.proxy.utils.asyncio.create_task", side_effect=_fake_create_task
):
with patch("litellm.proxy.utils.asyncio.create_task", side_effect=_fake_create_task):
await client.start_db_health_watchdog_task()
assert client._db_health_watchdog_task is dummy_task
@ -369,12 +333,8 @@ async def test_recreate_prisma_client_kills_old_engine_without_disconnect(
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
)
disconnect_mock = AsyncMock(
side_effect=AssertionError("disconnect must not be called on reconnect path")
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
disconnect_mock = AsyncMock(side_effect=AssertionError("disconnect must not be called on reconnect path"))
client.db._original_prisma.disconnect = disconnect_mock
with (
@ -410,9 +370,7 @@ async def test_get_generic_data_retries_on_transport_error_for_config_table(
four concurrent `get_generic_data` calls, so a single transport flap used
to surface as four `db_exceptions` alerts and a stale config window.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
invocations: list[None] = []
@ -449,13 +407,9 @@ async def test_get_generic_data_retries_on_transport_error_for_config_table(
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
"""If reconnect itself does not succeed, propagate the original transport
error and let the existing failure_handler / db_exceptions telemetry fire."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client.db.litellm_config.find_first = AsyncMock(
side_effect=httpx.ReadError("simulated transport blip")
)
client.db.litellm_config.find_first = AsyncMock(side_effect=httpx.ReadError("simulated transport blip"))
client.attempt_db_reconnect = AsyncMock(return_value=False)
with pytest.raises(httpx.ReadError):
@ -492,16 +446,12 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
The fix moves the reset into the success branch — the flag must stay True
when heavy reconnect raises.
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._engine_confirmed_dead = True
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
# Make the heavy reconnect path raise.
client.db.recreate_prisma_client = AsyncMock(
side_effect=RuntimeError("simulated heavy reconnect failure")
)
client.db.recreate_prisma_client = AsyncMock(side_effect=RuntimeError("simulated heavy reconnect failure"))
client._start_engine_watcher = AsyncMock()
client._cleanup_engine_watcher = MagicMock()
client._reap_all_zombies = MagicMock()
@ -513,3 +463,109 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
# The flag must STILL be True so the next attempt re-enters the heavy
# branch instead of silently demoting to the lightweight path.
assert client._engine_confirmed_dead is True
# ---------------------------------------------------------------------------
# Windows: _poll_engine_proc must not use os.kill(pid, 0) (regression)
#
# On Windows, Python's os.kill(pid, 0) calls TerminateProcess(pid, 0) rather
# than performing a no-op existence check as it does on POSIX. This caused the
# engine-watchdog poll loop to kill the Prisma query-engine sidecar on the
# very first tick, immediately after start_db_health_watchdog_task() yielded
# control back to the asyncio event loop.
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_poll_engine_proc_keeps_alive_engine_running(mock_proxy_logging):
"""_poll_engine_proc must not terminate a living engine process.
Regression: before the fix, _poll_engine_proc called os.kill(pid, 0) which
on Windows routes to TerminateProcess and kills the sidecar on the first
poll tick. The fix delegates to _is_engine_alive(), which on Windows uses
Popen.poll() instead.
"""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._engine_pid = 12345
client._watching_engine = True
client.attempt_db_reconnect = AsyncMock()
alive_process = MagicMock()
alive_process.poll.return_value = None # process still running
with (
patch.object(client, "_is_engine_alive", return_value=True) as mock_alive,
patch("litellm.proxy.utils.asyncio.sleep", AsyncMock(side_effect=asyncio.CancelledError())),
):
try:
await client._poll_engine_proc()
except asyncio.CancelledError:
pass
mock_alive.assert_called_once()
client.attempt_db_reconnect.assert_not_awaited()
@pytest.mark.asyncio
async def test_poll_engine_proc_triggers_reconnect_when_engine_dead(mock_proxy_logging):
"""_poll_engine_proc must trigger reconnect when the engine is no longer alive."""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._engine_pid = 12345
client._watching_engine = True
client._engine_confirmed_dead = False
client._reap_all_zombies = MagicMock(return_value=set())
client._cleanup_engine_watcher = MagicMock()
client.attempt_db_reconnect = AsyncMock(return_value=True)
with patch.object(client, "_is_engine_alive", return_value=False):
await client._poll_engine_proc()
client.attempt_db_reconnect.assert_awaited_once_with(
reason="engine_process_death",
force=True,
)
assert client._engine_confirmed_dead is True
def test_is_engine_alive_windows_living_process(mock_proxy_logging):
"""_is_engine_alive_windows returns True when Popen.poll() is None."""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._engine_pid = 99999
alive_process = MagicMock()
alive_process.poll.return_value = None
engine_mock = MagicMock()
engine_mock.process = alive_process
client.writer_db._original_prisma._engine = engine_mock
assert client._is_engine_alive_windows() is True
alive_process.poll.assert_called_once()
def test_is_engine_alive_windows_dead_process(mock_proxy_logging):
"""_is_engine_alive_windows returns False when Popen.poll() returns an exit code."""
client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging)
client._engine_pid = 99999
dead_process = MagicMock()
dead_process.poll.return_value = 0 # exited with code 0
engine_mock = MagicMock()
engine_mock.process = dead_process
client.writer_db._original_prisma._engine = engine_mock
assert client._is_engine_alive_windows() is False
dead_process.poll.assert_called_once()
def test_poll_engine_proc_does_not_call_os_kill(mock_proxy_logging):
"""_poll_engine_proc must not call os.kill — on Windows os.kill(pid, 0)
calls TerminateProcess and kills the engine instead of checking liveness."""
import inspect
import litellm.proxy.utils as proxy_utils
source = inspect.getsource(proxy_utils.PrismaClient._poll_engine_proc)
assert "os.kill" not in source, (
"_poll_engine_proc must not call os.kill directly; use _is_engine_alive() which is Windows-safe"
)