scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge e07785aa76 into eb0e3f8c18

Browse source
This commit is contained in:
Stephen Gran 2026-08-27 19:15:28 -05:00 • committed by GitHub
commit 3894dcec03
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 228 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
litellm/integrations/prometheus.py
View file

@ -2649,15 +2649,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:
@ -2883,13 +2882,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(
@ -3246,55 +3244,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,
) -> Mapping[str, str]:
"""
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,

15
litellm/router_utils/cooldown_callbacks.py
View file

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

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

@ -625,7 +625,6 @@ class PrometheusMetricLabels:
litellm_deployment_state = [
UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value,
UserAPIKeyLabelNames.MODEL_ID.value,
UserAPIKeyLabelNames.API_BASE.value,
UserAPIKeyLabelNames.API_PROVIDER.value,
]

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

@ -685,7 +685,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",
)
@ -1002,7 +1001,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",
)
@ -1150,7 +1148,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",
}
@ -1159,7 +1156,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)

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

@ -211,11 +211,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(
@ -286,7 +285,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] == [

170
tests/test_litellm/router_utils/test_cooldown_callbacks.py Normal file
View file

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