mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat(prometheus): add caching metrics (#18755)
This commit is contained in:
parent
91b5c66cf2
commit
92f7789f10
3 changed files with 318 additions and 15 deletions
|
|
@ -14,6 +14,7 @@ from typing import (
|
|||
Literal,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
|
|
@ -44,6 +45,7 @@ def _get_cached_end_user_id_for_cost_tracking():
|
|||
global _get_end_user_id_for_cost_tracking
|
||||
if _get_end_user_id_for_cost_tracking is None:
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
_get_end_user_id_for_cost_tracking = get_end_user_id_for_cost_tracking
|
||||
return _get_end_user_id_for_cost_tracking
|
||||
|
||||
|
|
@ -329,6 +331,25 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_requests_metric"),
|
||||
)
|
||||
|
||||
# Cache metrics
|
||||
self.litellm_cache_hits_metric = self._counter_factory(
|
||||
name="litellm_cache_hits_metric",
|
||||
documentation="Total number of LiteLLM cache hits",
|
||||
labelnames=self.get_labels_for_metric("litellm_cache_hits_metric"),
|
||||
)
|
||||
|
||||
self.litellm_cache_misses_metric = self._counter_factory(
|
||||
name="litellm_cache_misses_metric",
|
||||
documentation="Total number of LiteLLM cache misses",
|
||||
labelnames=self.get_labels_for_metric("litellm_cache_misses_metric"),
|
||||
)
|
||||
|
||||
self.litellm_cached_tokens_metric = self._counter_factory(
|
||||
name="litellm_cached_tokens_metric",
|
||||
documentation="Total tokens served from LiteLLM cache",
|
||||
labelnames=self.get_labels_for_metric("litellm_cached_tokens_metric"),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print_verbose(f"Got exception on init prometheus client {str(e)}")
|
||||
raise e
|
||||
|
|
@ -795,7 +816,7 @@ class PrometheusLogger(CustomLogger):
|
|||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
_metadata = litellm_params.get("metadata", {})
|
||||
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
|
||||
|
||||
|
||||
end_user_id = get_end_user_id_for_cost_tracking(
|
||||
litellm_params, service_type="prometheus"
|
||||
)
|
||||
|
|
@ -815,7 +836,7 @@ class PrometheusLogger(CustomLogger):
|
|||
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
|
||||
"metadata"
|
||||
].get("user_api_key_auth_metadata")
|
||||
|
||||
|
||||
# Include top-level metadata fields (excluding nested dictionaries)
|
||||
# This allows accessing fields like requester_ip_address from top-level metadata
|
||||
top_level_metadata = standard_logging_payload.get("metadata", {})
|
||||
|
|
@ -826,7 +847,7 @@ class PrometheusLogger(CustomLogger):
|
|||
for k, v in top_level_metadata.items()
|
||||
if not isinstance(v, dict) # Exclude nested dicts to avoid conflicts
|
||||
}
|
||||
|
||||
|
||||
combined_metadata: Dict[str, Any] = {
|
||||
**top_level_fields, # Include top-level fields first
|
||||
**(_requester_metadata if _requester_metadata else {}),
|
||||
|
|
@ -945,6 +966,12 @@ class PrometheusLogger(CustomLogger):
|
|||
kwargs, start_time, end_time, enum_values, output_tokens
|
||||
)
|
||||
|
||||
# cache metrics
|
||||
self._increment_cache_metrics(
|
||||
standard_logging_payload=standard_logging_payload, # type: ignore
|
||||
enum_values=enum_values,
|
||||
)
|
||||
|
||||
if (
|
||||
standard_logging_payload["stream"] is True
|
||||
): # log successful streaming requests from logging event hook.
|
||||
|
|
@ -1014,6 +1041,54 @@ class PrometheusLogger(CustomLogger):
|
|||
standard_logging_payload["completion_tokens"]
|
||||
)
|
||||
|
||||
def _increment_cache_metrics(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
):
|
||||
"""
|
||||
Increment cache-related Prometheus metrics based on cache hit/miss status.
|
||||
|
||||
Args:
|
||||
standard_logging_payload: Contains cache_hit field (True/False/None)
|
||||
enum_values: Label values for Prometheus metrics
|
||||
"""
|
||||
cache_hit = standard_logging_payload.get("cache_hit")
|
||||
|
||||
# Only track if cache_hit has a definite value (True or False)
|
||||
if cache_hit is None:
|
||||
return
|
||||
|
||||
if cache_hit is True:
|
||||
# Increment cache hits counter
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_cache_hits_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_cache_hits_metric.labels(**_labels).inc()
|
||||
|
||||
# Increment cached tokens counter
|
||||
total_tokens = standard_logging_payload.get("total_tokens", 0)
|
||||
if total_tokens > 0:
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_cached_tokens_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_cached_tokens_metric.labels(**_labels).inc(total_tokens)
|
||||
else:
|
||||
# cache_hit is False - increment cache misses counter
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_cache_misses_metric"
|
||||
),
|
||||
enum_values=enum_values,
|
||||
)
|
||||
self.litellm_cache_misses_metric.labels(**_labels).inc()
|
||||
|
||||
async def _increment_remaining_budget_metrics(
|
||||
self,
|
||||
user_api_team: Optional[str],
|
||||
|
|
@ -1196,7 +1271,7 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
|
||||
|
||||
|
||||
end_user_id = get_end_user_id_for_cost_tracking(
|
||||
litellm_params, service_type="prometheus"
|
||||
)
|
||||
|
|
@ -1398,7 +1473,6 @@ class PrometheusLogger(CustomLogger):
|
|||
api_provider=llm_provider or "",
|
||||
)
|
||||
if exception is not None:
|
||||
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name="litellm_deployment_failure_responses"
|
||||
|
|
@ -1431,12 +1505,11 @@ class PrometheusLogger(CustomLogger):
|
|||
enum_values: UserAPIKeyLabelValues,
|
||||
output_tokens: float = 1.0,
|
||||
):
|
||||
|
||||
try:
|
||||
verbose_logger.debug("setting remaining tokens requests metric")
|
||||
standard_logging_payload: Optional[StandardLoggingPayload] = (
|
||||
request_kwargs.get("standard_logging_object")
|
||||
)
|
||||
standard_logging_payload: Optional[
|
||||
StandardLoggingPayload
|
||||
] = request_kwargs.get("standard_logging_object")
|
||||
|
||||
if standard_logging_payload is None:
|
||||
return
|
||||
|
|
@ -2208,10 +2281,10 @@ class PrometheusLogger(CustomLogger):
|
|||
from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
prometheus_loggers: List[CustomLogger] = (
|
||||
litellm.logging_callback_manager.get_custom_loggers_for_type(
|
||||
callback_type=PrometheusLogger
|
||||
)
|
||||
prometheus_loggers: List[
|
||||
CustomLogger
|
||||
] = litellm.logging_callback_manager.get_custom_loggers_for_type(
|
||||
callback_type=PrometheusLogger
|
||||
)
|
||||
# we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them
|
||||
verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers))
|
||||
|
|
@ -2283,7 +2356,7 @@ def prometheus_label_factory(
|
|||
|
||||
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
|
||||
get_end_user_id_for_cost_tracking = _get_cached_end_user_id_for_cost_tracking()
|
||||
|
||||
|
||||
filtered_labels["end_user"] = get_end_user_id_for_cost_tracking(
|
||||
litellm_params={"user_api_key_end_user_id": enum_values.end_user},
|
||||
service_type="prometheus",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import re
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Dict, List, Literal, Optional, Tuple, Union
|
||||
from typing import Dict, List, Literal, Optional, Tuple
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import Annotated
|
||||
|
|
@ -185,6 +185,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
"litellm_redis_daily_spend_update_queue_size",
|
||||
"litellm_in_memory_spend_update_queue_size",
|
||||
"litellm_redis_spend_update_queue_size",
|
||||
# Cache metrics
|
||||
"litellm_cache_hits_metric",
|
||||
"litellm_cache_misses_metric",
|
||||
"litellm_cached_tokens_metric",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -436,6 +440,21 @@ class PrometheusMetricLabels:
|
|||
|
||||
litellm_redis_spend_update_queue_size: List[str] = []
|
||||
|
||||
# Cache metrics - track cache hits, misses, and tokens served from cache
|
||||
_cache_metric_labels = [
|
||||
UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value,
|
||||
UserAPIKeyLabelNames.API_KEY_HASH.value,
|
||||
UserAPIKeyLabelNames.API_KEY_ALIAS.value,
|
||||
UserAPIKeyLabelNames.TEAM.value,
|
||||
UserAPIKeyLabelNames.TEAM_ALIAS.value,
|
||||
UserAPIKeyLabelNames.END_USER.value,
|
||||
UserAPIKeyLabelNames.USER.value,
|
||||
]
|
||||
|
||||
litellm_cache_hits_metric = _cache_metric_labels
|
||||
litellm_cache_misses_metric = _cache_metric_labels
|
||||
litellm_cached_tokens_metric = _cache_metric_labels
|
||||
|
||||
@staticmethod
|
||||
def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]:
|
||||
default_labels = getattr(PrometheusMetricLabels, label_name)
|
||||
|
|
|
|||
211
tests/test_litellm/integrations/test_prometheus_cache_metrics.py
Normal file
211
tests/test_litellm/integrations/test_prometheus_cache_metrics.py
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
"""
|
||||
Unit tests for cache Prometheus metrics.
|
||||
|
||||
Run with: poetry run pytest tests/test_litellm/integrations/test_prometheus_cache_metrics.py -v
|
||||
"""
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch
|
||||
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
|
||||
|
||||
|
||||
class TestPrometheusCacheMetrics:
|
||||
"""Tests for cache-related Prometheus metrics"""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_enum_values(self):
|
||||
"""Create sample enum values for labels"""
|
||||
return UserAPIKeyLabelValues(
|
||||
end_user="test-end-user",
|
||||
hashed_api_key="test-key-hash",
|
||||
api_key_alias="test-key-alias",
|
||||
team="test-team",
|
||||
team_alias="test-team-alias",
|
||||
user="test-user",
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
|
||||
def test_cache_metrics_defined_in_types(self):
|
||||
"""Test that cache metrics are defined in DEFINED_PROMETHEUS_METRICS"""
|
||||
from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS
|
||||
from typing import get_args
|
||||
|
||||
defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS)
|
||||
|
||||
assert "litellm_cache_hits_metric" in defined_metrics
|
||||
assert "litellm_cache_misses_metric" in defined_metrics
|
||||
assert "litellm_cached_tokens_metric" in defined_metrics
|
||||
|
||||
def test_cache_metric_labels_defined(self):
|
||||
"""Test that cache metric labels are properly defined"""
|
||||
from litellm.types.integrations.prometheus import PrometheusMetricLabels
|
||||
|
||||
# Verify labels are defined for each cache metric
|
||||
assert hasattr(PrometheusMetricLabels, "litellm_cache_hits_metric")
|
||||
assert hasattr(PrometheusMetricLabels, "litellm_cache_misses_metric")
|
||||
assert hasattr(PrometheusMetricLabels, "litellm_cached_tokens_metric")
|
||||
|
||||
# Verify labels include expected keys
|
||||
expected_labels = [
|
||||
"model",
|
||||
"hashed_api_key",
|
||||
"api_key_alias",
|
||||
"team",
|
||||
"team_alias",
|
||||
"end_user",
|
||||
"user",
|
||||
]
|
||||
for label in expected_labels:
|
||||
assert label in PrometheusMetricLabels.litellm_cache_hits_metric
|
||||
assert label in PrometheusMetricLabels.litellm_cache_misses_metric
|
||||
assert label in PrometheusMetricLabels.litellm_cached_tokens_metric
|
||||
|
||||
def test_increment_cache_metrics_on_cache_hit(self, sample_enum_values):
|
||||
"""Test that cache hit increments the correct metrics"""
|
||||
# Create mock for PrometheusLogger instance
|
||||
mock_logger = MagicMock()
|
||||
|
||||
# Import the method directly and bind it to our mock
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
# Create a mock standard logging payload with cache_hit=True
|
||||
standard_logging_payload = {
|
||||
"cache_hit": True,
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
"model_group": "openai",
|
||||
"request_tags": [],
|
||||
}
|
||||
|
||||
# Create mock metrics
|
||||
mock_logger.litellm_cache_hits_metric = MagicMock()
|
||||
mock_logger.litellm_cache_misses_metric = MagicMock()
|
||||
mock_logger.litellm_cached_tokens_metric = MagicMock()
|
||||
mock_logger.get_labels_for_metric = MagicMock(
|
||||
return_value=[
|
||||
"model",
|
||||
"hashed_api_key",
|
||||
"api_key_alias",
|
||||
"team",
|
||||
"team_alias",
|
||||
"end_user",
|
||||
"user",
|
||||
]
|
||||
)
|
||||
|
||||
# Call the method using unbound method approach
|
||||
PrometheusLogger._increment_cache_metrics(
|
||||
mock_logger,
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=sample_enum_values,
|
||||
)
|
||||
|
||||
# Verify cache hits metric was incremented
|
||||
mock_logger.litellm_cache_hits_metric.labels.assert_called()
|
||||
mock_logger.litellm_cache_hits_metric.labels().inc.assert_called_once()
|
||||
|
||||
# Verify cached tokens metric was incremented with total_tokens
|
||||
mock_logger.litellm_cached_tokens_metric.labels.assert_called()
|
||||
mock_logger.litellm_cached_tokens_metric.labels().inc.assert_called_once_with(
|
||||
100
|
||||
)
|
||||
|
||||
# Verify cache misses metric was NOT called
|
||||
mock_logger.litellm_cache_misses_metric.labels.assert_not_called()
|
||||
|
||||
def test_increment_cache_metrics_on_cache_miss(self, sample_enum_values):
|
||||
"""Test that cache miss increments the correct metrics"""
|
||||
# Create mock for PrometheusLogger instance
|
||||
mock_logger = MagicMock()
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
# Create a mock standard logging payload with cache_hit=False
|
||||
standard_logging_payload = {
|
||||
"cache_hit": False,
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
"model_group": "openai",
|
||||
"request_tags": [],
|
||||
}
|
||||
|
||||
# Create mock metrics
|
||||
mock_logger.litellm_cache_hits_metric = MagicMock()
|
||||
mock_logger.litellm_cache_misses_metric = MagicMock()
|
||||
mock_logger.litellm_cached_tokens_metric = MagicMock()
|
||||
mock_logger.get_labels_for_metric = MagicMock(
|
||||
return_value=[
|
||||
"model",
|
||||
"hashed_api_key",
|
||||
"api_key_alias",
|
||||
"team",
|
||||
"team_alias",
|
||||
"end_user",
|
||||
"user",
|
||||
]
|
||||
)
|
||||
|
||||
# Call the method
|
||||
PrometheusLogger._increment_cache_metrics(
|
||||
mock_logger,
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=sample_enum_values,
|
||||
)
|
||||
|
||||
# Verify cache misses metric was incremented
|
||||
mock_logger.litellm_cache_misses_metric.labels.assert_called()
|
||||
mock_logger.litellm_cache_misses_metric.labels().inc.assert_called_once()
|
||||
|
||||
# Verify cache hits and cached tokens metrics were NOT called
|
||||
mock_logger.litellm_cache_hits_metric.labels.assert_not_called()
|
||||
mock_logger.litellm_cached_tokens_metric.labels.assert_not_called()
|
||||
|
||||
def test_increment_cache_metrics_when_cache_hit_is_none(self, sample_enum_values):
|
||||
"""Test that no metrics are incremented when cache_hit is None"""
|
||||
# Create mock for PrometheusLogger instance
|
||||
mock_logger = MagicMock()
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
# Create a mock standard logging payload with cache_hit=None
|
||||
standard_logging_payload = {
|
||||
"cache_hit": None,
|
||||
"total_tokens": 100,
|
||||
"prompt_tokens": 50,
|
||||
"completion_tokens": 50,
|
||||
"model_group": "openai",
|
||||
"request_tags": [],
|
||||
}
|
||||
|
||||
# Create mock metrics
|
||||
mock_logger.litellm_cache_hits_metric = MagicMock()
|
||||
mock_logger.litellm_cache_misses_metric = MagicMock()
|
||||
mock_logger.litellm_cached_tokens_metric = MagicMock()
|
||||
mock_logger.get_labels_for_metric = MagicMock(
|
||||
return_value=[
|
||||
"model",
|
||||
"hashed_api_key",
|
||||
"api_key_alias",
|
||||
"team",
|
||||
"team_alias",
|
||||
"end_user",
|
||||
"user",
|
||||
]
|
||||
)
|
||||
|
||||
# Call the method
|
||||
PrometheusLogger._increment_cache_metrics(
|
||||
mock_logger,
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=sample_enum_values,
|
||||
)
|
||||
|
||||
# Verify NO metrics were called
|
||||
mock_logger.litellm_cache_hits_metric.labels.assert_not_called()
|
||||
mock_logger.litellm_cache_misses_metric.labels.assert_not_called()
|
||||
mock_logger.litellm_cached_tokens_metric.labels.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
Loading…
Add table
Reference in a new issue