mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
Merge 8070294564 into c2c2a623c0
This commit is contained in:
commit
015f8775d4
2 changed files with 328 additions and 6 deletions
|
|
@ -8,6 +8,7 @@ import subprocess
|
|||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
|
@ -138,7 +139,93 @@ def clear_pid_record(path: Path | None = None) -> None:
|
|||
resolved_path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
_PROCESS_QUERY_LIMITED_INFORMATION: Final = 0x1000
|
||||
_STILL_ACTIVE: Final = 259
|
||||
_ERROR_ACCESS_DENIED: Final = 5
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _WindowsProcessApi:
|
||||
"""The kernel32 surface ``_windows_pid_exists`` needs, injected so the probe is testable off Windows."""
|
||||
|
||||
open_query_handle: Callable[[int], int]
|
||||
exit_code: Callable[[int], int | None]
|
||||
close_handle: Callable[[int], None]
|
||||
last_error: Callable[[], int]
|
||||
|
||||
|
||||
def _load_windows_process_api() -> _WindowsProcessApi: # pragma: no cover - kernel32 is Windows-only; CI is Linux
|
||||
import ctypes
|
||||
from ctypes import wintypes
|
||||
|
||||
kernel32: Final = ctypes.WinDLL("kernel32", use_last_error=True)
|
||||
kernel32.OpenProcess.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD)
|
||||
kernel32.OpenProcess.restype = wintypes.HANDLE
|
||||
kernel32.GetExitCodeProcess.argtypes = (wintypes.HANDLE, ctypes.POINTER(wintypes.DWORD))
|
||||
kernel32.GetExitCodeProcess.restype = wintypes.BOOL
|
||||
kernel32.CloseHandle.argtypes = (wintypes.HANDLE,)
|
||||
kernel32.CloseHandle.restype = wintypes.BOOL
|
||||
|
||||
def exit_code(handle: int) -> int | None:
|
||||
code = wintypes.DWORD()
|
||||
if not kernel32.GetExitCodeProcess(handle, ctypes.byref(code)):
|
||||
return None
|
||||
return code.value
|
||||
|
||||
return _WindowsProcessApi(
|
||||
open_query_handle=lambda pid: kernel32.OpenProcess(_PROCESS_QUERY_LIMITED_INFORMATION, False, pid),
|
||||
exit_code=exit_code,
|
||||
close_handle=kernel32.CloseHandle,
|
||||
last_error=ctypes.get_last_error,
|
||||
)
|
||||
|
||||
|
||||
def _windows_pid_exists(pid: int, api: _WindowsProcessApi | None = None) -> bool:
|
||||
"""Liveness check for Windows that does not signal the process.
|
||||
|
||||
Opens a query-only handle and asks for the exit code. ``OpenProcess``
|
||||
failing with ``ERROR_ACCESS_DENIED`` means the pid exists but is not ours
|
||||
to open, which is the case the ``PermissionError`` arm covers on POSIX.
|
||||
An unreadable exit code is reported as alive rather than guessed away.
|
||||
|
||||
Caveat kept deliberately: a process whose real exit code is 259 reads as
|
||||
alive, because ``GetExitCodeProcess`` reports ``STILL_ACTIVE`` (259) for a
|
||||
running process and cannot distinguish the two. That is the standard
|
||||
trade-off for this API and it is strictly better than the previous
|
||||
behaviour, which killed the process it was asked about.
|
||||
"""
|
||||
resolved: Final = api if api is not None else _load_windows_process_api()
|
||||
handle: Final = resolved.open_query_handle(pid)
|
||||
if not handle:
|
||||
return resolved.last_error() == _ERROR_ACCESS_DENIED
|
||||
try:
|
||||
code: Final = resolved.exit_code(handle)
|
||||
return code is None or code == _STILL_ACTIVE
|
||||
finally:
|
||||
resolved.close_handle(handle)
|
||||
|
||||
|
||||
def is_running(pid: int) -> bool:
|
||||
"""Report whether ``pid`` names a live process, without signalling it.
|
||||
|
||||
``os.kill(pid, 0)`` is the POSIX idiom and is a genuine no-op there, but it
|
||||
is not a probe on Windows. ``signal.CTRL_C_EVENT`` is 0, so ``os.kill(pid,
|
||||
0)`` *is* ``os.kill(pid, CTRL_C_EVENT)`` and reaches
|
||||
``GenerateConsoleCtrlEvent``. That API's second argument is a process
|
||||
*group*, and ``launch_proxy`` above starts the proxy through
|
||||
``subprocess.Popen`` with no ``CREATE_NEW_PROCESS_GROUP``, so the child
|
||||
shares this console's group: the Ctrl-C is delivered to the proxy, to the
|
||||
``lite`` process asking the question, and to anything else on the console.
|
||||
|
||||
Measured on Windows 11, Python 3.12.10. Probing a sleeping child returns
|
||||
``True`` with no exception raised; the child then exits with 3221225786
|
||||
(``0xC000013A``, ``STATUS_CONTROL_C_EXIT``) and a ``KeyboardInterrupt``
|
||||
arrives in the caller a moment later -- so the traceback does not point at
|
||||
the probe. ``KeyboardInterrupt`` is a ``BaseException``, so neither the
|
||||
``ProcessLookupError`` nor the ``PermissionError`` arm below ever sees it.
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
return _windows_pid_exists(pid)
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
except ProcessLookupError:
|
||||
|
|
@ -149,17 +236,31 @@ def is_running(pid: int) -> bool:
|
|||
|
||||
|
||||
def terminate(pid: int, grace_period: float = 5.0) -> None:
|
||||
"""Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed."""
|
||||
"""Terminate a process by pid, escalating from SIGTERM to SIGKILL if needed.
|
||||
|
||||
Both kills tolerate the same errors ``prisma_client`` already tolerates for
|
||||
this exact case -- "already dead or inaccessible". ``ProcessLookupError``
|
||||
alone is not enough: on Windows ``os.kill`` routes to ``TerminateProcess``,
|
||||
and a process that has already exited answers ``ERROR_ACCESS_DENIED``, so
|
||||
the call raises ``PermissionError`` and escapes the suppression.
|
||||
"""
|
||||
if not is_running(pid):
|
||||
return
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError, OSError):
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
deadline: Final = time.monotonic() + grace_period
|
||||
while time.monotonic() < deadline and is_running(pid):
|
||||
time.sleep(0.2)
|
||||
if is_running(pid):
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.kill(pid, signal.SIGKILL)
|
||||
# signal.SIGKILL does not exist on Windows. Referencing it raises
|
||||
# AttributeError, which contextlib.suppress(ProcessLookupError) does
|
||||
# not catch, so this escalation path crashed instead of hard-killing
|
||||
# a proxy that ignored SIGTERM. os.kill with any signal other than 0
|
||||
# or 1 routes to TerminateProcess on Windows, so SIGTERM is a real
|
||||
# kill there. litellm/proxy/db/prisma_client.py already resolves the
|
||||
# signal this way.
|
||||
with contextlib.suppress(ProcessLookupError, PermissionError, OSError):
|
||||
os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM))
|
||||
|
||||
|
||||
def stream_log(log_path: Path, stop_event: threading.Event) -> None:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
import os
|
||||
import signal
|
||||
import socket
|
||||
from typing import Optional
|
||||
from unittest.mock import patch
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -10,6 +14,8 @@ from litellm.proxy.client.cli.commands.autoroute.process import (
|
|||
PidRecord,
|
||||
ProcessLaunchError,
|
||||
UpError,
|
||||
_windows_pid_exists,
|
||||
_WindowsProcessApi,
|
||||
clear_pid_record,
|
||||
is_port_available,
|
||||
is_running,
|
||||
|
|
@ -17,6 +23,7 @@ from litellm.proxy.client.cli.commands.autoroute.process import (
|
|||
missing_proxy_runtime_modules,
|
||||
poll_liveliness,
|
||||
read_pid_record,
|
||||
terminate,
|
||||
write_pid_record,
|
||||
)
|
||||
|
||||
|
|
@ -34,6 +41,23 @@ class FakeResponse:
|
|||
self.status_code = status_code
|
||||
|
||||
|
||||
def _signals_sent(fake_kill: Mock) -> tuple[int, ...]:
|
||||
"""Signal numbers ``os.kill`` was called with, in call order."""
|
||||
return tuple(call.args[1] for call in fake_kill.call_args_list)
|
||||
|
||||
|
||||
def _windows_api(handle: int, exit_code: int | None = None, last_error: int = 0) -> tuple[_WindowsProcessApi, Mock]:
|
||||
"""A kernel32 stand-in, plus the Mock that recorded ``CloseHandle``."""
|
||||
close_handle: Final = Mock()
|
||||
api: Final = _WindowsProcessApi(
|
||||
open_query_handle=Mock(return_value=handle),
|
||||
exit_code=Mock(return_value=exit_code),
|
||||
close_handle=close_handle,
|
||||
last_error=Mock(return_value=last_error),
|
||||
)
|
||||
return api, close_handle
|
||||
|
||||
|
||||
class TestIsPortAvailable:
|
||||
def test_true_for_a_free_port(self):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
|
|
@ -111,6 +135,10 @@ class TestIsRunning:
|
|||
def test_huge_unlikely_pid_is_not_running(self):
|
||||
assert is_running(2**30) is False
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform == "win32",
|
||||
reason="POSIX branch only; win32 answers with a query-only process handle, not os.kill",
|
||||
)
|
||||
def test_permission_error_from_kill_is_treated_as_running(self, monkeypatch):
|
||||
def fake_kill(pid: int, sig: int) -> None:
|
||||
raise PermissionError("not permitted to signal this pid")
|
||||
|
|
@ -120,6 +148,199 @@ class TestIsRunning:
|
|||
assert is_running(999) is True
|
||||
|
||||
|
||||
class TestIsRunningDoesNotSignal:
|
||||
"""``os.kill(pid, 0)`` is not a liveness probe on Windows.
|
||||
|
||||
``signal.CTRL_C_EVENT`` is 0, so the call reaches
|
||||
``GenerateConsoleCtrlEvent``, whose second argument is a process *group*.
|
||||
``launch_proxy`` spawns the proxy without ``CREATE_NEW_PROCESS_GROUP``, so
|
||||
the child shares this console's group and the Ctrl-C hits the child, the
|
||||
caller, and anything else attached to the console.
|
||||
|
||||
Measured on Windows 11, Python 3.12.10: probing a sleeping child returns
|
||||
``True`` with no exception, the child then exits 3221225786
|
||||
(``0xC000013A`` / ``STATUS_CONTROL_C_EXIT``), and a ``KeyboardInterrupt``
|
||||
lands in the caller afterwards -- a ``BaseException``, so the handlers in
|
||||
``is_running`` never see it.
|
||||
"""
|
||||
|
||||
def test_probing_a_child_does_not_terminate_it(self):
|
||||
"""The behavioural check, with nothing mocked."""
|
||||
child: Final = subprocess.Popen((sys.executable, "-c", "import time; time.sleep(30)"))
|
||||
try:
|
||||
assert is_running(child.pid) is True
|
||||
time.sleep(1.0)
|
||||
assert child.poll() is None, (
|
||||
f"the liveness probe terminated the process it was asked about (exit code {child.returncode})"
|
||||
)
|
||||
assert is_running(child.pid) is True
|
||||
finally:
|
||||
child.kill()
|
||||
child.wait(timeout=10)
|
||||
|
||||
def test_probing_our_own_pid_leaves_this_process_alone(self):
|
||||
"""What ``TestIsRunning.test_current_process_is_running`` above relies on.
|
||||
|
||||
The assertion there passes either way; it is the interpreter surviving
|
||||
the next statement that separates a probe from a signal.
|
||||
"""
|
||||
assert is_running(os.getpid()) is True
|
||||
time.sleep(0.5)
|
||||
assert is_running(os.getpid()) is True
|
||||
|
||||
def test_win32_answers_without_calling_os_kill(self, monkeypatch):
|
||||
fake_kill: Final = Mock()
|
||||
monkeypatch.setattr(process_module.sys, "platform", "win32")
|
||||
monkeypatch.setattr(process_module.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(process_module, "_windows_pid_exists", lambda pid: True)
|
||||
|
||||
assert is_running(1234) is True
|
||||
assert _signals_sent(fake_kill) == (), "os.kill reached on win32; signal 0 is CTRL_C_EVENT there"
|
||||
|
||||
|
||||
class TestWindowsPidExists:
|
||||
"""The Windows probe, driven through an injected kernel32 so it runs on any platform.
|
||||
|
||||
``OpenProcess`` / ``GetExitCodeProcess`` only exist on Windows, so CI could
|
||||
never execute this decision logic before. The four calls the probe needs
|
||||
now arrive as a ``_WindowsProcessApi``; only the loader that builds the real
|
||||
one out of ``ctypes`` stays Windows-only.
|
||||
"""
|
||||
|
||||
_ERROR_ACCESS_DENIED: Final = 5
|
||||
_ERROR_INVALID_PARAMETER: Final = 87
|
||||
_STILL_ACTIVE: Final = 259
|
||||
|
||||
def test_access_denied_on_open_means_the_pid_exists(self):
|
||||
"""A pid we may not open is still a live pid; POSIX answers the same way via PermissionError."""
|
||||
api, close_handle = _windows_api(handle=0, last_error=self._ERROR_ACCESS_DENIED)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is True
|
||||
close_handle.assert_not_called()
|
||||
|
||||
def test_any_other_open_failure_means_the_pid_is_gone(self):
|
||||
api, close_handle = _windows_api(handle=0, last_error=self._ERROR_INVALID_PARAMETER)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is False
|
||||
close_handle.assert_not_called()
|
||||
|
||||
def test_still_active_means_running(self):
|
||||
api, close_handle = _windows_api(handle=99, exit_code=self._STILL_ACTIVE)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is True
|
||||
close_handle.assert_called_once_with(99)
|
||||
|
||||
def test_a_real_exit_code_means_the_process_finished(self):
|
||||
api, close_handle = _windows_api(handle=99, exit_code=0)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is False
|
||||
close_handle.assert_called_once_with(99)
|
||||
|
||||
def test_an_unreadable_exit_code_is_reported_as_running(self):
|
||||
"""``GetExitCodeProcess`` failing is not evidence the process died, so do not claim it did."""
|
||||
api, close_handle = _windows_api(handle=99, exit_code=None)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is True
|
||||
close_handle.assert_called_once_with(99)
|
||||
|
||||
def test_the_handle_is_released_even_when_the_query_raises(self):
|
||||
"""A leaked query handle would keep the exited process object alive for the life of the CLI."""
|
||||
close_handle: Final = Mock()
|
||||
api: Final = _WindowsProcessApi(
|
||||
open_query_handle=Mock(return_value=99),
|
||||
exit_code=Mock(side_effect=OSError(22, "Invalid argument")),
|
||||
close_handle=close_handle,
|
||||
last_error=Mock(return_value=0),
|
||||
)
|
||||
|
||||
with pytest.raises(OSError, match="Invalid argument"):
|
||||
_windows_pid_exists(4321, api)
|
||||
|
||||
close_handle.assert_called_once_with(99)
|
||||
|
||||
def test_the_probe_asks_about_the_pid_it_was_given(self):
|
||||
api, _ = _windows_api(handle=99, exit_code=self._STILL_ACTIVE)
|
||||
|
||||
_windows_pid_exists(4321, api)
|
||||
|
||||
api.open_query_handle.assert_called_once_with(4321)
|
||||
|
||||
def test_the_probe_never_signals_the_process(self, monkeypatch):
|
||||
"""The bug this replaces: ``os.kill(pid, 0)`` is ``CTRL_C_EVENT`` on Windows."""
|
||||
fake_kill: Final = Mock()
|
||||
monkeypatch.setattr(process_module.os, "kill", fake_kill)
|
||||
api, _ = _windows_api(handle=99, exit_code=self._STILL_ACTIVE)
|
||||
|
||||
assert _windows_pid_exists(4321, api) is True
|
||||
assert _signals_sent(fake_kill) == ()
|
||||
|
||||
|
||||
class TestTerminateHardKillSignal:
|
||||
def test_escalation_resolves_a_signal_when_sigkill_is_absent(self, monkeypatch):
|
||||
"""``signal.SIGKILL`` does not exist on Windows.
|
||||
|
||||
Referencing it raises ``AttributeError``, and ``terminate`` wraps the
|
||||
call in ``contextlib.suppress(ProcessLookupError)``, which does not
|
||||
catch that -- so the branch that exists for a proxy ignoring SIGTERM
|
||||
crashed instead of hard-killing it. ``os.kill`` with any signal other
|
||||
than 0 or 1 routes to ``TerminateProcess`` on Windows, so SIGTERM is a
|
||||
real kill there.
|
||||
"""
|
||||
fake_kill: Final = Mock()
|
||||
monkeypatch.setattr(process_module.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(process_module, "is_running", lambda pid: True)
|
||||
monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None)
|
||||
monkeypatch.delattr(process_module.signal, "SIGKILL", raising=False)
|
||||
|
||||
terminate(4242, grace_period=0.0)
|
||||
|
||||
assert _signals_sent(fake_kill) == (signal.SIGTERM, signal.SIGTERM)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not hasattr(signal, "SIGKILL"),
|
||||
reason="POSIX escalation path; SIGKILL is absent on this platform",
|
||||
)
|
||||
def test_escalation_still_uses_sigkill_where_it_exists(self, monkeypatch):
|
||||
fake_kill: Final = Mock()
|
||||
monkeypatch.setattr(process_module.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(process_module, "is_running", lambda pid: True)
|
||||
monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None)
|
||||
|
||||
terminate(4242, grace_period=0.0)
|
||||
|
||||
assert _signals_sent(fake_kill) == (signal.SIGTERM, signal.SIGKILL)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raised",
|
||||
(
|
||||
PermissionError(5, "Access is denied"),
|
||||
OSError(22, "Invalid argument"),
|
||||
ProcessLookupError(3, "No such process"),
|
||||
),
|
||||
ids=("permission-denied", "oserror", "already-gone"),
|
||||
)
|
||||
def test_a_kill_that_cannot_land_does_not_escape(self, monkeypatch, raised):
|
||||
"""A process that exited between the probe and the kill must not crash ``down``.
|
||||
|
||||
On Windows ``os.kill`` routes to ``TerminateProcess``, and a process
|
||||
that has already exited answers ``ERROR_ACCESS_DENIED`` -- so the call
|
||||
raises ``PermissionError``, not ``ProcessLookupError``. Measured with
|
||||
the real CLI on Windows 11: ``lite autoroute down`` ended in
|
||||
``PermissionError: [WinError 5]`` out of ``terminate``.
|
||||
``litellm/proxy/db/prisma_client.py`` already tolerates all three for
|
||||
this same "already dead or inaccessible" case.
|
||||
"""
|
||||
|
||||
fake_kill: Final = Mock(side_effect=raised)
|
||||
|
||||
monkeypatch.setattr(process_module.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(process_module, "is_running", lambda pid: True)
|
||||
monkeypatch.setattr(process_module.time, "sleep", lambda seconds: None)
|
||||
|
||||
terminate(4242, grace_period=0.0)
|
||||
|
||||
assert _signals_sent(fake_kill) == (signal.SIGTERM, getattr(signal, "SIGKILL", signal.SIGTERM))
|
||||
|
||||
class TestPollLiveliness:
|
||||
def test_succeeds_when_health_check_returns_200_quickly(self, monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(process_module.requests, "get", lambda url, timeout: FakeResponse(200))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue