mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-14 05:06:33 +00:00
Ryan Crabbe
4fa6742b01
When a gunicorn worker exits (e.g. from max_requests recycling), its per-process prometheus .db files remain on disk. For gauges using livesum/liveall mode, this means the dead worker's last-known values persist as if the process were still alive. Wire gunicorn's child_exit hook to call mark_process_dead() so live-tracking gauges accurately reflect only running workers.
114 lines
4.1 KiB
Python
114 lines
4.1 KiB
Python
"""
|
|
Tests for litellm.proxy.prometheus_cleanup.wipe_directory and
|
|
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.prometheus_cleanup import mark_worker_exit, wipe_directory
|
|
from litellm.proxy.proxy_cli import ProxyInitializationHelpers
|
|
|
|
|
|
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"))
|
|
|
|
|
|
class TestMarkWorkerExit:
|
|
def test_calls_mark_process_dead_when_env_set(self, tmp_path):
|
|
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
|
with patch(
|
|
"prometheus_client.multiprocess.mark_process_dead"
|
|
) as mock_mark:
|
|
mark_worker_exit(12345)
|
|
mock_mark.assert_called_once_with(12345)
|
|
|
|
def test_noop_when_env_not_set(self):
|
|
with patch.dict(os.environ, {}, clear=False):
|
|
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
|
with patch(
|
|
"prometheus_client.multiprocess.mark_process_dead"
|
|
) as mock_mark:
|
|
mark_worker_exit(12345)
|
|
mock_mark.assert_not_called()
|
|
|
|
def test_exception_is_caught_and_logged(self, tmp_path):
|
|
with patch.dict(os.environ, {"PROMETHEUS_MULTIPROC_DIR": str(tmp_path)}):
|
|
with patch(
|
|
"prometheus_client.multiprocess.mark_process_dead",
|
|
side_effect=FileNotFoundError("gone"),
|
|
) as mock_mark:
|
|
# Should not raise
|
|
mark_worker_exit(99)
|
|
mock_mark.assert_called_once_with(99)
|
|
|
|
|
|
class TestMaybeSetupPrometheusMultiprocDir:
|
|
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)
|
|
|
|
@pytest.mark.parametrize(
|
|
"num_workers, litellm_settings",
|
|
[
|
|
(1, {"callbacks": ["prometheus"]}),
|
|
(4, {"callbacks": ["langfuse"]}),
|
|
(4, None),
|
|
],
|
|
)
|
|
def test_noop_when_setup_not_needed(self, num_workers, litellm_settings):
|
|
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=num_workers,
|
|
litellm_settings=litellm_settings,
|
|
)
|
|
|
|
assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") is None
|
|
|
|
@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=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)
|
|
|
|
# Cleanup
|
|
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|