From e960d96ae98ea69181a598ef3e6d681e56d3000f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 05:20:00 +0000 Subject: [PATCH] fix(prometheus): set the remaining rate-limit gauges from decimal headers again (#47039) * test(integration): wait for both workers to load each metered model before sending * fix(prometheus): read remaining rate-limit headers as numbers and let the router fill unreadable ones * fix(prometheus): keep deployment metrics on non-string header keys and cover decimal and router-filled gauges Accept any mapping key when reading additional headers, so a non-string key no longer drops every success metric. The router fallback test now uses a real Router and covers inf, nan and unreadable requests values, the decimal test uses a fractional requests value, and the integration cell sends 12.5 and an unreadable tokens value to deployments with tpm set * fix(prometheus): leave the router fallback unchanged and keep only the decimal header fix * fix(prometheus): read remaining rate-limit headers with float() like every release did * test(integration): check the remaining rate-limit gauges on /v1/responses and /v1/messages too --- litellm/integrations/prometheus.py | 25 ++++++--- .../test_openai_like_response_headers_wire.py | 45 ++++++++++------ ...etheus_remaining_tokens_router_fallback.py | 51 ++++++++++++++++++- 3 files changed, 96 insertions(+), 25 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 517a28cf44a..43250fd0fd4 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -259,6 +259,7 @@ _POSITIVE_SERIES_CAP: Final[TypeAdapter[int]] = TypeAdapter(Annotated[int, Field _POSITIVE_SERIES_TTL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(gt=0)]) _SERIES_CLEANUP_INTERVAL: Final[TypeAdapter[float]] = TypeAdapter(Annotated[float, Field(ge=0)]) _DEFAULT_SERIES_CLEANUP_INTERVAL_SECONDS: Final = 60.0 +_ADDITIONAL_HEADERS: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) def _number_or_none(value: object, limit: TypeAdapter[_SeriesLimitT]) -> _SeriesLimitT | None: @@ -270,6 +271,15 @@ def _number_or_none(value: object, limit: TypeAdapter[_SeriesLimitT]) -> _Series return None +def _gauge_reading_or_none(value: object) -> float | None: + if not isinstance(value, str | int | float): + return None + try: + return float(value) + except (ValueError, OverflowError): + return None + + def _positive_or_ignored(setting: str, value: object, limit: TypeAdapter[_SeriesLimitT]) -> _SeriesLimitT | None: if value is None: return None @@ -3350,12 +3360,11 @@ class PrometheusLogger(CustomLogger): model_group=enum_values.model_group, ) - remaining_requests: int | None = None - remaining_tokens: int | None = None - if additional_headers := standard_logging_payload["hidden_params"]["additional_headers"]: - # OpenAI / OpenAI Compatible headers - remaining_requests = additional_headers.get("x_ratelimit_remaining_requests", None) - remaining_tokens = additional_headers.get("x_ratelimit_remaining_tokens", None) + additional_headers: Final = _ADDITIONAL_HEADERS.validate_python( + standard_logging_payload["hidden_params"]["additional_headers"] or {} + ) + remaining_requests: Final = _gauge_reading_or_none(additional_headers.get("x_ratelimit_remaining_requests")) + remaining_tokens: Final = _gauge_reading_or_none(additional_headers.get("x_ratelimit_remaining_tokens")) if litellm_overhead_time_ms := standard_logging_payload["hidden_params"].get("litellm_overhead_time_ms"): _labels = prometheus_label_factory( @@ -3373,7 +3382,7 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) - if isinstance(remaining_requests, int): + if remaining_requests is not None: """ "model_group", "api_provider", @@ -3387,7 +3396,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) - if isinstance(remaining_tokens, int): + if remaining_tokens is not None: _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, diff --git a/tests/integration/providers/test_openai_like_response_headers_wire.py b/tests/integration/providers/test_openai_like_response_headers_wire.py index 1ffa77cf326..f4d3965fb0f 100644 --- a/tests/integration/providers/test_openai_like_response_headers_wire.py +++ b/tests/integration/providers/test_openai_like_response_headers_wire.py @@ -16,6 +16,8 @@ from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.prometheus_series import Sample, scrape, series_cap_config from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter Endpoint: TypeAlias = Literal["chat", "responses", "messages"] @@ -32,13 +34,18 @@ _PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType( {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} ) _RATE_LIMIT_HEADERS: Final[Mapping[str, str]] = MappingProxyType( - {"x-ratelimit-remaining-requests": "0", "x-ratelimit-remaining-tokens": "856, abc"} + {"x-ratelimit-remaining-requests": "12.5", "x-ratelimit-remaining-tokens": "856, abc"} ) -_METERED_CALLS: Final[tuple[tuple[str, bool], ...]] = ( - ("groq", False), - ("groq", True), - ("openai", False), - ("openai", True), +_METERED_WORKERS: Final = 2 +_METERED_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( + ("chat", "groq", False), + ("chat", "groq", True), + ("chat", "openai", False), + ("chat", "openai", True), + ("responses", "groq", False), + ("responses", "groq", True), + ("messages", "groq", False), + ("messages", "groq", True), ) _LATIN1_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( ("chat", "bedrock_mantle", False), @@ -399,8 +406,10 @@ def test_provider_header_outside_latin1_does_not_fail_the_call(gateway: Gateway, assert len(wire.drain()) == len(_LATIN1_CALLS) -@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -def test_an_unparseable_rate_limit_header_keeps_the_deployment_metrics(gateway: Gateway, tmp_path: Path) -> None: +@pytest.mark.timeout(2 * graceful_stop_seconds() + 2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + 120) +def test_rate_limit_headers_set_the_remaining_gauges_and_keep_the_deployment_metrics( + gateway: Gateway, tmp_path: Path +) -> None: multiproc_dir: Final = tmp_path / "prometheus" multiproc_dir.mkdir() with ( @@ -410,19 +419,25 @@ def test_an_unparseable_rate_limit_header_keeps_the_deployment_metrics(gateway: tmp_path, {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)}, config=series_cap_config(tmp_path, {}), - workers=2, + workers=_METERED_WORKERS, ) as metered, metered.scenario() as scenario, ): calls: Final = { - f"{member}-{'stream' if stream else 'whole'}": (_MEMBERS[member](metered, scenario, wire), stream) - for member, stream in _METERED_CALLS + f"{endpoint}-{member}-{'stream' if stream else 'whole'}": ( + endpoint, + _MEMBERS[member](metered, scenario, wire), + stream, + ) + for endpoint, member, stream in _METERED_CALLS } - models: Final = {call: model for call, (model, _) in calls.items()} + models: Final = {call: model for call, (_, model, _) in calls.items()} + for model in models.values(): + wait_for_worker_readiness(metered, model, 1, _METERED_WORKERS) key: Final = scenario.key(models=list(models.values())) statuses: Final = { - call: _send(metered, "chat", model, key, stream=stream).status_code - for call, (model, stream) in calls.items() + call: _send(metered, endpoint, model, key, stream=stream).status_code + for call, (endpoint, model, stream) in calls.items() } assert statuses == dict.fromkeys(calls, 200) observed: Final = eventually( @@ -431,5 +446,5 @@ def test_an_unparseable_rate_limit_header_keeps_the_deployment_metrics(gateway: 30, return_last_on_timeout=True, ) - assert observed == dict.fromkeys(calls, (1.0, (0.0,), ())), observed + assert observed == dict.fromkeys(calls, (1.0, (12.5,), ())), observed assert len(wire.drain()) == len(calls) diff --git a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py index 28d744cdb06..5e5a32575b9 100644 --- a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py +++ b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py @@ -308,15 +308,16 @@ def _sample_values(sample_name: str) -> tuple[float, ...]: return tuple(sample.value for sample in _registry_samples() if sample.name == sample_name) +@pytest.mark.parametrize("remaining_tokens", ["856, abc", 10**400], ids=["garbled", "too-large-for-a-float"]) def test_unparseable_provider_ratelimit_header_skips_only_its_own_gauge( - prometheus_logger: PrometheusLogger, + prometheus_logger: PrometheusLogger, remaining_tokens: str | int ): payload: Final = _build_payload( model_group="groq-group", custom_llm_provider="groq", additional_headers={ "x_ratelimit_remaining_requests": 0, - "x_ratelimit_remaining_tokens": "856, abc", + "x_ratelimit_remaining_tokens": remaining_tokens, }, ) request_kwargs: Final = { @@ -336,3 +337,49 @@ def test_unparseable_provider_ratelimit_header_skips_only_its_own_gauge( assert _sample_values("litellm_remaining_requests_metric") == (0.0,) assert _sample_values("litellm_remaining_tokens_metric") == () assert _sample_values("litellm_deployment_success_responses_total") == (1.0,) + + +@pytest.mark.parametrize( + ("additional_headers", "remaining_requests", "remaining_tokens"), + [ + ({"x_ratelimit_remaining_requests": "12.5", "x_ratelimit_remaining_tokens": "856.5"}, 12.5, 856.5), + ( + {1: "non-string key", "x_ratelimit_remaining_requests": "12.5", "x_ratelimit_remaining_tokens": "856.5"}, + 12.5, + 856.5, + ), + ( + {"x_ratelimit_remaining_requests": "inf", "x_ratelimit_remaining_tokens": "\u0661\u0662.\u0665"}, + float("inf"), + 12.5, + ), + ], + ids=["decimal", "decimal-with-a-non-string-key", "infinity-and-arabic-indic-digits"], +) +def test_provider_ratelimit_headers_set_both_gauges_to_the_value_float_reads( + prometheus_logger: PrometheusLogger, + additional_headers: dict[object, object], + remaining_requests: float, + remaining_tokens: float, +): + payload: Final = _build_payload( + model_group="groq-group", + custom_llm_provider="groq", + additional_headers=additional_headers, + ) + + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs={ + "litellm_params": {"custom_llm_provider": "groq", "metadata": {}}, + "model": "llama-3.1-8b-instant", + "standard_logging_object": payload, + }, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1), + enum_values=_enum_values(model_group="groq-group"), + output_tokens=2.0, + ) + + assert _sample_values("litellm_remaining_requests_metric") == (remaining_requests,) + assert _sample_values("litellm_remaining_tokens_metric") == (remaining_tokens,) + assert _sample_values("litellm_deployment_success_responses_total") == (1.0,)