Merge 3dc66d7f06 into cfe0100e1f

litellm/integrations/prometheus.py

@@ -43,6 +43,10 @@ from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker i
     BoundedPrometheusSeriesTracker,
     PrometheusSeriesLimits,
 )
+from litellm.integrations.prometheus_helpers.deployment_state_series import (
+    DeploymentStateSeries,
+    DeploymentStateSeriesTracker,
+)
 from litellm.integrations.prometheus_helpers.shared_prometheus_series_admissions import (
     SharedPrometheusSeriesAdmissions,
 )
@@ -532,6 +536,7 @@ class PrometheusLogger(CustomLogger):
             _custom_buckets: Final = litellm.prometheus_latency_buckets
             self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS
             self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker()
+            self._deployment_state_tracker = DeploymentStateSeriesTracker(write=self.set_litellm_deployment_state)
             _multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR")
             self._series_cap_tracker = (
                 BoundedPrometheusSeriesTracker()
@@ -3952,7 +3957,9 @@ class PrometheusLogger(CustomLogger):
         api_base: str,
         api_provider: str,
     ):
-        self.set_litellm_deployment_state(0, litellm_model_name, model_id, api_base, api_provider)
+        self._deployment_state_tracker.healthy(
+            DeploymentStateSeries(litellm_model_name, model_id, api_base, api_provider)
+        )
 
     def set_deployment_partial_outage(
         self,
@@ -3961,7 +3968,9 @@ class PrometheusLogger(CustomLogger):
         api_base: str | None,
         api_provider: str,
     ):
-        self.set_litellm_deployment_state(1, litellm_model_name, model_id, api_base, api_provider)
+        self._deployment_state_tracker.partial_outage(
+            DeploymentStateSeries(litellm_model_name, model_id, api_base, api_provider)
+        )
 
     def set_deployment_complete_outage(
         self,
@@ -3970,7 +3979,9 @@ class PrometheusLogger(CustomLogger):
         api_base: str | None,
         api_provider: str,
     ):
