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(