scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge pull request #14530 from BerriAI/litellm_dev_09_10_2025_p1

Browse source 
fix(prometheus.py): make prometheus work for multiple workers
This commit is contained in:
Krish Dholakia 2025-09-19 22:26:56 -07:00 • committed by GitHub
commit d6553045a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 498 additions and 141 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

271
enterprise/litellm_enterprise/integrations/prometheus.py
View file

@ -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}"
@ -134,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(
@ -183,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(
@ -221,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(
@ -231,18 +314,20 @@ 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"),
multiprocess_mode="mostrecent",
)
self.litellm_deployment_cooled_down = self._counter_factory(
@ -320,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
@ -729,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()
@ -847,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,
)
@ -1020,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,
):
@ -1045,7 +1224,6 @@ class PrometheusLogger(CustomLogger):
),
enum_values=enum_values,
)
self.litellm_spend_metric.labels(**_labels).inc(response_cost)
def _set_virtual_key_rate_limit_metrics(
@ -2173,13 +2351,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
@ -2190,14 +2369,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(
@ -2328,9 +2527,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
@ -2338,7 +2534,6 @@ 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}")

2
litellm/proxy/_new_secret_config.yaml
View file

@ -14,4 +14,4 @@ model_list:
litellm_params:
model: hosted_vllm/whisper-v3
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
api_key: dummy

31
litellm/proxy/common_utils/callback_utils.py
View file

@ -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 = [

53
litellm/proxy/proxy_cli.py
View file

@ -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 (

46
litellm/proxy/proxy_server.py
View file

@ -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

236
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
View file

@ -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