fix(BA-4550): use persistent directory for Prometheus multiprocess metrics (#9114)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: HyeockJinKim

changes/9114.fix.md

@@ -0,0 +1 @@
+Fix Prometheus metrics scrape crash when `/tmp` multiprocess directory is cleaned by OS, by using a persistent directory (`/var/run/backendai/prometheus/`) outside `/tmp`.

src/ai/backend/common/metrics/multiprocess.py

@@ -29,14 +29,18 @@
 
 _multiprocess_dir: Path | None = None
 
+_DEFAULT_BASE_DIR = Path(tempfile.gettempdir()) / "backendai" / "prometheus"
 
-def setup_prometheus_multiprocess_dir(component: str = "manager") -> Path:
+
+def setup_prometheus_multiprocess_dir(
+    component: str = "manager",
+) -> Path:
     """
     Set up the prometheus multiprocess directory and environment variable.
 
     MUST be called before any prometheus_client import.
 
-    Creates a temporary directory for prometheus multiprocess files and sets
+    Creates a directory for prometheus multiprocess files and sets
     the PROMETHEUS_MULTIPROC_DIR environment variable.
 
     Args:
@@ -50,15 +54,15 @@ def setup_prometheus_multiprocess_dir(component: str = "manager") -> Path:
     if _multiprocess_dir is not None:
         return _multiprocess_dir
 
-    base_dir = Path(tempfile.gettempdir()) / "backendai-prometheus"
-    base_dir.mkdir(parents=True, exist_ok=True)
+    multiprocess_dir = _DEFAULT_BASE_DIR / component
+    multiprocess_dir.mkdir(parents=True, exist_ok=True)
 
-    multiprocess_dir = Path(
-        tempfile.mkdtemp(
-            prefix=f"{component}-",
-            dir=base_dir,
-        )
-    )
+    # Clean stale .db files from previous runs
+    for db_file in multiprocess_dir.glob("*.db"):
+        try:
+            db_file.unlink()
+        except OSError:
+            pass
 
     os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(multiprocess_dir)
     _multiprocess_dir = multiprocess_dir
@@ -75,9 +79,31 @@ def generate_latest_multiprocess() -> bytes:
 
     This should be used by multi-worker components (manager, agent, storage, etc.).
     """
-    registry = CollectorRegistry()
-    MultiProcessCollector(registry)  # type: ignore[no-untyped-call]
-    return generate_latest(registry)
+    try:
+        registry = CollectorRegistry()
+        MultiProcessCollector(registry)  # type: ignore[no-untyped-call]
+        return generate_latest(registry)
+    except ValueError:
+        # Directory may have been deleted (e.g., by systemd-tmpfiles-clean).
+        # Attempt to recreate it and retry once.
+        if _multiprocess_dir is not None:
+            try:
+                _multiprocess_dir.mkdir(parents=True, exist_ok=True)
+                os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(_multiprocess_dir)
+                registry = CollectorRegistry()
+                MultiProcessCollector(registry)  # type: ignore[no-untyped-call]
+                log.warning(
+                    "Prometheus multiprocess dir was missing and has been recreated: %s",
+                    _multiprocess_dir,
+                )
+                return generate_latest(registry)
+            except Exception:
+                log.error(
+                    "Failed to recover prometheus multiprocess dir: %s",
+                    _multiprocess_dir,
+                    exc_info=True,
+                )
+        return b""
 
 
 def generate_latest_singleprocess() -> bytes:

tests/unit/common/test_prometheus_multiprocess.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+import ai.backend.common.metrics.multiprocess as mp_mod
+from ai.backend.common.metrics.multiprocess import (
+    cleanup_prometheus_multiprocess_dir,
+    generate_latest_multiprocess,
+    setup_prometheus_multiprocess_dir,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_multiprocess_state() -> Iterator[None]:
+    """Reset the module-level global state before each test."""
+    original_dir = mp_mod._multiprocess_dir
+    original_env = os.environ.get("PROMETHEUS_MULTIPROC_DIR")
+    yield
+    mp_mod._multiprocess_dir = original_dir
+    if original_env is not None:
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = original_env
+    elif "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+        del os.environ["PROMETHEUS_MULTIPROC_DIR"]
+
+
+class TestSetupPrometheusMultiprocDir:
+    def test_creates_directory_with_default_base(self, tmp_path: Path) -> None:
+        with patch.object(mp_mod, "_DEFAULT_BASE_DIR", tmp_path):
+            result = setup_prometheus_multiprocess_dir("manager")
+
+        assert result == tmp_path / "manager"
+        assert result.is_dir()
+        assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == str(result)
+
+    def test_cleans_stale_db_files(self, tmp_path: Path) -> None:
+        prom_dir = tmp_path / "manager"
+        prom_dir.mkdir(parents=True)
+        (prom_dir / "gauge_liveall_123.db").touch()
+        (prom_dir / "counter_456.db").touch()
+        (prom_dir / "keep_this.txt").touch()
+
+        with patch.object(mp_mod, "_DEFAULT_BASE_DIR", tmp_path):
+            setup_prometheus_multiprocess_dir("manager")
+
+        assert not (prom_dir / "gauge_liveall_123.db").exists()
+        assert not (prom_dir / "counter_456.db").exists()
+        assert (prom_dir / "keep_this.txt").exists()
+
+    def test_idempotent_returns_same_path(self, tmp_path: Path) -> None:
+        with patch.object(mp_mod, "_DEFAULT_BASE_DIR", tmp_path):
+            first = setup_prometheus_multiprocess_dir("manager")
+            second = setup_prometheus_multiprocess_dir("manager")
+
+        assert first == second
+
+    def test_idempotent_ignores_different_component(self, tmp_path: Path) -> None:
+        """Once setup, calling again with different component still returns the first path."""
+        with patch.object(mp_mod, "_DEFAULT_BASE_DIR", tmp_path):
+            first = setup_prometheus_multiprocess_dir("manager")
+            second = setup_prometheus_multiprocess_dir("agent")
+
+        assert first == second  # idempotent, returns first result
+
+
+class TestGenerateLatestMultiprocess:
+    def test_returns_bytes_normally(self, tmp_path: Path) -> None:
+        prom_dir = tmp_path / "test-component"
+        prom_dir.mkdir(parents=True)
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(prom_dir)
+        mp_mod._multiprocess_dir = prom_dir
+
+        result = generate_latest_multiprocess()
+        assert isinstance(result, bytes)
+
+    def test_recovers_from_missing_directory(self, tmp_path: Path) -> None:
+        prom_dir = tmp_path / "test-component"
+        prom_dir.mkdir(parents=True)
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(prom_dir)
+        mp_mod._multiprocess_dir = prom_dir
+
+        # Simulate systemd-tmpfiles-clean deleting the directory
+        prom_dir.rmdir()
+        assert not prom_dir.exists()
+
+        result = generate_latest_multiprocess()
+        # Should recover by recreating the directory
+        assert isinstance(result, bytes)
+        assert prom_dir.exists()
+
+    def test_returns_empty_bytes_on_unrecoverable_failure(self) -> None:
+        # Set an invalid path that can't be recreated
+        mp_mod._multiprocess_dir = Path("/nonexistent/impossible/path")
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = "/nonexistent/impossible/path"
+
+        result = generate_latest_multiprocess()
+        assert result == b""
+
+
+class TestCleanupPrometheusMultiprocDir:
+    def test_removes_directory_and_env(self, tmp_path: Path) -> None:
+        prom_dir = tmp_path / "manager"
+        prom_dir.mkdir(parents=True)
+        (prom_dir / "test.db").touch()
+        os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(prom_dir)
+        mp_mod._multiprocess_dir = prom_dir
+
+        cleanup_prometheus_multiprocess_dir()
+
+        assert not prom_dir.exists()
+        assert "PROMETHEUS_MULTIPROC_DIR" not in os.environ
+        assert mp_mod._multiprocess_dir is None
+
+    def test_noop_when_not_initialized(self) -> None:
+        mp_mod._multiprocess_dir = None
+        if "PROMETHEUS_MULTIPROC_DIR" in os.environ:
+            del os.environ["PROMETHEUS_MULTIPROC_DIR"]
+
+        # Should not raise
+        cleanup_prometheus_multiprocess_dir()