From 0ee8cb5f021cb2d02149c63bdb0274c46462cb47 Mon Sep 17 00:00:00 2001
From: Ryan Crabbe <rcrabbe@berkeley.edu>
Date: Thu, 26 Feb 2026 12:17:31 -0800
Subject: [PATCH] refactor: remove shutdown cleanup, rely solely on startup
 wipe
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Phase 2 (per-worker mark_process_dead on shutdown) only ever fired
when all workers shut down together, making it redundant — Phase 1
wipes everything on next startup anyway. This aligns with the
prometheus_client docs: just wipe the directory between runs.
---
 litellm/proxy/prometheus_cleanup.py           | 31 +------------------
 litellm/proxy/proxy_server.py                 |  8 -----
 .../proxy/test_prometheus_cleanup.py          | 31 ++-----------------
 3 files changed, 3 insertions(+), 67 deletions(-)

diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py
index 8c0abb6386..6907752e80 100644
--- a/litellm/proxy/prometheus_cleanup.py
+++ b/litellm/proxy/prometheus_cleanup.py
@@ -1,25 +1,16 @@
 """
 Prometheus multiprocess directory cleanup utilities.
 
-mark_process_dead() only removes gauge_live* files — counter/histogram
-files are kept for correct aggregation and wiped in bulk on startup.
+Wipes all .db files on startup so workers start with a clean slate.
 """
 
 from __future__ import annotations
 
 import glob
 import os
-from typing import Optional
 
 from litellm._logging import verbose_proxy_logger
 
-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 wipe_directory(directory: str) -> None:
     """Delete all .db files in the directory. Called once before workers fork."""
@@ -35,23 +26,3 @@ def wipe_directory(directory: str) -> None:
         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 via mark_process_dead() (worker shutdown)."""
-    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}"
-        )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 4becffa857..6def660632 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -709,14 +709,6 @@ 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()
 
diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py
index 81ef1dd449..276f2b592d 100644
--- a/tests/test_litellm/proxy/test_prometheus_cleanup.py
+++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py
@@ -1,5 +1,5 @@
 """
-Tests for litellm.proxy.prometheus_cleanup module and
+Tests for litellm.proxy.prometheus_cleanup.wipe_directory and
 ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir.
 """
 
@@ -10,10 +10,7 @@ from unittest.mock import patch
 
 import pytest
 
-from litellm.proxy.prometheus_cleanup import (
-    cleanup_own_pid_files,
-    wipe_directory,
-)
+from litellm.proxy.prometheus_cleanup import wipe_directory
 from litellm.proxy.proxy_cli import ProxyInitializationHelpers
 
 
@@ -26,30 +23,6 @@ class TestWipeDirectory:
         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 TestMaybeSetupPrometheusMultiprocDir:
     def test_respects_existing_env_var(self, tmp_path):
         """When PROMETHEUS_MULTIPROC_DIR is already set, don't override it."""