From 3dc66d7f066b0b187af071a63edca812e1806723 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 24 Oct 2026 21:09:35 +0000 Subject: [PATCH] fix(prometheus): stop litellm_deployment_state sticking at 2 after a cooldown The cooldown callback labelled the deployment state with the router group name and the config api_base, while requests label it with the resolved model name and the request api_base. A cooldown therefore wrote 2 to a second series that nothing ever reset, so it stayed at 2 after the provider recovered and alerts on it kept firing A new DeploymentStateSeriesTracker remembers the series each deployment wrote since its last success. A cooldown marks those series at 2, falling back to the series derived from config when there is no history. Failures logged while the deployment cools down stay at 2, and the next success clears every marked series back to 0. The cooldown callback now labels its state with the resolved model name and provider, and litellm_deployment_cooled_down keeps its labels The tracker holds at most 10000 deployments and 8 series per deployment, and it writes 0 to any marked series it stops tracking so none is left at 2. It takes its lock as an argument so the concurrency test can watch for contention instead of timing it. set_litellm_deployment_state keeps its five argument signature and the gauge states are a DeploymentState enum The tests check the final gauge state per series in the tracker, the logger setters against the real registry, and the cooldown callback against a real Router. The registry isolation and deployment state readers live in the unit conftest --- litellm/integrations/prometheus.py | 17 +- .../deployment_state_series.py | 114 +++++++ litellm/router_utils/cooldown_callbacks.py | 28 +- tests/unit/conftest.py | 40 ++- .../test_prometheus_logging_callbacks.py | 138 ++++++++- .../prometheus_helpers/__init__.py | 0 .../test_deployment_state_series.py | 283 ++++++++++++++++++ .../router_utils/test_cooldown_handlers.py | 110 ++++++- 8 files changed, 707 insertions(+), 23 deletions(-) create mode 100644 litellm/integrations/prometheus_helpers/deployment_state_series.py create mode 100644 tests/unit/integrations/prometheus_helpers/__init__.py create mode 100644 tests/unit/integrations/prometheus_helpers/test_deployment_state_series.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index bc8344fa686..f9301ece536 100644 --- a/litellm/integrations/prometheus.py +++ b/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, diff --git a/litellm/integrations/prometheus_helpers/deployment_state_series.py b/litellm/integrations/prometheus_helpers/deployment_state_series.py new file mode 100644 index 00000000000..13b51f9d3ac --- /dev/null +++ b/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)) diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py index 65aa01750ac..5a73e227710 100644 --- a/litellm/router_utils/cooldown_callbacks.py +++ b/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, diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index e64f5cca5b8..5e9b0a2996d 100644 --- a/tests/unit/conftest.py +++ b/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)) diff --git a/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index de6a5601d5c..9972c79021a 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/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 diff --git a/tests/unit/integrations/prometheus_helpers/__init__.py b/tests/unit/integrations/prometheus_helpers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/prometheus_helpers/test_deployment_state_series.py b/tests/unit/integrations/prometheus_helpers/test_deployment_state_series.py new file mode 100644 index 00000000000..abd00cfcee1 --- /dev/null +++ b/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 diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0c9d7ab9f85..450d48d9ef9 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/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(