From 6f356cd920dee3f0c36bfa61cfa6c215662b5203 Mon Sep 17 00:00:00 2001
From: xuan <118178092+yuee3@users.noreply.github.com>
Date: Mon, 12 Oct 2026 13:45:12 +0800
Subject: [PATCH] fix: clean up auto-created prometheus multiproc dirs

Auto-created mkdtemp dirs were never removed. Register an atexit hook
that rmtree's the dir only from the creating PID so forked workers
sharing it do not delete samples their siblings are still writing.
Close the remaining test leaks: test_handles_string_callbacks and
test_separate_metrics_port_forces_dir_for_single_worker now rmtree
the dirs they create.
---
 litellm/proxy/proxy_cli.py                  | 13 +++++++++++++
 tests/unit/proxy/test_prometheus_cleanup.py |  6 +++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

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."""