From fcdd92c2bcde031bce78668470559d5fce7aa449 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 10 Oct 2026 17:29:17 +0200 Subject: [PATCH] fix(proxy): measure months from timezone-aware datetimes and assert whole usage values datetime.fromtimestamp without a tz tripped DTZ006. The month arithmetic now runs on an aware local datetime, which keeps the same wall-clock result. The duration tests move to parametrized functions, since subTest payloads carrying datetimes cannot cross the xdist worker boundary, and the new usage tests compare the whole entry --- litellm/litellm_core_utils/duration_parser.py | 4 +- .../proxy/hooks/model_max_budget_limiter.py | 2 +- .../test_duration_parser.py | 51 ++++++++------- ...test_unit_test_max_model_budget_limiter.py | 65 +++++++++++++------ 4 files changed, 76 insertions(+), 46 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 9e6ad7493f6..d52716951df 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -74,7 +74,9 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - return _months_in_seconds(start=datetime.fromtimestamp(time_module.time()), months=value) + return _months_in_seconds( + start=datetime.fromtimestamp(time_module.time(), tz=timezone.utc).astimezone(), months=value + ) else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index a518d60de05..f562956f77a 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -304,7 +304,7 @@ def _window_reset_at(started_at: float | None, budget_duration: str | None, now: if started_at is None or budget_duration is None: return None ends_at: Final = started_at + duration_in_seconds_from( - start=datetime.fromtimestamp(started_at), duration=budget_duration + start=datetime.fromtimestamp(started_at, tz=timezone.utc).astimezone(), duration=budget_duration ) if ends_at <= now: return None diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index 9212da633e0..f14015effe6 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -4,6 +4,8 @@ from typing import Final from unittest.mock import patch from zoneinfo import ZoneInfo +import pytest + import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, @@ -374,32 +376,33 @@ class TestGetBudgetWindowStart(unittest.TestCase): ) -class TestDurationInSecondsFrom(unittest.TestCase): - def test_a_month_lasts_as_long_as_the_month_it_starts_in(self): - cases: Final = ( - (datetime(2026, 1, 15, 12), "1mo", 31), - (datetime(2026, 2, 10, 12), "1mo", 28), - (datetime(2026, 1, 31, 12), "1mo", 28), - (datetime(2028, 1, 31, 12), "1mo", 29), - (datetime(2026, 12, 15, 12), "1mo", 31), - (datetime(2026, 1, 15, 12), "2mo", 59), - ) - for start, duration, days in cases: - with self.subTest(start=start, duration=duration): - self.assertEqual(duration_in_seconds_from(start, duration), days * 86400) +@pytest.mark.parametrize( + "start,duration,days", + [ + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 2, 10, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2028, 1, 31, 12, tzinfo=timezone.utc), "1mo", 29), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "2mo", 59), + ], + ids=["january", "february", "clamped_to_february", "leap_february", "across_new_year", "two_months"], +) +def test_a_month_lasts_as_long_as_the_month_it_starts_in(start, duration, days): + assert duration_in_seconds_from(start, duration) == days * 86400 - def test_fixed_length_units_do_not_depend_on_the_start(self): - for duration in ("45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"): - with self.subTest(duration=duration): - self.assertEqual( - duration_in_seconds_from(datetime(2026, 2, 10, 12), duration), - duration_in_seconds(duration), - ) - def test_a_month_starting_now_matches_duration_in_seconds(self): - now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000) - with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): - self.assertEqual(duration_in_seconds("1mo"), duration_in_seconds_from(now, "1mo")) +@pytest.mark.parametrize("duration", ["45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"]) +def test_fixed_length_units_do_not_depend_on_the_start(duration): + assert duration_in_seconds_from(datetime(2026, 2, 10, 12, tzinfo=timezone.utc), duration) == duration_in_seconds( + duration + ) + + +def test_a_month_starting_now_matches_duration_in_seconds(): + now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000, tzinfo=timezone.utc).astimezone() + with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): + assert duration_in_seconds("1mo") == duration_in_seconds_from(now, "1mo") if __name__ == "__main__": diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index b523e5fced0..98170a406ea 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -1141,7 +1141,9 @@ async def test_usage_reports_when_the_open_window_resets(seconds_since_window_op now=_WINDOW_OPENED_AT + timedelta(seconds=seconds_since_window_opened), ) - assert usage["gpt-4"]["budget_reset_at"] == expected_reset_at + assert usage == { + "gpt-4": {"current_spend": 0.4, "budget_limit": 1.0, "time_period": "1h", "budget_reset_at": expected_reset_at} + } _DAY: Final = 86400 @@ -1151,10 +1153,10 @@ _DAY: Final = 86400 @pytest.mark.parametrize( "opened_at,read_after_days,window_days", [ - (datetime(2026, 1, 15, 12), 26, 31), - (datetime(2026, 1, 15, 12), 29, 31), - (datetime(2026, 1, 31, 12), 20, 28), - (datetime(2026, 12, 15, 12), 30, 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 26, 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 29, 31), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), 20, 28), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), 30, 31), ], ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"], ) @@ -1179,23 +1181,28 @@ async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_a entity_id="user-1", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, cache=dual_cache, - now=datetime.fromtimestamp(started_at + read_after_days * _DAY, tz=timezone.utc), + now=opened_at + timedelta(days=read_after_days), ) - assert usage["gpt-4"]["budget_reset_at"] == ( - datetime.fromtimestamp(started_at + window_days * _DAY, tz=timezone.utc).isoformat() - ) + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": (opened_at + timedelta(days=window_days)).isoformat(), + } + } @pytest.mark.asyncio async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): dual_cache = DualCache() - started_at: Final = datetime(2026, 1, 15, 12).timestamp() + opened_at: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) await dual_cache.async_set_cache( key=model_budget_start_time_cache_key( entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" ), - value=started_at, + value=opened_at.timestamp(), ) usage = await build_model_max_budget_usage( @@ -1203,10 +1210,12 @@ async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): entity_id="user-1", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, cache=dual_cache, - now=datetime.fromtimestamp(started_at + 31 * _DAY, tz=timezone.utc), + now=opened_at + timedelta(days=31), ) - assert usage["gpt-4"]["budget_reset_at"] is None + assert usage == { + "gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1mo", "budget_reset_at": None} + } @pytest.mark.asyncio @@ -1235,9 +1244,21 @@ async def test_each_budgeted_model_reports_its_own_window_reset(): now=_WINDOW_OPENED_AT, ) - assert usage["gpt-4"]["budget_reset_at"] == "2026-10-10T13:00:00+00:00" - assert usage["claude-3"]["budget_reset_at"] == "2026-10-11T07:00:00+00:00" - assert usage["never-used"]["budget_reset_at"] is None + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": "2026-10-10T13:00:00+00:00", + }, + "claude-3": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1d", + "budget_reset_at": "2026-10-11T07:00:00+00:00", + }, + "never-used": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}, + } @pytest.mark.asyncio @@ -1649,10 +1670,14 @@ async def test_a_replica_reports_the_reset_of_a_window_another_replica_opened(): entity_type=Litellm_EntityType.KEY, entity_id="vk-shared", budget_model="gpt-4", budget_duration="1h" ) ] - assert ( - usage_on_other["gpt-4"]["budget_reset_at"] - == (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat() - ) + assert usage_on_other == { + "gpt-4": { + "current_spend": 0.5, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat(), + } + } def _log_success(limiter, **kwargs):