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"