From a9fddbf4add08865bd671a5e940146a931f8a0d3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 10 Sep 2025 16:20:17 -0700 Subject: [PATCH 1/2] fix(prometheus.py): make prometheus work for multiple workers --- .../integrations/prometheus.py | 298 +++++++++++++++--- litellm/proxy/_new_secret_config.yaml | 28 +- litellm/proxy/common_utils/callback_utils.py | 31 ++ litellm/proxy/proxy_cli.py | 53 ++++ litellm/proxy/proxy_server.py | 46 +++ 5 files changed, 403 insertions(+), 53 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index efee1a7783e..836c79abb44 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -1,8 +1,11 @@ # used for /metrics endpoint on LiteLLM Proxy #### What this does #### # On success, log events to Prometheus +import os import sys +import tempfile from datetime import datetime, timedelta +from pathlib import Path from typing import ( TYPE_CHECKING, Any, @@ -16,6 +19,64 @@ from typing import ( cast, ) + +# CRITICAL: Set up multiprocess mode BEFORE importing prometheus_client +# This must happen at module import time, not at class instantiation time +def _setup_early_multiprocess_mode(): + """Setup multiprocess mode at import time if needed.""" + try: + # Check if we're in a multiprocess environment + num_workers = os.environ.get("NUM_WORKERS", "1") + is_multiprocess = False + + try: + if int(num_workers) > 1: + is_multiprocess = True + except (ValueError, TypeError): + pass + + # Check for gunicorn worker environment variables + if os.environ.get("GUNICORN_CMD_ARGS") or os.environ.get("GUNICORN_WORKER_ID"): + is_multiprocess = True + + # Check if PROMETHEUS_MULTIPROC_DIR is explicitly set (admin override) + if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + is_multiprocess = True + + if is_multiprocess: + existing_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if not existing_dir: + # Set up multiprocess directory + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + + # Ensure the directory exists + Path(multiproc_dir).mkdir(parents=True, exist_ok=True) + + verbose_logger.info( + f"Prometheus multiprocess mode auto-enabled with directory: {multiproc_dir}" + ) + else: + # Directory already set, just ensure it exists + Path(existing_dir).mkdir(parents=True, exist_ok=True) + verbose_logger.info( + f"Using existing Prometheus multiprocess directory: {existing_dir}" + ) + + except PermissionError as e: + verbose_logger.warning( + f"Warning: Unable to create Prometheus multiprocess directory due to permission error. " + f"Running in non-root environment. Prometheus metrics may not work correctly in multiprocess mode. Error: {e}" + ) + except Exception as e: + verbose_logger.warning(f"Warning: Failed to setup early multiprocess mode: {e}") + + +# Set up multiprocess mode before any prometheus imports +_setup_early_multiprocess_mode() + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -44,6 +105,18 @@ class PrometheusLogger(CustomLogger): # Always initialize label_filters, even for non-premium users self.label_filters = self._parse_prometheus_config() + # Initialize multiprocess mode for Prometheus metrics to handle multiple workers + self._setup_multiprocess_mode() + + # Debug: Check if multiprocess mode is active + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if multiproc_dir: + verbose_logger.info( + f"Prometheus multiprocess mode active with directory: {multiproc_dir}" + ) + else: + verbose_logger.info("Prometheus running in single-process mode") + if premium_user is not True: verbose_logger.warning( f"šŸšØšŸšØšŸšØ Prometheus Metrics is on LiteLLM Enterprise\nšŸšØ {CommonProxyErrors.not_premium_user.value}" @@ -102,7 +175,9 @@ class PrometheusLogger(CustomLogger): # "team", # "team_alias", # ], - labelnames=self.get_labels_for_metric("litellm_llm_api_time_to_first_token_metric"), + labelnames=self.get_labels_for_metric( + "litellm_llm_api_time_to_first_token_metric" + ), buckets=LATENCY_BUCKETS, ) @@ -132,47 +207,52 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_output_tokens_metric"), ) - # Remaining Budget for Team + # Remaining Budget for Team (use 'mostrecent' for multiprocess mode) self.litellm_remaining_team_budget_metric = self._gauge_factory( "litellm_remaining_team_budget_metric", "Remaining budget for team", labelnames=self.get_labels_for_metric( "litellm_remaining_team_budget_metric" ), + multiprocess_mode="mostrecent", ) - # Max Budget for Team + # Max Budget for Team (use 'mostrecent' for multiprocess mode) self.litellm_team_max_budget_metric = self._gauge_factory( "litellm_team_max_budget_metric", "Maximum budget set for team", labelnames=self.get_labels_for_metric("litellm_team_max_budget_metric"), + multiprocess_mode="mostrecent", ) - # Team Budget Reset At + # Team Budget Reset At (use 'mostrecent' for multiprocess mode) self.litellm_team_budget_remaining_hours_metric = self._gauge_factory( "litellm_team_budget_remaining_hours_metric", "Remaining days for team budget to be reset", labelnames=self.get_labels_for_metric( "litellm_team_budget_remaining_hours_metric" ), + multiprocess_mode="mostrecent", ) - # Remaining Budget for API Key + # Remaining Budget for API Key (use 'mostrecent' for multiprocess mode) self.litellm_remaining_api_key_budget_metric = self._gauge_factory( "litellm_remaining_api_key_budget_metric", "Remaining budget for api key", labelnames=self.get_labels_for_metric( "litellm_remaining_api_key_budget_metric" ), + multiprocess_mode="mostrecent", ) - # Max Budget for API Key + # Max Budget for API Key (use 'mostrecent' for multiprocess mode) self.litellm_api_key_max_budget_metric = self._gauge_factory( "litellm_api_key_max_budget_metric", "Maximum budget set for api key", labelnames=self.get_labels_for_metric( "litellm_api_key_max_budget_metric" ), + multiprocess_mode="mostrecent", ) self.litellm_api_key_budget_remaining_hours_metric = self._gauge_factory( @@ -181,36 +261,40 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric( "litellm_api_key_budget_remaining_hours_metric" ), + multiprocess_mode="mostrecent", ) ######################################## # LiteLLM Virtual API KEY metrics ######################################## - # Remaining MODEL RPM limit for API Key + # Remaining MODEL RPM limit for API Key (use 'mostrecent' for multiprocess mode) self.litellm_remaining_api_key_requests_for_model = self._gauge_factory( "litellm_remaining_api_key_requests_for_model", "Remaining Requests API Key can make for model (model based rpm limit on key)", labelnames=["hashed_api_key", "api_key_alias", "model"], + multiprocess_mode="mostrecent", ) - # Remaining MODEL TPM limit for API Key + # Remaining MODEL TPM limit for API Key (use 'mostrecent' for multiprocess mode) self.litellm_remaining_api_key_tokens_for_model = self._gauge_factory( "litellm_remaining_api_key_tokens_for_model", "Remaining Tokens API Key can make for model (model based tpm limit on key)", labelnames=["hashed_api_key", "api_key_alias", "model"], + multiprocess_mode="mostrecent", ) ######################################## # LLM API Deployment Metrics / analytics ######################################## - # Remaining Rate Limit for model + # Remaining Rate Limit for model (use 'mostrecent' for multiprocess mode) self.litellm_remaining_requests_metric = self._gauge_factory( "litellm_remaining_requests", "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_requests_metric" ), + multiprocess_mode="mostrecent", ) self.litellm_remaining_tokens_metric = self._gauge_factory( @@ -219,6 +303,7 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric( "litellm_remaining_tokens_metric" ), + multiprocess_mode="mostrecent", ) self.litellm_overhead_latency_metric = self._histogram_factory( @@ -229,25 +314,27 @@ class PrometheusLogger(CustomLogger): ), buckets=LATENCY_BUCKETS, ) - # llm api provider budget metrics + # llm api provider budget metrics (use 'mostrecent' for multiprocess mode) self.litellm_provider_remaining_budget_metric = self._gauge_factory( "litellm_provider_remaining_budget_metric", "Remaining budget for provider - used when you set provider budget limits", labelnames=["api_provider"], + multiprocess_mode="mostrecent", ) - # Metric for deployment state + # Metric for deployment state (use 'mostrecent' for multiprocess mode) self.litellm_deployment_state = self._gauge_factory( "litellm_deployment_state", "LLM Deployment Analytics - The state of the deployment: 0 = healthy, 1 = partial outage, 2 = complete outage", - labelnames=self.get_labels_for_metric("litellm_deployment_state") + labelnames=self.get_labels_for_metric("litellm_deployment_state"), + multiprocess_mode="mostrecent", ) self.litellm_deployment_cooled_down = self._counter_factory( "litellm_deployment_cooled_down", "LLM Deployment Analytics - Number of times a deployment has been cooled down by LiteLLM load balancing logic. exception_status is the status of the exception that caused the deployment to be cooled down", # labelnames=_logged_llm_labels + [EXCEPTION_STATUS], - labelnames=self.get_labels_for_metric("litellm_deployment_cooled_down") + labelnames=self.get_labels_for_metric("litellm_deployment_cooled_down"), ) self.litellm_deployment_success_responses = self._counter_factory( @@ -318,6 +405,105 @@ class PrometheusLogger(CustomLogger): print_verbose(f"Got exception on init prometheus client {str(e)}") raise e + def _setup_multiprocess_mode(self): + """ + Setup Prometheus multiprocess mode to handle multiple workers properly. + This ensures that metrics are aggregated correctly across all worker processes. + """ + import os + import tempfile + from pathlib import Path + + try: + # Check if we're in a multiprocess environment (multiple workers) + if not self._is_multiprocess_environment(): + verbose_logger.debug( + "Single process environment detected, skipping multiprocess setup" + ) + return + + # Set up multiprocess directory if not already configured + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if not multiproc_dir: + # Create a temp directory for multiprocess metrics + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + verbose_logger.debug(f"Set PROMETHEUS_MULTIPROC_DIR to {multiproc_dir}") + + # Ensure the directory exists + Path(multiproc_dir).mkdir(parents=True, exist_ok=True) + + # Force the prometheus_client to recognize multiprocess mode + # This is important because the environment variable must be set BEFORE importing prometheus_client + try: + from prometheus_client import multiprocess + + # This will trigger the multiprocess mode if the env var is set + verbose_logger.debug( + "Prometheus multiprocess module imported successfully" + ) + except Exception as e: + verbose_logger.warning( + f"Failed to import prometheus multiprocess module: {e}" + ) + + verbose_logger.info( + f"Prometheus multiprocess mode enabled with directory: {multiproc_dir}" + ) + + except Exception as e: + verbose_logger.warning(f"Failed to setup Prometheus multiprocess mode: {e}") + + def _is_multiprocess_environment(self) -> bool: + """ + Detect if we're running in a multiprocess environment (uvicorn/gunicorn with multiple workers). + """ + import os + + # Check for common environment variables that indicate multiple workers + num_workers = os.environ.get("NUM_WORKERS", "1") + try: + if int(num_workers) > 1: + return True + except (ValueError, TypeError): + pass + + # Check for gunicorn worker environment variables + if os.environ.get("GUNICORN_CMD_ARGS") or os.environ.get("GUNICORN_WORKER_ID"): + return True + + # Check if PROMETHEUS_MULTIPROC_DIR is explicitly set (admin override) + if os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + return True + + return False + + @staticmethod + def cleanup_multiprocess_metrics(): + """ + Clean up multiprocess metrics directory on startup. + This should be called once during application startup to prevent stale metrics. + """ + import os + from pathlib import Path + + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if multiproc_dir and os.path.exists(multiproc_dir): + try: + # Remove all files in the directory but keep the directory itself + for file_path in Path(multiproc_dir).glob("*"): + if file_path.is_file(): + file_path.unlink() + verbose_logger.info( + f"Cleaned up Prometheus multiprocess metrics directory: {multiproc_dir}" + ) + except Exception as e: + verbose_logger.warning( + f"Failed to cleanup Prometheus multiprocess directory: {e}" + ) + def _parse_prometheus_config(self) -> Dict[str, List[str]]: """Parse prometheus metrics configuration for label filtering and enabled metrics""" import litellm @@ -727,7 +913,16 @@ class PrometheusLogger(CustomLogger): metric_name = args[0] if args else kwargs.get("name", "") if self._is_metric_enabled(metric_name): - return metric_class(*args, **kwargs) + # Handle multiprocess_mode parameter for Gauge metrics + if metric_class.__name__ == "Gauge" and "multiprocess_mode" in kwargs: + # Pass through multiprocess_mode to the Gauge constructor + return metric_class(*args, **kwargs) + else: + # For Counter and Histogram, remove multiprocess_mode if present + filtered_kwargs = { + k: v for k, v in kwargs.items() if k != "multiprocess_mode" + } + return metric_class(*args, **filtered_kwargs) else: return NoOpMetric() @@ -1039,20 +1234,11 @@ class PrometheusLogger(CustomLogger): _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_proxy_total_requests_metric" + metric_name="litellm_spend_metric" ), enum_values=enum_values, ) - - self.litellm_spend_metric.labels( - end_user_id, - user_api_key, - user_api_key_alias, - model, - user_api_team, - user_api_team_alias, - user_id, - ).inc(response_cost) + self.litellm_spend_metric.labels(**_labels).inc(response_cost) def _set_virtual_key_rate_limit_metrics( self, @@ -2179,13 +2365,14 @@ class PrometheusLogger(CustomLogger): def _mount_metrics_endpoint(premium_user: bool): """ Mount the Prometheus metrics endpoint with optional authentication. + Uses multiprocess collector when running with multiple workers. Args: premium_user (bool): Whether the user is a premium user - require_auth (bool, optional): Whether to require authentication for the metrics endpoint. - Defaults to False. """ - from prometheus_client import make_asgi_app + import os + + from prometheus_client import CollectorRegistry, make_asgi_app from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors @@ -2196,14 +2383,34 @@ class PrometheusLogger(CustomLogger): f"Prometheus metrics are only available for premium users. {CommonProxyErrors.not_premium_user.value}" ) - # Create metrics ASGI app - metrics_app = make_asgi_app() + # Check if we're in multiprocess mode + multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + + if multiproc_dir: + # Use multiprocess collector for worker aggregation + try: + from prometheus_client import multiprocess + + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + metrics_app = make_asgi_app(registry) + verbose_proxy_logger.info( + f"Starting Prometheus Metrics on /metrics with multiprocess collector (directory: {multiproc_dir})" + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to setup multiprocess collector, falling back to default: {e}" + ) + metrics_app = make_asgi_app() + else: + # Use default single-process collector + metrics_app = make_asgi_app() + verbose_proxy_logger.debug( + "Starting Prometheus Metrics on /metrics (single process mode)" + ) # Mount the metrics app to the app app.mount("/metrics", metrics_app) - verbose_proxy_logger.debug( - "Starting Prometheus Metrics on /metrics (no authentication)" - ) def prometheus_label_factory( @@ -2280,7 +2487,9 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]: return result -def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: str) -> bool: +def _tag_matches_wildcard_configured_pattern( + tags: List[str], configured_tag: str +) -> bool: """ Check if any of the request tags matches a wildcard configured pattern @@ -2305,6 +2514,7 @@ def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: st import re from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + pattern_router = PatternMatchRouter() regex_pattern = pattern_router._pattern_to_regex(configured_tag) return any(re.match(pattern=regex_pattern, string=tag) for tag in tags) @@ -2313,11 +2523,11 @@ def _tag_matches_wildcard_configured_pattern(tags: List[str], configured_tag: st def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]: """ Get custom labels from tags based on admin configuration. - + Supports both exact matches and wildcard patterns: - Exact match: "prod" matches "prod" exactly - - Wildcard pattern: "User-Agent: curl/*" matches "User-Agent: curl/7.68.0" - + - Wildcard pattern: "User-Agent: curl/*" matches "User-Agent: curl/7.68.0" + Reuses PatternMatchRouter for wildcard pattern matching. Returns dict of label_name: "true" if the tag matches the configured tag, "false" otherwise @@ -2331,9 +2541,6 @@ def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]: "tag_Service_web_app_v1": "false", } """ - import re - - from litellm.router_utils.pattern_match_deployments import PatternMatchRouter from litellm.types.integrations.prometheus import _sanitize_prometheus_label_name configured_tags = litellm.custom_prometheus_tags @@ -2341,21 +2548,22 @@ def get_custom_labels_from_tags(tags: List[str]) -> Dict[str, str]: return {} result: Dict[str, str] = {} - pattern_router = PatternMatchRouter() for configured_tag in configured_tags: label_name = _sanitize_prometheus_label_name(f"tag_{configured_tag}") - + # Check for exact match first (backwards compatibility) if configured_tag in tags: result[label_name] = "true" continue - + # Use PatternMatchRouter for wildcard pattern matching - if "*" in configured_tag and _tag_matches_wildcard_configured_pattern(tags=tags, configured_tag=configured_tag): + if "*" in configured_tag and _tag_matches_wildcard_configured_pattern( + tags=tags, configured_tag=configured_tag + ): result[label_name] = "true" continue - + # No match found result[label_name] = "false" diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c785dd05c40..07e04c3211d 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -1,9 +1,21 @@ model_list: - - model_name: fake-openai-endpoint - litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ - - model_name: wildcard_models/* - litellm_params: - model: openai/* + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ + - model_name: gpt-5-mini + litellm_params: + model: azure/gpt-5-mini + api_base: os.environ/AZURE_GPT_5_MINI_API_BASE # runs os.getenv("AZURE_API_BASE") + api_key: os.environ/AZURE_GPT_5_MINI_API_KEY # runs os.getenv("AZURE_API_KEY") + stream_timeout: 60 + merge_reasoning_content_in_choices: true + model_info: + mode: chat + +router_settings: + model_group_alias: {"my-fake-gpt-4": "fake-openai-endpoint"} + +litellm_settings: + callbacks: ["prometheus"] \ No newline at end of file diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d52592952bc..697ecb467df 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -268,12 +268,43 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 litellm.callbacks = imported_list # type: ignore if "prometheus" in value: + # CRITICAL: Set up prometheus multiprocess mode BEFORE importing + import os + import tempfile + from pathlib import Path + + # Check if we're in a multiprocess environment + num_workers = os.environ.get('NUM_WORKERS', '1') + is_multiprocess = False + + try: + if int(num_workers) > 1: + is_multiprocess = True + except (ValueError, TypeError): + pass + + # Check for gunicorn worker environment variables + if os.environ.get('GUNICORN_CMD_ARGS') or os.environ.get('GUNICORN_WORKER_ID'): + is_multiprocess = True + + if is_multiprocess and not os.environ.get('PROMETHEUS_MULTIPROC_DIR'): + # Set up multiprocess directory + multiproc_dir = os.path.join(tempfile.gettempdir(), 'litellm_prometheus_multiproc') + os.environ['PROMETHEUS_MULTIPROC_DIR'] = multiproc_dir + + # Ensure the directory exists + Path(multiproc_dir).mkdir(parents=True, exist_ok=True) + + verbose_proxy_logger.info(f"Prometheus multiprocess mode enabled with directory: {multiproc_dir}") + try: from litellm_enterprise.integrations.prometheus import PrometheusLogger except Exception: PrometheusLogger = None if PrometheusLogger: + # Clean up any existing multiprocess metrics before mounting + PrometheusLogger.cleanup_multiprocess_metrics() PrometheusLogger._mount_metrics_endpoint(premium_user) else: litellm.callbacks = [ diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 867a395d627..0009b0e37b5 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -186,6 +186,39 @@ class ProxyInitializationHelpers: ssl_certfile_path: str, ssl_keyfile_path: str, ): + # Set up Prometheus multiprocess mode for gunicorn workers + import os + import tempfile + from pathlib import Path + + from litellm._logging import verbose_proxy_logger + + if num_workers > 1 and not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + + try: + Path(multiproc_dir).mkdir(parents=True, exist_ok=True) + + # Clean up any stale files from previous runs + for file in Path(multiproc_dir).glob("*"): + if file.is_file(): + try: + file.unlink() + except Exception: + pass # Ignore errors if file is in use + + except PermissionError: + verbose_proxy_logger.warning( + f"Warning: Unable to create Prometheus multiprocess directory at {multiproc_dir} due to permission error. " + f"Running in non-root environment. Prometheus metrics may not work correctly in multiprocess mode." + ) + except Exception as e: + verbose_proxy_logger.warning( + f"Warning: Failed to create Prometheus multiprocess directory at {multiproc_dir}: {e}" + ) """ Run litellm with `gunicorn` """ @@ -525,6 +558,26 @@ def run_server( # noqa: PLR0915 skip_server_startup, keepalive_timeout, ): + # CRITICAL: Set up Prometheus multiprocess mode BEFORE any imports + # This ensures all worker processes will use multiprocess mode + import os + import tempfile + from pathlib import Path + + if num_workers > 1 and not os.environ.get("PROMETHEUS_MULTIPROC_DIR"): + multiproc_dir = os.path.join( + tempfile.gettempdir(), "litellm_prometheus_multiproc" + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + Path(multiproc_dir).mkdir(parents=True, exist_ok=True) + + # Clean up any stale files from previous runs + for file in Path(multiproc_dir).glob("*"): + if file.is_file(): + try: + file.unlink() + except Exception: + pass # Ignore errors if file is in use args = locals() if local: from proxy_server import ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ed6569f85f..198070156d3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1910,6 +1910,48 @@ class ProxyConfig: callback ) if "prometheus" in callback: + # CRITICAL: Set up prometheus multiprocess mode BEFORE importing + import os + import tempfile + from pathlib import Path + + # Check if we're in a multiprocess environment + num_workers = os.environ.get("NUM_WORKERS", "1") + is_multiprocess = False + + try: + if int(num_workers) > 1: + is_multiprocess = True + except (ValueError, TypeError): + pass + + # Check for gunicorn worker environment variables + if os.environ.get( + "GUNICORN_CMD_ARGS" + ) or os.environ.get("GUNICORN_WORKER_ID"): + is_multiprocess = True + + if is_multiprocess and not os.environ.get( + "PROMETHEUS_MULTIPROC_DIR" + ): + # Set up multiprocess directory + multiproc_dir = os.path.join( + tempfile.gettempdir(), + "litellm_prometheus_multiproc", + ) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = ( + multiproc_dir + ) + + # Ensure the directory exists + Path(multiproc_dir).mkdir( + parents=True, exist_ok=True + ) + + verbose_proxy_logger.info( + f"Prometheus multiprocess mode enabled with directory: {multiproc_dir}" + ) + try: from litellm_enterprise.integrations.prometheus import ( PrometheusLogger, @@ -1921,6 +1963,8 @@ class ProxyConfig: verbose_proxy_logger.debug( "mounting metrics endpoint" ) + # Clean up any existing multiprocess metrics before mounting + PrometheusLogger.cleanup_multiprocess_metrics() PrometheusLogger._mount_metrics_endpoint( premium_user ) @@ -2165,6 +2209,8 @@ class ProxyConfig: if assistant_settings: for k, v in assistant_settings["litellm_params"].items(): if isinstance(v, str) and v.startswith("os.environ/"): + import os + _v = v.replace("os.environ/", "") v = os.getenv(_v) assistant_settings["litellm_params"][k] = v From 459e66a9cdbb65512816e493a44cc9a41fd984c2 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 13 Sep 2025 13:36:05 -0700 Subject: [PATCH 2/2] fix: fix test --- .../integrations/prometheus.py | 14 -- .../test_prometheus_logging_callbacks.py | 236 ++++++++++-------- 2 files changed, 134 insertions(+), 116 deletions(-) diff --git a/enterprise/litellm_enterprise/integrations/prometheus.py b/enterprise/litellm_enterprise/integrations/prometheus.py index 836c79abb44..58943dc2dd4 100644 --- a/enterprise/litellm_enterprise/integrations/prometheus.py +++ b/enterprise/litellm_enterprise/integrations/prometheus.py @@ -1040,13 +1040,6 @@ class PrometheusLogger(CustomLogger): # increment total LLM requests and spend metric self._increment_top_level_request_and_spend_metrics( - end_user_id=end_user_id, - user_api_key=user_api_key, - user_api_key_alias=user_api_key_alias, - model=model, - user_api_team=user_api_team, - user_api_team_alias=user_api_team_alias, - user_id=user_id, response_cost=response_cost, enum_values=enum_values, ) @@ -1213,13 +1206,6 @@ class PrometheusLogger(CustomLogger): def _increment_top_level_request_and_spend_metrics( self, - end_user_id: Optional[str], - user_api_key: Optional[str], - user_api_key_alias: Optional[str], - model: Optional[str], - user_api_team: Optional[str], - user_api_team_alias: Optional[str], - user_id: Optional[str], response_cost: float, enum_values: UserAPIKeyLabelValues, ): diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 936d93c4417..2b23aab6285 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -560,13 +560,6 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger): prometheus_logger.litellm_spend_metric = MagicMock() prometheus_logger._increment_top_level_request_and_spend_metrics( - end_user_id="user1", - user_api_key="key1", - user_api_key_alias="alias1", - model="gpt-3.5-turbo", - user_api_team="team1", - user_api_team_alias="team_alias1", - user_id="user1", response_cost=0.1, enum_values=enum_values, ) @@ -584,7 +577,13 @@ def test_increment_top_level_request_and_spend_metrics(prometheus_logger): prometheus_logger.litellm_requests_metric.labels().inc.assert_called_once() prometheus_logger.litellm_spend_metric.labels.assert_called_once_with( - "user1", "key1", "alias1", "gpt-3.5-turbo", "team1", "team_alias1", "user1" + end_user=None, + hashed_api_key="test_hash", + api_key_alias="test_alias", + model="gpt-3.5-turbo", + team="test_team", + team_alias="test_team_alias", + user=None, ) prometheus_logger.litellm_spend_metric.labels().inc.assert_called_once_with(0.1) @@ -1141,22 +1140,28 @@ def test_get_custom_labels_from_tags_wildcard_patterns(monkeypatch): # Configure tags with wildcard patterns monkeypatch.setattr( - "litellm.custom_prometheus_tags", - ["User-Agent: curl/*", "User-Agent: python-requests/*", "Environment: prod*", "Service: api-gateway*", "exact-match"] + "litellm.custom_prometheus_tags", + [ + "User-Agent: curl/*", + "User-Agent: python-requests/*", + "Environment: prod*", + "Service: api-gateway*", + "exact-match", + ], ) - + # Test tags that should match the wildcard patterns tags = [ - "User-Agent: curl/7.68.0", - "User-Agent: python-requests/2.28.1", + "User-Agent: curl/7.68.0", + "User-Agent: python-requests/2.28.1", "Environment: production", "Service: api-gateway-v2", "exact-match", - "other-tag" + "other-tag", ] - + result = get_custom_labels_from_tags(tags) - + expected = { "tag_User_Agent__curl__": "true", # matches "User-Agent: curl/*" "tag_User_Agent__python_requests__": "true", # matches "User-Agent: python-requests/*" @@ -1164,7 +1169,7 @@ def test_get_custom_labels_from_tags_wildcard_patterns(monkeypatch): "tag_Service__api_gateway_": "true", # matches "Service: api-gateway*" "tag_exact_match": "true", # exact match } - + assert result == expected @@ -1174,26 +1179,26 @@ def test_get_custom_labels_from_tags_wildcard_no_matches(monkeypatch): # Configure tags with wildcard patterns monkeypatch.setattr( - "litellm.custom_prometheus_tags", - ["User-Agent: firefox/*", "Environment: dev*", "Service: web-app*"] + "litellm.custom_prometheus_tags", + ["User-Agent: firefox/*", "Environment: dev*", "Service: web-app*"], ) - + # Test tags that should NOT match the wildcard patterns tags = [ "User-Agent: curl/7.68.0", # doesn't match "User-Agent: firefox/*" - "Environment: production", # doesn't match "Environment: dev*" + "Environment: production", # doesn't match "Environment: dev*" "Service: api-gateway-v2", # doesn't match "Service: web-app*" - "other-tag" + "other-tag", ] - + result = get_custom_labels_from_tags(tags) - + expected = { "tag_User_Agent__firefox__": "false", # no match for "User-Agent: firefox/*" "tag_Environment__dev_": "false", # no match for "Environment: dev*" "tag_Service__web_app_": "false", # no match for "Service: web-app*" } - + assert result == expected @@ -1204,48 +1209,69 @@ def test_tag_matches_wildcard_configured_pattern(): ) # Test cases that should match - assert _tag_matches_wildcard_configured_pattern( - tags=["User-Agent: curl/7.68.0", "prod", "other"], - configured_tag="User-Agent: curl/*" - ) is True - - assert _tag_matches_wildcard_configured_pattern( - tags=["User-Agent: python-requests/2.28.1", "test"], - configured_tag="User-Agent: python-requests/*" - ) is True - - assert _tag_matches_wildcard_configured_pattern( - tags=["Environment: production", "debug"], - configured_tag="Environment: prod*" - ) is True - + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["User-Agent: curl/7.68.0", "prod", "other"], + configured_tag="User-Agent: curl/*", + ) + is True + ) + + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["User-Agent: python-requests/2.28.1", "test"], + configured_tag="User-Agent: python-requests/*", + ) + is True + ) + + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["Environment: production", "debug"], + configured_tag="Environment: prod*", + ) + is True + ) + # Test exact match (no wildcard) - assert _tag_matches_wildcard_configured_pattern( - tags=["prod", "test"], - configured_tag="prod" - ) is True - + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["prod", "test"], configured_tag="prod" + ) + is True + ) + # Test cases that should NOT match - assert _tag_matches_wildcard_configured_pattern( - tags=["User-Agent: firefox/98.0", "prod"], - configured_tag="User-Agent: curl/*" - ) is False - - assert _tag_matches_wildcard_configured_pattern( - tags=["Environment: development", "test"], - configured_tag="Environment: prod*" - ) is False - - assert _tag_matches_wildcard_configured_pattern( - tags=["staging", "test"], - configured_tag="prod" - ) is False - + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["User-Agent: firefox/98.0", "prod"], + configured_tag="User-Agent: curl/*", + ) + is False + ) + + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["Environment: development", "test"], + configured_tag="Environment: prod*", + ) + is False + ) + + assert ( + _tag_matches_wildcard_configured_pattern( + tags=["staging", "test"], configured_tag="prod" + ) + is False + ) + # Test with empty tags - assert _tag_matches_wildcard_configured_pattern( - tags=[], - configured_tag="User-Agent: curl/*" - ) is False + assert ( + _tag_matches_wildcard_configured_pattern( + tags=[], configured_tag="User-Agent: curl/*" + ) + is False + ) @pytest.mark.asyncio(scope="session") @@ -1908,12 +1934,12 @@ def test_set_llm_deployment_success_metrics_with_label_filtering(): async def test_prometheus_token_metrics_with_prometheus_config(): """ Test that validates the renamed token metrics are incremented correctly with a prometheus config. - + This test ensures that after the metric renaming (git diff): - litellm_total_tokens -> litellm_total_tokens_metric - - litellm_input_tokens -> litellm_input_tokens_metric + - litellm_input_tokens -> litellm_input_tokens_metric - litellm_output_tokens -> litellm_output_tokens_metric - + All three metrics should be properly incremented when making a successful completion request. """ from prometheus_client import CollectorRegistry, Counter @@ -1925,39 +1951,39 @@ async def test_prometheus_token_metrics_with_prometheus_config(): collectors = list(REGISTRY._collector_to_names.keys()) for collector in collectors: REGISTRY.unregister(collector) - + # Set up prometheus configuration that includes the token metrics config = [ PrometheusMetricsConfig( group="token_metrics_test", metrics=[ "litellm_total_tokens_metric", - "litellm_input_tokens_metric", + "litellm_input_tokens_metric", "litellm_output_tokens_metric", - "litellm_requests_metric" + "litellm_requests_metric", ], include_labels=[ "model", - "hashed_api_key", + "hashed_api_key", "api_key_alias", "team", - "team_alias" + "team_alias", ], ) ] - + # Mock litellm.prometheus_metrics_config with patch("litellm.prometheus_metrics_config", config): # Create PrometheusLogger with the configuration prometheus_logger = PrometheusLogger() - + # Test data with specific token counts standard_logging_payload = create_standard_logging_payload() standard_logging_payload["total_tokens"] = 1500 standard_logging_payload["prompt_tokens"] = 900 standard_logging_payload["completion_tokens"] = 600 standard_logging_payload["response_cost"] = 0.075 - + kwargs = { "model": "gpt-3.5-turbo", "stream": False, @@ -1971,7 +1997,7 @@ async def test_prometheus_token_metrics_with_prometheus_config(): } }, "start_time": datetime.now() - timedelta(seconds=2), - "completion_start_time": datetime.now() - timedelta(seconds=1), + "completion_start_time": datetime.now() - timedelta(seconds=1), "api_call_start_time": datetime.now() - timedelta(seconds=1.5), "end_time": datetime.now(), "standard_logging_object": standard_logging_payload, @@ -1987,69 +2013,75 @@ async def test_prometheus_token_metrics_with_prometheus_config(): print("final registry values", REGISTRY._collector_to_names) - # Get metric collectors directly from registry + # Get metric collectors directly from registry metric_collectors = {} for collector, names in REGISTRY._collector_to_names.items(): metric_name = names[0] # First name is the base metric name metric_collectors[metric_name] = collector print("=== Final Metric Values (Direct Access) ===") - - # Expected values + + # Expected values expected_values = { "litellm_total_tokens_metric": 1500.0, "litellm_input_tokens_metric": 900.0, "litellm_output_tokens_metric": 600.0, - "litellm_requests_metric": 1.0 + "litellm_requests_metric": 1.0, } - + expected_label_values = { - 'api_key_alias': 'test_alias', - 'hashed_api_key': 'test_hash', - 'model': 'gpt-3.5-turbo', - 'team': 'test_team', - 'team_alias': 'test_team_alias' + "api_key_alias": "test_alias", + "hashed_api_key": "test_hash", + "model": "gpt-3.5-turbo", + "team": "test_team", + "team_alias": "test_team_alias", } # Validate each metric directly for metric_name, expected_value in expected_values.items(): if metric_name in metric_collectors: collector = metric_collectors[metric_name] - + # Get all samples for this metric samples = list(collector.collect())[0].samples - + # Find the _total sample (the actual counter value) total_sample = None for sample in samples: - if sample.name.endswith('_total'): + if sample.name.endswith("_total"): total_sample = sample break - + if total_sample: actual_value = total_sample.value actual_labels = total_sample.labels - - print(f"āœ“ {metric_name}: expected={expected_value}, actual={actual_value}") + + print( + f"āœ“ {metric_name}: expected={expected_value}, actual={actual_value}" + ) print(f" Labels: {actual_labels}") - + # Validate the value - assert actual_value == expected_value, f"Expected {expected_value}, got {actual_value} for {metric_name}" - + assert ( + actual_value == expected_value + ), f"Expected {expected_value}, got {actual_value} for {metric_name}" + # Validate the labels - for label_key, expected_label_value in expected_label_values.items(): + for ( + label_key, + expected_label_value, + ) in expected_label_values.items(): actual_label_value = actual_labels.get(label_key) - assert actual_label_value == expected_label_value, f"Expected label {label_key}={expected_label_value}, got {actual_label_value}" - + assert ( + actual_label_value == expected_label_value + ), f"Expected label {label_key}={expected_label_value}, got {actual_label_value}" + print(f" āœ“ {metric_name} VALIDATED") else: raise AssertionError(f"No _total sample found for {metric_name}") else: raise AssertionError(f"Metric {metric_name} not found in registry") - + print("āœ“ All token metrics validated successfully!") # check final value of metrics in registry - - -