fix(prometheus): retain deployment attribution on proxy failures (#49717)

This commit is contained in:
yuneng-jiang 2026-11-02 12:56:53 -08:00 • committed by GitHub
parent 044bd824c0
commit 1563776cd8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 50 additions and 4 deletions

View file

@ -3036,7 +3036,7 @@ class PrometheusLogger(CustomLogger):
async def async_post_call_failure_hook(
self,
request_data: dict,
request_data: dict[str, object],
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: str | None = None,
@ -3059,10 +3059,14 @@ class PrometheusLogger(CustomLogger):
try:
_tags: Final = StandardLoggingPayloadSetup.get_request_tags(
litellm_params=request_data,
proxy_server_request=request_data.get("proxy_server_request", {}),
proxy_server_request=_logged_fields(request_data.get("proxy_server_request")),
)
_metadata: Final = _logged_fields(request_data.get("metadata"))
model_id: Final = (
_logged_fields(_metadata.get("model_info")).get("id")
or _logged_fields(request_data.get("model_info")).get("id")
or _logged_fields(request_data.get("standard_logging_object")).get("model_id")
)
_metadata: Final = request_data.get("metadata", {}) or {}
model_id: Final = _metadata.get("model_info", {}).get("id") or request_data.get("model_info", {}).get("id")
rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(original_exception)
api_provider: Final = self._extract_api_provider_from_request_data(
request_data

View file

@ -5,6 +5,7 @@ Unit tests for prometheus metric labels configuration
from typing import Final, Literal
import pytest
from pydantic import JsonValue
import litellm
from litellm.types.integrations.prometheus import (
@ -802,6 +803,47 @@ async def test_failure_hook_emits_api_provider_value_on_failed_requests_metric()
_clear_prometheus_registry()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("metadata_id", "request_id", "logged_id", "expected"),
(
(None, None, "logged-deployment", "logged-deployment"),
("", None, "logged-deployment", "logged-deployment"),
("metadata-deployment", "request-deployment", "logged-deployment", "metadata-deployment"),
(None, "request-deployment", "logged-deployment", "request-deployment"),
(None, None, None, "None"),
),
)
async def test_failure_counters_keep_logged_deployment_attribution(
metadata_id: str | None, request_id: str | None, logged_id: str | None, expected: str
) -> None:
from fastapi import HTTPException
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import UserAPIKeyAuth
request_data: Final[dict[str, JsonValue]] = {
"model": "metric-test-alias",
"metadata": {"model_info": {"id": metadata_id}},
"model_info": {"id": request_id},
"standard_logging_object": {"model_id": logged_id, "custom_llm_provider": "openai"},
}
_clear_prometheus_registry()
try:
await PrometheusLogger().async_post_call_failure_hook(
request_data=request_data,
original_exception=HTTPException(status_code=400, detail="guardrail rejected"),
user_api_key_dict=UserAPIKeyAuth(token="metric-test-token"),
)
for metric_name in ("litellm_proxy_failed_requests_metric_total", "litellm_proxy_total_requests_metric_total"):
samples: Final = _collected_samples(metric_name)
assert [(sample.labels["model_id"], sample.value) for sample in samples] == [
(expected, 1.0)
]
finally:
_clear_prometheus_registry()
@pytest.mark.asyncio
async def test_failure_hook_labels_a_cli_session_with_the_per_user_alias_not_the_login_token():
from litellm.integrations.prometheus import PrometheusLogger