scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

feat(prometheus): add opt-in model_group label to deployment metrics (#31025)

Browse source 
Deployment-level metrics carried only model_id, so a model group spread
across several deployments showed up as repeated model_id series with no
way to tell which configured group each belonged to. litellm_deployment_state,
litellm_deployment_tpm_limit, litellm_deployment_rpm_limit,
litellm_deployment_cooled_down and litellm_deployment_latency_per_output_token
can now emit model_group alongside model_id.

Adding a label changes a metric's time-series identity, so the label is
opt-in behind litellm.prometheus_emit_deployment_model_group_label (default
False), mirroring prometheus_emit_rate_limit_labels. Off by default preserves
each metric's historical label set across upgrade; enable it once downstream
dashboards and recording rules account for the new dimension. The label is
appended in PrometheusMetricLabels.get_labels when the flag is set, so it also
respects the include_labels filter.

The cooldown callback previously passed the deployment alias as
litellm_model_name, which disagreed with the success/failure logging paths and
fragmented litellm_deployment_state into two series per deployment. It now
reports the prefix-stripped underlying model as litellm_model_name and the
alias as model_group, and resolves api_base from the underlying model.
increment_deployment_cooled_down was moved off positional label args onto
prometheus_label_factory so it respects the label config like every other
deployment metric.

Fixes #30748
This commit is contained in:
Ali Khan 2026-06-23 09:04:19 -04:00 • committed by Sameer Kankute
parent 5c525eb086
commit 0e59c3bb17
No known key found for this signature in database
7 changed files with 411 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/__init__.py
View file

@ -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

46
litellm/integrations/prometheus.py
View file

@ -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,

16
litellm/router_utils/cooldown_callbacks.py
View file

@ -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

18
litellm/types/integrations/prometheus.py
View file

@ -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,

70
tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py
View file

@ -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

64
tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py
View file

@ -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():
"""

211
tests/test_litellm/integrations/test_prometheus_model_group_labels.py Normal file
View file

@ -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"