diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 2e677857573..242b8ba0c68 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -712,6 +712,7 @@ class ProxyInitializationHelpers: with prometheus configured as a callback in config.yaml, or the separate metrics server (always, since callbacks may also be enabled from the DB after startup). """ + import shutil import tempfile from litellm.proxy.prometheus_cleanup import wipe_directory @@ -731,8 +732,20 @@ class ProxyInitializationHelpers: multiproc_dir, action = configured_dir, "Using existing" else: # Unique per boot: a fixed shared name lets a second proxy on this host wipe the first one's counters. + import atexit + multiproc_dir, action = tempfile.mkdtemp(prefix="litellm_prometheus_multiproc_"), "Auto-created" + # Only the creating process removes the dir on exit. Forked workers share it and + # must not delete samples their siblings are still writing. + creator_pid = os.getpid() + + def _cleanup_auto_created_dir() -> None: + if os.getpid() == creator_pid: + shutil.rmtree(multiproc_dir, ignore_errors=True) + + atexit.register(_cleanup_auto_created_dir) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") return multiproc_dir diff --git a/tests/unit/proxy/test_prometheus_cleanup.py b/tests/unit/proxy/test_prometheus_cleanup.py index 95903540c47..65163e7eb65 100644 --- a/tests/unit/proxy/test_prometheus_cleanup.py +++ b/tests/unit/proxy/test_prometheus_cleanup.py @@ -146,13 +146,15 @@ class TestMaybeSetupPrometheusMultiprocDir: os.environ.pop("prometheus_multiproc_dir", None) # Should not raise TypeError - ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( num_workers=4, litellm_settings=litellm_settings, ) # Cleanup os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + if result_dir is not None: + shutil.rmtree(result_dir, ignore_errors=True) @pytest.mark.parametrize( "num_workers, litellm_settings", @@ -200,6 +202,7 @@ class TestMaybeSetupPrometheusMultiprocDir: # Cleanup os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + shutil.rmtree(result_dir, ignore_errors=True) @pytest.mark.parametrize( "litellm_settings", @@ -227,6 +230,7 @@ class TestMaybeSetupPrometheusMultiprocDir: assert os.path.isdir(result_dir) os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + shutil.rmtree(result_dir, ignore_errors=True) 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."""