mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-15 08:27:36 +00:00
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:
@@ -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}"
|
||||
)
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user