-        self.set_litellm_deployment_state(2, litellm_model_name, model_id, api_base, api_provider)
+        self._deployment_state_tracker.complete_outage(
+            DeploymentStateSeries(litellm_model_name, model_id, api_base, api_provider)
+        )
 
     def increment_deployment_cooled_down(
         self,

litellm/integrations/prometheus_helpers/deployment_state_series.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Callable, Iterable
+from contextlib import AbstractContextManager
+from dataclasses import dataclass
+from enum import IntEnum
+from threading import Lock
+from typing import Final, TypeAlias
+
+from typing_extensions import assert_never
+
+MAX_TRACKED_DEPLOYMENTS: Final = 10_000
+MAX_SERIES_PER_DEPLOYMENT: Final = 8
+
+
+class DeploymentState(IntEnum):
+    HEALTHY = 0
+    PARTIAL_OUTAGE = 1
+    COMPLETE_OUTAGE = 2
+
+
+@dataclass(frozen=True, slots=True)
+class DeploymentStateSeries:
+    litellm_model_name: str
+    model_id: str | None
+    api_base: str | None
+    api_provider: str
+
+
+@dataclass(frozen=True, slots=True)
+class Serving:
+    written: tuple[DeploymentStateSeries, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class CoolingDown:
+    written: tuple[DeploymentStateSeries, ...]
+
+
+SeriesHistory: TypeAlias = Serving | CoolingDown
+
+
+def including(
+    written: tuple[DeploymentStateSeries, ...], series: DeploymentStateSeries
+) -> tuple[DeploymentStateSeries, ...]:
+    return (*(known for known in written if known != series), series)
+
+
+def stuck_at_two(history: SeriesHistory) -> tuple[DeploymentStateSeries, ...]:
+    match history:
+        case Serving():
+            return ()
+        case CoolingDown(written):
+            return written
+        case _:
+            assert_never(history)
+
+
+class DeploymentStateSeriesTracker:
+    """
+    A cooldown marks each series written since the last success, as its config labels can differ from the request
+    labels. The next success clears them.
+    """
+
+    def __init__(
+        self,
+        write: Callable[[DeploymentState, str, str | None, str | None, str], None],
+        lock: AbstractContextManager[object] | None = None,
+    ) -> None:
+        self._write = write
+        self._by_id: OrderedDict[str | None, SeriesHistory] = OrderedDict()  # mutable-ok: O(1) LRU
+        self._lock: AbstractContextManager[object] = Lock() if lock is None else lock
+
+    def healthy(self, series: DeploymentStateSeries) -> None:
+        with self._lock:
+            history: Final = self._by_id.get(series.model_id, Serving())
+            self._write_all(DeploymentState.HEALTHY, including(stuck_at_two(history), series))
+            self._remember(series.model_id, Serving(written=(series,)))
+
+    def partial_outage(self, series: DeploymentStateSeries) -> None:
+        with self._lock:
+            history: Final = self._by_id.get(series.model_id, Serving())
+            listed: Final = including(history.written, series)
+            tracked: Final = listed[-MAX_SERIES_PER_DEPLOYMENT:]
+            match history:
+                case Serving():
+                    self._write_all(DeploymentState.PARTIAL_OUTAGE, (series,))
+                    self._remember(series.model_id, Serving(written=tracked))
+                case CoolingDown():
+                    self._write_all(DeploymentState.HEALTHY, listed[:-MAX_SERIES_PER_DEPLOYMENT])
+                    self._write_all(DeploymentState.COMPLETE_OUTAGE, (series,))
+                    self._remember(series.model_id, CoolingDown(written=tracked))
+                case _:
+                    assert_never(history)
+
+    def complete_outage(self, derived_from_config: DeploymentStateSeries) -> None:
+        with self._lock:
+            history: Final = self._by_id.get(derived_from_config.model_id, Serving())
+            marked: Final = history.written or (derived_from_config,)
+            self._write_all(DeploymentState.COMPLETE_OUTAGE, marked)
+            self._remember(derived_from_config.model_id, CoolingDown(written=marked))
+
+    def _write_all(self, state: DeploymentState, series: Iterable[DeploymentStateSeries]) -> None:
+        for item in series:
+            self._write(state, item.litellm_model_name, item.model_id, item.api_base, item.api_provider)
+
+    def _remember(self, model_id: str | None, history: SeriesHistory) -> None:
+        self._by_id[model_id] = history
+        self._by_id.move_to_end(model_id)
+        if len(self._by_id) <= MAX_TRACKED_DEPLOYMENTS:
+            return
+        _, evicted = self._by_id.popitem(last=False)
+        self._write_all(DeploymentState.HEALTHY, stuck_at_two(evicted))

litellm/router_utils/cooldown_callbacks.py

@@ -18,6 +18,14 @@ else:
     PrometheusLogger = Any
 
 
+def resolved_model_and_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]:
+    try:
+        resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
+    except Exception:
+        return model, ""
+    return resolved_model, provider
+
+
 async def router_cooldown_event_callback(
     litellm_router_instance: LitellmRouter,
     deployment_id: str,
@@ -41,34 +49,28 @@ async def router_cooldown_event_callback(
     _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_group_name: Final = _deployment.get("model_name", None) or ""
+    _api_base: Final = litellm.get_api_base(model=model_group_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 ""
-    llm_provider = ""
-    try:
-        _, llm_provider, _, _ = litellm.get_llm_provider(
-            model=litellm_model_name,
-            custom_llm_provider=temp_litellm_params.get("custom_llm_provider"),
-        )
-    except Exception:
-        pass
+    litellm_model_name, llm_provider = resolved_model_and_provider(
+        temp_litellm_params.get("model") or "", temp_litellm_params.get("custom_llm_provider")
+    )
 
     # get the prometheus logger from in memory loggers
     prometheusLogger: Final[PrometheusLogger | None] = 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,
         )
 
         prometheusLogger.increment_deployment_cooled_down(
-            litellm_model_name=_model_name,
+            litellm_model_name=model_group_name,
             model_id=model_id,
             api_base=_api_base,
             api_provider=llm_provider,

tests/unit/conftest.py

@@ -4,15 +4,19 @@ import importlib
 import os
 from collections.abc import Callable, Coroutine, Iterator
 from dataclasses import dataclass, field
+from itertools import chain
 from pathlib import Path
 from types import MappingProxyType
-from typing import Final, Literal, TypeAlias
+from typing import TYPE_CHECKING, Final, Literal, TypeAlias
 
 import boto3
 import httpx
 import pytest
 from pytest_socket import enable_socket, socket_allow_hosts
 
+if TYPE_CHECKING:
+    from prometheus_client.samples import Sample
+
 HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset(
     (
         "PATH",
@@ -436,6 +440,40 @@ async def drained_logging_worker() -> None:
     await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10)
 
 
+@pytest.fixture
+def isolated_prometheus_registry() -> Iterator[None]:
+    from prometheus_client import REGISTRY
+
+    registered: Final = tuple(REGISTRY._collector_to_names)
+    for collector in registered:
+        REGISTRY.unregister(collector)
+    yield
+    for collector in tuple(REGISTRY._collector_to_names):
+        REGISTRY.unregister(collector)
+    for collector in registered:
+        REGISTRY.register(collector)
+
+
+@pytest.fixture
+def deployment_state_samples() -> Callable[[str], list["Sample"]]:
+    from prometheus_client import REGISTRY
+
+    def read(model_id: str) -> list["Sample"]:
+        families: Final = (family for family in REGISTRY.collect() if family.name == "litellm_deployment_state")
+        samples: Final = chain.from_iterable(family.samples for family in families)
+        return [sample for sample in samples if sample.labels["model_id"] == model_id]
+
+    return read
+
+
+@pytest.fixture
+def deployment_state_values(deployment_state_samples: Callable[[str], list["Sample"]]) -> Callable[[str], list[float]]:
+    def read(model_id: str) -> list[float]:
+        return sorted(sample.value for sample in deployment_state_samples(model_id))
+
+    return read
+
+
 def pytest_sessionfinish() -> None:
     for name in MODULE_LEVEL_CLIENTS:
         _close_handler_if_needed(litellm.__dict__.pop(name, None))

tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py

@@ -2,6 +2,7 @@
 
 import asyncio
 from datetime import datetime, timedelta, timezone
+from typing import Final
 from unittest.mock import MagicMock, call, patch
 
 import pytest
@@ -25,12 +26,17 @@ except Exception:
     PrometheusLogger = None
 from litellm.proxy._types import UserAPIKeyAuth
 
+REQUEST_SERIES: Final = {
+    "litellm_model_name": "gpt-5.4-mini",
+    "model_id": "dep-1",
+    "api_base": "https://api.openai.com/v1",
+    "api_provider": "openai",
+}
+LABELS_DERIVED_FROM_CONFIG: Final = {**REQUEST_SERIES, "litellm_model_name": "prod-chat", "api_base": ""}
+
 
 @pytest.fixture
-def prometheus_logger() -> PrometheusLogger:
-    collectors = list(REGISTRY._collector_to_names.keys())
-    for collector in collectors:
-        REGISTRY.unregister(collector)
+def prometheus_logger(isolated_prometheus_registry: None) -> PrometheusLogger:
     return PrometheusLogger()
 
 
@@ -1191,6 +1197,130 @@ def test_deployment_state_management(prometheus_logger):
     prometheus_logger.litellm_deployment_state.labels().set.assert_called_with(2)
 
 
+def test_the_public_state_setter_writes_the_given_labels_straight_to_the_gauge(
+    prometheus_logger, deployment_state_values
+):
+    prometheus_logger.set_litellm_deployment_state(2, "gpt-5.4-mini", "dep-1", "https://api.openai.com/v1", "openai")
+
+    assert deployment_state_values("dep-1") == [2.0]
+
+
+@pytest.mark.parametrize("record_request_outcome", ["set_deployment_healthy", "set_deployment_partial_outage"])
+def test_cooldown_marks_the_series_the_request_path_wrote_and_the_next_success_clears_it(
+    prometheus_logger, deployment_state_values, record_request_outcome
+):
+    prometheus_logger.set_deployment_healthy(**{**REQUEST_SERIES, "model_id": "dep-2"})
+    getattr(prometheus_logger, record_request_outcome)(**REQUEST_SERIES)
+
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    assert deployment_state_values("dep-1") == [2.0]
+    assert deployment_state_values("dep-2") == [0.0]
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0]
+
+
+def test_cooldown_before_any_request_series_is_cleared_by_the_next_success(prometheus_logger, deployment_state_values):
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [2.0, 2.0]
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0, 0.0]
+
+
+def test_a_failure_logged_after_the_cooldown_keeps_the_series_at_two_until_the_next_success(
+    prometheus_logger, deployment_state_values
+):
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [2.0]
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0]
+
+
+def test_a_second_cooldown_leaves_no_series_stuck_at_two_after_the_next_success(
+    prometheus_logger, deployment_state_values
+):
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0, 0.0]
+
+
+def test_cooldown_marks_every_series_written_since_the_last_success_and_the_next_success_clears_them(
+    prometheus_logger, deployment_state_values
+):
+    streaming_series = {**REQUEST_SERIES, "api_base": "https://api.openai.com/v1/streaming"}
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_partial_outage(**streaming_series)
+
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    assert deployment_state_values("dep-1") == [2.0, 2.0]
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0, 0.0]
+
+
+def test_a_failure_after_the_deployment_recovered_from_a_cooldown_is_a_partial_outage_again(
+    prometheus_logger, deployment_state_values
+):
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [1.0]
+
+
+def test_every_failure_logged_after_a_cooldown_stays_at_two(prometheus_logger, deployment_state_values):
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+    prometheus_logger.set_deployment_partial_outage(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [2.0]
+
+
+def test_a_cooldown_leaves_a_series_the_deployment_stopped_reporting_before_its_last_success_alone(
+    prometheus_logger, deployment_state_values
+):
+    stale_series = {**REQUEST_SERIES, "api_base": "https://api.openai.com/v1/stale"}
+    prometheus_logger.set_deployment_partial_outage(**stale_series)
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    prometheus_logger.set_deployment_complete_outage(**LABELS_DERIVED_FROM_CONFIG)
+
+    assert deployment_state_values("dep-1") == [1.0, 2.0]
+
+
+def test_a_success_leaves_the_partial_outage_of_another_series_alone_when_no_cooldown_happened(
+    prometheus_logger, deployment_state_values
+):
+    other_series = {**REQUEST_SERIES, "litellm_model_name": "gpt-5.6"}
+    prometheus_logger.set_deployment_partial_outage(**other_series)
+
+    prometheus_logger.set_deployment_healthy(**REQUEST_SERIES)
+
+    assert deployment_state_values("dep-1") == [0.0, 1.0]
+
+
 def test_increment_deployment_cooled_down(prometheus_logger):
     import inspect
 

tests/unit/integrations/prometheus_helpers/test_deployment_state_series.py

@@ -0,0 +1,283 @@
+import threading
+from collections.abc import Callable
+from typing import Final, TypeAlias
+
+import pytest
+
+from litellm.integrations.prometheus_helpers.deployment_state_series import (
+    MAX_SERIES_PER_DEPLOYMENT,
+    MAX_TRACKED_DEPLOYMENTS,
+    DeploymentState,
+    DeploymentStateSeries,
+    DeploymentStateSeriesTracker,
+)
+
+HEALTHY: Final = DeploymentState.HEALTHY
+PARTIAL_OUTAGE: Final = DeploymentState.PARTIAL_OUTAGE
+COMPLETE_OUTAGE: Final = DeploymentState.COMPLETE_OUTAGE
+
+Gauge: TypeAlias = dict[DeploymentStateSeries, DeploymentState]
+StateChange: TypeAlias = Callable[[DeploymentStateSeriesTracker, DeploymentStateSeries], None]
+
+
+def gauge_writer(
+    gauge: Gauge, before_write: Callable[[DeploymentState], None] = lambda state: None
+) -> Callable[[DeploymentState, str, str | None, str | None, str], None]:
+    def write(
+        state: DeploymentState, litellm_model_name: str, model_id: str | None, api_base: str | None, api_provider: str
+    ) -> None:
+        before_write(state)
+        gauge[DeploymentStateSeries(litellm_model_name, model_id, api_base, api_provider)] = state
+
+    return write
+
+
+def tracker_and_gauge() -> tuple[DeploymentStateSeriesTracker, Gauge]:
+    gauge: Final[Gauge] = {}
+    return DeploymentStateSeriesTracker(write=gauge_writer(gauge)), gauge
+
+
+def request_series(model_id: str, api_base: str = "https://api.example/v1") -> DeploymentStateSeries:
+    return DeploymentStateSeries("gpt", model_id, api_base, "openai")
+
+
+def config_series(model_id: str) -> DeploymentStateSeries:
+    return DeploymentStateSeries("group", model_id, "", "openai")
+
+
+def distinct_series(count: int) -> list[DeploymentStateSeries]:
+    return [request_series("dep-1", f"https://api.example/{index}") for index in range(count)]
+
+
+class ContentionRecordingLock:
+    def __init__(self, on_contention: threading.Event) -> None:
+        self.contended: Final = threading.Event()
+        self._on_contention = on_contention
+        self._lock = threading.Lock()
+
+    def __enter__(self) -> None:
+        if not self._lock.acquire(blocking=False):
+            self.contended.set()
+            self._on_contention.set()
+            self._lock.acquire()
+
+    def __exit__(self, *exc_info: object) -> None:
+        self._lock.release()
+
+
+def test_a_deployment_is_remembered_until_more_than_the_cap_of_deployments_were_seen():
+    tracker, gauge = tracker_and_gauge()
+    for index in range(MAX_TRACKED_DEPLOYMENTS):
+        tracker.healthy(request_series(f"dep-{index}"))
+
+    tracker.complete_outage(config_series("dep-0"))
+
+    assert gauge == {
+        **{request_series(f"dep-{index}"): HEALTHY for index in range(1, MAX_TRACKED_DEPLOYMENTS)},
+        request_series("dep-0"): COMPLETE_OUTAGE,
+    }
+
+
+def test_the_deployment_unused_for_longest_is_forgotten_first():
+    tracker, gauge = tracker_and_gauge()
+    tracker.healthy(request_series("busy"))
+    for index in range(MAX_TRACKED_DEPLOYMENTS - 1):
+        tracker.healthy(request_series(f"dep-{index}"))
+    tracker.healthy(request_series("busy"))
+    tracker.healthy(request_series("newest"))
+
+    tracker.complete_outage(config_series("busy"))
+    tracker.complete_outage(config_series("dep-0"))
+
+    assert gauge == {
+        **{request_series(f"dep-{index}"): HEALTHY for index in range(MAX_TRACKED_DEPLOYMENTS - 1)},
+        request_series("busy"): COMPLETE_OUTAGE,
+        request_series("newest"): HEALTHY,
+        config_series("dep-0"): COMPLETE_OUTAGE,
+    }
+
+
+def test_a_cooldown_marks_only_the_most_recent_series_of_a_deployment():
+    tracker, gauge = tracker_and_gauge()
+    failing: Final = distinct_series(MAX_SERIES_PER_DEPLOYMENT + 3)
+    for series in failing:
+        tracker.partial_outage(series)
+
+    tracker.complete_outage(config_series("dep-1"))
+
+    assert gauge == {
+        **{series: PARTIAL_OUTAGE for series in failing[:-MAX_SERIES_PER_DEPLOYMENT]},
+        **{series: COMPLETE_OUTAGE for series in failing[-MAX_SERIES_PER_DEPLOYMENT:]},
+    }
+
+
+def test_repeated_failures_of_one_series_do_not_push_another_series_out_of_the_cooldown():
+    tracker, gauge = tracker_and_gauge()
+    first, repeating = request_series("dep-1", "first"), request_series("dep-1", "repeating")
+    tracker.partial_outage(first)
+    for _ in range(MAX_SERIES_PER_DEPLOYMENT + 2):
+        tracker.partial_outage(repeating)
+
+    tracker.complete_outage(config_series("dep-1"))
+
+    assert gauge == {first: COMPLETE_OUTAGE, repeating: COMPLETE_OUTAGE}
+
+
+def test_a_success_under_a_new_label_clears_every_series_a_full_cooldown_marked():
+    tracker, gauge = tracker_and_gauge()
+    failing: Final = distinct_series(MAX_SERIES_PER_DEPLOYMENT)
+    for series in failing:
+        tracker.partial_outage(series)
+    tracker.complete_outage(config_series("dep-1"))
+    recovered: Final = request_series("dep-1", "https://api.example/recovered")
+
+    tracker.healthy(recovered)
+
+    assert gauge == {**{series: HEALTHY for series in failing}, recovered: HEALTHY}
+
+
+def test_no_series_stays_at_two_after_recovery_when_a_cooldown_saw_more_labels_than_it_tracks():
+    tracker, gauge = tracker_and_gauge()
+    labels: Final = distinct_series(MAX_SERIES_PER_DEPLOYMENT + 1)
+    for series in labels[:MAX_SERIES_PER_DEPLOYMENT]:
+        tracker.partial_outage(series)
+    tracker.complete_outage(config_series("dep-1"))
+    tracker.partial_outage(labels[-1])
+
+    tracker.healthy(labels[-1])
+
+    assert gauge == {series: HEALTHY for series in labels}
+
+
+def test_a_series_dropped_during_a_cooldown_stays_cleared_through_a_second_cooldown():
+    tracker, gauge = tracker_and_gauge()
+    labels: Final = distinct_series(MAX_SERIES_PER_DEPLOYMENT + 1)
+    for series in labels[:MAX_SERIES_PER_DEPLOYMENT]:
+        tracker.partial_outage(series)
+    tracker.complete_outage(config_series("dep-1"))
+    tracker.partial_outage(labels[-1])
+
+    tracker.complete_outage(config_series("dep-1"))
+
+    assert gauge == {labels[0]: HEALTHY, **{series: COMPLETE_OUTAGE for series in labels[1:]}}
+
+
+def test_failures_beyond_the_series_cap_outside_a_cooldown_leave_the_older_series_at_one():
+    tracker, gauge = tracker_and_gauge()
+    labels: Final = distinct_series(MAX_SERIES_PER_DEPLOYMENT + 1)
+    for series in labels:
+        tracker.partial_outage(series)
+
+    assert gauge == {series: PARTIAL_OUTAGE for series in labels}
+
+
+def test_a_deployment_pushed_out_of_the_tracker_keeps_the_state_its_requests_last_reported():
+    tracker, gauge = tracker_and_gauge()
+    failing: Final = request_series("failing")
+    tracker.partial_outage(failing)
+    for index in range(MAX_TRACKED_DEPLOYMENTS):
+        tracker.healthy(request_series(f"dep-{index}"))
+
+    assert gauge == {
+        failing: PARTIAL_OUTAGE,
+        **{request_series(f"dep-{index}"): HEALTHY for index in range(MAX_TRACKED_DEPLOYMENTS)},
+    }
+
+
+def test_a_cooling_deployment_pushed_out_of_the_tracker_is_cleared_instead_of_left_at_two():
+    tracker, gauge = tracker_and_gauge()
+    cooling: Final = request_series("cooling")
+    tracker.healthy(cooling)
+    tracker.complete_outage(config_series("cooling"))
+    for index in range(MAX_TRACKED_DEPLOYMENTS):
+        tracker.healthy(request_series(f"dep-{index}"))
+
+    assert gauge == {
+        cooling: HEALTHY,
+        **{request_series(f"dep-{index}"): HEALTHY for index in range(MAX_TRACKED_DEPLOYMENTS)},
+    }
+
+
+@pytest.mark.timeout(60)
+@pytest.mark.parametrize(
+    ("blocked_state", "in_flight_change", "in_flight_series", "waiting_change", "waiting_series", "expected_gauge"),
+    [
+        pytest.param(
+            HEALTHY,
+            DeploymentStateSeriesTracker.healthy,
+            request_series("dep-1"),
+            DeploymentStateSeriesTracker.complete_outage,
+            config_series("dep-1"),
+            {request_series("dep-1"): COMPLETE_OUTAGE},
+            id="cooldown_waits_for_a_success_write",
+        ),
+        pytest.param(
+            COMPLETE_OUTAGE,
+            DeploymentStateSeriesTracker.complete_outage,
+            config_series("dep-1"),
+            DeploymentStateSeriesTracker.healthy,
+            request_series("dep-1"),
+            {config_series("dep-1"): HEALTHY, request_series("dep-1"): HEALTHY},
+            id="success_waits_for_a_cooldown_write",
+        ),
+        pytest.param(
+            PARTIAL_OUTAGE,
+            DeploymentStateSeriesTracker.partial_outage,
+            request_series("dep-1"),
+            DeploymentStateSeriesTracker.complete_outage,
+            config_series("dep-1"),
+            {request_series("dep-1"): COMPLETE_OUTAGE},
+            id="cooldown_waits_for_a_failure_write",
+        ),
+        pytest.param(
+            COMPLETE_OUTAGE,
+            DeploymentStateSeriesTracker.complete_outage,
+            config_series("dep-1"),
+            DeploymentStateSeriesTracker.partial_outage,
+            request_series("dep-1"),
+            {config_series("dep-1"): COMPLETE_OUTAGE, request_series("dep-1"): COMPLETE_OUTAGE},
+            id="failure_waits_for_a_cooldown_write",
+        ),
+    ],
+)
+def test_a_write_in_flight_holds_back_the_next_state_change_of_the_deployment(
+    blocked_state: DeploymentState,
+    in_flight_change: StateChange,
+    in_flight_series: DeploymentStateSeries,
+    waiting_change: StateChange,
+    waiting_series: DeploymentStateSeries,
+    expected_gauge: Gauge,
+):
+    in_flight: Final = threading.Event()
+    release: Final = threading.Event()
+    second_settled: Final = threading.Event()
+    gauge: Final[Gauge] = {}
+
+    def block_on_the_blocked_state(state: DeploymentState) -> None:
+        if state == blocked_state:
+            in_flight.set()
+            release.wait()
+
+    def change_state_of_waiting_call() -> None:
+        try:
+            waiting_change(tracker, waiting_series)
+        finally:
+            second_settled.set()
+
+    lock: Final = ContentionRecordingLock(second_settled)
+    tracker: Final = DeploymentStateSeriesTracker(write=gauge_writer(gauge, block_on_the_blocked_state), lock=lock)
+    first: Final = threading.Thread(target=in_flight_change, args=(tracker, in_flight_series), daemon=True)
+    second: Final = threading.Thread(target=change_state_of_waiting_call, daemon=True)
+
+    try:
+        first.start()
+        in_flight.wait()
+        second.start()
+        second_settled.wait()
+    finally:
+        release.set()
+    first.join()
+    second.join()
+
+    assert lock.contended.is_set()
+    assert gauge == expected_gauge

tests/unit/router_utils/test_cooldown_handlers.py

@@ -4,22 +4,25 @@ import importlib
 import itertools
 import random
 import time
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from typing import Final
 from unittest.mock import MagicMock, patch
 
 import httpx
 import pytest
 import respx
+from prometheus_client.samples import Sample
 
 import litellm
 from litellm import Router
 from litellm._internal_context import current_service_target
 from litellm.caching.dual_cache import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.integrations.prometheus import PrometheusLogger
 from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
 from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
-from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
+from litellm.router_utils.cooldown_callbacks import resolved_model_and_provider, router_cooldown_event_callback
 from litellm.router_utils.cooldown_handlers import (
     _get_deployment_cooldown_policy,
     _has_explicit_allowed_fails_policy_for_exception,
@@ -52,6 +55,15 @@ from litellm.types.router import (
 from litellm.types.utils import CooldownDeploymentInfo, CredentialItem
 from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
 
+CHAT_COMPLETION: Final = {
+    "id": "chatcmpl-1",
+    "object": "chat.completion",
+    "created": 1700000000,
+    "model": "gpt-5.4-mini",
+    "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+    "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
+}
+
 
 class TestGetDeploymentCooldownPolicy:
     def _make_router(self, deployment_id: str, model_info: dict | None = None):
@@ -1752,6 +1764,100 @@ async def test_router_cooldown_event_callback_no_deployment():
     mock_router.get_deployment.assert_called_once_with(model_id="test-deployment")
 
 
+async def await_scheduled_success_logging() -> None:
+    await asyncio.sleep(0)
+    await GLOBAL_LOGGING_WORKER.flush()
+
+
+@pytest.mark.parametrize(
+    "endpoint_params",
+    [
+        pytest.param({}, id="provider_default_endpoint"),
+        pytest.param({"api_base": "https://upstream.example/v1/"}, id="api_base_with_trailing_slash"),
+    ],
+)
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "isolated_prometheus_registry")
+@pytest.mark.asyncio
+async def test_deployment_state_returns_to_healthy_after_cooldown_when_group_name_differs_from_model(
+    endpoint_params: dict[str, str],
+    monkeypatch: pytest.MonkeyPatch,
+    respx_mock: respx.MockRouter,
+    deployment_state_values: Callable[[str], list[float]],
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "callbacks", [PrometheusLogger()])
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "prod-chat",
+                "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake", **endpoint_params},
+                "model_info": {"id": "dep-1"},
+            }
+        ]
+    )
+
+    respx_mock.post(url__regex=r".*/chat/completions$").mock(return_value=httpx.Response(200, json=CHAT_COMPLETION))
+
+    await router.acompletion(model="prod-chat", messages=[{"role": "user", "content": "hi"}])
+    await await_scheduled_success_logging()
+    assert deployment_state_values("dep-1") == [0.0]
+
+    await router_cooldown_event_callback(
+        litellm_router_instance=router, deployment_id="dep-1", exception_status=404, cooldown_time=5.0
+    )
+    assert deployment_state_values("dep-1") == [2.0]
+
+    await router.acompletion(model="prod-chat", messages=[{"role": "user", "content": "hi"}])
+    await await_scheduled_success_logging()
+
+    assert deployment_state_values("dep-1") == [0.0]
+
+
+@pytest.mark.parametrize(
+    ("model_params", "expected_model_and_provider"),
+    [
+        pytest.param({"model": "openai/gpt-5.4-mini"}, ("gpt-5.4-mini", "openai"), id="provider_prefix_is_stripped"),
+        pytest.param(
+            {"model": "private-model", "custom_llm_provider": "openai"},
+            ("private-model", "openai"),
+            id="explicit_provider",
+        ),
+    ],
+)
+@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown", "isolated_prometheus_registry")
+@pytest.mark.asyncio
+async def test_cooldown_without_request_history_labels_the_state_with_the_resolved_model_and_provider(
+    model_params: dict[str, str],
+    expected_model_and_provider: tuple[str, str],
+    monkeypatch: pytest.MonkeyPatch,
+    deployment_state_samples: Callable[[str], list[Sample]],
+) -> None:
+    monkeypatch.setattr(litellm, "callbacks", [PrometheusLogger()])
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "prod-chat",
+                "litellm_params": {"api_key": "fake", "api_base": "https://upstream.example/v1", **model_params},
+                "model_info": {"id": "dep-1"},
+            }
+        ]
+    )
+
+    await router_cooldown_event_callback(
+        litellm_router_instance=router, deployment_id="dep-1", exception_status=404, cooldown_time=5.0
+    )
+
+    labelled: Final = {
+        (sample.labels["litellm_model_name"], sample.labels["api_provider"])
+        for sample in deployment_state_samples("dep-1")
+    }
+    assert labelled == {expected_model_and_provider}
+
+
+def test_a_model_litellm_cannot_resolve_keeps_its_configured_name_and_has_no_provider() -> None:
+    assert resolved_model_and_provider("private-model", None) == ("private-model", "")
+
+
 @pytest.fixture
 def testing_litellm_router():
     return Router(