refactor: remove shutdown cleanup, rely solely on startup wipe

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.
This commit is contained in:
Ryan Crabbe
2026-02-26 12:17:31 -08:00
parent acfb5ade97
commit 0ee8cb5f02
3 changed files with 3 additions and 67 deletions
+1 -30
View File
@@ -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}"
)
-8
View File
@@ -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()
@@ -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."""