mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-23 20:26:28 +00:00
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.
This commit is contained in:
Ryan Crabbe
@@ -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}"
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user