mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
refactor: remove periodic dead PID cleanup, trim redundant tests
Startup wipe + graceful shutdown cleanup are sufficient. Remove hourly mark_dead_pids scan, its helpers, and redundant test cases.
This commit is contained in:
Ryan Crabbe
Sameer Kankute
parent
d10b94ecb1
commit
02d5b7bde5
3 changed files with 11 additions and 272 deletions
|
|
@ -9,14 +9,10 @@ from __future__ import annotations
|
|||
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
from typing import Optional, Set
|
||||
from typing import Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
_PID_PATTERN = re.compile(r"_(\d+)\.db$")
|
||||
|
||||
|
||||
def _get_multiproc_dir() -> Optional[str]:
|
||||
"""Return the PROMETHEUS_MULTIPROC_DIR env var value, or None."""
|
||||
return os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get(
|
||||
|
|
@ -24,31 +20,6 @@ def _get_multiproc_dir() -> Optional[str]:
|
|||
)
|
||||
|
||||
|
||||
def _is_pid_alive(pid: int) -> bool:
|
||||
"""Check if a process is alive using signal 0 (conservative: unknown = alive)."""
|
||||
try:
|
||||
os.kill(pid, 0)
|
||||
return True
|
||||
except ProcessLookupError:
|
||||
return False
|
||||
except (PermissionError, OSError):
|
||||
return True
|
||||
|
||||
|
||||
def _extract_pids_from_dir(directory: str) -> Set[int]:
|
||||
"""Scan .db filenames in a directory and extract PIDs (e.g. counter_1234.db -> 1234)."""
|
||||
pids: Set[int] = set()
|
||||
try:
|
||||
for filename in os.listdir(directory):
|
||||
if not filename.endswith(".db"):
|
||||
continue
|
||||
match = _PID_PATTERN.search(filename)
|
||||
if match:
|
||||
pids.add(int(match.group(1)))
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
return pids
|
||||
|
||||
|
||||
def wipe_directory(directory: str) -> None:
|
||||
"""Delete all .db files in the directory. Called once before workers fork."""
|
||||
|
|
@ -84,35 +55,3 @@ def cleanup_own_pid_files() -> None:
|
|||
verbose_proxy_logger.warning(
|
||||
f"Failed to mark worker PID {pid} as dead: {e}"
|
||||
)
|
||||
|
||||
|
||||
def mark_dead_pids(skip_pid: Optional[int] = None) -> None:
|
||||
"""Scan for dead PIDs and call mark_process_dead() for each. Skips current process by default."""
|
||||
directory = _get_multiproc_dir()
|
||||
if not directory or not os.path.isdir(directory):
|
||||
return
|
||||
|
||||
if skip_pid is None:
|
||||
skip_pid = os.getpid()
|
||||
|
||||
pids = _extract_pids_from_dir(directory)
|
||||
|
||||
from prometheus_client import multiprocess
|
||||
|
||||
dead_pids = []
|
||||
for pid in pids:
|
||||
if pid == skip_pid:
|
||||
continue
|
||||
if not _is_pid_alive(pid):
|
||||
try:
|
||||
multiprocess.mark_process_dead(pid)
|
||||
dead_pids.append(pid)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to mark PID {pid} as dead: {e}"
|
||||
)
|
||||
|
||||
if dead_pids:
|
||||
verbose_proxy_logger.info(
|
||||
f"Prometheus cleanup: marked {len(dead_pids)} dead PIDs: {dead_pids}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -901,12 +901,6 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
|
|||
## Initialize shared aiohttp session for connection reuse
|
||||
shared_aiohttp_session = await _initialize_shared_aiohttp_session()
|
||||
|
||||
## Start periodic prometheus multiproc directory cleanup
|
||||
from litellm.proxy.prometheus_cleanup import _get_multiproc_dir
|
||||
|
||||
if _get_multiproc_dir():
|
||||
asyncio.create_task(_periodic_prometheus_cleanup())
|
||||
|
||||
# End of startup event
|
||||
yield
|
||||
|
||||
|
|
@ -2061,21 +2055,6 @@ def _schedule_background_health_check_db_save(
|
|||
)
|
||||
|
||||
|
||||
async def _periodic_prometheus_cleanup():
|
||||
"""
|
||||
Periodically mark dead worker PIDs in the prometheus multiproc directory.
|
||||
"""
|
||||
from litellm.proxy.prometheus_cleanup import mark_dead_pids
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(3600) # 1 hour
|
||||
try:
|
||||
mark_dead_pids()
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Error in periodic prometheus cleanup: {e}"
|
||||
)
|
||||
|
||||
|
||||
async def _run_background_health_check():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,74 +6,17 @@ ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir.
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.prometheus_cleanup import (
|
||||
_extract_pids_from_dir,
|
||||
_is_pid_alive,
|
||||
cleanup_own_pid_files,
|
||||
mark_dead_pids,
|
||||
wipe_directory,
|
||||
)
|
||||
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
|
||||
|
||||
|
||||
class TestExtractPidsFromDir:
|
||||
def test_counter_files(self, tmp_path):
|
||||
(tmp_path / "counter_1234.db").touch()
|
||||
(tmp_path / "counter_5678.db").touch()
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == {1234, 5678}
|
||||
|
||||
def test_histogram_files(self, tmp_path):
|
||||
(tmp_path / "histogram_1234.db").touch()
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == {1234}
|
||||
|
||||
def test_gauge_files(self, tmp_path):
|
||||
(tmp_path / "gauge_livesum_1234.db").touch()
|
||||
(tmp_path / "gauge_liveall_5678.db").touch()
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == {1234, 5678}
|
||||
|
||||
def test_multiple_pids(self, tmp_path):
|
||||
(tmp_path / "counter_100.db").touch()
|
||||
(tmp_path / "histogram_200.db").touch()
|
||||
(tmp_path / "gauge_livesum_300.db").touch()
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == {100, 200, 300}
|
||||
|
||||
def test_non_db_files_ignored(self, tmp_path):
|
||||
(tmp_path / "counter_1234.db").touch()
|
||||
(tmp_path / "readme.txt").touch()
|
||||
(tmp_path / "data.json").touch()
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == {1234}
|
||||
|
||||
def test_empty_directory(self, tmp_path):
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == set()
|
||||
|
||||
def test_nonexistent_directory(self):
|
||||
assert _extract_pids_from_dir("/nonexistent/path/abc123") == set()
|
||||
|
||||
def test_malformed_filenames(self, tmp_path):
|
||||
(tmp_path / "counter_.db").touch() # no PID
|
||||
(tmp_path / "random.db").touch() # no underscore+PID pattern
|
||||
(tmp_path / "counter_abc.db").touch() # non-numeric PID
|
||||
assert _extract_pids_from_dir(str(tmp_path)) == set()
|
||||
|
||||
|
||||
class TestIsPidAlive:
|
||||
def test_own_pid_is_alive(self):
|
||||
assert _is_pid_alive(os.getpid()) is True
|
||||
|
||||
def test_dead_pid(self):
|
||||
# A very high PID is almost certainly dead
|
||||
assert _is_pid_alive(4_000_000) is False
|
||||
|
||||
def test_permission_error_treated_as_alive(self):
|
||||
with patch("os.kill", side_effect=PermissionError):
|
||||
assert _is_pid_alive(99999) is True
|
||||
|
||||
|
||||
class TestWipeDirectory:
|
||||
def test_deletes_all_db_files(self, tmp_path):
|
||||
(tmp_path / "counter_1234.db").touch()
|
||||
|
|
@ -82,19 +25,6 @@ class TestWipeDirectory:
|
|||
wipe_directory(str(tmp_path))
|
||||
assert not list(tmp_path.glob("*.db"))
|
||||
|
||||
def test_preserves_non_db_files(self, tmp_path):
|
||||
(tmp_path / "counter_1234.db").touch()
|
||||
(tmp_path / "readme.txt").touch()
|
||||
(tmp_path / "config.json").touch()
|
||||
wipe_directory(str(tmp_path))
|
||||
assert not list(tmp_path.glob("*.db"))
|
||||
assert (tmp_path / "readme.txt").exists()
|
||||
assert (tmp_path / "config.json").exists()
|
||||
|
||||
def test_empty_directory(self, tmp_path):
|
||||
wipe_directory(str(tmp_path))
|
||||
assert not list(tmp_path.glob("*.db"))
|
||||
|
||||
|
||||
class TestCleanupOwnPidFiles:
|
||||
def test_calls_mark_process_dead_for_own_pid(self, tmp_path):
|
||||
|
|
@ -120,99 +50,7 @@ class TestCleanupOwnPidFiles:
|
|||
mock_mark_dead.assert_not_called()
|
||||
|
||||
|
||||
class TestMarkDeadPids:
|
||||
def test_calls_mark_process_dead_for_dead_pids(self, tmp_path):
|
||||
"""mark_dead_pids should call mark_process_dead() for dead PIDs."""
|
||||
dead_pid = 4_000_000
|
||||
(tmp_path / f"counter_{dead_pid}.db").touch()
|
||||
(tmp_path / f"gauge_livesum_{dead_pid}.db").touch()
|
||||
|
||||
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
||||
with patch(
|
||||
"prometheus_client.multiprocess.mark_process_dead"
|
||||
) as mock_mark_dead:
|
||||
mark_dead_pids()
|
||||
mock_mark_dead.assert_called_once_with(dead_pid)
|
||||
|
||||
def test_skips_own_pid(self, tmp_path):
|
||||
"""Should not call mark_process_dead for the current process."""
|
||||
own_pid = os.getpid()
|
||||
(tmp_path / f"counter_{own_pid}.db").touch()
|
||||
|
||||
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
||||
with patch(
|
||||
"prometheus_client.multiprocess.mark_process_dead"
|
||||
) as mock_mark_dead:
|
||||
mark_dead_pids()
|
||||
mock_mark_dead.assert_not_called()
|
||||
|
||||
def test_skips_alive_pids(self, tmp_path):
|
||||
"""Should not call mark_process_dead for alive PIDs."""
|
||||
alive_pid = 99999
|
||||
(tmp_path / f"counter_{alive_pid}.db").touch()
|
||||
|
||||
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
||||
with patch(
|
||||
"litellm.proxy.prometheus_cleanup._is_pid_alive",
|
||||
return_value=True,
|
||||
):
|
||||
with patch(
|
||||
"prometheus_client.multiprocess.mark_process_dead"
|
||||
) as mock_mark_dead:
|
||||
mark_dead_pids()
|
||||
mock_mark_dead.assert_not_called()
|
||||
|
||||
def test_handles_mixed_alive_and_dead(self, tmp_path):
|
||||
"""Should only call mark_process_dead for dead PIDs."""
|
||||
dead_pid = 4_000_000
|
||||
own_pid = os.getpid()
|
||||
(tmp_path / f"counter_{dead_pid}.db").touch()
|
||||
(tmp_path / f"counter_{own_pid}.db").touch()
|
||||
|
||||
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
||||
with patch(
|
||||
"prometheus_client.multiprocess.mark_process_dead"
|
||||
) as mock_mark_dead:
|
||||
mark_dead_pids()
|
||||
mock_mark_dead.assert_called_once_with(dead_pid)
|
||||
|
||||
def test_noop_when_not_configured(self, tmp_path):
|
||||
"""Should do nothing when PROMETHEUS_MULTIPROC_DIR is not set."""
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
os.environ.pop("prometheus_multiproc_dir", None)
|
||||
with patch(
|
||||
"prometheus_client.multiprocess.mark_process_dead"
|
||||
) as mock_mark_dead:
|
||||
mark_dead_pids()
|
||||
mock_mark_dead.assert_not_called()
|
||||
|
||||
|
||||
class TestMaybeSetupPrometheusMultiprocDir:
|
||||
def test_auto_creates_dir_when_prometheus_configured(self):
|
||||
"""When multiple workers + prometheus callback, auto-creates temp dir."""
|
||||
litellm_settings = {"callbacks": ["prometheus"]}
|
||||
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
os.environ.pop("prometheus_multiproc_dir", None)
|
||||
|
||||
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=4,
|
||||
litellm_settings=litellm_settings,
|
||||
)
|
||||
|
||||
result_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR")
|
||||
assert result_dir is not None
|
||||
assert os.path.isdir(result_dir)
|
||||
expected = os.path.join(
|
||||
tempfile.gettempdir(), "litellm_prometheus_multiproc"
|
||||
)
|
||||
assert result_dir == expected
|
||||
|
||||
# Cleanup
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
|
||||
def test_respects_existing_env_var(self, tmp_path):
|
||||
"""When PROMETHEUS_MULTIPROC_DIR is already set, don't override it."""
|
||||
custom_dir = str(tmp_path / "custom_prom")
|
||||
|
|
@ -227,28 +65,6 @@ class TestMaybeSetupPrometheusMultiprocDir:
|
|||
assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == custom_dir
|
||||
assert os.path.isdir(custom_dir)
|
||||
|
||||
def test_wipes_stale_files_on_setup(self, tmp_path):
|
||||
"""Should wipe existing .db files from a previous run."""
|
||||
custom_dir = str(tmp_path / "prom_dir")
|
||||
os.makedirs(custom_dir)
|
||||
# Simulate stale files from a previous run
|
||||
for name in ["counter_9999.db", "histogram_9999.db", "gauge_livesum_9999.db"]:
|
||||
with open(os.path.join(custom_dir, name), "w") as f:
|
||||
f.write("stale")
|
||||
|
||||
litellm_settings = {"callbacks": ["prometheus"]}
|
||||
|
||||
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": custom_dir}):
|
||||
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=4,
|
||||
litellm_settings=litellm_settings,
|
||||
)
|
||||
|
||||
# All .db files should be wiped
|
||||
import glob
|
||||
remaining = glob.glob(os.path.join(custom_dir, "*.db"))
|
||||
assert remaining == []
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_workers, litellm_settings",
|
||||
[
|
||||
|
|
@ -269,16 +85,21 @@ class TestMaybeSetupPrometheusMultiprocDir:
|
|||
|
||||
assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") is None
|
||||
|
||||
def test_prometheus_in_success_callback(self):
|
||||
"""Prometheus in success_callback should also trigger setup."""
|
||||
litellm_settings = {"success_callback": ["prometheus"]}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_settings",
|
||||
[
|
||||
{"callbacks": ["prometheus"]},
|
||||
{"success_callback": ["prometheus"]},
|
||||
],
|
||||
)
|
||||
def test_auto_creates_dir_when_prometheus_configured(self, litellm_settings):
|
||||
"""When multiple workers + prometheus callback, auto-creates temp dir."""
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
os.environ.pop("prometheus_multiproc_dir", None)
|
||||
|
||||
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=2,
|
||||
num_workers=4,
|
||||
litellm_settings=litellm_settings,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue