Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/prometheus_cleanup.py: 20%

1""" 

2Prometheus multiprocess directory cleanup utilities. 

3 

4Wipes all .db files on startup so workers start with a clean slate. 

5""" 

6 

7from __future__ import annotations 

8 

9import glob 

10import os 

11import re 

12from typing import Final 

13 

14from litellm._logging import verbose_proxy_logger 

15 

16_LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$") 

17 

18 

19def wipe_directory(directory: str) -> None: 

20 """Delete all .db files in the directory. Called once before workers fork.""" 

21 files: Final = glob.glob(os.path.join(directory, "*.db")) 

22 deleted = 0 

23 for filepath in files: 

24 try: 

25 os.remove(filepath) 

26 deleted += 1 

27 except OSError as e: 

28 verbose_proxy_logger.warning("Failed to delete stale prometheus file %s: %s", filepath, e) 

29 if deleted: 

30 verbose_proxy_logger.info("Prometheus cleanup: wiped %s stale .db files from %s", deleted, directory) 

31 

32 

33def mark_worker_exit(worker_pid: int) -> None: 

34 """Remove prometheus .db files for a dead worker. Called by gunicorn child_exit hook.""" 

35 if not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): 

36 return 

37 try: 

38 from prometheus_client import multiprocess 

39 

40 multiprocess.mark_process_dead(worker_pid) 

41 verbose_proxy_logger.info("Prometheus cleanup: marked worker %s as dead", worker_pid) 

42 except Exception as e: 

43 verbose_proxy_logger.warning("Failed to mark prometheus worker %s as dead: %s", worker_pid, e) 

44 

45 

46def _is_running(pid: int) -> bool: 

47 try: 

48 os.kill(pid, 0) 

49 except ProcessLookupError: 

50 return False 

51 except PermissionError: 

52 return True 

53 return True 

54 

55 

56def mark_dead_workers(directory: str) -> tuple[int, ...]: 

57 """Drop the live-gauge files of workers that no longer exist and return their pids. 

58 

59 Uvicorn's multi-worker supervisor has no exit hook, so a replacement worker calls this at startup; without it 

60 a crashed worker's in-flight gauges stay in the aggregate forever. 

61 """ 

62 owners: Final = frozenset( 

63 int(match.group(1)) 

64 for match in map(_LIVE_GAUGE_PID.search, glob.glob(os.path.join(directory, "gauge_live*_*.db"))) 

65 if match is not None 

66 ) 

67 dead: Final = tuple(sorted(pid for pid in owners if pid != os.getpid() and not _is_running(pid))) 

68 if not dead: 

69 return dead 

70 from prometheus_client import multiprocess 

71 

72 for pid in dead: 

73 multiprocess.mark_process_dead(pid, path=directory) 

74 verbose_proxy_logger.info("Prometheus cleanup: marked dead workers %s in %s", dead, directory) 

75 return dead