From 026beed0015e8c11755c0515d21c6cb96cb7bfd8 Mon Sep 17 00:00:00 2001 From: NTZ_ARM Date: Fri, 19 Jun 2026 17:50:05 +0700 Subject: [PATCH] fix(proxy): use Popen.poll() instead of os.kill(pid, 0) for engine liveness on Windows 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. --- litellm/proxy/utils.py | 30 ++- .../proxy/db/test_prisma_self_heal.py | 206 +++++++++++------- 2 files changed, 149 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 705690c3294..b573ce89dd5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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: diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index 265940e51ed..1cf1a0c930c 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -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" + )