scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-13 03:59:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(proxy): measure months from timezone-aware datetimes and assert whole usage values

Browse source 
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
This commit is contained in:
Thijmen Stavenuiter 2026-10-10 17:29:17 +02:00
parent 25d52df97b
commit fcdd92c2bc
4 changed files with 76 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/duration_parser.py
View file

@ -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}")

2
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -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

51
tests/unit/litellm_core_utils/test_duration_parser.py
View file

@ -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__":

65
tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
View file

@ -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):