mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 02:34:29 +00:00
281 lines
11 KiB
Python
281 lines
11 KiB
Python
# used for monitoring litellm services health on `/metrics` endpoint on LiteLLM Proxy
|
|
#### What this does ####
|
|
# On success + failure, log events to Prometheus for litellm / adjacent services (litellm, redis, postgres, llm api providers)
|
|
|
|
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
from litellm._logging import print_verbose, verbose_logger
|
|
from litellm.types.integrations.prometheus import LATENCY_BUCKETS
|
|
from litellm.types.services import (
|
|
DEFAULT_SERVICE_CONFIGS,
|
|
ServiceLoggerPayload,
|
|
ServiceMetrics,
|
|
ServiceTypes,
|
|
)
|
|
|
|
FAILED_REQUESTS_LABELS = ["error_class", "function_name"]
|
|
|
|
|
|
class PrometheusServicesLogger:
|
|
# Class variables or attributes
|
|
litellm_service_latency = None # Class-level attribute to store the Histogram
|
|
|
|
def __init__(
|
|
self,
|
|
mock_testing: bool = False,
|
|
**kwargs,
|
|
):
|
|
try:
|
|
try:
|
|
from prometheus_client import REGISTRY, Counter, Gauge, Histogram
|
|
from prometheus_client.gc_collector import Collector
|
|
except ImportError:
|
|
raise Exception(
|
|
"Missing prometheus_client. Run `pip install prometheus-client`"
|
|
)
|
|
|
|
self.Histogram = Histogram
|
|
self.Counter = Counter
|
|
self.Gauge = Gauge
|
|
self.REGISTRY = REGISTRY
|
|
|
|
verbose_logger.debug("in init prometheus services metrics")
|
|
|
|
self.payload_to_prometheus_map: Dict[
|
|
str, List[Union[Histogram, Counter, Gauge, Collector]]
|
|
] = {}
|
|
|
|
for service in ServiceTypes:
|
|
service_metrics: List[Union[Histogram, Counter, Gauge, Collector]] = []
|
|
|
|
metrics_to_initialize = self._get_service_metrics_initialize(service)
|
|
|
|
# Initialize only the configured metrics for each service
|
|
if ServiceMetrics.HISTOGRAM in metrics_to_initialize:
|
|
histogram = self.create_histogram(
|
|
service.value, type_of_request="latency"
|
|
)
|
|
if histogram:
|
|
service_metrics.append(histogram)
|
|
|
|
if ServiceMetrics.COUNTER in metrics_to_initialize:
|
|
counter_failed_request = self.create_counter(
|
|
service.value,
|
|
type_of_request="failed_requests",
|
|
additional_labels=FAILED_REQUESTS_LABELS,
|
|
)
|
|
if counter_failed_request:
|
|
service_metrics.append(counter_failed_request)
|
|
counter_total_requests = self.create_counter(
|
|
service.value, type_of_request="total_requests"
|
|
)
|
|
if counter_total_requests:
|
|
service_metrics.append(counter_total_requests)
|
|
|
|
if ServiceMetrics.GAUGE in metrics_to_initialize:
|
|
gauge = self.create_gauge(service.value, type_of_request="size")
|
|
if gauge:
|
|
service_metrics.append(gauge)
|
|
|
|
if service_metrics:
|
|
self.payload_to_prometheus_map[service.value] = service_metrics
|
|
|
|
self.prometheus_to_amount_map: dict = {}
|
|
### MOCK TESTING ###
|
|
self.mock_testing = mock_testing
|
|
self.mock_testing_success_calls = 0
|
|
self.mock_testing_failure_calls = 0
|
|
|
|
except Exception as e:
|
|
print_verbose(f"Got exception on init prometheus client {str(e)}")
|
|
raise e
|
|
|
|
def _get_service_metrics_initialize(
|
|
self, service: ServiceTypes
|
|
) -> List[ServiceMetrics]:
|
|
DEFAULT_METRICS = [ServiceMetrics.COUNTER, ServiceMetrics.HISTOGRAM]
|
|
if service not in DEFAULT_SERVICE_CONFIGS:
|
|
return DEFAULT_METRICS
|
|
|
|
metrics = DEFAULT_SERVICE_CONFIGS.get(service, {}).get("metrics", [])
|
|
if not metrics:
|
|
verbose_logger.debug(f"No metrics found for service {service}")
|
|
return DEFAULT_METRICS
|
|
return metrics
|
|
|
|
def is_metric_registered(self, metric_name) -> bool:
|
|
for metric in self.REGISTRY.collect():
|
|
if metric_name == metric.name:
|
|
return True
|
|
return False
|
|
|
|
def _get_metric(self, metric_name):
|
|
"""
|
|
Helper function to get a metric from the registry by name.
|
|
"""
|
|
return self.REGISTRY._names_to_collectors.get(metric_name)
|
|
|
|
def create_histogram(self, service: str, type_of_request: str):
|
|
metric_name = "litellm_{}_{}".format(service, type_of_request)
|
|
is_registered = self.is_metric_registered(metric_name)
|
|
if is_registered:
|
|
return self._get_metric(metric_name)
|
|
return self.Histogram(
|
|
metric_name,
|
|
"Latency for {} service".format(service),
|
|
labelnames=[service],
|
|
buckets=LATENCY_BUCKETS,
|
|
)
|
|
|
|
def create_gauge(self, service: str, type_of_request: str):
|
|
metric_name = "litellm_{}_{}".format(service, type_of_request)
|
|
is_registered = self.is_metric_registered(metric_name)
|
|
if is_registered:
|
|
return self._get_metric(metric_name)
|
|
return self.Gauge(
|
|
metric_name, "Gauge for {} service".format(service), labelnames=[service]
|
|
)
|
|
|
|
def create_counter(
|
|
self,
|
|
service: str,
|
|
type_of_request: str,
|
|
additional_labels: Optional[List[str]] = None,
|
|
):
|
|
metric_name = "litellm_{}_{}".format(service, type_of_request)
|
|
is_registered = self.is_metric_registered(metric_name)
|
|
if is_registered:
|
|
return self._get_metric(metric_name)
|
|
return self.Counter(
|
|
metric_name,
|
|
"Total {} for {} service".format(type_of_request, service),
|
|
labelnames=[service] + (additional_labels or []),
|
|
)
|
|
|
|
def observe_histogram(
|
|
self,
|
|
histogram,
|
|
labels: str,
|
|
amount: float,
|
|
):
|
|
assert isinstance(histogram, self.Histogram)
|
|
|
|
histogram.labels(labels).observe(amount)
|
|
|
|
def update_gauge(
|
|
self,
|
|
gauge,
|
|
labels: str,
|
|
amount: float,
|
|
):
|
|
assert isinstance(gauge, self.Gauge)
|
|
gauge.labels(labels).set(amount)
|
|
|
|
def increment_counter(
|
|
self,
|
|
counter,
|
|
labels: str,
|
|
amount: float,
|
|
additional_labels: Optional[List[str]] = [],
|
|
):
|
|
assert isinstance(counter, self.Counter)
|
|
|
|
if additional_labels:
|
|
counter.labels(labels, *additional_labels).inc(amount)
|
|
else:
|
|
counter.labels(labels).inc(amount)
|
|
|
|
def service_success_hook(self, payload: ServiceLoggerPayload):
|
|
if self.mock_testing:
|
|
self.mock_testing_success_calls += 1
|
|
|
|
if payload.service.value in self.payload_to_prometheus_map:
|
|
prom_objects = self.payload_to_prometheus_map[payload.service.value]
|
|
for obj in prom_objects:
|
|
if isinstance(obj, self.Histogram):
|
|
self.observe_histogram(
|
|
histogram=obj,
|
|
labels=payload.service.value,
|
|
amount=payload.duration,
|
|
)
|
|
elif isinstance(obj, self.Counter) and "total_requests" in obj._name:
|
|
self.increment_counter(
|
|
counter=obj,
|
|
labels=payload.service.value,
|
|
amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS
|
|
)
|
|
|
|
def service_failure_hook(self, payload: ServiceLoggerPayload):
|
|
if self.mock_testing:
|
|
self.mock_testing_failure_calls += 1
|
|
|
|
if payload.service.value in self.payload_to_prometheus_map:
|
|
prom_objects = self.payload_to_prometheus_map[payload.service.value]
|
|
for obj in prom_objects:
|
|
if isinstance(obj, self.Counter):
|
|
self.increment_counter(
|
|
counter=obj,
|
|
labels=payload.service.value,
|
|
amount=1, # LOG ERROR COUNT / TOTAL REQUESTS TO PROMETHEUS
|
|
)
|
|
|
|
async def async_service_success_hook(self, payload: ServiceLoggerPayload):
|
|
"""
|
|
Log successful call to prometheus
|
|
"""
|
|
if self.mock_testing:
|
|
self.mock_testing_success_calls += 1
|
|
|
|
if payload.service.value in self.payload_to_prometheus_map:
|
|
prom_objects = self.payload_to_prometheus_map[payload.service.value]
|
|
for obj in prom_objects:
|
|
if isinstance(obj, self.Histogram):
|
|
self.observe_histogram(
|
|
histogram=obj,
|
|
labels=payload.service.value,
|
|
amount=payload.duration,
|
|
)
|
|
elif isinstance(obj, self.Counter) and "total_requests" in obj._name:
|
|
self.increment_counter(
|
|
counter=obj,
|
|
labels=payload.service.value,
|
|
amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS
|
|
)
|
|
elif isinstance(obj, self.Gauge):
|
|
if payload.event_metadata:
|
|
self.update_gauge(
|
|
gauge=obj,
|
|
labels=payload.event_metadata.get("gauge_labels") or "",
|
|
amount=payload.event_metadata.get("gauge_value") or 0,
|
|
)
|
|
|
|
async def async_service_failure_hook(
|
|
self,
|
|
payload: ServiceLoggerPayload,
|
|
error: Union[str, Exception],
|
|
):
|
|
if self.mock_testing:
|
|
self.mock_testing_failure_calls += 1
|
|
error_class = error.__class__.__name__
|
|
function_name = payload.call_type
|
|
|
|
if payload.service.value in self.payload_to_prometheus_map:
|
|
prom_objects = self.payload_to_prometheus_map[payload.service.value]
|
|
for obj in prom_objects:
|
|
# increment both failed and total requests
|
|
if isinstance(obj, self.Counter):
|
|
if "failed_requests" in obj._name:
|
|
self.increment_counter(
|
|
counter=obj,
|
|
labels=payload.service.value,
|
|
# log additional_labels=["error_class", "function_name"], used for debugging what's going wrong with the DB
|
|
additional_labels=[error_class, function_name],
|
|
amount=1, # LOG ERROR COUNT TO PROMETHEUS
|
|
)
|
|
else:
|
|
self.increment_counter(
|
|
counter=obj,
|
|
labels=payload.service.value,
|
|
amount=1, # LOG TOTAL REQUESTS TO PROMETHEUS
|
|
)
|