From 359aa247742b079a842bb8d762480938665d9110 Mon Sep 17 00:00:00 2001 From: DanBrima <40828002+DanBrima@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:17:50 +0000 Subject: [PATCH] perf(prometheus): track the last team labelset instead of scanning Retiring a renamed team's old series scanned the metric's children on every team request, once per gauge. That cost work proportional to the total number of team series ever emitted, and any authenticated caller could amplify it with ordinary traffic. Remember the last labelset emitted per (metric, team, model) and retire that one directly, which is O(1) and drops the registry scan entirely. The map is also the reason the gauges no longer need collect(), so the label-exclusion wrapper is no longer involved in cleanup at all. --- litellm/integrations/prometheus.py | 69 ++++++++++++------- ...test_prometheus_team_rate_limit_metrics.py | 1 + 2 files changed, 46 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 4425bfcecd9..f3a55a5204d 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -7,7 +7,7 @@ import asyncio import math import os import sys -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast @@ -228,6 +228,11 @@ class PrometheusLogger(CustomLogger): _custom_buckets: Final = litellm.prometheus_latency_buckets self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker() + # Last labelset emitted per (metric, team, model), so a renamed team's + # previous series can be retired without scanning the registry. + self._team_series_label_values: MutableMapping[ # mutable-ok: per-process emission state, rewritten as teams are renamed + tuple[str, str, str], tuple[str, ...] + ] = {} # Create metric factory functions self._counter_factory = self._create_metric_factory(Counter) @@ -2176,11 +2181,14 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) label_values: Final = tuple(labels.get(name, "") for name in labelnames) - self._drop_superseded_team_series(gauge=gauge, labelnames=labelnames, labels=labels) + self._drop_superseded_team_series( + gauge=gauge, metric_name=metric_name, labels=labels, label_values=label_values + ) if value is not None: gauge.labels(*label_values).set(value) return + self._forget_team_series(metric_name=metric_name, labels=labels) try: gauge.remove(*label_values) except KeyError: @@ -2191,36 +2199,49 @@ class PrometheusLogger(CustomLogger): def _drop_superseded_team_series( self, gauge: _LabeledGauge, - labelnames: Sequence[str], + metric_name: DEFINED_PROMETHEUS_METRICS, labels: Mapping[str, str], + label_values: tuple[str, ...], ) -> None: """ - Retire child series that describe this same team and model under a - different alias. Renaming a team changes ``team_alias``, which starts a - new series, and the old one would otherwise keep publishing the values - it held at rename time, double counting the team on any sum over - ``team``. - """ - collect: Final = getattr(gauge, "collect", None) - if collect is None: - return + Retire the series this team and model last published under a different + alias. Renaming a team changes ``team_alias``, which starts a new + series, and the old one would otherwise keep publishing the values it + held at rename time, double counting the team on any sum over ``team``. - team_label: Final = UserAPIKeyLabelNames.TEAM.value - alias_label: Final = UserAPIKeyLabelNames.TEAM_ALIAS.value - model_label: Final = UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value - superseded: Final = tuple( - tuple(sample.labels.get(name, "") for name in labelnames) - for metric in collect() - for sample in metric.samples - if sample.labels.get(team_label) == labels.get(team_label) - and sample.labels.get(model_label) == labels.get(model_label) - and sample.labels.get(alias_label) != labels.get(alias_label) + The previously emitted labelset is remembered per (metric, team, model) + rather than found by scanning the registry. A scan would cost every + team request work proportional to the total number of team series ever + emitted, which any authenticated caller could amplify by sending + ordinary traffic. + """ + identity: Final = ( + metric_name, + labels.get(UserAPIKeyLabelNames.TEAM.value, ""), + labels.get(UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, ""), ) - for label_values in superseded: + previous: Final = self._team_series_label_values.get(identity) + if previous is not None and previous != label_values: try: - gauge.remove(*label_values) + gauge.remove(*previous) except KeyError: pass + self._team_series_label_values[identity] = label_values + + def _forget_team_series( + self, + metric_name: DEFINED_PROMETHEUS_METRICS, + labels: Mapping[str, str], + ) -> None: + """Stop tracking a (metric, team, model) whose series has been dropped.""" + self._team_series_label_values.pop( + ( + metric_name, + labels.get(UserAPIKeyLabelNames.TEAM.value, ""), + labels.get(UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, ""), + ), + None, + ) def _set_latency_metrics( self, diff --git a/tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py b/tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py index fbee2ba58c4..63d70a6def5 100644 --- a/tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py +++ b/tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py @@ -48,6 +48,7 @@ def _logger_with_mock_team_gauges() -> PrometheusLogger: for metric_name in TEAM_RATE_LIMIT_METRICS: setattr(logger, metric_name, MagicMock()) logger.get_labels_for_metric = MagicMock(side_effect=PrometheusMetricLabels.get_labels) + logger._team_series_label_values = {} return logger