diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index bc8344fa686..5fdbc7b3ba3 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -9,7 +9,7 @@ import os import sys from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import replace -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast @@ -4921,9 +4921,12 @@ class PrometheusLogger(CustomLogger): def _get_remaining_hours_for_budget_reset(self, budget_reset_at: datetime) -> float: """ - Get remaining hours for budget reset + Get remaining hours for budget reset. A naive ``budget_reset_at`` is UTC, as the DB stores it """ - return (budget_reset_at - datetime.now(budget_reset_at.tzinfo)).total_seconds() / 3600 + reset_at: Final = ( + budget_reset_at if budget_reset_at.tzinfo is not None else budget_reset_at.replace(tzinfo=timezone.utc) + ) + return (reset_at - datetime.now(timezone.utc)).total_seconds() / 3600 def _safe_duration_seconds( self, 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..4998853d470 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -298,8 +298,8 @@ async def test_increment_remaining_budget_metrics(prometheus_logger): prometheus_logger.litellm_api_key_budget_remaining_hours_metric = MagicMock() # Create a future budget reset time for testing - future_reset_time_team = datetime.now() + timedelta(hours=10) - future_reset_time_key = datetime.now() + timedelta(hours=12) + future_reset_time_team = datetime.now(timezone.utc) + timedelta(hours=10) + future_reset_time_key = datetime.now(timezone.utc) + timedelta(hours=12) # Mock the get_team_object and get_key_object functions to return objects with budget reset times with ( patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team, @@ -1614,7 +1614,7 @@ async def test_initialize_remaining_budget_metrics(prometheus_logger): ) as mock_get_teams, ): # Create mock team data with proper datetime objects for budget_reset_at - future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now + future_reset = datetime.now(timezone.utc) + timedelta(hours=24) # Reset 24 hours from now mock_teams = [ MagicMock( team_id="team1", @@ -1790,7 +1790,7 @@ async def test_initialize_api_key_budget_metrics(prometheus_logger): ) as mock_list_keys, ): # Create mock key data with proper datetime objects for budget_reset_at - future_reset = datetime.now() + timedelta(hours=24) # Reset 24 hours from now + future_reset = datetime.now(timezone.utc) + timedelta(hours=24) # Reset 24 hours from now key1 = UserAPIKeyAuth( api_key="key1_hash", key_alias="alias1", diff --git a/tests/unit/integrations/test_prometheus.py b/tests/unit/integrations/test_prometheus.py new file mode 100644 index 00000000000..25b4b1aafc4 --- /dev/null +++ b/tests/unit/integrations/test_prometheus.py @@ -0,0 +1,211 @@ +import time +from collections.abc import Callable, Iterator +from datetime import datetime, timedelta, timezone, tzinfo +from typing import Final + +import pytest +from prometheus_client import REGISTRY +from typing_extensions import Self + +import litellm +import litellm.integrations.prometheus as prometheus_module +from litellm.integrations.prometheus import PrometheusLogger +from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec + +_HOURS_TO_RESET: Final = 10.0 +_HOURS_SINCE_EXPIRY: Final = 3.0 +# One fixed UTC instant for every test, so the remaining hours are exact rather than approximately right +_NOW_UTC: Final = datetime(2026, 10, 20, 12, 0, 0, 250000, tzinfo=timezone.utc) +_RESET_AT_UTC: Final = _NOW_UTC + timedelta(hours=_HOURS_TO_RESET) +# What Postgres to_jsonb renders for a timestamp(3) column: ISO 8601, no offset, trailing zeros dropped +_RESET_AT_TO_JSONB: Final = "2026-10-20T22:00:00.25" + +_PLUS_0530: Final = timezone(timedelta(hours=5, minutes=30)) +_MINUS_0700: Final = timezone(timedelta(hours=-7)) + +SetGauge = Callable[[PrometheusLogger, datetime], None] + + +class _FrozenDatetime(datetime): + @classmethod + def now(cls, tz: tzinfo | None = None) -> Self: + return cls.fromtimestamp(_NOW_UTC.timestamp(), tz) + + +def _unregister_all_collectors() -> None: + for collector in tuple(REGISTRY._collector_to_names): + REGISTRY.unregister(collector) + + +@pytest.fixture(autouse=True) +def fresh_registry() -> Iterator[None]: + _unregister_all_collectors() + yield + _unregister_all_collectors() + + +@pytest.fixture(autouse=True) +def frozen_now(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(prometheus_module, "datetime", _FrozenDatetime) + + +@pytest.fixture(params=("UTC", "Asia/Kolkata", "America/Los_Angeles")) +def host_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """Run under a zero, a positive and a negative UTC offset; a naive reset time read as local shifts by it.""" + if not hasattr(time, "tzset"): + pytest.skip("time.tzset is POSIX only") + zone: Final[str] = request.param + monkeypatch.setenv("TZ", zone) + time.tzset() + yield zone + monkeypatch.undo() + time.tzset() + + +@pytest.fixture +def prometheus_logger(monkeypatch: pytest.MonkeyPatch) -> PrometheusLogger: + monkeypatch.setattr(litellm, "enable_end_user_cost_tracking_prometheus_only", True) + return PrometheusLogger() + + +def _set_key(logger: PrometheusLogger, reset_at: datetime) -> None: + logger._set_key_budget_metrics(UserAPIKeyAuth(token="hashed-key", max_budget=10.0, budget_reset_at=reset_at)) + + +def _set_team(logger: PrometheusLogger, reset_at: datetime) -> None: + logger._set_team_budget_metrics(LiteLLM_TeamTable(team_id="team", max_budget=10.0, budget_reset_at=reset_at)) + + +def _set_user(logger: PrometheusLogger, reset_at: datetime) -> None: + logger._set_user_budget_metrics(LiteLLM_UserTable(user_id="user", max_budget=10.0, budget_reset_at=reset_at)) + + +def _set_org(logger: PrometheusLogger, reset_at: datetime) -> None: + logger._set_org_budget_metrics(org_id="org", org_alias="", spend=0.0, max_budget=10.0, budget_reset_at=reset_at) + + +def _set_customer(logger: PrometheusLogger, reset_at: datetime) -> None: + logger._set_customer_budget_metrics(end_user_id="customer", spend=0.0, max_budget=10.0, budget_reset_at=reset_at) + + +_REMAINING_HOURS_GAUGES: Final = ( + pytest.param(_set_key, "litellm_api_key_budget_remaining_hours_metric", id="key"), + pytest.param(_set_team, "litellm_team_budget_remaining_hours_metric", id="team"), + pytest.param(_set_user, "litellm_user_budget_remaining_hours_metric", id="user"), + pytest.param(_set_org, "litellm_org_budget_remaining_hours_metric", id="org"), + pytest.param(_set_customer, "litellm_customer_budget_remaining_hours_metric", id="customer"), +) + + +def _same_instant_in_each_shape(reset_at_utc: datetime) -> tuple[datetime, ...]: + """The naive UTC wall time the prefetch cache holds, then the instant carrying UTC, +05:30 and -07:00 offsets.""" + return ( + reset_at_utc.replace(tzinfo=None), + reset_at_utc, + reset_at_utc.astimezone(_PLUS_0530), + reset_at_utc.astimezone(_MINUS_0700), + ) + + +def _sample_values(name: str) -> Iterator[float]: + for metric in REGISTRY.collect(): + yield from (sample.value for sample in metric.samples if sample.name == name) + + +def _only_sample(name: str) -> float: + values: Final = tuple(_sample_values(name)) + assert len(values) == 1, values + return values[0] + + +def _gauge_after(logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str, reset_at: datetime) -> float: + set_gauge(logger, reset_at) + return _only_sample(metric_name) + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("host_timezone") +def test_naive_budget_reset_at_is_read_as_utc( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + set_gauge(prometheus_logger, _RESET_AT_UTC.replace(tzinfo=None)) + + assert _only_sample(metric_name) == _HOURS_TO_RESET + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("host_timezone") +def test_aware_budget_reset_at_in_another_zone_keeps_its_offset( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + set_gauge(prometheus_logger, _RESET_AT_UTC.astimezone(_MINUS_0700)) + + assert _only_sample(metric_name) == _HOURS_TO_RESET + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("host_timezone") +def test_naive_and_aware_shapes_of_one_instant_give_identical_hours( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + hours: Final = tuple( + _gauge_after(prometheus_logger, set_gauge, metric_name, reset_at) + for reset_at in _same_instant_in_each_shape(_RESET_AT_UTC) + ) + + assert hours == (_HOURS_TO_RESET,) * len(hours) + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("host_timezone") +def test_expired_budget_reports_negative_hours_unclamped( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + expired_at_utc: Final = _NOW_UTC - timedelta(hours=_HOURS_SINCE_EXPIRY) + + hours: Final = tuple( + _gauge_after(prometheus_logger, set_gauge, metric_name, reset_at) + for reset_at in _same_instant_in_each_shape(expired_at_utc) + ) + + assert hours == (-_HOURS_SINCE_EXPIRY,) * len(hours) + + +@pytest.mark.usefixtures("host_timezone") +def test_cold_prefetch_rows_and_warm_lookups_report_the_same_hours(prometheus_logger: PrometheusLogger) -> None: + """A cold read caches the to_jsonb row, which parses naive; a warm Prisma lookup carries the offset.""" + cold_team: Final = CacheCodec.deserialize( + CacheCodec.serialize( + {"team_id": "team", "max_budget": 10.0, "budget_reset_at": _RESET_AT_TO_JSONB}, + model_type=LiteLLM_TeamTableCachedObj, + ), + LiteLLM_TeamTableCachedObj, + ) + cold_user: Final = CacheCodec.deserialize( + CacheCodec.serialize( + {"user_id": "user", "max_budget": 10.0, "budget_reset_at": _RESET_AT_TO_JSONB}, + model_type=LiteLLM_UserTable, + ), + LiteLLM_UserTable, + ) + assert cold_team is not None and cold_user is not None + assert cold_team.budget_reset_at is not None and cold_team.budget_reset_at.tzinfo is None + assert cold_user.budget_reset_at is not None and cold_user.budget_reset_at.tzinfo is None + + prometheus_logger._set_team_budget_metrics(cold_team) + prometheus_logger._set_user_budget_metrics(cold_user) + cold_hours: Final = ( + _only_sample("litellm_team_budget_remaining_hours_metric"), + _only_sample("litellm_user_budget_remaining_hours_metric"), + ) + _set_team(prometheus_logger, _RESET_AT_UTC) + _set_user(prometheus_logger, _RESET_AT_UTC) + warm_hours: Final = ( + _only_sample("litellm_team_budget_remaining_hours_metric"), + _only_sample("litellm_user_budget_remaining_hours_metric"), + ) + + assert cold_hours == warm_hours == (_HOURS_TO_RESET, _HOURS_TO_RESET)