fix: Final-annotate the per-boot dir pick and test the exit hook

main's type-discipline gate counts assignments without a Final declaration
against the merge-base ceiling, and unpacking targets cannot carry Final:
the second `multiproc_dir, action = ...` pick re-binds both names (LIT010).
Declare both as Final and defer the assignment to the branch that picks the
value. The creator-PID exit hook becomes a staticmethod so the
creator-only-removes branch is unit-covered, which also lifts patch
coverage back above target.
This commit is contained in:
xuan 2026-10-12 14:20:40 +08:00
parent 6f356cd920
commit 327c956c46
2 changed files with 23 additions and 11 deletions

View file

@ -4,6 +4,7 @@ import json
import os
import random
import re
import shutil
import subprocess
import sys
import urllib.parse as urlparse
@ -712,7 +713,6 @@ 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
@ -725,26 +725,24 @@ class ProxyInitializationHelpers:
wipe_directory(configured_dir)
return None
multiproc_dir: Final[str]
action: Final[str]
if configured_dir:
# An operator-set directory is theirs to manage; keep the historical wipe-on-boot.
os.makedirs(configured_dir, exist_ok=True)
wipe_directory(configured_dir)
multiproc_dir, action = configured_dir, "Using existing"
multiproc_dir = configured_dir
action = "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"
multiproc_dir = tempfile.mkdtemp(prefix="litellm_prometheus_multiproc_")
action = "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)
atexit.register(ProxyInitializationHelpers._rmtree_auto_created_dir, multiproc_dir, os.getpid())
os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir
print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}")
@ -752,6 +750,13 @@ class ProxyInitializationHelpers:
_maybe_setup_prometheus_multiproc_dir = maybe_setup_prometheus_multiproc_dir
@staticmethod
def _rmtree_auto_created_dir(multiproc_dir: str, creator_pid: int) -> None:
"""atexit hook for a boot-auto-created PROMETHEUS_MULTIPROC_DIR: only the creating process
removes the dir, so forked workers sharing it never delete samples their siblings still write."""
if os.getpid() == creator_pid:
shutil.rmtree(multiproc_dir, ignore_errors=True)
@click.command()
@click.argument("cli_args", nargs=-1)

View file

@ -292,8 +292,15 @@ class TestMaybeSetupPrometheusMultiprocDir:
assert os.path.exists(stale)
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
# The exit hook removes the dir only from the creating process: a forked worker
# (different PID) sharing the dir must leave its siblings' samples alone.
ProxyInitializationHelpers._rmtree_auto_created_dir(first, 999999)
assert os.path.exists(first)
ProxyInitializationHelpers._rmtree_auto_created_dir(second, os.getpid())
assert not os.path.exists(second)
shutil.rmtree(first, ignore_errors=True)
shutil.rmtree(second, ignore_errors=True)
def test_configured_dir_is_still_wiped_on_boot(self, tmp_path):
"""A directory the operator set keeps its historical wipe-on-boot lifecycle."""