From f9b4b153bce0f0c9da4821f10ebefa40c67f293f Mon Sep 17 00:00:00 2001 From: agnij-dutta Date: Sun, 18 Oct 2026 02:00:46 +0530 Subject: [PATCH 1/2] 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 19d161b8aa4..b9435e090f3 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 @@ -4826,9 +4826,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) From e0d98055408d71808290421ce43838d3b0559a72 Mon Sep 17 00:00:00 2001 From: agnij-dutta Date: Tue, 20 Oct 2026 20:15:37 +0530 Subject: [PATCH 2/2] test(prometheus): cover remaining hours across UTC, positive and negative offsets Freeze one UTC instant through the module's datetime and run the remaining hours cases under TZ=UTC, Asia/Kolkata and America/Los_Angeles, so the assertions are exact and a naive value read as local time shows up under both a positive and a negative offset Assert that the naive UTC wall time the prefetch cache holds and the same instant carrying UTC, +05:30 and -07:00 offsets give identical hours for the key, team, user, org and customer gauges, that an expired budget reports negative hours rather than being clamped, and that a team and user row shaped like the to_jsonb prefetch output, round-tripped through CacheCodec, reports the same hours as a warm lookup carrying the offset --- tests/unit/integrations/test_prometheus.py | 133 ++++++++++++++++++--- 1 file changed, 116 insertions(+), 17 deletions(-) diff --git a/tests/unit/integrations/test_prometheus.py b/tests/unit/integrations/test_prometheus.py index 325f02ea1ba..25b4b1aafc4 100644 --- a/tests/unit/integrations/test_prometheus.py +++ b/tests/unit/integrations/test_prometheus.py @@ -1,23 +1,40 @@ import time from collections.abc import Callable, Iterator -from datetime import datetime, timedelta, timezone +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 +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 -_ONE_MINUTE_IN_HOURS: Final = 1 / 60 +_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) @@ -30,13 +47,20 @@ def fresh_registry() -> Iterator[None]: _unregister_all_collectors() -@pytest.fixture -def local_clock_ahead_of_utc(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: +@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") - monkeypatch.setenv("TZ", "Asia/Kolkata") + zone: Final[str] = request.param + monkeypatch.setenv("TZ", zone) time.tzset() - yield + yield zone monkeypatch.undo() time.tzset() @@ -76,6 +100,16 @@ _REMAINING_HOURS_GAUGES: Final = ( ) +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) @@ -87,26 +121,91 @@ def _only_sample(name: str) -> float: 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("local_clock_ahead_of_utc") +@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: - reset_at_utc: Final = datetime.now(timezone.utc) + timedelta(hours=_HOURS_TO_RESET) + set_gauge(prometheus_logger, _RESET_AT_UTC.replace(tzinfo=None)) - 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) + assert _only_sample(metric_name) == _HOURS_TO_RESET @pytest.mark.parametrize(("set_gauge", "metric_name"), _REMAINING_HOURS_GAUGES) -@pytest.mark.usefixtures("local_clock_ahead_of_utc") +@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: - 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_UTC.astimezone(_MINUS_0700)) - set_gauge(prometheus_logger, reset_at) + assert _only_sample(metric_name) == _HOURS_TO_RESET - 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("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)