From 9ea100a9f73654f087e60c6e408a76e3559566b5 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Wed, 25 Feb 2026 17:20:41 -0800 Subject: [PATCH] feat: add prometheus multiprocess directory cleanup Adds cleanup utilities for PROMETHEUS_MULTIPROC_DIR to prevent unbounded RAM/disk growth from stale .db files in multi-worker setups. Three-part lifecycle aligned with upstream prometheus_client docs: 1. Startup: wipe entire directory before workers fork (clean slate) 2. Shutdown: mark_process_dead() for own PID (removes gauge_live* only) 3. Periodic (hourly): scan for dead PIDs and call mark_process_dead() Counter/histogram files are never individually deleted at runtime to avoid partial counter resets that cause false spikes in rate()/increase(). Also auto-creates PROMETHEUS_MULTIPROC_DIR when prometheus callback is configured with multiple workers and the env var is not already set. --- litellm/proxy/prometheus_cleanup.py | 168 +++++++++ litellm/proxy/proxy_cli.py | 54 +++ litellm/proxy/proxy_server.py | 34 ++ .../proxy/test_prometheus_cleanup.py | 339 ++++++++++++++++++ 4 files changed, 595 insertions(+) create mode 100644 litellm/proxy/prometheus_cleanup.py create mode 100644 tests/test_litellm/proxy/test_prometheus_cleanup.py diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py new file mode 100644 index 0000000000..34c48de74f --- /dev/null +++ b/litellm/proxy/prometheus_cleanup.py @@ -0,0 +1,168 @@ +""" +Prometheus multiprocess directory cleanup utilities. + +When running with multiple workers and PROMETHEUS_MULTIPROC_DIR set, +each worker creates memory-mapped .db files (e.g., counter_1234.db). +When workers die or restart, gauge_live* files for dead PIDs must be +cleaned up via mark_process_dead(). Counter and histogram files are +kept since they contain cumulative data needed for correct aggregation. +""" + +from __future__ import annotations + +import glob +import os +import re +from typing import Optional, Set + +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( + "prometheus_multiproc_dir" + ) + + +def _is_pid_alive(pid: int) -> bool: + """ + Check if a process with the given PID is alive. + + Uses os.kill(pid, 0) which doesn't send a signal but checks existence. + - ProcessLookupError: process does not exist (dead) + - PermissionError: process exists but we can't signal it (alive, conservative) + - OSError: other error, treat as alive (conservative) + """ + try: + os.kill(pid, 0) + return True + except ProcessLookupError: + return False + except PermissionError: + # Process exists but we don't have permission to signal it + return True + except OSError: + # Conservative: treat unknown errors as alive + return True + + +def _extract_pids_from_dir(directory: str) -> Set[int]: + """ + Scan .db filenames in a directory and extract PIDs. + + Prometheus client creates files like: + - counter_1234.db + - histogram_1234.db + - gauge_livesum_1234.db + - gauge_liveall_1234.db + + Returns a set of integer PIDs found. + """ + 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 prometheus multiproc directory. + + Called once in the master process before workers fork. Per the + prometheus_client docs: "This directory must be wiped between + process runs (before startup is recommended)." + + Any .db files present at this point are stale from a previous run. + """ + files = glob.glob(os.path.join(directory, "*.db")) + for filepath in files: + try: + os.remove(filepath) + except OSError as e: + verbose_proxy_logger.warning( + f"Failed to delete stale prometheus file {filepath}: {e}" + ) + if files: + verbose_proxy_logger.info( + f"Prometheus cleanup: wiped {len(files)} stale .db files from {directory}" + ) + + +def cleanup_own_pid_files() -> None: + """ + Mark the current process as dead for prometheus multiproc cleanup. + + Called during per-worker shutdown. Uses mark_process_dead() which + only removes gauge_live* files — counter and histogram files are + preserved since they contain cumulative data needed for correct + aggregation until the directory is wiped on next startup. + """ + directory = _get_multiproc_dir() + if not directory or not os.path.isdir(directory): + return + + from prometheus_client import multiprocess + + pid = os.getpid() + try: + multiprocess.mark_process_dead(pid) + verbose_proxy_logger.info( + f"Prometheus cleanup: marked worker PID {pid} as dead" + ) + except Exception as e: + 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 the prometheus multiproc directory and call mark_process_dead() + for PIDs that no longer exist. + + Uses prometheus_client.multiprocess.mark_process_dead() which only + removes gauge_live* files — counter and histogram files are preserved + since they contain cumulative data needed for correct aggregation. + + Args: + skip_pid: PID to skip (typically os.getpid()). If None, skips + the current process's PID. + """ + 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}" + ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e91447af89..1be4931aee 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -314,6 +314,54 @@ class ProxyInitializationHelpers: return None # Let uvicorn choose the default loop on Windows return "uvloop" + @staticmethod + def _maybe_setup_prometheus_multiproc_dir( + num_workers: int, + litellm_settings: Optional[dict], + ) -> None: + """ + Auto-create PROMETHEUS_MULTIPROC_DIR when running with multiple workers + and prometheus is configured as a callback. + + If the env var is already set by the user, just ensure the directory exists. + Otherwise, create a temp directory and set the env var. + """ + import tempfile + + if num_workers <= 1 or litellm_settings is None: + return + + # Check if prometheus is in any callback list + callbacks = litellm_settings.get("callbacks") or [] + success_callbacks = litellm_settings.get("success_callback") or [] + failure_callbacks = litellm_settings.get("failure_callback") or [] + all_callbacks = callbacks + success_callbacks + failure_callbacks + if "prometheus" not in all_callbacks: + return + + from litellm.proxy.prometheus_cleanup import wipe_directory + + existing_dir = os.environ.get( + "PROMETHEUS_MULTIPROC_DIR" + ) or os.environ.get("prometheus_multiproc_dir") + if existing_dir: + os.makedirs(existing_dir, exist_ok=True) + wipe_directory(existing_dir) + print( # noqa + f"LiteLLM: Using existing PROMETHEUS_MULTIPROC_DIR={existing_dir}" + ) + return + + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.makedirs(multiproc_dir, exist_ok=True) + wipe_directory(multiproc_dir) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + print( # noqa + f"LiteLLM: Auto-created PROMETHEUS_MULTIPROC_DIR={multiproc_dir}" + ) + @click.command() @click.option( @@ -819,6 +867,12 @@ def run_server( # noqa: PLR0915 # DO NOT DELETE - enables global variables to work across files from litellm.proxy.proxy_server import app # noqa + # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=num_workers, + litellm_settings=litellm_settings if config else None, + ) + # --- SEPARATE HEALTH APP LOGIC --- # To run the health app separately, use: # uvicorn litellm.proxy.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=4001 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f0b1e66818..c24942e135 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -709,6 +709,14 @@ async def proxy_shutdown_event(): # [DO NOT BLOCK shutdown events for this] pass + # Clean up this worker's prometheus multiproc .db files + try: + from litellm.proxy.prometheus_cleanup import cleanup_own_pid_files + + cleanup_own_pid_files() + except Exception as e: + verbose_proxy_logger.warning(f"Error cleaning up prometheus files: {e}") + ## RESET CUSTOM VARIABLES ## cleanup_router_config_variables() @@ -891,6 +899,12 @@ 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 @@ -2045,6 +2059,26 @@ def _schedule_background_health_check_db_save( ) +async def _periodic_prometheus_cleanup(): + """ + Periodically mark dead worker PIDs in the prometheus multiproc directory. + + Uses mark_process_dead() which only removes gauge_live* files, preserving + counter/histogram data for correct aggregation. First run is 1 hour after + startup (startup wipe handles stale files), then every hour thereafter. + """ + 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(): """ Periodically run health checks in the background on the endpoints. diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py new file mode 100644 index 0000000000..9f69402b5b --- /dev/null +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -0,0 +1,339 @@ +""" +Tests for litellm.proxy.prometheus_cleanup module and +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() + (tmp_path / "histogram_5678.db").touch() + (tmp_path / "gauge_livesum_9999.db").touch() + 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): + """Should call mark_process_dead with current PID on shutdown.""" + pid = os.getpid() + + with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}): + with patch( + "prometheus_client.multiprocess.mark_process_dead" + ) as mock_mark_dead: + cleanup_own_pid_files() + mock_mark_dead.assert_called_once_with(pid) + + def test_noop_when_not_configured(self, tmp_path): + """Should not call mark_process_dead when env var 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: + cleanup_own_pid_files() + 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_handles_malformed_filenames(self, tmp_path): + """Malformed filenames should be ignored, dead PIDs still cleaned.""" + dead_pid = 4_000_000 + (tmp_path / "counter_.db").touch() + (tmp_path / "random.db").touch() + (tmp_path / f"counter_{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_permission_error_treated_as_alive(self, tmp_path): + """PermissionError from os.kill means process is alive, skip it.""" + target_pid = 99999 + (tmp_path / f"counter_{target_pid}.db").touch() + + with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}): + with patch("os.kill", side_effect=PermissionError): + with patch( + "prometheus_client.multiprocess.mark_process_dead" + ) as mock_mark_dead: + mark_dead_pids() + mock_mark_dead.assert_not_called() + + 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") + 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, + ) + + 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 == [] + + def test_noop_for_single_worker(self): + """Single worker doesn't need multiproc 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=1, + litellm_settings=litellm_settings, + ) + + assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") is None + + def test_noop_without_prometheus_callback(self): + """No prometheus callback = no setup needed.""" + litellm_settings = {"callbacks": ["langfuse"]} + + 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, + ) + + assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") is None + + def test_noop_with_none_litellm_settings(self): + """None litellm_settings = no setup needed.""" + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + + ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=4, + litellm_settings=None, + ) + + 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"]} + + 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, + litellm_settings=litellm_settings, + ) + + result_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + assert result_dir is not None + assert os.path.isdir(result_dir) + + # Cleanup + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)