From 465480af41200231505926f97eb9d1aca73230cd Mon Sep 17 00:00:00 2001 From: DanBrima <40828002+DanBrima@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:44:57 +0000 Subject: [PATCH] feat(prometheus): expose team-scoped rate limit gauges Configured and remaining rate limits were only observable at virtual key scope (litellm_remaining_api_key_*_for_model) and deployment scope (litellm_deployment_{tpm,rpm}_limit). At team scope the only gauges were dollar-denominated budgets, so there was no way to alert on a team approaching the model_tpm_limit / model_rpm_limit set on its team object. The v3 rate limiter already computes current_limit and limit_remaining for its model_per_team descriptor and publishes them as x-ratelimit-model_per_team-{remaining,limit}-{requests,tokens}, which land in the standard logging payload. Read those existing values in async_log_success_event and set four new gauges labeled by team, team_alias and model. No enforcement change and no extra Redis round trips. Scope is the per-model team limits only. Every gauge carries a model label, and a model-agnostic team-wide limit has no single model to attribute to. A team with no limit configured for the requested model produces no header, and therefore no series. --- litellm/integrations/prometheus.py | 168 +++++++++++- litellm/types/integrations/prometheus.py | 18 +- ...test_prometheus_team_rate_limit_metrics.py | 256 ++++++++++++++++++ 3 files changed, 427 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index a9056aaf4e1..42a9148066f 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -9,7 +9,7 @@ import os import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -114,6 +114,31 @@ def _get_budget_metrics_per_request_timeout() -> float: return parsed +class _LabeledGauge(Protocol): + """Structural type shared by ``prometheus_client.Gauge`` and the no-op / label-excluding wrappers above.""" + + def labels(self, **labels: str) -> _LabeledGauge: ... + + def set(self, value: float) -> None: ... + + +_TEAM_RATE_LIMIT_GAUGE_SPECS: Final[ + tuple[ + tuple[ + DEFINED_PROMETHEUS_METRICS, + Literal["remaining", "limit"], + Literal["requests", "tokens"], + ], + ..., + ] +] = ( + ("litellm_remaining_team_requests_for_model", "remaining", "requests"), + ("litellm_remaining_team_tokens_for_model", "remaining", "tokens"), + ("litellm_team_rpm_limit", "limit", "requests"), + ("litellm_team_tpm_limit", "limit", "tokens"), +) + + class PrometheusLogger(CustomLogger): # Class variables or attributes @@ -387,6 +412,34 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"), ) + ######################################## + # LiteLLM Team rate limit metrics + ######################################## + + self.litellm_remaining_team_requests_for_model = self._gauge_factory( + "litellm_remaining_team_requests_for_model", + "Remaining Requests team can make for model (model based rpm limit on team)", + labelnames=self.get_labels_for_metric("litellm_remaining_team_requests_for_model"), + ) + + self.litellm_remaining_team_tokens_for_model = self._gauge_factory( + "litellm_remaining_team_tokens_for_model", + "Remaining Tokens team can make for model (model based tpm limit on team)", + labelnames=self.get_labels_for_metric("litellm_remaining_team_tokens_for_model"), + ) + + self.litellm_team_rpm_limit = self._gauge_factory( + "litellm_team_rpm_limit", + "Configured RPM limit for team + model (model based rpm limit on team)", + labelnames=self.get_labels_for_metric("litellm_team_rpm_limit"), + ) + + self.litellm_team_tpm_limit = self._gauge_factory( + "litellm_team_tpm_limit", + "Configured TPM limit for team + model (model based tpm limit on team)", + labelnames=self.get_labels_for_metric("litellm_team_tpm_limit"), + ) + ######################################## # LLM API Deployment Metrics / analytics ######################################## @@ -1385,6 +1438,14 @@ class PrometheusLogger(CustomLogger): model_id=enum_values.model_id, ) + # set team rpm/tpm metrics for the requested model + self._set_team_rate_limit_metrics( + user_api_team=user_api_team, + user_api_team_alias=user_api_team_alias, + model_group=standard_logging_payload["model_group"], + standard_logging_payload=standard_logging_payload, + ) + # set latency metrics self._set_latency_metrics( kwargs=kwargs, @@ -1895,18 +1956,21 @@ class PrometheusLogger(CustomLogger): ) @staticmethod - def _get_remaining_from_v3_rate_limit_headers( + def _get_v3_rate_limit_header( standard_logging_payload: StandardLoggingPayload | None, + descriptor_key: Literal["model_per_key", "model_per_team"], + value_type: Literal["remaining", "limit"], rate_limit_type: Literal["requests", "tokens"], ) -> int | None: """ - Read the per-(key, model) remaining value emitted by the v3 rate - limiter (``parallel_request_limiter_v3.py``), which writes - ``x-ratelimit-model_per_key-remaining-{requests,tokens}`` into - ``standard_logging_object.hidden_params.additional_headers`` instead - of the ``litellm-key-remaining-*`` metadata keys the legacy limiter - sets. The header carries no model group; it always refers to this - request's model group, which is what the gauges are labeled with. + Read a per-(scope, model) value emitted by the v3 rate limiter + (``parallel_request_limiter_v3.py``), which writes + ``x-ratelimit-{descriptor_key}-{remaining,limit}-{requests,tokens}`` + into ``standard_logging_object.hidden_params.additional_headers`` + instead of the ``litellm-key-remaining-*`` metadata keys the legacy + limiter sets. The header carries no model group; it always refers to + this request's model group, which is what the gauges are labeled + with. A scope with no configured limit produces no header at all. Values are written in-process as plain ints (never HTTP-serialized strings), so anything else is rejected rather than coerced. """ @@ -1918,7 +1982,7 @@ class PrometheusLogger(CustomLogger): additional_headers: Final = hidden_params.get("additional_headers") if additional_headers is None: return None - value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}") + value: Final = dict(additional_headers).get(f"x-ratelimit-{descriptor_key}-{value_type}-{rate_limit_type}") if isinstance(value, bool) or not isinstance(value, int): return None return value @@ -1944,15 +2008,21 @@ class PrometheusLogger(CustomLogger): remaining_requests = metadata.get(remaining_requests_variable_name) if remaining_requests is None: - remaining_requests = self._get_remaining_from_v3_rate_limit_headers( - standard_logging_payload=standard_logging_payload, rate_limit_type="requests" + remaining_requests = self._get_v3_rate_limit_header( + standard_logging_payload=standard_logging_payload, + descriptor_key="model_per_key", + value_type="remaining", + rate_limit_type="requests", ) if remaining_requests is None: remaining_requests = sys.maxsize remaining_tokens = metadata.get(remaining_tokens_variable_name) if remaining_tokens is None: - remaining_tokens = self._get_remaining_from_v3_rate_limit_headers( - standard_logging_payload=standard_logging_payload, rate_limit_type="tokens" + remaining_tokens = self._get_v3_rate_limit_header( + standard_logging_payload=standard_logging_payload, + descriptor_key="model_per_key", + value_type="remaining", + rate_limit_type="tokens", ) if remaining_tokens is None: remaining_tokens = sys.maxsize @@ -1983,6 +2053,76 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(remaining_tokens) + def _set_team_rate_limit_metrics( + self, + user_api_team: str | None, + user_api_team_alias: str | None, + model_group: str | None, + standard_logging_payload: StandardLoggingPayload | None, + ) -> None: + """ + Emit the per-(team, model) rate limit gauges from the values the v3 + rate limiter already computed for its ``model_per_team`` descriptor + and shipped to the client as ``x-ratelimit-model_per_team-*`` + headers. A team with no per-model limit configured for the requested + model produces no header, and therefore no series, which matches how + the per-key gauges behave. + """ + if user_api_team is None: + return + + configured: Final = tuple( + (metric_name, value) + for metric_name, value_type, rate_limit_type in _TEAM_RATE_LIMIT_GAUGE_SPECS + if ( + value := self._get_v3_rate_limit_header( + standard_logging_payload=standard_logging_payload, + descriptor_key="model_per_team", + value_type=value_type, + rate_limit_type=rate_limit_type, + ) + ) + is not None + ) + if not configured: + return + + enum_values: Final = UserAPIKeyLabelValues( + team=user_api_team, + team_alias=user_api_team_alias, + model=model_group, + custom_metadata_labels=get_custom_labels_from_metadata( + metadata=_get_combined_custom_metadata_from_standard_logging_payload( + standard_logging_payload=standard_logging_payload + ) + ), + ) + label_context: Final = PrometheusLabelFactoryContext(enum_values) + + for metric_name, value in configured: + self._set_team_rate_limit_gauge( + gauge=getattr(self, metric_name), + metric_name=metric_name, + value=value, + enum_values=enum_values, + label_context=label_context, + ) + + def _set_team_rate_limit_gauge( + self, + gauge: _LabeledGauge, + metric_name: DEFINED_PROMETHEUS_METRICS, + value: int, + enum_values: UserAPIKeyLabelValues, + label_context: PrometheusLabelFactoryContext, + ) -> None: + labels: Final = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric(metric_name), + enum_values=enum_values, + label_context=label_context, + ) + gauge.labels(**labels).set(value) + def _set_latency_metrics( self, kwargs: dict, diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index ebec5df55fa..9fe51f026e9 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -1,5 +1,5 @@ import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import MISSING, dataclass, field, fields from enum import Enum from types import MappingProxyType @@ -257,6 +257,10 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_deployment_rpm_limit", "litellm_remaining_api_key_requests_for_model", "litellm_remaining_api_key_tokens_for_model", + "litellm_remaining_team_requests_for_model", + "litellm_remaining_team_tokens_for_model", + "litellm_team_rpm_limit", + "litellm_team_tpm_limit", "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", @@ -670,6 +674,18 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.MODEL_ID.value, ] + litellm_remaining_team_requests_for_model: ClassVar[Sequence[str]] = [ + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + litellm_remaining_team_tokens_for_model = litellm_remaining_team_requests_for_model + + litellm_team_rpm_limit = litellm_remaining_team_requests_for_model + + litellm_team_tpm_limit = litellm_remaining_team_requests_for_model + litellm_llm_api_failed_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, 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 new file mode 100644 index 00000000000..820c4b75ea1 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_team_rate_limit_metrics.py @@ -0,0 +1,256 @@ +""" +Tests for the team-scoped rate limit Prometheus gauges. + +LiteLLM exposed configured/remaining rate limits at virtual key scope +(``litellm_remaining_api_key_*_for_model``) and deployment scope +(``litellm_deployment_{tpm,rpm}_limit``) but not at team scope, so there was +no way to alert on a team approaching the ``model_tpm_limit`` / +``model_rpm_limit`` configured on its team object. + +The v3 rate limiter already computes those numbers for its ``model_per_team`` +descriptor and ships them to clients as +``x-ratelimit-model_per_team-{remaining,limit}-{requests,tokens}``. These +tests cover routing those already-computed values to Prometheus. +""" + +from typing import get_args +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.types.integrations.prometheus import ( + DEFINED_PROMETHEUS_METRICS, + PrometheusMetricLabels, + UserAPIKeyLabelNames, +) + +TEAM_RATE_LIMIT_METRICS = ( + "litellm_remaining_team_requests_for_model", + "litellm_remaining_team_tokens_for_model", + "litellm_team_rpm_limit", + "litellm_team_tpm_limit", +) + + +def _logger_with_mock_team_gauges(labels_are_real: bool = False) -> PrometheusLogger: + with patch("litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None): + logger = PrometheusLogger() + for metric_name in TEAM_RATE_LIMIT_METRICS: + setattr(logger, metric_name, MagicMock()) + if labels_are_real: + logger.get_labels_for_metric = MagicMock(side_effect=PrometheusMetricLabels.get_labels) + else: + logger.get_labels_for_metric = MagicMock(return_value=[]) + return logger + + +def _payload_with_headers(additional_headers: dict) -> dict: + return { + "metadata": {}, + "hidden_params": {"additional_headers": additional_headers}, + } + + +def _set_team_metrics(logger: PrometheusLogger, standard_logging_payload: dict) -> None: + logger._set_team_rate_limit_metrics( + user_api_team="team-abc", + user_api_team_alias="research", + model_group="gpt-4o-mini", + standard_logging_payload=standard_logging_payload, + ) + + +def _assert_set_once(logger: PrometheusLogger, metric_name: str, value: int) -> None: + getattr(logger, metric_name).labels.return_value.set.assert_called_once_with(value) + + +ALL_TEAM_HEADERS = { + "x-ratelimit-model_per_team-remaining-requests": 42, + "x-ratelimit-model_per_team-remaining-tokens": 900, + "x-ratelimit-model_per_team-limit-requests": 100, + "x-ratelimit-model_per_team-limit-tokens": 1000, +} + + +def test_team_metrics_are_defined_with_team_and_model_labels(): + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + expected_labels = [ + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + ] + + for metric_name in TEAM_RATE_LIMIT_METRICS: + assert metric_name in defined_metrics + labels = PrometheusMetricLabels.get_labels(metric_name) + for expected_label in expected_labels: + assert expected_label in labels + + +def test_every_logger_owned_metric_resolves_labels(): + """ + ``PrometheusMetricLabels.get_labels`` resolves a metric name to a label + list via ``getattr``, so a metric added to the literal without a matching + label attribute fails at logger construction time in production rather + than at lint time. + + ``litellm_in_flight_requests`` is excluded because it is a label-free + gauge registered by the in-flight middleware, not by ``PrometheusLogger``; + it appears in the literal only so ``prometheus_metrics_config`` can name it. + """ + for metric_name in get_args(DEFINED_PROMETHEUS_METRICS): + if metric_name == "litellm_in_flight_requests": + continue + assert isinstance(PrometheusMetricLabels.get_labels(metric_name), list) + + +def test_sets_every_team_gauge_from_v3_headers(): + logger = _logger_with_mock_team_gauges() + + _set_team_metrics(logger, _payload_with_headers(dict(ALL_TEAM_HEADERS))) + + _assert_set_once(logger, "litellm_remaining_team_requests_for_model", 42) + _assert_set_once(logger, "litellm_remaining_team_tokens_for_model", 900) + _assert_set_once(logger, "litellm_team_rpm_limit", 100) + _assert_set_once(logger, "litellm_team_tpm_limit", 1000) + + +def test_labels_carry_team_and_requested_model(): + logger = _logger_with_mock_team_gauges(labels_are_real=True) + + _set_team_metrics(logger, _payload_with_headers(dict(ALL_TEAM_HEADERS))) + + for metric_name in TEAM_RATE_LIMIT_METRICS: + labels_kwargs = getattr(logger, metric_name).labels.call_args.kwargs + assert labels_kwargs["team"] == "team-abc" + assert labels_kwargs["team_alias"] == "research" + assert labels_kwargs["model"] == "gpt-4o-mini" + + +def test_emits_nothing_when_team_has_no_configured_limits(): + """A team without per-model limits gets no descriptor, so no header, so no series.""" + logger = _logger_with_mock_team_gauges() + + _set_team_metrics( + logger, + _payload_with_headers( + { + "x-ratelimit-model_per_key-remaining-requests": 42, + "x-ratelimit-model_per_key-limit-requests": 100, + } + ), + ) + + for metric_name in TEAM_RATE_LIMIT_METRICS: + getattr(logger, metric_name).labels.assert_not_called() + + +def test_emits_only_the_dimension_the_team_configured(): + """A team with only an RPM limit must not get a fabricated TPM series.""" + logger = _logger_with_mock_team_gauges() + + _set_team_metrics( + logger, + _payload_with_headers( + { + "x-ratelimit-model_per_team-remaining-requests": 7, + "x-ratelimit-model_per_team-limit-requests": 60, + } + ), + ) + + _assert_set_once(logger, "litellm_remaining_team_requests_for_model", 7) + _assert_set_once(logger, "litellm_team_rpm_limit", 60) + logger.litellm_remaining_team_tokens_for_model.labels.assert_not_called() + logger.litellm_team_tpm_limit.labels.assert_not_called() + + +def test_emits_zero_remaining_rather_than_skipping_it(): + """An exhausted team is the case operators alert on, so 0 must be a real sample.""" + logger = _logger_with_mock_team_gauges() + + _set_team_metrics( + logger, + _payload_with_headers( + { + "x-ratelimit-model_per_team-remaining-requests": 0, + "x-ratelimit-model_per_team-remaining-tokens": 0, + } + ), + ) + + _assert_set_once(logger, "litellm_remaining_team_requests_for_model", 0) + _assert_set_once(logger, "litellm_remaining_team_tokens_for_model", 0) + + +def test_emits_nothing_for_a_request_with_no_team(): + logger = _logger_with_mock_team_gauges() + + logger._set_team_rate_limit_metrics( + user_api_team=None, + user_api_team_alias=None, + model_group="gpt-4o-mini", + standard_logging_payload=_payload_with_headers(dict(ALL_TEAM_HEADERS)), + ) + + for metric_name in TEAM_RATE_LIMIT_METRICS: + getattr(logger, metric_name).labels.assert_not_called() + + +@pytest.mark.parametrize("bad_value", ["100", None, True, 12.5]) +def test_ignores_non_int_header_values(bad_value): + logger = _logger_with_mock_team_gauges() + + _set_team_metrics( + logger, + _payload_with_headers({"x-ratelimit-model_per_team-remaining-requests": bad_value}), + ) + + logger.litellm_remaining_team_requests_for_model.labels.assert_not_called() + + +def test_raises_nothing_when_payload_has_no_hidden_params(): + logger = _logger_with_mock_team_gauges() + + _set_team_metrics(logger, {"metadata": {}}) + + for metric_name in TEAM_RATE_LIMIT_METRICS: + getattr(logger, metric_name).labels.assert_not_called() + + +def test_limiter_publishes_team_headers_in_the_shape_the_gauges_read(): + """ + Pins the producer/consumer contract: the gauges read header names the v3 + limiter builds from ``descriptor_key`` + ``rate_limit_type``, so a change + to that format would otherwise silently stop the team series. + """ + headers = _PROXY_MaxParallelRequestsHandler_v3._merge_ratelimit_statuses_into_additional_headers( + additional_headers={}, + statuses=[ + { + "code": "OK", + "current_limit": 100, + "limit_remaining": 42, + "rate_limit_type": "requests", + "descriptor_key": "model_per_team", + }, + { + "code": "OK", + "current_limit": 1000, + "limit_remaining": 900, + "rate_limit_type": "tokens", + "descriptor_key": "model_per_team", + }, + ], + ) + + assert headers == { + "x-ratelimit-model_per_team-remaining-requests": 42, + "x-ratelimit-model_per_team-limit-requests": 100, + "x-ratelimit-model_per_team-remaining-tokens": 900, + "x-ratelimit-model_per_team-limit-tokens": 1000, + }