From 303f699bef42580be2b18397eba09837847a8fe2 Mon Sep 17 00:00:00 2001 From: agnij-dutta Date: Sun, 18 Oct 2026 02:00:46 +0530 Subject: [PATCH] fix(prometheus): treat naive budget_reset_at as UTC for remaining hours gauges The remaining hours gauges subtracted datetime.now(budget_reset_at.tzinfo), which is local wall time when budget_reset_at is naive. The auth object prefetch builds its cached user, team and org rows with to_jsonb, and Postgres renders timestamp(3) without an offset, so those rows parse naive while they hold UTC. On a host whose TZ is not UTC every gauge fed from that cache was off by the host's UTC offset Read a naive budget_reset_at as UTC and compare against UTC now. The existing remaining hours tests built their reset times from local datetime.now(), which only matched under TZ=UTC, so they now use UTC --- litellm/integrations/prometheus.py | 9 +- .../test_prometheus_logging_callbacks.py | 8 +- tests/unit/integrations/test_prometheus.py | 112 ++++++++++++++++++ 3 files changed, 122 insertions(+), 7 deletions(-) create mode 100644 tests/unit/integrations/test_prometheus.py diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 43250fd0fd4..c286f6bfffa 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 @@ -4768,9 +4768,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..325f02ea1ba --- /dev/null +++ b/tests/unit/integrations/test_prometheus.py @@ -0,0 +1,112 @@ +import time +from collections.abc import Callable, Iterator +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest +from prometheus_client import REGISTRY + +import litellm +from litellm.integrations.prometheus import PrometheusLogger +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import UserAPIKeyAuth + +_HOURS_TO_RESET: Final = 10.0 +_ONE_MINUTE_IN_HOURS: Final = 1 / 60 + +SetGauge = Callable[[PrometheusLogger, datetime], None] + + +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 +def local_clock_ahead_of_utc(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + if not hasattr(time, "tzset"): + pytest.skip("time.tzset is POSIX only") + monkeypatch.setenv("TZ", "Asia/Kolkata") + time.tzset() + yield + 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 _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] + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("local_clock_ahead_of_utc") +def test_naive_budget_reset_at_is_read_as_utc( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + reset_at_utc: Final = datetime.now(timezone.utc) + timedelta(hours=_HOURS_TO_RESET) + + set_gauge(prometheus_logger, reset_at_utc.replace(tzinfo=None)) + + assert _only_sample(metric_name) == pytest.approx(_HOURS_TO_RESET, abs=_ONE_MINUTE_IN_HOURS) + + +@pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) +@pytest.mark.usefixtures("local_clock_ahead_of_utc") +def test_aware_budget_reset_at_in_another_zone_keeps_its_offset( + prometheus_logger: PrometheusLogger, set_gauge: SetGauge, metric_name: str +) -> None: + pacific: Final = timezone(timedelta(hours=-7)) + reset_at: Final = (datetime.now(timezone.utc) + timedelta(hours=_HOURS_TO_RESET)).astimezone(pacific) + + set_gauge(prometheus_logger, reset_at) + + assert _only_sample(metric_name) == pytest.approx(_HOURS_TO_RESET, abs=_ONE_MINUTE_IN_HOURS)