From 6cb570da3b0ff180c9270e4b62c43d31718c81d5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 17:47:40 -0700 Subject: [PATCH] fix(prometheus): address greptile review - Sequence type, is-not-None check, drop hasattr guards, add exclude tests, fix registry assertions for labeled histograms --- litellm/__init__.py | 5 +- litellm/integrations/prometheus.py | 20 ++- .../test_prometheus_latency_buckets.py | 136 ++++++++++++++---- 3 files changed, 129 insertions(+), 32 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 72dd135b386..39eb53f42a6 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -24,6 +24,7 @@ from typing import ( List, Optional, Dict, + Sequence, Union, Any, Literal, @@ -398,8 +399,8 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None custom_prometheus_metadata_labels: List[str] = [] custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None -prometheus_latency_buckets: Optional[Tuple] = ( - None # override default LATENCY_BUCKETS for histogram metrics +prometheus_latency_buckets: Optional[Sequence[float]] = ( + None # override default LATENCY_BUCKETS for histogram metrics; must be set before PrometheusLogger is first instantiated ) prometheus_exclude_metrics: Optional[List[str]] = ( None # metric names to disable entirely diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index adb5e8bfc56..a3cc7b6fd32 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -14,6 +14,7 @@ from typing import ( List, Literal, Optional, + Sequence, Tuple, Union, cast, @@ -445,15 +446,26 @@ class PrometheusLogger(CustomLogger): print_verbose(f"Got exception on init prometheus client {str(e)}") raise e - def _get_latency_buckets(self) -> tuple: + def _get_latency_buckets(self) -> Sequence[float]: """Return latency buckets to use for histogram metrics. Uses ``litellm.prometheus_latency_buckets`` when set, falling back to the module-level ``LATENCY_BUCKETS`` constant. + + .. note:: + This value is read once at ``PrometheusLogger`` instantiation time. + ``litellm.prometheus_latency_buckets`` must be configured **before** the + first request is processed (i.e. before the logger is created). + Changes made after initialisation are silently ignored because Prometheus + histograms cannot be re-registered mid-run. """ import litellm - return litellm.prometheus_latency_buckets or LATENCY_BUCKETS + return ( + litellm.prometheus_latency_buckets + if litellm.prometheus_latency_buckets is not None + else LATENCY_BUCKETS + ) def _parse_exclude_config(self) -> None: """Populate self.excluded_metrics and self.excluded_labels from litellm module vars.""" @@ -857,7 +869,7 @@ class PrometheusLogger(CustomLogger): def _is_metric_enabled(self, metric_name: str) -> bool: """Check if a metric is enabled based on configuration""" # Check exclude list first — excluded metrics are always disabled - if hasattr(self, "excluded_metrics") and metric_name in self.excluded_metrics: + if metric_name in self.excluded_metrics: return False # If no specific include configuration is provided, enable all metrics (default behavior) @@ -902,7 +914,7 @@ class PrometheusLogger(CustomLogger): labels = [label for label in default_labels if label in configured_labels] # Strip globally excluded labels - if hasattr(self, "excluded_labels") and self.excluded_labels: + if self.excluded_labels: labels = [label for label in labels if label not in self.excluded_labels] return labels diff --git a/tests/test_litellm/integrations/test_prometheus_latency_buckets.py b/tests/test_litellm/integrations/test_prometheus_latency_buckets.py index 3160bb5a9de..de0e6849fc0 100644 --- a/tests/test_litellm/integrations/test_prometheus_latency_buckets.py +++ b/tests/test_litellm/integrations/test_prometheus_latency_buckets.py @@ -1,7 +1,10 @@ """ -Tests that prometheus_latency_buckets is respected end-to-end: +Tests that prometheus_latency_buckets, prometheus_exclude_metrics, and +prometheus_exclude_labels are respected end-to-end: - histograms are registered with the custom boundaries - observed values land in the expected bucket +- excluded metrics are replaced by NoOpMetric +- excluded labels are stripped from all metrics """ from datetime import datetime, timedelta @@ -10,26 +13,32 @@ import pytest from prometheus_client import REGISTRY import litellm -from litellm.integrations.prometheus import PrometheusLogger +from litellm.integrations.prometheus import NoOpMetric, PrometheusLogger from litellm.types.integrations.prometheus import UserAPIKeyLabelValues @pytest.fixture(autouse=True) def reset_prometheus_registry(): - collectors = list(REGISTRY._collector_to_names.keys()) - for c in collectors: - REGISTRY.unregister(c) + """Unregister only collectors added during the test, avoiding private-API churn.""" + before = set(REGISTRY._names_to_collectors.keys()) yield - collectors = list(REGISTRY._collector_to_names.keys()) - for c in collectors: - REGISTRY.unregister(c) + after = set(REGISTRY._names_to_collectors.keys()) + for name in after - before: + try: + REGISTRY.unregister(REGISTRY._names_to_collectors[name]) + except Exception: + pass @pytest.fixture(autouse=True) -def reset_litellm_buckets(): - original = litellm.prometheus_latency_buckets +def reset_litellm_prometheus_settings(): + original_buckets = litellm.prometheus_latency_buckets + original_exclude_metrics = litellm.prometheus_exclude_metrics + original_exclude_labels = litellm.prometheus_exclude_labels yield - litellm.prometheus_latency_buckets = original + litellm.prometheus_latency_buckets = original_buckets + litellm.prometheus_exclude_metrics = original_exclude_metrics + litellm.prometheus_exclude_labels = original_exclude_labels def _make_enum_values() -> UserAPIKeyLabelValues: @@ -70,23 +79,61 @@ def _get_bucket_count(metric_name: str, le: str) -> float: raise AssertionError(f"Bucket le={le} not found in metric {metric_name}") +def _get_registered_bucket_les(metric_name: str) -> set: + """Return the set of 'le' label values registered for a histogram metric.""" + return { + sample.labels["le"] + for family in REGISTRY.collect() + if family.name == metric_name + for sample in family.samples + if sample.name.endswith("_bucket") + } + + +def _observe_once(logger: PrometheusLogger) -> None: + """Make one zero-latency observation so that labeled histograms emit samples. + + prometheus_client only populates the registry for labeled metric instances + that have received at least one ``observe()`` call. Without this, REGISTRY.collect() + returns an empty sample list for any labeled Histogram, making bucket-boundary + inspection via the registry impossible. + """ + now = datetime.now() + logger._set_latency_metrics( + kwargs={ + "start_time": now, + "end_time": now, + "api_call_start_time": now, + "litellm_params": {"metadata": {}}, + "model": "gpt-4", + }, + model="gpt-4", + user_api_key="test-key", + user_api_key_alias="test-alias", + user_api_team=None, + user_api_team_alias=None, + enum_values=_make_enum_values(), + ) + + def test_custom_latency_buckets_registered(): """Histograms use prometheus_latency_buckets when set.""" litellm.prometheus_latency_buckets = (1.0, 5.0, float("inf")) logger = PrometheusLogger() - assert logger._get_latency_buckets() == (1.0, 5.0, float("inf")) - assert logger.litellm_request_total_latency_metric._upper_bounds == [ - 1.0, - 5.0, - float("inf"), - ] - assert logger.litellm_llm_api_latency_metric._upper_bounds == [ - 1.0, - 5.0, - float("inf"), - ] + assert list(logger._get_latency_buckets()) == [1.0, 5.0, float("inf")] + + # prometheus_client only emits samples for labeled metrics that have been observed; + # trigger a zero-latency observation so the registry is populated before we inspect it. + _observe_once(logger) + + # Verify via the public registry text format, not private _upper_bounds + les = _get_registered_bucket_les("litellm_request_total_latency_metric") + assert les == {"1.0", "5.0", "+Inf"} + + les_api = _get_registered_bucket_les("litellm_llm_api_latency_metric") + assert les_api == {"1.0", "5.0", "+Inf"} def test_observation_lands_in_correct_custom_bucket(): @@ -138,7 +185,44 @@ def test_default_buckets_used_when_not_set(): logger = PrometheusLogger() - assert logger._get_latency_buckets() == LATENCY_BUCKETS - assert logger.litellm_request_total_latency_metric._upper_bounds == list( - LATENCY_BUCKETS - ) + assert list(logger._get_latency_buckets()) == list(LATENCY_BUCKETS) + + # Trigger an observation so the labeled histogram emits samples to the registry. + _observe_once(logger) + + # Verify via the registry + les = _get_registered_bucket_les("litellm_request_total_latency_metric") + expected = {str(b) if b != float("inf") else "+Inf" for b in LATENCY_BUCKETS} + assert les == expected + + +def test_exclude_metrics_replaces_with_noop(): + """Excluded metrics are replaced by NoOpMetric and never registered in Prometheus.""" + litellm.prometheus_exclude_metrics = ["litellm_overhead_latency_metric"] + + logger = PrometheusLogger() + + assert isinstance(logger.litellm_overhead_latency_metric, NoOpMetric) + + # The metric must not appear in the registry at all + registered_names = {f.name for f in REGISTRY.collect()} + assert "litellm_overhead_latency_metric" not in registered_names + + # Other metrics are real + assert not isinstance(logger.litellm_request_total_latency_metric, NoOpMetric) + + +def test_exclude_labels_strips_label_from_metrics(): + """Excluded labels are absent from the label set of all affected metrics.""" + litellm.prometheus_exclude_labels = ["end_user"] + + logger = PrometheusLogger() + + # Verify via the registry: no sample for litellm_request_total_latency_metric + # should carry the end_user label key + for family in REGISTRY.collect(): + if family.name == "litellm_request_total_latency_metric": + for sample in family.samples: + assert ( + "end_user" not in sample.labels + ), f"end_user label found in sample {sample}"