litellm/tests/test_litellm/proxy/test_prometheus_cleanup.py

"""
Tests for litellm.proxy.prometheus_cleanup.wipe_directory and
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir.
"""

from __future__ import annotations

import os
import subprocess
import sys
from pathlib import Path
from typing import Final
from unittest.mock import patch

import pytest
from prometheus_client import CollectorRegistry, multiprocess

from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit, wipe_directory
from litellm.proxy.proxy_cli import ProxyInitializationHelpers

_WORKER: Final = """
import sys, time
from prometheus_client import Gauge
Gauge("litellm_in_flight", "", multiprocess_mode="livesum").set(float(sys.argv[1]))
print("ready", flush=True)
if sys.argv[2] == "stay":
    time.sleep(120)
"""


def _spawn_worker(directory: Path, in_flight: str, lifetime: str) -> subprocess.Popen[str]:
    env = {**os.environ, "PROMETHEUS_MULTIPROC_DIR": str(directory)}
    worker = subprocess.Popen(
        [sys.executable, "-c", _WORKER, in_flight, lifetime], env=env, stdout=subprocess.PIPE, text=True
    )
    assert worker.stdout is not None and worker.stdout.readline() == "ready\n"
    return worker


def _livesum(directory: Path) -> float:
    registry = CollectorRegistry()
    multiprocess.MultiProcessCollector(registry, path=str(directory))
    value = registry.get_sample_value("litellm_in_flight")
    return 0.0 if value is None else value


class TestMarkDeadWorkers:
    def test_drops_live_gauges_of_exited_workers_and_keeps_running_ones(self, tmp_path: Path) -> None:
        """A worker that died mid-request leaves its livesum file behind; the replacement worker's startup prune
        must remove exactly that file so the aggregate stops counting requests nobody is serving."""
        dead = _spawn_worker(tmp_path, "3", "exit")
        assert dead.wait(timeout=30) == 0
        alive = _spawn_worker(tmp_path, "2", "stay")
        try:
            assert (tmp_path / f"gauge_livesum_{dead.pid}.db").exists()
            assert _livesum(tmp_path) == 5.0

            assert mark_dead_workers(str(tmp_path)) == (dead.pid,)

            assert not (tmp_path / f"gauge_livesum_{dead.pid}.db").exists()
            assert (tmp_path / f"gauge_livesum_{alive.pid}.db").exists()
            assert _livesum(tmp_path) == 2.0
            assert mark_dead_workers(str(tmp_path)) == ()
        finally:
            alive.kill()
            alive.wait(timeout=30)

    def test_leaves_counters_of_exited_workers_alone(self, tmp_path: Path) -> None:
        (tmp_path / "counter_424242.db").touch()
        (tmp_path / "histogram_424242.db").touch()
        assert mark_dead_workers(str(tmp_path)) == ()
        assert sorted(p.name for p in tmp_path.glob("*.db")) == ["counter_424242.db", "histogram_424242.db"]

    def test_keeps_live_gauges_of_workers_it_may_not_signal(self, tmp_path: Path) -> None:
        """Signal 0 to pid 1 raises PermissionError for an unprivileged proxy; that pid is alive, not dead."""
        (tmp_path / "gauge_livesum_1.db").touch()
        assert mark_dead_workers(str(tmp_path)) == ()
        assert (tmp_path / "gauge_livesum_1.db").exists()


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(
        "litellm_settings",
        [
            {"callbacks": "prometheus"},
            {"success_callback": "prometheus"},
            {"failure_callback": "prometheus"},
            {"callbacks": "custom_callback"},  # string but not prometheus
        ],
    )
    def test_handles_string_callbacks(self, litellm_settings):
        """When callbacks are specified as a string instead of a list, should not crash."""
        with patch.dict(os.environ, {}, clear=False):
            os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
            os.environ.pop("prometheus_multiproc_dir", None)

            # Should not raise TypeError
            ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
                num_workers=4,
                litellm_settings=litellm_settings,
            )

            # Cleanup
            os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)

    @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)

    @pytest.mark.parametrize(
        "litellm_settings",
        [
            {"callbacks": ["prometheus"]},
            {"callbacks": ["langfuse"]},
            None,
        ],
    )
    def test_separate_metrics_port_forces_dir_for_single_worker(self, litellm_settings):
        """The separate metrics process reads the samples, so one worker still needs the shared dir, even when
        prometheus is not in config.yaml (callbacks can be turned on from the DB after startup)."""
        with patch.dict(os.environ, {}, clear=False):
            os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
            os.environ.pop("prometheus_multiproc_dir", None)

            result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
                num_workers=1,
                litellm_settings=litellm_settings,
                prometheus_metrics_port=4001,
            )

            assert result_dir is not None
            assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") == result_dir
            assert os.path.isdir(result_dir)

            os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)

    def test_lowercase_env_var_is_reused_and_exported_uppercase(self, tmp_path):
        """prometheus_client honours both spellings; the metrics server only reads the uppercase one."""
        with patch.dict(os.environ, {"prometheus_multiproc_dir": str(tmp_path)}, clear=False):
            os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)

            result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
                num_workers=4,
                litellm_settings={"callbacks": "prometheus"},
            )

            assert result_dir == str(tmp_path)
            assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == str(tmp_path)