From 2b9be52a6216f842e2610838f1619a7a47378cee Mon Sep 17 00:00:00 2001 From: Steve Gran Date: Mon, 3 Aug 2026 12:57:08 +0100 Subject: [PATCH 1/2] Canonicalise prometheus labels for the metric litellm_deployment_state These labels are being constructed in 3 different places, leading to drift when the values change. This results in multiple series for each model if a model ever goes unhealthy and then becomes healthy again, making alerting impractical. This simplifies the logic by removing api_base from the set of labels and centralising the construction of the label set. --- litellm/integrations/prometheus.py | 81 +++++---- litellm/router_utils/cooldown_callbacks.py | 15 +- litellm/types/integrations/prometheus.py | 1 - .../test_prometheus_logging_callbacks.py | 2 - .../test_prometheus_unit_tests.py | 4 +- .../router_utils/test_cooldown_callbacks.py | 170 ++++++++++++++++++ 6 files changed, 228 insertions(+), 45 deletions(-) create mode 100644 tests/test_litellm/router_utils/test_cooldown_callbacks.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index a9056aaf4e1..f50059a1201 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2576,15 +2576,14 @@ class PrometheusLogger(CustomLogger): """ log these labels - ["litellm_model_name", "model_id", "api_base", "api_provider"] + ["litellm_model_name", "model_id", "api_provider"] """ # Only mark a deployment outage when one was actually picked. if deployment_selected: self.set_deployment_partial_outage( - litellm_model_name=litellm_model_name or "", + litellm_model_name=litellm_model_name, model_id=model_id, - api_base=api_base, - api_provider=llm_provider or "", + api_provider=llm_provider or self._extract_api_provider_from_request_data(request_kwargs), ) _deployment_label_ctx: Final = PrometheusLabelFactoryContext(enum_values) if exception is not None: @@ -2810,13 +2809,12 @@ class PrometheusLogger(CustomLogger): """ log these labels - ["litellm_model_name", "requested_model", model_id", "api_base", "api_provider"] + ["litellm_model_name", "model_id", "api_provider"] """ self.set_deployment_healthy( - litellm_model_name=litellm_model_name or "", - model_id=model_id or "", - api_base=api_base or "", - api_provider=llm_provider or "", + litellm_model_name=litellm_model_name, + model_id=model_id, + api_provider=llm_provider or self._extract_api_provider_from_request_data(request_kwargs), ) PrometheusLogger._inc_labeled_counter( @@ -3173,55 +3171,70 @@ class PrometheusLogger(CustomLogger): label_context=PrometheusLabelFactoryContext(enum_values), ) + def get_deployment_state_labels( + self, + litellm_model_name: str | None, + model_id: str | None, + api_provider: str | None, + ) -> dict: + """ + Returns the complete label set for the litellm_deployment_state gauge. + + Every writer of the gauge (success, failure, and cooldown paths) must + derive its labels through this helper so a deployment always maps to + exactly one time series; a second labelset would leave a stale state + value exported forever. Values are coerced to non-empty-or-"" strings + so every label is always present. + """ + return prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_state"), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name or "", + model_id=model_id or "", + api_provider=api_provider or "", + ), + ) + def set_litellm_deployment_state( self, state: int, - litellm_model_name: str, + litellm_model_name: str | None, model_id: str | None, - api_base: str | None, - api_provider: str, + api_provider: str | None, ): """ Set the deployment state. """ - ### get labels - _labels: Final = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_state"), - enum_values=UserAPIKeyLabelValues( - litellm_model_name=litellm_model_name, - model_id=model_id, - api_base=api_base, - api_provider=api_provider, - ), + _labels: Final = self.get_deployment_state_labels( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_provider=api_provider, ) self.litellm_deployment_state.labels(**_labels).set(state) def set_deployment_healthy( self, - litellm_model_name: str, - model_id: str, - api_base: str, - api_provider: str, + litellm_model_name: str | None, + model_id: str | None, + api_provider: str | None, ): - self.set_litellm_deployment_state(0, litellm_model_name, model_id, api_base, api_provider) + self.set_litellm_deployment_state(0, litellm_model_name, model_id, api_provider) def set_deployment_partial_outage( self, - litellm_model_name: str, + litellm_model_name: str | None, model_id: str | None, - api_base: str | None, - api_provider: str, + api_provider: str | None, ): - self.set_litellm_deployment_state(1, litellm_model_name, model_id, api_base, api_provider) + self.set_litellm_deployment_state(1, litellm_model_name, model_id, api_provider) def set_deployment_complete_outage( self, - litellm_model_name: str, + litellm_model_name: str | None, model_id: str | None, - api_base: str | None, - api_provider: str, + api_provider: str | None, ): - self.set_litellm_deployment_state(2, litellm_model_name, model_id, api_base, api_provider) + self.set_litellm_deployment_state(2, litellm_model_name, model_id, api_provider) def increment_deployment_cooled_down( self, diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 9773964461b..24695168785 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -38,15 +38,21 @@ async def router_cooldown_event_callback( deployment_id, ) return + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + _litellm_params: Final = _deployment["litellm_params"] temp_litellm_params = copy.deepcopy(_litellm_params) temp_litellm_params = dict(temp_litellm_params) - _model_name: Final = _deployment.get("model_name", None) or "" - _api_base: Final = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" model_info: Final = _deployment["model_info"] model_id: Final = model_info.id litellm_model_name: Final = temp_litellm_params.get("model") or "" + _api_base: Final = ( + StandardLoggingPayloadSetup.strip_trailing_slash( + litellm.get_api_base(model=litellm_model_name, optional_params=temp_litellm_params) + ) + or "" + ) llm_provider = "" try: _, llm_provider, _, _ = litellm.get_llm_provider( @@ -61,14 +67,13 @@ async def router_cooldown_event_callback( if prometheusLogger is not None: prometheusLogger.set_deployment_complete_outage( - litellm_model_name=_model_name, + litellm_model_name=litellm_model_name, model_id=model_id, - api_base=_api_base, api_provider=llm_provider, ) prometheusLogger.increment_deployment_cooled_down( - litellm_model_name=_model_name, + litellm_model_name=litellm_model_name, model_id=model_id, api_base=_api_base, api_provider=llm_provider, diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index ebec5df55fa..5531c147eef 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -520,7 +520,6 @@ class PrometheusMetricLabels: litellm_deployment_state = [ UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, - UserAPIKeyLabelNames.API_BASE.value, UserAPIKeyLabelNames.API_PROVIDER.value, ] diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index b6c9cd0294b..309f2527d30 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -1153,7 +1153,6 @@ def test_deployment_state_management(prometheus_logger): test_params = { "litellm_model_name": "gpt-5-mini", "model_id": "model-123", - "api_base": "https://api.openai.com", "api_provider": "openai", } @@ -1162,7 +1161,6 @@ def test_deployment_state_management(prometheus_logger): prometheus_logger.litellm_deployment_state.labels.assert_called_with( litellm_model_name=test_params["litellm_model_name"], model_id=test_params["model_id"], - api_base=test_params["api_base"], api_provider=test_params["api_provider"], ) prometheus_logger.litellm_deployment_state.labels().set.assert_called_with(0) diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index 55c4cbae821..a35dd48827a 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -215,11 +215,10 @@ class CustomPrometheusLogger(PrometheusLogger): self, litellm_model_name: str, model_id: str, - api_base: str, api_provider: str, ): self.deployment_complete_outages.append( - [litellm_model_name, model_id, api_base, api_provider] + [litellm_model_name, model_id, api_provider] ) def increment_deployment_cooled_down( @@ -290,7 +289,6 @@ async def test_router_cooldown_event_callback(): assert prometheus_logger.deployment_complete_outages[0] == [ "gpt-5-mini", "test-model-id", - "https://api.openai.com", "openai", ] assert prometheus_logger.deployment_cooled_downs[0] == [ diff --git a/tests/test_litellm/router_utils/test_cooldown_callbacks.py b/tests/test_litellm/router_utils/test_cooldown_callbacks.py new file mode 100644 index 00000000000..5e3ccb61943 --- /dev/null +++ b/tests/test_litellm/router_utils/test_cooldown_callbacks.py @@ -0,0 +1,170 @@ +""" +Tests for litellm/router_utils/cooldown_callbacks.py +""" + +import datetime + +import pytest + +import litellm +from litellm import Router +from litellm.integrations.prometheus import PrometheusLogger +from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback + + +def _clear_prometheus_registry() -> None: + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names.keys()): + try: + REGISTRY.unregister(collector) + except Exception: + pass + + +def _collected_samples(metric_name: str): + from prometheus_client import REGISTRY + + return [ + sample + for metric in REGISTRY.collect() + for sample in metric.samples + if sample.name == metric_name + ] + + +def _standard_logging_payload(litellm_model_name: str, model_id: str) -> dict: + return { + "id": "t", + "call_type": "completion", + "response_cost": 0.001, + "status": "success", + "total_tokens": 30, + "prompt_tokens": 20, + "completion_tokens": 10, + "startTime": 1.0, + "endTime": 2.0, + "completionStartTime": 1.5, + "model": litellm_model_name, + "model_id": model_id, + "model_group": "claude-opus-4-8", + "api_base": "https://api.anthropic.com/v1/messages", + "custom_llm_provider": "anthropic", + "request_tags": [], + "end_user": None, + "cache_hit": False, + "metadata": { + "user_api_key_hash": "h", + "user_api_key_alias": "a", + "user_api_key_team_id": "t", + "user_api_key_team_alias": "ta", + "user_api_key_user_id": "u", + "user_api_key_user_email": "e@x.com", + "user_api_key_org_id": None, + "user_api_key_org_alias": None, + "requester_metadata": None, + "user_api_key_end_user_id": None, + }, + "hidden_params": {"litellm_overhead_time_ms": None, "additional_headers": None}, + } + + +@pytest.mark.asyncio +async def test_deployment_state_is_single_series_across_cooldown_failure_and_recovery(monkeypatch): + """ + Regression test for litellm_deployment_state series fragmentation. + + The gauge's series identity is (litellm_model_name, model_id, + api_provider). router_cooldown_event_callback used to label it with the + deployment's public model_name group alias, and the gauge additionally + carried an api_base label whose value differed between the config-driven + cooldown path and the request-driven success/failure paths. Each writer + therefore created its own series: the cooldown's state=2 was never reset by + recovery, so dashboards reported the deployment as permanently unhealthy. + + Drives the real cooldown callback and the real success/failure logging + paths for one deployment with no configured api_base (mirroring production + anthropic/bedrock configs) and asserts every write lands on a single fully + labeled time series that follows the deployment's actual state. + """ + model_id = "deployment-state-regression-id" + litellm_model_name = "anthropic/claude-opus-4-8" + + router = Router( + model_list=[ + { + "model_name": "claude-opus-4-8", + "litellm_params": { + "model": litellm_model_name, + "api_key": "fake-key", + }, + "model_info": {"id": model_id}, + } + ] + ) + + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + + async def cooldown(): + await router_cooldown_event_callback( + litellm_router_instance=router, + deployment_id=model_id, + exception_status="429", + cooldown_time=60.0, + ) + + def single_deployment_state_sample(): + samples = _collected_samples("litellm_deployment_state") + assert len(samples) == 1, ( + f"expected a single litellm_deployment_state series, got {[s.labels for s in samples]}" + ) + assert samples[0].labels == { + "litellm_model_name": litellm_model_name, + "model_id": model_id, + "api_provider": "anthropic", + } + return samples[0] + + await cooldown() + assert single_deployment_state_sample().value == 2 + + cooled_down_samples = _collected_samples("litellm_deployment_cooled_down_total") + assert len(cooled_down_samples) == 1 + assert cooled_down_samples[0].labels["litellm_model_name"] == litellm_model_name + + now = datetime.datetime.now() + await logger.async_log_success_event( + { + "model": litellm_model_name, + "litellm_params": { + "custom_llm_provider": "anthropic", + "metadata": {"model_info": {"id": model_id}}, + }, + "standard_logging_object": _standard_logging_payload(litellm_model_name, model_id), + }, + None, + now, + now, + ) + assert single_deployment_state_sample().value == 0 + + logger.set_llm_deployment_failure_metrics( + { + "model": litellm_model_name, + "exception": litellm.exceptions.RateLimitError("rate limited", "anthropic", litellm_model_name), + "litellm_params": { + "custom_llm_provider": "anthropic", + "metadata": {"model_info": {"id": model_id}}, + }, + "standard_logging_object": _standard_logging_payload(litellm_model_name, model_id), + } + ) + assert single_deployment_state_sample().value == 1 + + await cooldown() + assert single_deployment_state_sample().value == 2 + finally: + _clear_prometheus_registry() From e07785aa762556ab3d5242d789f4161d23c997e9 Mon Sep 17 00:00:00 2001 From: Steve Gran Date: Mon, 3 Aug 2026 13:29:14 +0100 Subject: [PATCH 2/2] Fix failing tests --- litellm/integrations/prometheus.py | 2 +- .../enterprise_callbacks/test_prometheus_logging_callbacks.py | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f50059a1201..a43a3cc1b54 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3176,7 +3176,7 @@ class PrometheusLogger(CustomLogger): litellm_model_name: str | None, model_id: str | None, api_provider: str | None, - ) -> dict: + ) -> Mapping[str, str]: """ Returns the complete label set for the litellm_deployment_state gauge. diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 309f2527d30..f875d46142f 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -688,7 +688,6 @@ async def test_async_log_failure_event(prometheus_logger): prometheus_logger.set_deployment_partial_outage.assert_called_once_with( litellm_model_name="gpt-5-mini", model_id="model-123", - api_base="https://api.openai.com", api_provider="openai", ) @@ -1005,7 +1004,6 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): prometheus_logger.set_deployment_healthy.assert_called_once_with( litellm_model_name="gpt-5-mini", model_id="model-123", - api_base="https://api.openai.com", api_provider="openai", )