scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(dynamic_rate_limiter_v3): prometheus share, borrowed and throttle metrics per team and model (#51417)

Browse source 
* 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>
This commit is contained in:
Kerry Lu 2026-11-12 15:28:51 -08:00 • committed by GitHub
parent f87d0a98df
commit 7130a5b4fd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 781 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

87
litellm/integrations/prometheus.py
View file

@ -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.

184
litellm/integrations/prometheus_helpers/grafana_dashboards/dashboard_all_metrics/grafana_dashboard.json
View file

@ -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,

115
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -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

1
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -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)

22
litellm/types/integrations/prometheus.py
View file

@ -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

151
tests/integration/routing/test_dynamic_limiter_prometheus_metrics.py Normal file
View file

@ -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)

4
tests/unit/integrations/test_prometheus_labels.py
View file

@ -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")

230
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -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"