From 7130a5b4fdefb6bd6b2c9fd7b720d8f908bb1e78 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 15:28:51 -0800 Subject: [PATCH] feat(dynamic_rate_limiter_v3): prometheus share, borrowed and throttle metrics per team and model (#51417) * feat(dynamic_rate_limiter_v3): Prometheus share, borrowed and throttle metrics per team and model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): name the BLE001 suppressions on the metrics guards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): count only final refusals, publish the served model's share, join the priority label Review round 1: a steered alternate's refusal no longer counts as a throttle when the preferred model then admits the request, a router fallback's settlement publishes the fallback target's own share under its model group, the gauges carry the same priority label the 429 carries, the gauges aggregate with livemostrecent so a new window replaces an old worker's value, the share read is skipped when no Prometheus logger is configured, the unit tests restore the process registry, and the integration test polls for the settled totals. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prometheus): type the gauge setter and the label factory so the type gate holds _set_labeled_gauge takes a _MetricLike and sets the series only when it is a real Gauge, prometheus_label_factory declares the dict[str, str | None] it returns, the priority resolver goes back to its original shape, and the label test tolerates the now visible None values. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(prometheus): keep the label test diff to the two None-tolerant assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/prometheus.py | 87 ++++++- .../grafana_dashboard.json | 184 ++++++++++++++ .../proxy/hooks/dynamic_rate_limiter_v3.py | 115 ++++++++- .../hooks/parallel_request_limiter_v3.py | 1 + litellm/types/integrations/prometheus.py | 22 ++ ...test_dynamic_limiter_prometheus_metrics.py | 151 ++++++++++++ .../integrations/test_prometheus_labels.py | 4 +- .../hooks/test_dynamic_rate_limiter_v3.py | 230 ++++++++++++++++++ 8 files changed, 781 insertions(+), 13 deletions(-) create mode 100644 tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 7984c7d04f3..6f56693b77e 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1061,18 +1061,39 @@ class PrometheusLogger(CustomLogger): labelnames=self.get_labels_for_metric("litellm_deployment_latency_per_output_token"), ) - self.litellm_deployment_successful_fallbacks = self._counter_factory( + self.litellm_deployment_successful_fallbacks: _MetricLike = self._counter_factory( "litellm_deployment_successful_fallbacks", "LLM Deployment Analytics - Number of successful fallback requests from primary model -> fallback model", labelnames=self.get_labels_for_metric("litellm_deployment_successful_fallbacks"), ) - self.litellm_deployment_failed_fallbacks = self._counter_factory( + self.litellm_deployment_failed_fallbacks: _MetricLike = self._counter_factory( "litellm_deployment_failed_fallbacks", "LLM Deployment Analytics - Number of failed fallback requests from primary model -> fallback model", labelnames=self.get_labels_for_metric("litellm_deployment_failed_fallbacks"), ) + self.litellm_dynamic_limiter_share_tokens: _MetricLike = self._gauge_factory( + "litellm_dynamic_limiter_share_tokens", + "Dynamic rate limiter: the team's share of the model's TPM, as of the team's last settled request", + labelnames=self.get_labels_for_metric("litellm_dynamic_limiter_share_tokens"), + multiprocess_mode="livemostrecent", + ) + + self.litellm_dynamic_limiter_borrowed_tokens: _MetricLike = self._gauge_factory( + "litellm_dynamic_limiter_borrowed_tokens", + "Dynamic rate limiter: tokens the team has used above its share this window, 0 while under its share", + labelnames=self.get_labels_for_metric("litellm_dynamic_limiter_borrowed_tokens"), + multiprocess_mode="livemostrecent", + ) + + self.litellm_dynamic_limiter_throttled_total: _MetricLike = self._counter_factory( + "litellm_dynamic_limiter_throttled_total", + "Dynamic rate limiter: requests refused per team, model and reason, " + "the reason being the value of the response's rate_limited_by header", + labelnames=self.get_labels_for_metric("litellm_dynamic_limiter_throttled_total"), + ) + # Callback Logging Failure Metrics self.litellm_callback_logging_failures_metric = self._counter_factory( name="litellm_callback_logging_failures_metric", @@ -4172,6 +4193,66 @@ class PrometheusLogger(CustomLogger): sanitize_prometheus_label_value(exception_status), ).inc() + def set_dynamic_limiter_share( + self, + model_group: str, + team: str, + team_alias: str | None, + priority: str, + share_tokens: int, + borrowed_tokens: int, + ) -> None: + enum_values: Final = UserAPIKeyLabelValues( + model_group=model_group, team=team, team_alias=team_alias, priority=priority + ) + self._set_labeled_gauge( + self.litellm_dynamic_limiter_share_tokens, "litellm_dynamic_limiter_share_tokens", enum_values, share_tokens + ) + self._set_labeled_gauge( + self.litellm_dynamic_limiter_borrowed_tokens, + "litellm_dynamic_limiter_borrowed_tokens", + enum_values, + borrowed_tokens, + ) + + def _set_labeled_gauge( + self, + gauge: _MetricLike, + metric_name: DEFINED_PROMETHEUS_METRICS, + enum_values: UserAPIKeyLabelValues, + value: float, + ) -> None: + from prometheus_client import Gauge + + labels: Final = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name), enum_values=enum_values + ) + series: Final = gauge.labels(**labels) + if isinstance(series, Gauge): + series.set(value) + + def increment_dynamic_limiter_throttled( + self, + model_group: str, + team: str, + team_alias: str | None, + priority: str, + rate_limit_type: str, + reason: str, + ) -> None: + self._inc_labeled_counter( + self.litellm_dynamic_limiter_throttled_total, + "litellm_dynamic_limiter_throttled_total", + UserAPIKeyLabelValues( + model_group=model_group, + team=team, + team_alias=team_alias, + priority=priority, + rate_limit_type=rate_limit_type, + reason=reason, + ), + ) + def increment_callback_logging_failure( self, callback_name: str, @@ -5221,7 +5302,7 @@ def prometheus_label_factory( tag: str | None = None, *, label_context: PrometheusLabelFactoryContext | None = None, -) -> dict: +) -> dict[str, str | None]: """ Returns a dictionary of label + values for prometheus. diff --git a/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json b/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json index cedbbd0f07b..59861fcf462 100644 --- a/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json +++ b/litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json @@ -6644,6 +6644,190 @@ ], "title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate", "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { + "h": 1, + "w": 24, + "x": 0, + "y": 456 + }, + "id": 117, + "panels": [], + "title": "Dynamic rate limiter (needs dynamic_rate_limiter_v3 and prometheus callbacks)", + "type": "row" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "The team's token share on the model this window, per team and model group", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 457 + }, + "id": 118, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "litellm_dynamic_limiter_share_tokens", + "legendFormat": "__auto", + "range": true, + "refId": "A" + } + ], + "title": "litellm_dynamic_limiter_share_tokens", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Tokens the team has used above its share this window, 0 while under it", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 457 + }, + "id": 119, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "litellm_dynamic_limiter_borrowed_tokens", + "legendFormat": "__auto", + "range": true, + "refId": "A" + } + ], + "title": "litellm_dynamic_limiter_borrowed_tokens", + "type": "timeseries" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Requests the dynamic limiter refused per second, by team, model and the 429's rate_limited_by reason", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "reqps" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 24, + "x": 0, + "y": 465 + }, + "id": 120, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum by (team, team_alias, model_group, reason) (rate(litellm_dynamic_limiter_throttled_total[5m]))", + "legendFormat": "__auto", + "range": true, + "refId": "A" + } + ], + "title": "litellm_dynamic_limiter_throttled_total", + "type": "timeseries" } ], "preload": false, diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 89818c21850..2873121cf56 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import json import os -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType @@ -79,6 +79,7 @@ from litellm.types.router import ModelGroupInfo from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, Usage if TYPE_CHECKING: + from litellm.integrations.prometheus import PrometheusLogger from litellm.types.utils import PriorityReservationSettings @@ -353,6 +354,7 @@ class _SettlementMetadata(BaseModel): user_api_key_auth_metadata: Mapping[str, object] | None = None user_api_key_team_id: str | None = None + user_api_key_team_alias: str | None = None @property def priority(self) -> str | None: @@ -368,6 +370,26 @@ class _SettlementLoggingObject(BaseModel): metadata: _SettlementMetadata | None = None +def _prometheus_loggers() -> "Sequence[PrometheusLogger]": + from litellm.integrations.prometheus import PrometheusLogger + + return tuple( + logger + for logger in litellm.logging_callback_manager.get_custom_loggers_for_type(PrometheusLogger) + if isinstance(logger, PrometheusLogger) + ) + + +def _tightest_team_share(descriptors: Iterable[RateLimitDescriptor | None]) -> tuple[RateLimitDescriptor, int] | None: + """The team share descriptor with the smallest TPM share, the one that binds first at saturation.""" + shares: Final = tuple( + (descriptor, tokens) + for descriptor in descriptors + if descriptor is not None and (tokens := (descriptor["rate_limit"] or {}).get("tokens_per_unit")) is not None + ) + return min(shares, key=lambda share: share[1], default=None) + + def _usage_of(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: """Usage of a modeled response, a Responses WebSocket turn's combined usage, or its raw event list.""" usage: Final = getattr(response_obj, "usage", None) @@ -865,6 +887,63 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): model, status, saturation, priority, user_api_key_dict, team_share.limits ) + def _throttled( + self, refusal: ProxyRateLimitError, model: str, user_api_key_dict: UserAPIKeyAuth + ) -> ProxyRateLimitError: + """Count a refusal that reaches the caller under the reason its ``rate_limited_by`` header carries.""" + headers: Final = refusal.headers or {} + reason: Final = headers.get("rate_limited_by") + if user_api_key_dict.team_id is None or reason is None: + return refusal + for logger in _prometheus_loggers(): + try: + logger.increment_dynamic_limiter_throttled( + model_group=model, + team=user_api_key_dict.team_id, + team_alias=user_api_key_dict.team_alias, + priority=headers.get("x-litellm-priority") or "default", + rate_limit_type=headers.get("rate_limit_type") or "tokens", + reason=reason, + ) + except Exception as metrics_error: # noqa: BLE001 # a failing metrics sink must not change the 429 + verbose_proxy_logger.warning("Dynamic rate limiter throttle metric failed: %s", metrics_error) + return refusal + + async def _record_team_share( + self, + stash: RequestRateLimiterStash | None, + model_group: str | None, + settlement: _SettlementMetadata, + hold: BorrowingTeamModelHold | None, + ) -> None: + """Publish the team's share on the model that served the request, the fallback target's own hold when a + router fallback served it, and the tokens settled above that share so far this window.""" + try: + loggers: Final = _prometheus_loggers() + if not loggers or stash is None or model_group is None or settlement.user_api_key_team_id is None: + return + share: Final = _tightest_team_share( + (stash.borrowing_team_model_limit, stash.priority_class_share) if hold is None else (hold.descriptor,) + ) + if share is None: + return + descriptor, share_tokens = share + used: Final = await self._get_saturation_value_from_cache( + RateLimitScope(descriptor["key"], descriptor["value"]).counter("tokens") + ) + used_tokens: Final = 0 if used is None else int(float(used)) + for logger in loggers: + logger.set_dynamic_limiter_share( + model_group=model_group, + team=settlement.user_api_key_team_id, + team_alias=settlement.user_api_key_team_alias, + priority=stash.dynamic_limiter_priority or "default", + share_tokens=share_tokens, + borrowed_tokens=max(0, used_tokens - share_tokens), + ) + except Exception as metrics_error: # noqa: BLE001 # a failing metrics sink must not change settlement + verbose_proxy_logger.warning("Dynamic rate limiter share metrics failed: %s", metrics_error) + def _team_borrow_ceiling_error( self, model: str, @@ -940,7 +1019,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): share: TeamShare, ) -> ProxyRateLimitError: resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) - return ProxyRateLimitError( + error: Final = ProxyRateLimitError( detail={ "error": f"Priority class share reached for {model} and the model is saturated. " f"Team: {user_api_key_dict.team_id}, " @@ -963,6 +1042,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): model=resolved_model, llm_provider=llm_provider, ) + return error async def enforce_token_limits_for_frame( self, user_api_key_dict: UserAPIKeyAuth, requested_model: str | None @@ -1020,10 +1100,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return over_class_share: Final = over_limit["descriptor_key"] == PRIORITY_CLASS_SHARE_AT_SATURATION if share is not None and (over_class_share or borrowing is None): - raise self._class_share_error(requested_model, over_limit, saturation, priority, user_api_key_dict, share) + raise self._throttled( + self._class_share_error(requested_model, over_limit, saturation, priority, user_api_key_dict, share), + requested_model, + user_api_key_dict, + ) if borrowing is not None: - raise self._team_share_error( - requested_model, over_limit, saturation, priority, user_api_key_dict, borrowing + raise self._throttled( + self._team_share_error(requested_model, over_limit, saturation, priority, user_api_key_dict, borrowing), + requested_model, + user_api_key_dict, ) def _create_model_tracking_descriptor( @@ -1122,6 +1208,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): get_or_create_request_stash().priority_class_share = next( (d for d in priority_descriptors if d["key"] == PRIORITY_CLASS_SHARE_AT_SATURATION), None ) + get_or_create_request_stash().dynamic_limiter_priority = priority if priority_descriptors: descriptors_to_check.extend(priority_descriptors) @@ -1481,6 +1568,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): admitted: Final = await self._admit_steered_or_preferred( model, requested_group_info, steered, user_api_key_dict, data, call_type ) + except ProxyRateLimitError as refusal: + raise self._throttled(refusal, model, user_api_key_dict) from None except HTTPException: raise except Exception as e: @@ -1728,10 +1817,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): team_model_tokens: Final = _team_model_settlement( stash, model_group, total_tokens, self.v3_limiter.window_size ) + served: Final = ( + None + if stash is None or model_group is None + else stash.borrowing_team_model_fallback_limits.get(model_group) + ) if stash is not None: - served: Final = ( - stash.borrowing_team_model_fallback_limits.get(model_group) if model_group is not None else None - ) stash.borrowing_team_model_reserved_tokens = 0 stash.borrowing_team_model_reserved_window = None stash.borrowing_team_model_fallback_limits = MappingProxyType( @@ -1801,6 +1892,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): pipeline_operations=pipeline_operations, parent_otel_span=litellm_parent_otel_span, ) + await self._record_team_share(stash, model_group, settlement_metadata, served) # Only log 'priority' if it's known safe; otherwise, redact. SAFE_PRIORITIES: Final = {"low", "medium", "high", "default"} @@ -1864,6 +1956,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) if refund: await self.v3_limiter.async_increment_reservation_aware_tokens(pipeline_operations=refund) + await self._record_team_share( + stash, + failed_group, + _SettlementLoggingObject.model_validate(callback_kwargs.get("standard_logging_object") or {}).metadata + or _SettlementMetadata(), + failed_hold, + ) PROXY_DynamicRateLimitHandlerV3: Final = _PROXY_DynamicRateLimitHandlerV3 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 602f955d8a3..7d5daff71cd 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -631,6 +631,7 @@ class RequestRateLimiterStash: default_factory=lambda: MappingProxyType({}) ) priority_class_share: RateLimitDescriptor | None = None + dynamic_limiter_priority: str | None = None reservation_released: bool = False io_reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index ca084bb358c..d1415f583d2 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -217,6 +217,8 @@ class UserAPIKeyLabelNames(Enum): EXCEPTION_CLASS = EXCEPTION_CLASS RATE_LIMIT_CATEGORY = RATE_LIMIT_CATEGORY RATE_LIMIT_TYPE = RATE_LIMIT_TYPE + REASON = "reason" + PRIORITY = "priority" STATUS_CODE = "status_code" FALLBACK_MODEL = "fallback_model" ROUTE = "route" @@ -258,6 +260,9 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_images_generated_metric", "litellm_deployment_successful_fallbacks", "litellm_deployment_failed_fallbacks", + "litellm_dynamic_limiter_share_tokens", + "litellm_dynamic_limiter_borrowed_tokens", + "litellm_dynamic_limiter_throttled_total", "litellm_remaining_team_budget_metric", "litellm_team_max_budget_metric", "litellm_team_budget_remaining_hours_metric", @@ -716,6 +721,21 @@ class PrometheusMetricLabels: litellm_deployment_failed_fallbacks = litellm_deployment_successful_fallbacks + litellm_dynamic_limiter_share_tokens: ClassVar[tuple[str, ...]] = ( + UserAPIKeyLabelNames.MODEL_GROUP.value, + UserAPIKeyLabelNames.TEAM.value, + UserAPIKeyLabelNames.TEAM_ALIAS.value, + UserAPIKeyLabelNames.PRIORITY.value, + ) + + litellm_dynamic_limiter_borrowed_tokens = litellm_dynamic_limiter_share_tokens + + litellm_dynamic_limiter_throttled_total: ClassVar[tuple[str, ...]] = ( + *litellm_dynamic_limiter_share_tokens, + UserAPIKeyLabelNames.RATE_LIMIT_TYPE.value, + UserAPIKeyLabelNames.REASON.value, + ) + litellm_remaining_team_budget_metric = [ UserAPIKeyLabelNames.TEAM.value, UserAPIKeyLabelNames.TEAM_ALIAS.value, @@ -1095,6 +1115,8 @@ class UserAPIKeyLabelValues: exception_class: str | None = None rate_limit_category: str | None = None rate_limit_type: str | None = None + reason: str | None = None + priority: str | None = None status_code: str | None = None fallback_model: str | None = None route: str | None = None diff --git a/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py b/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py new file mode 100644 index 00000000000..f5bdbb6acb8 --- /dev/null +++ b/tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py @@ -0,0 +1,151 @@ +"""The dynamic rate limiter publishes each team's share of a model, what the team has used above that share, and +every refusal with the same reason the 429 carries in its rate_limited_by header.""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator, Mapping, Sequence +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.prometheus_series import Sample, scrape +from integration._support.wire import Reply, Request, wire_server + +OPENAI_MODEL: Final = "gpt-4.1-mini" +MODEL_TPM: Final = 600 +TEAM_TPM: Final = 100 +REPLY_TOKENS: Final = 30 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl_dynamic_limiter_metrics", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "metered"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": REPLY_TOKENS}, + } +).encode() + + +def _respond(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/v1/chat/completions"), request.target + return Reply(body=UPSTREAM_REPLY) + + +@pytest.fixture(scope="module") +def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp("dynamic-limiter-metrics") + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["litellm_settings"] = { + **configuration["litellm_settings"], + "callbacks": ["dynamic_rate_limiter_v3", "prometheus"], + "priority_reservation": {"interactive": 0.6, "batch": 0.4}, + "priority_reservation_settings": { + "default_priority": 1.0, + "saturation_threshold": 0.6, + "saturation_check_cache_ttl": 0, + "team_model_limit_borrows_idle_capacity_from_other_teams": True, + }, + } + path: Final = directory / "dynamic_limiter_metrics.yaml" + path.write_text(yaml.safe_dump(configuration)) + with ( + gateway_from_environment() as rig, + wire_server(_respond) as wire, + owned_proxy(rig, directory, {}, config=path) as owned, + ): + yield owned, wire.url + + +def _team_series(samples: Sequence[Sample], name: str, team: str) -> Mapping[str, float]: + """The metric's samples for one team, keyed by the label values that are not the team itself.""" + return { + json.dumps({k: v for k, v in sample.labels.items() if k != "team"}, sort_keys=True): sample.value + for sample in samples + if sample.name == name and sample.labels.get("team") == team + } + + +def test_a_team_over_its_share_shows_up_as_borrowed_tokens_and_a_matching_throttle_reason( + saturated_model: tuple[Gateway, str], +) -> None: + owned, upstream = saturated_model + with owned.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM + ) + team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}}) + key: Final = scenario.key(team_id=team) + responses: Final = tuple( + owned.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"metered {uuid.uuid4().hex}"}]}, + key=key, + ) + for _ in range(16) + ) + blocked: Final = tuple(response for response in responses if response.status_code == 429) + expected_borrowed: Final = max(0, (len(responses) - len(blocked)) * REPLY_TOKENS - TEAM_TPM) + samples: Final = eventually( + lambda: scrape(owned), + lambda scraped: ( + tuple(_team_series(scraped, "litellm_dynamic_limiter_borrowed_tokens", team).values()) + == (expected_borrowed,) + and tuple(_team_series(scraped, "litellm_dynamic_limiter_throttled_total", team).values()) + == (len(blocked),) + ), + seconds=30, + return_last_on_timeout=True, + ) + + assert responses[0].status_code == 200, responses[0].text + assert blocked, [r.text for r in responses] + reason: Final = blocked[0].headers.get("rate_limited_by") + assert reason == "model_per_team_saturated", dict(blocked[0].headers) + + share: Final = _team_series(samples, "litellm_dynamic_limiter_share_tokens", team) + borrowed: Final = _team_series(samples, "litellm_dynamic_limiter_borrowed_tokens", team) + throttled: Final = _team_series(samples, "litellm_dynamic_limiter_throttled_total", team) + assert share, [s for s in samples if s.name.startswith("litellm_dynamic_limiter")] + ((share_labels, share_tokens),) = share.items() + labels: Final[dict[str, str]] = json.loads(share_labels) + assert labels["model_group"] == model and labels["priority"] == "default", labels + assert share_tokens == TEAM_TPM + assert borrowed[share_labels] == expected_borrowed, borrowed + assert borrowed[share_labels] > 0, borrowed + ((throttle_labels, refused),) = throttled.items() + throttle: Final[dict[str, str]] = json.loads(throttle_labels) + assert throttle["reason"] == reason and throttle["rate_limit_type"] == "tokens", throttle + assert refused == len(blocked) + + +def test_a_team_under_its_share_borrows_nothing(saturated_model: tuple[Gateway, str]) -> None: + owned, upstream = saturated_model + with owned.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM + ) + team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}}) + response: Final[httpx.Response] = owned.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"one call {uuid.uuid4().hex}"}]}, + key=scenario.key(team_id=team), + ) + samples: Final = eventually( + lambda: scrape(owned), + lambda scraped: bool(_team_series(scraped, "litellm_dynamic_limiter_share_tokens", team)), + seconds=30, + return_last_on_timeout=True, + ) + assert response.status_code == 200, response.text + assert tuple(_team_series(samples, "litellm_dynamic_limiter_share_tokens", team).values()) == (TEAM_TPM,) + assert tuple(_team_series(samples, "litellm_dynamic_limiter_borrowed_tokens", team).values()) == (0,) + assert not _team_series(samples, "litellm_dynamic_limiter_throttled_total", team) diff --git a/tests/unit/integrations/test_prometheus_labels.py b/tests/unit/integrations/test_prometheus_labels.py index 04eb7c3a5a4..b3ed782f127 100644 --- a/tests/unit/integrations/test_prometheus_labels.py +++ b/tests/unit/integrations/test_prometheus_labels.py @@ -619,12 +619,12 @@ def test_prometheus_label_value_sanitization(): # U+2028 must be stripped assert ( - "\u2028" not in labels["requested_model"] + "\u2028" not in (labels["requested_model"] or "") ), f"U+2028 should be removed from label value, got: {repr(labels['requested_model'])}" assert labels["requested_model"] == "claude-haiku-4-5-20251001" # Newlines must be replaced with spaces, quotes must be escaped - assert "\n" not in labels["api_key_alias"] + assert "\n" not in (labels["api_key_alias"] or "") assert labels["api_key_alias"] == 'My Key \\"test\\" with newline' print("✅ Prometheus label values are properly sanitized") diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index efd2e9e7e18..1e4d89b011a 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -13,6 +13,7 @@ from typing import Final from unittest.mock import AsyncMock, patch import pytest +from prometheus_client import CollectorRegistry import litellm from litellm import DualCache, Router @@ -29,6 +30,7 @@ from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3, + claim_request_stash_for_data, ) from litellm.proxy.hooks.priority_class_share import pool_member from litellm.redis.stores.rate_limits import RateLimitScope @@ -3359,3 +3361,231 @@ async def test_tokens_served_by_a_fallback_model_land_on_the_teams_counter_for_t with pytest.raises(ProxyRateLimitError) as exc: await _check(handler, _class_team("team-a"), saturation=0.62, priority="interactive", model="other-model") _held_to_class_share(exc.value) + + +def _prometheus_sink(monkeypatch: pytest.MonkeyPatch, sink_class: type | None = None) -> CollectorRegistry: + """A PrometheusLogger registered as a callback next to the dynamic rate limiter, on the process registry emptied + for the test and restored with its original collectors at teardown.""" + from prometheus_client import REGISTRY + + from litellm.integrations.prometheus import PrometheusLogger + + monkeypatch.setattr(REGISTRY, "_collector_to_names", {}) # pyright: ignore[reportPrivateUsage] # the registry's whole state + monkeypatch.setattr(REGISTRY, "_names_to_collectors", {}) # pyright: ignore[reportPrivateUsage] # the registry's whole state + prometheus: Final = (sink_class or PrometheusLogger)() + monkeypatch.setattr(litellm, "callbacks", [prometheus, "dynamic_rate_limiter_v3"]) + registry: CollectorRegistry = REGISTRY + return registry + + +async def _admit_at_the_hook(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None: + assert await handler.async_pre_call_hook(auth, DualCache(), {"model": "shared-model"}, "completion") is None + await _settle_tokens(handler, tokens) + + +async def _refused_at_the_hook( + handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, model: str = "shared-model" +) -> ProxyRateLimitError: + with pytest.raises(ProxyRateLimitError) as exc: + await handler.async_pre_call_hook(auth, DualCache(), {"model": model}, "completion") + return exc.value + + +def _sample(registry: CollectorRegistry, name: str, **labels: str) -> float | None: + matches: Final = tuple( + sample.value + for metric in registry.collect() + for sample in metric.samples + if sample.name == name and all(sample.labels.get(label) == value for label, value in labels.items()) + ) # comprehension-ok: one metric family holds the series of all teams + assert len(matches) <= 1, matches + return matches[0] if matches else None + + +TEAM_LABELS: Final = {"model_group": "shared-model", "team": "team-a", "priority": "default"} + + +@pytest.mark.asyncio +async def test_team_share_gauges_follow_settled_usage_above_the_share(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) is None + + await _admit_then_settle(handler, _team_key(), tokens=60) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 0 + + await _admit_then_settle(handler, _team_key(), tokens=70) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 30 + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", **TEAM_LABELS) is None + + +@pytest.mark.asyncio +async def test_a_failed_request_settles_the_team_share_gauges_too(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + await _admit_then_settle(handler, _team_key(), tokens=90) + claim_request_stash_for_data({"litellm_call_id": "failed-call"}) + await _check(handler, _team_key(), saturation=0.0, data=_prompt(words=40), priority=None) + await handler.async_log_failure_event( + kwargs={ + "litellm_params": {"metadata": {"model_group": "shared-model"}}, + "litellm_call_id": "failed-call", + "standard_logging_object": {"metadata": {"user_api_key_team_id": "team-a"}}, + }, + response_obj=litellm.ModelResponse(usage=litellm.Usage(prompt_tokens=50, completion_tokens=0, total_tokens=50)), + start_time=datetime.now(), + end_time=datetime.now(), + ) + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 40 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth, burns, saturation_threshold, reason", + [ + (_team_key(), (60, 70), 0.1, "model_per_team_saturated"), + (_capped_team_key(2), (150, 60), 0.5, "model_per_team_max_borrow_multiple"), + ], + ids=["saturated", "borrow_cap"], +) +async def test_a_refused_request_is_counted_under_the_reason_its_429_header_carries( + monkeypatch: pytest.MonkeyPatch, + auth: UserAPIKeyAuth, + burns: tuple[int, ...], + saturation_threshold: float, + reason: str, +): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=saturation_threshold) + registry = _prometheus_sink(monkeypatch) + for tokens in burns: + await _admit_at_the_hook(handler, auth, tokens=tokens) + refusal = await _refused_at_the_hook(handler, auth) + assert refusal.headers is not None + assert refusal.headers["rate_limited_by"] == reason + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", **TEAM_LABELS) == 1 + assert ( + _sample( + registry, "litellm_dynamic_limiter_throttled_total", **TEAM_LABELS, rate_limit_type="tokens", reason=reason + ) + == 1 + ) + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == sum(burns) - 100 + + +@pytest.mark.asyncio +async def test_a_class_share_refusal_is_counted_under_the_teams_priority(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=False, splits=True) + registry = _prometheus_sink(monkeypatch) + await _burn(handler, _class_team("team-b"), "interactive", 20) + await _burn(handler, _class_team("team-a", "interactive"), "interactive", 300) + await _burn(handler, _class_team("team-a", "batch"), "batch", 300) + _held_to_class_share(await _refused_at_the_hook(handler, _class_team("team-a", "batch"))) + assert ( + _sample( + registry, + "litellm_dynamic_limiter_throttled_total", + model_group="shared-model", + team="team-a", + priority="batch", + reason="priority_class_share_saturated", + ) + == 1 + ) + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="team-b") is None + + +@pytest.mark.asyncio +async def test_an_alternates_refusal_is_not_counted_while_the_preferred_model_admits(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "priority_reservation", None) + monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) + monkeypatch.setattr( + litellm.priority_reservation_settings, "team_model_limit_borrows_idle_capacity_from_other_teams", True + ) + registry = _prometheus_sink(monkeypatch) + handler: Final[_PROXY_DynamicRateLimitHandlerV3] = _steering_handler() + router: Final = _steering_router("luna", "opus", preferred_rpm=4) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + handler.update_variables(llm_router=router) + flexible: Final = UserAPIKeyAuth( + team_id="flexible-team", + team_models=["luna", "opus"], + team_metadata={ + "model_preferences": [{"model": "luna", "alternates": ["opus"]}], + "model_rpm_limit": {"opus": 1, "luna": 1}, + "team_model_limit_max_borrow_multiple": 1, + }, + ) + await _saturate(handler, "luna", requests=2) + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") == { + "model": "opus" + } + assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team") is None + + refusal: Final = await _refused_at_the_hook(handler, flexible, model="luna") + assert refusal.headers is not None + assert refusal.headers["rate_limited_by"] == "model_per_team_saturated" + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="luna") == 1 + assert ( + _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="opus") is None + ) + + +@pytest.mark.asyncio +async def test_share_gauges_and_throttles_carry_the_teams_priority_so_the_series_join(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + registry = _prometheus_sink(monkeypatch) + team: Final = UserAPIKeyAuth( + api_key="sk-interactive", + team_id="team-a", + team_metadata={"model_tpm_limit": {"shared-model": 100}, "priority": "interactive"}, + ) + await _admit_at_the_hook(handler, team, tokens=60) + await _admit_at_the_hook(handler, team, tokens=70) + labels: Final = {"model_group": "shared-model", "team": "team-a", "priority": "interactive"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **labels) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **labels) == 30 + refusal: Final = await _refused_at_the_hook(handler, team) + assert refusal.headers is not None + assert refusal.headers["x-litellm-priority"] == "interactive" + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", **labels) == 1 + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", team="team-a", priority="default") is None + + +@pytest.mark.asyncio +async def test_a_fallback_served_request_publishes_the_teams_share_on_the_fallback_model( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + team: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 40, "shared-model": 100}}) + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None + assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True + await _settle_tokens(handler, 150, model_group="shared-model") + shared: Final = {"model_group": "shared-model", "team": "team-a", "priority": "default"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **shared) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **shared) == 50 + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", model_group="other-model") is None + + +@pytest.mark.asyncio +async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.prometheus import PrometheusLogger + + class BrokenSink(PrometheusLogger): + def set_dynamic_limiter_share(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("metrics backend down") + + def increment_dynamic_limiter_throttled(self, *args: object, **kwargs: object) -> None: + raise RuntimeError("metrics backend down") + + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + _prometheus_sink(monkeypatch, BrokenSink) + await _admit_at_the_hook(handler, _team_key(), tokens=60) + await _admit_at_the_hook(handler, _team_key(), tokens=70) + assert await _team_tokens_remaining(handler) == -30 + refusal: Final = await _refused_at_the_hook(handler, _team_key()) + assert refusal.headers is not None + assert refusal.headers["rate_limited_by"] == "model_per_team_saturated"