diff --git a/litellm/__init__.py b/litellm/__init__.py index 15e95ded906..885fde0d913 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -444,6 +444,12 @@ prometheus_emit_stream_label: bool = False # are ready to split 429s by source (vendor vs. litellm) and dimension # (RPM/TPM/concurrent/budget). prometheus_emit_rate_limit_labels: bool = False +# Opt-in: emit the `model_group` label on the deployment-level metrics +# (litellm_deployment_state / _tpm_limit / _rpm_limit / _cooled_down / +# _latency_per_output_token). Off by default so each metric's historical label +# set is preserved across upgrade; enable once downstream dashboards / recording +# rules account for the added label dimension. +prometheus_emit_deployment_model_group_label: bool = False prometheus_user_budget_label_include_email_alias: bool = False prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000 prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0 diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index b517cb0c38d..76fa8a9b4c2 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2140,6 +2140,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider or "", + model_group=model_group, ) _deployment_label_ctx = PrometheusLabelFactoryContext(enum_values) if exception is not None: @@ -2173,6 +2174,7 @@ class PrometheusLogger(CustomLogger): model_id: Optional[str], api_base: Optional[str], llm_provider: Optional[str], + model_group: Optional[str] = None, ): """ Set the deployment TPM and RPM limits metrics @@ -2185,6 +2187,7 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_tpm_limit"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, + model_group=model_group, model_id=model_id, api_base=api_base, api_provider=llm_provider, @@ -2197,6 +2200,7 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_rpm_limit"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, + model_group=model_group, model_id=model_id, api_base=api_base, api_provider=llm_provider, @@ -2309,6 +2313,7 @@ class PrometheusLogger(CustomLogger): _litellm_params = request_kwargs.get("litellm_params", {}) or {} _metadata = get_litellm_metadata_from_kwargs(request_kwargs) litellm_model_name = request_kwargs.get("model", None) + model_group = standard_logging_payload.get("model_group") llm_provider = _litellm_params.get("custom_llm_provider", None) _model_info = _metadata.get("model_info") or {} model_id = _model_info.get("id", None) @@ -2321,6 +2326,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, llm_provider=llm_provider, + model_group=model_group, ) remaining_requests: Optional[int] = None @@ -2371,6 +2377,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id or "", api_base=api_base or "", api_provider=llm_provider or "", + model_group=model_group, ) PrometheusLogger._inc_labeled_counter( @@ -2736,6 +2743,7 @@ class PrometheusLogger(CustomLogger): model_id: Optional[str], api_base: Optional[str], api_provider: str, + model_group: Optional[str] = None, ): """ Set the deployment state. @@ -2745,6 +2753,7 @@ class PrometheusLogger(CustomLogger): supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_deployment_state"), enum_values=UserAPIKeyLabelValues( litellm_model_name=litellm_model_name, + model_group=model_group, model_id=model_id, api_base=api_base, api_provider=api_provider, @@ -2758,8 +2767,11 @@ class PrometheusLogger(CustomLogger): model_id: str, api_base: str, api_provider: str, + model_group: Optional[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_base, api_provider, model_group + ) def set_deployment_partial_outage( self, @@ -2767,8 +2779,11 @@ class PrometheusLogger(CustomLogger): model_id: Optional[str], api_base: Optional[str], api_provider: str, + model_group: Optional[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_base, api_provider, model_group + ) def set_deployment_complete_outage( self, @@ -2776,8 +2791,11 @@ class PrometheusLogger(CustomLogger): model_id: Optional[str], api_base: Optional[str], api_provider: str, + model_group: Optional[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_base, api_provider, model_group + ) def increment_deployment_cooled_down( self, @@ -2786,17 +2804,25 @@ class PrometheusLogger(CustomLogger): api_base: str, api_provider: str, exception_status: str, + model_group: Optional[str] = None, ): """ increment metric when litellm.Router / load balancing logic places a deployment in cool down """ - self.litellm_deployment_cooled_down.labels( - _sanitize_prometheus_label_value(litellm_model_name), - _sanitize_prometheus_label_value(model_id), - _sanitize_prometheus_label_value(api_base), - _sanitize_prometheus_label_value(api_provider), - _sanitize_prometheus_label_value(exception_status), - ).inc() + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_cooled_down" + ), + enum_values=UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_group=model_group, + model_id=model_id, + api_base=api_base, + api_provider=api_provider, + exception_status=exception_status, + ), + ) + self.litellm_deployment_cooled_down.labels(**_labels).inc() def increment_callback_logging_failure( self, diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 313037a6364..8ba073df31a 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/litellm/router_utils/cooldown_callbacks.py @@ -41,37 +41,45 @@ async def router_cooldown_event_callback( temp_litellm_params = copy.deepcopy(_litellm_params) temp_litellm_params = dict(temp_litellm_params) _model_name = _deployment.get("model_name", None) or "" - _api_base = litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params) or "" model_info = _deployment["model_info"] model_id = model_info.id litellm_model_name = temp_litellm_params.get("model") or "" llm_provider = "" try: - _, llm_provider, _, _ = litellm.get_llm_provider( + litellm_model_name, llm_provider, _, _ = litellm.get_llm_provider( model=litellm_model_name, custom_llm_provider=temp_litellm_params.get("custom_llm_provider"), ) except Exception: pass + _api_base = ( + litellm.get_api_base( + model=litellm_model_name, optional_params=temp_litellm_params + ) + or "" + ) + # get the prometheus logger from in memory loggers prometheusLogger: Optional[PrometheusLogger] = _get_prometheus_logger_from_callbacks() 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, + model_group=_model_name, ) 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, exception_status=str(exception_status), + model_group=_model_name, ) return diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 8f460b79955..dc99b9b9499 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -773,6 +773,24 @@ class PrometheusMetricLabels: if label not in default_labels and label not in custom_labels: custom_labels.append(label) + # Conditionally add model_group to the deployment-level metrics. Off by + # default so each metric's historical label set is preserved across + # upgrade; enable via ``litellm.prometheus_emit_deployment_model_group_label`` + # once downstream dashboards / recording rules account for the new label. + _model_group_deployment_metrics = { + "litellm_deployment_state", + "litellm_deployment_tpm_limit", + "litellm_deployment_rpm_limit", + "litellm_deployment_cooled_down", + "litellm_deployment_latency_per_output_token", + } + if ( + label_name in _model_group_deployment_metrics + and litellm.prometheus_emit_deployment_model_group_label is True + and UserAPIKeyLabelNames.MODEL_GROUP.value not in default_labels + ): + custom_labels.append(UserAPIKeyLabelNames.MODEL_GROUP.value) + if label_name in PrometheusMetricLabels._org_label_metrics: for label in [ UserAPIKeyLabelNames.ORG_ID.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 be199c19149..3c4b5c0b21c 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 @@ -26,6 +26,7 @@ try: UserAPIKeyLabelValues, get_custom_labels_from_metadata, ) + from litellm.types.integrations.prometheus import PrometheusMetricLabels except Exception: PrometheusLogger = None from litellm.proxy._types import UserAPIKeyAuth @@ -680,6 +681,7 @@ async def test_async_log_failure_event(prometheus_logger): model_id="model-123", api_base="https://api.openai.com", api_provider="openai", + model_group="openai-gpt", ) # deployment failure responses incremented - verify key labels are populated @@ -906,7 +908,11 @@ async def test_async_post_call_success_hook(prometheus_logger): prometheus_logger.litellm_proxy_total_requests_metric.labels.assert_not_called() -def test_set_llm_deployment_success_metrics(prometheus_logger): +def test_set_llm_deployment_success_metrics(prometheus_logger, monkeypatch): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + prometheus_logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) # Mock all the metrics used in the method prometheus_logger.litellm_remaining_requests_metric = MagicMock() prometheus_logger.litellm_remaining_tokens_metric = MagicMock() @@ -995,6 +1001,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): model_id="model-123", api_base="https://api.openai.com", api_provider="openai", + model_group="my_custom_model_group", ) # Verify success responses metric @@ -1032,6 +1039,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): # Verify latency per output token metric prometheus_logger.litellm_deployment_latency_per_output_token.labels.assert_called_once_with( litellm_model_name="gpt-5-mini", + model_group="my_custom_model_group", model_id="model-123", api_base="https://api.openai.com", api_provider="openai", @@ -1135,7 +1143,11 @@ async def test_log_failure_fallback_event(prometheus_logger): prometheus_logger.litellm_deployment_failed_fallbacks.labels().inc.assert_called_once() -def test_deployment_state_management(prometheus_logger): +def test_deployment_state_management(prometheus_logger, monkeypatch): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + prometheus_logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) prometheus_logger.litellm_deployment_state = MagicMock() test_params = { @@ -1143,12 +1155,14 @@ def test_deployment_state_management(prometheus_logger): "model_id": "model-123", "api_base": "https://api.openai.com", "api_provider": "openai", + "model_group": "openai-gpt", } # Test set_deployment_healthy (state=0) prometheus_logger.set_deployment_healthy(**test_params) prometheus_logger.litellm_deployment_state.labels.assert_called_with( litellm_model_name=test_params["litellm_model_name"], + model_group=test_params["model_group"], model_id=test_params["model_id"], api_base=test_params["api_base"], api_provider=test_params["api_provider"], @@ -1164,9 +1178,14 @@ def test_deployment_state_management(prometheus_logger): prometheus_logger.litellm_deployment_state.labels().set.assert_called_with(2) -def test_increment_deployment_cooled_down(prometheus_logger): +def test_increment_deployment_cooled_down(prometheus_logger, monkeypatch): import inspect + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + prometheus_logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + method_sig = inspect.signature(prometheus_logger.increment_deployment_cooled_down) expected_label_count = len([p for p in method_sig.parameters.keys() if p != "self"]) @@ -1192,14 +1211,57 @@ def test_increment_deployment_cooled_down(prometheus_logger): api_base="https://api.openai.com", api_provider="openai", exception_status="429", + model_group="openai-gpt", ) prometheus_logger.litellm_deployment_cooled_down.labels.assert_called_once_with( - "gpt-5-mini", "model-123", "https://api.openai.com", "openai", "429" + litellm_model_name="gpt-5-mini", + model_group="openai-gpt", + model_id="model-123", + api_base="https://api.openai.com", + api_provider="openai", + exception_status="429", ) mock_chain.inc.assert_called_once() +def test_set_deployment_tpm_rpm_limit_metrics_includes_model_group( + prometheus_logger, monkeypatch +): + """ + Regression for https://github.com/BerriAI/litellm/issues/30748: the tpm/rpm + limit gauges must carry model_group alongside model_id so a limit can be + attributed to its configured model group, not just an opaque deployment id. + """ + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + prometheus_logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + prometheus_logger.litellm_deployment_tpm_limit = MagicMock() + prometheus_logger.litellm_deployment_rpm_limit = MagicMock() + + prometheus_logger._set_deployment_tpm_rpm_limit_metrics( + model_info={"tpm": 1000, "rpm": 60}, + litellm_params={}, + litellm_model_name="gpt-5-mini", + model_id="model-123", + api_base="https://api.openai.com", + llm_provider="openai", + model_group="openai-gpt", + ) + + tpm_labels = prometheus_logger.litellm_deployment_tpm_limit.labels.call_args.kwargs + rpm_labels = prometheus_logger.litellm_deployment_rpm_limit.labels.call_args.kwargs + + assert tpm_labels["model_group"] == "openai-gpt" + assert tpm_labels["litellm_model_name"] == "gpt-5-mini" + assert tpm_labels["model_id"] == "model-123" + assert rpm_labels["model_group"] == "openai-gpt" + + prometheus_logger.litellm_deployment_tpm_limit.labels().set.assert_called_with(1000) + prometheus_logger.litellm_deployment_rpm_limit.labels().set.assert_called_with(60) + + @pytest.mark.parametrize("enable_end_user_cost_tracking_prometheus_only", [True, False]) def test_prometheus_factory(monkeypatch, enable_end_user_cost_tracking_prometheus_only): from litellm.integrations.prometheus import prometheus_label_factory 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..9a253ba9e8b 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -10,6 +10,7 @@ except Exception: import asyncio import sys +from typing import Optional from dotenv import load_dotenv @@ -217,9 +218,10 @@ class CustomPrometheusLogger(PrometheusLogger): model_id: str, api_base: str, api_provider: str, + model_group: Optional[str] = None, ): self.deployment_complete_outages.append( - [litellm_model_name, model_id, api_base, api_provider] + [litellm_model_name, model_id, api_base, api_provider, model_group] ) def increment_deployment_cooled_down( @@ -229,9 +231,17 @@ class CustomPrometheusLogger(PrometheusLogger): api_base: str, api_provider: str, exception_status: str, + model_group: Optional[str] = None, ): self.deployment_cooled_downs.append( - [litellm_model_name, model_id, api_base, api_provider, exception_status] + [ + litellm_model_name, + model_id, + api_base, + api_provider, + exception_status, + model_group, + ] ) @@ -292,6 +302,7 @@ async def test_router_cooldown_event_callback(): "test-model-id", "https://api.openai.com", "openai", + "gpt-5-mini", ] assert prometheus_logger.deployment_cooled_downs[0] == [ "gpt-5-mini", @@ -299,9 +310,58 @@ async def test_router_cooldown_event_callback(): "https://api.openai.com", "openai", "429", + "gpt-5-mini", ] +@pytest.mark.asyncio +async def test_router_cooldown_event_callback_distinguishes_model_group_from_model(): + """ + Regression for https://github.com/BerriAI/litellm/issues/30748 + + The deployment's public alias (deployment["model_name"]) is the model_group, + while litellm_params["model"] is the underlying provider model. The cooldown + callback must report the underlying model as litellm_model_name and the alias + as model_group, matching the success path; otherwise litellm_deployment_state + fragments into two inconsistent series per deployment. Previously the alias + was passed as litellm_model_name and model_group was never emitted. + """ + from prometheus_client import REGISTRY + + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + + mock_router = MagicMock() + mock_deployment = { + "litellm_params": {"model": "gpt-5-mini"}, + "model_name": "my-gpt-group", + "model_info": ModelInfo(id="test-model-id"), + } + mock_router.get_deployment.return_value = mock_deployment + + prometheus_logger = CustomPrometheusLogger() + litellm.callbacks = [prometheus_logger] + + await router_cooldown_event_callback( + litellm_router_instance=mock_router, + deployment_id="test-deployment", + exception_status="429", + cooldown_time=60.0, + ) + + await asyncio.sleep(0.5) + + complete_outage = prometheus_logger.deployment_complete_outages[0] + cooled_down = prometheus_logger.deployment_cooled_downs[0] + + # index 0 = litellm_model_name (underlying model), last = model_group (alias) + assert complete_outage[0] == "gpt-5-mini" + assert complete_outage[-1] == "my-gpt-group" + assert cooled_down[0] == "gpt-5-mini" + assert cooled_down[-1] == "my-gpt-group" + + @pytest.mark.asyncio async def test_router_cooldown_event_callback_no_prometheus(): """ diff --git a/tests/test_litellm/integrations/test_prometheus_model_group_labels.py b/tests/test_litellm/integrations/test_prometheus_model_group_labels.py new file mode 100644 index 00000000000..f039fe5b81f --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_model_group_labels.py @@ -0,0 +1,211 @@ +""" +Tests for the opt-in `model_group` label on deployment-level Prometheus metrics +(issue #30748). + +The label is gated behind `litellm.prometheus_emit_deployment_model_group_label` +(default off) so the historical label set of each metric is preserved across +upgrade, mirroring `prometheus_emit_rate_limit_labels`. These tests assert both +the label-list wiring and that the label actually flows onto the emitted series +when the flag is enabled. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback +from litellm.types.integrations.prometheus import ( + PrometheusMetricLabels, + UserAPIKeyLabelNames, +) +from litellm.types.router import ModelInfo + +DEPLOYMENT_METRICS = [ + "litellm_deployment_state", + "litellm_deployment_tpm_limit", + "litellm_deployment_rpm_limit", + "litellm_deployment_cooled_down", + "litellm_deployment_latency_per_output_token", +] + + +def _logger_without_init() -> PrometheusLogger: + with patch( + "litellm.integrations.prometheus.PrometheusLogger.__init__", return_value=None + ): + return PrometheusLogger() + + +# --------------------------------------------------------------------------- +# Label-list wiring (flag on vs default off) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("metric_name", DEPLOYMENT_METRICS) +def test_model_group_included_when_flag_enabled(monkeypatch, metric_name): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + labels = PrometheusMetricLabels.get_labels(metric_name) + assert UserAPIKeyLabelNames.MODEL_GROUP.value in labels + # model_id must remain so a group can still be drilled down to a deployment. + assert UserAPIKeyLabelNames.MODEL_ID.value in labels + + +@pytest.mark.parametrize("metric_name", DEPLOYMENT_METRICS) +def test_model_group_omitted_by_default_for_back_compat(metric_name): + """Default-off preserves each metric's historical label set so existing + dashboards / recording rules keep matching after upgrade.""" + assert litellm.prometheus_emit_deployment_model_group_label is False + labels = PrometheusMetricLabels.get_labels(metric_name) + assert UserAPIKeyLabelNames.MODEL_GROUP.value not in labels + assert UserAPIKeyLabelNames.MODEL_ID.value in labels + + +# --------------------------------------------------------------------------- +# The label flows onto emitted series when enabled +# --------------------------------------------------------------------------- + + +def test_increment_deployment_cooled_down_emits_model_group(monkeypatch): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + logger = _logger_without_init() + logger.litellm_deployment_cooled_down = MagicMock() + logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + + logger.increment_deployment_cooled_down( + litellm_model_name="gpt-4o-mini", + model_id="model-123", + api_base="https://api.openai.com", + api_provider="openai", + exception_status="429", + model_group="gpt-group", + ) + + labels = logger.litellm_deployment_cooled_down.labels.call_args.kwargs + assert labels["model_group"] == "gpt-group" + assert labels["litellm_model_name"] == "gpt-4o-mini" + assert labels["model_id"] == "model-123" + assert labels["exception_status"] == "429" + logger.litellm_deployment_cooled_down.labels().inc.assert_called_once() + + +def test_set_litellm_deployment_state_emits_model_group(monkeypatch): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + logger = _logger_without_init() + logger.litellm_deployment_state = MagicMock() + logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + + logger.set_litellm_deployment_state( + state=2, + litellm_model_name="gpt-4o-mini", + model_id="model-123", + api_base="https://api.openai.com", + api_provider="openai", + model_group="gpt-group", + ) + + labels = logger.litellm_deployment_state.labels.call_args.kwargs + assert labels["model_group"] == "gpt-group" + assert labels["model_id"] == "model-123" + logger.litellm_deployment_state.labels().set.assert_called_with(2) + + +def test_set_deployment_tpm_rpm_limit_metrics_emit_model_group(monkeypatch): + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", True) + logger = _logger_without_init() + logger.litellm_deployment_tpm_limit = MagicMock() + logger.litellm_deployment_rpm_limit = MagicMock() + logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + + logger._set_deployment_tpm_rpm_limit_metrics( + model_info={"tpm": 1000, "rpm": 60}, + litellm_params={}, + litellm_model_name="gpt-4o-mini", + model_id="model-123", + api_base="https://api.openai.com", + llm_provider="openai", + model_group="gpt-group", + ) + + assert ( + logger.litellm_deployment_tpm_limit.labels.call_args.kwargs["model_group"] + == "gpt-group" + ) + assert ( + logger.litellm_deployment_rpm_limit.labels.call_args.kwargs["model_group"] + == "gpt-group" + ) + + +def test_deployment_metrics_omit_model_group_when_flag_disabled(monkeypatch): + """With the flag off (default), the factory must not emit model_group even + though the helper is handed a model_group value.""" + monkeypatch.setattr(litellm, "prometheus_emit_deployment_model_group_label", False) + logger = _logger_without_init() + logger.litellm_deployment_state = MagicMock() + logger.get_labels_for_metric = ( + lambda metric_name: PrometheusMetricLabels.get_labels(metric_name) + ) + + logger.set_litellm_deployment_state( + state=0, + litellm_model_name="gpt-4o-mini", + model_id="model-123", + api_base="https://api.openai.com", + api_provider="openai", + model_group="gpt-group", + ) + + assert "model_group" not in logger.litellm_deployment_state.labels.call_args.kwargs + + +# --------------------------------------------------------------------------- +# Cooldown callback: alias -> model_group, underlying model -> litellm_model_name +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_router_cooldown_callback_separates_alias_and_underlying_model( + monkeypatch, +): + """The deployment alias (model_name) becomes model_group while the + prefix-stripped underlying provider model becomes litellm_model_name, and + api_base is resolved from the underlying model rather than the alias.""" + mock_router = MagicMock() + mock_router.get_deployment.return_value = { + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_name": "my-gpt-group", + "model_info": ModelInfo(id="test-model-id"), + } + + logger = _logger_without_init() + logger.set_deployment_complete_outage = MagicMock() + logger.increment_deployment_cooled_down = MagicMock() + monkeypatch.setattr(litellm, "callbacks", [logger]) + + with patch("litellm.get_api_base", return_value="https://api.openai.com") as gab: + await router_cooldown_event_callback( + litellm_router_instance=mock_router, + deployment_id="test-deployment", + exception_status="429", + cooldown_time=60.0, + ) + + # api_base resolved from the underlying model, not the alias (P2). + assert gab.call_args.kwargs["model"] == "gpt-4o-mini" + + for mock in ( + logger.set_deployment_complete_outage, + logger.increment_deployment_cooled_down, + ): + kwargs = mock.call_args.kwargs + assert kwargs["litellm_model_name"] == "gpt-4o-mini" + assert kwargs["model_group"] == "my-gpt-group" + assert kwargs["model_id"] == "test-model-id"