scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(budget): fix timezone config lookup and replace hardcoded timezone map with ZoneInfo (#21754)

Browse source 
* fix(budget): fix timezone config lookup and replace hardcoded timezone map with ZoneInfo

* fix(budget): update stale docstring on get_budget_reset_time
This commit is contained in:
LeeJuOh 2026-02-22 12:35:06 +09:00 • committed by GitHub
parent ab032c292c
commit 50f36d9ca6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 154 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/my-website/docs/proxy/budget_reset_and_tz.md
View file

@ -22,6 +22,8 @@ litellm_settings:
This ensures that all budget resets happen at midnight in your specified timezone rather than in UTC.
If no timezone is specified, UTC will be used by default.
Any valid [IANA timezone string](https://en.wikipedia.org/wiki/List_of_tz_database_time_zones) is supported (powered by Python's `zoneinfo` module). DST transitions are handled automatically.
Common timezone values:
- `UTC` - Coordinated Universal Time

1
docs/my-website/sidebars.js
View file

@ -417,6 +417,7 @@ const sidebars = {
"proxy/dynamic_rate_limit",
"proxy/rate_limit_tiers",
"proxy/temporary_budget_increase",
"proxy/budget_reset_and_tz",
],
},
"proxy/caching",

28
litellm/litellm_core_utils/duration_parser.py
View file

@ -8,8 +8,9 @@ duration_in_seconds is used in diff parts of the code base, example
import re
import time
from datetime import datetime, timedelta, timezone
from datetime import datetime, timedelta, timezone, tzinfo
from typing import Optional, Tuple
from zoneinfo import ZoneInfo
def _extract_from_regex(duration: str) -> Tuple[int, str]:
@ -116,7 +117,7 @@ def get_next_standardized_reset_time(
- Next reset time at a standardized interval in the specified timezone
"""
# Set up timezone and normalize current time
current_time, timezone = _setup_timezone(current_time, timezone_str)
current_time, tz = _setup_timezone(current_time, timezone_str)
# Parse duration
value, unit = _parse_duration(duration)
@ -131,7 +132,7 @@ def get_next_standardized_reset_time(
# Handle different time units
if unit == "d":
return _handle_day_reset(current_time, base_midnight, value, timezone)
return _handle_day_reset(current_time, base_midnight, value, tz)
elif unit == "h":
return _handle_hour_reset(current_time, base_midnight, value)
elif unit == "m":
@ -147,22 +148,13 @@ def get_next_standardized_reset_time(
def _setup_timezone(
current_time: datetime, timezone_str: str = "UTC"
) -> Tuple[datetime, timezone]:
) -> Tuple[datetime, tzinfo]:
"""Set up timezone and normalize current time to that timezone."""
try:
if timezone_str is None:
tz = timezone.utc
tz: tzinfo = timezone.utc
else:
# Map common timezone strings to their UTC offsets
timezone_map = {
"US/Eastern": timezone(timedelta(hours=-4)), # EDT
"US/Pacific": timezone(timedelta(hours=-7)), # PDT
"Asia/Kolkata": timezone(timedelta(hours=5, minutes=30)), # IST
"Asia/Bangkok": timezone(timedelta(hours=7)), # ICT (Indochina Time)
"Europe/London": timezone(timedelta(hours=1)), # BST
"UTC": timezone.utc,
}
tz = timezone_map.get(timezone_str, timezone.utc)
tz = ZoneInfo(timezone_str)
except Exception:
# If timezone is invalid, fall back to UTC
tz = timezone.utc
@ -190,7 +182,7 @@ def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]:
def _handle_day_reset(
current_time: datetime, base_midnight: datetime, value: int, timezone: timezone
current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo
) -> datetime:
"""Handle day-based reset times."""
# Handle zero value - immediate expiration
@ -215,7 +207,7 @@ def _handle_day_reset(
minute=0,
second=0,
microsecond=0,
tzinfo=timezone,
tzinfo=tz,
)
else:
next_reset = datetime(
@ -226,7 +218,7 @@ def _handle_day_reset(
minute=0,
second=0,
microsecond=0,
tzinfo=timezone,
tzinfo=tz,
)
return next_reset
else: # Custom day value - next interval is value days from current

18
litellm/proxy/common_utils/timezone_utils.py
View file

@ -1,27 +1,23 @@
from datetime import datetime, timezone
import litellm
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
def get_budget_reset_timezone():
"""
Get the budget reset timezone from general_settings.
Get the budget reset timezone from litellm_settings.
Falls back to UTC if not specified.
litellm_settings values are set as attributes on the litellm module
by proxy_server.py at startup (via setattr(litellm, key, value)).
"""
# Import at function level to avoid circular imports
from litellm.proxy.proxy_server import general_settings
if general_settings:
litellm_settings = general_settings.get("litellm_settings", {})
if litellm_settings and "timezone" in litellm_settings:
return litellm_settings["timezone"]
return "UTC"
return getattr(litellm, "timezone", None) or "UTC"
def get_budget_reset_time(budget_duration: str):
"""
Get the budget reset time from general_settings.
Get the budget reset time based on the configured timezone.
Falls back to UTC if not specified.
"""

60
tests/test_litellm/litellm_core_utils/test_duration_parser.py
View file

@ -91,7 +91,9 @@ class TestStandardizedResetTime(unittest.TestCase):
# Test Bangkok timezone (UTC+7): 5:30 AM next day, so next reset is midnight the day after
bangkok = ZoneInfo("Asia/Bangkok")
bangkok_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=bangkok)
bangkok_result = get_next_standardized_reset_time("1d", base_time, "Asia/Bangkok")
bangkok_result = get_next_standardized_reset_time(
"1d", base_time, "Asia/Bangkok"
)
self.assertEqual(bangkok_result, bangkok_expected)
def test_edge_cases(self):
@ -125,6 +127,62 @@ class TestStandardizedResetTime(unittest.TestCase):
)
self.assertEqual(invalid_tz_result, invalid_tz_expected)
def test_iana_timezones_previously_unsupported(self):
"""Test IANA timezones that were previously unsupported by the hardcoded map."""
# Base time: 2023-05-15 15:00:00 UTC
base_time = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc)
# Asia/Tokyo (UTC+9): 15:00 UTC = 00:00 JST May 16, exactly on midnight boundary → next day
tokyo = ZoneInfo("Asia/Tokyo")
tokyo_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=tokyo)
tokyo_result = get_next_standardized_reset_time(
"1d", base_time, "Asia/Tokyo"
)
self.assertEqual(tokyo_result, tokyo_expected)
# Australia/Sydney (UTC+10): 2023-05-16 01:00 AEST
sydney = ZoneInfo("Australia/Sydney")
# At 15:00 UTC it's 01:00 AEST May 16 → next midnight is May 17 00:00 AEST
sydney_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=sydney)
sydney_result = get_next_standardized_reset_time(
"1d", base_time, "Australia/Sydney"
)
self.assertEqual(sydney_result, sydney_expected)
# America/Chicago (UTC-5): at 15:00 UTC it's 10:00 CDT → next midnight is May 16 00:00 CDT
chicago = ZoneInfo("America/Chicago")
chicago_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=chicago)
chicago_result = get_next_standardized_reset_time(
"1d", base_time, "America/Chicago"
)
self.assertEqual(chicago_result, chicago_expected)
def test_dst_fall_back(self):
"""Test DST fall-back transition (clocks go back 1 hour)."""
# US/Eastern DST ends first Sunday of November 2023 (Nov 5)
# At 2023-11-05 05:30 UTC = 01:30 EDT (before fall-back)
# After fall-back at 06:00 UTC = 01:00 EST
pre_fallback = datetime(2023, 11, 5, 5, 30, 0, tzinfo=timezone.utc)
eastern = ZoneInfo("US/Eastern")
# Daily reset: next midnight should be Nov 6 00:00 EST
expected = datetime(2023, 11, 6, 0, 0, 0, tzinfo=eastern)
result = get_next_standardized_reset_time("1d", pre_fallback, "US/Eastern")
self.assertEqual(result, expected)
def test_dst_spring_forward(self):
"""Test DST spring-forward transition (clocks go forward 1 hour)."""
# US/Eastern DST starts second Sunday of March 2023 (Mar 12)
# At 2023-03-12 06:30 UTC = 01:30 EST (before spring-forward)
# After spring-forward at 07:00 UTC = 03:00 EDT
pre_spring = datetime(2023, 3, 12, 6, 30, 0, tzinfo=timezone.utc)
eastern = ZoneInfo("US/Eastern")
# Daily reset: next midnight should be Mar 13 00:00 EDT
expected = datetime(2023, 3, 13, 0, 0, 0, tzinfo=eastern)
result = get_next_standardized_reset_time("1d", pre_spring, "US/Eastern")
self.assertEqual(result, expected)
if __name__ == "__main__":
unittest.main()

83
tests/test_litellm/proxy/common_utils/test_timezone_utils.py
View file

@ -1,18 +1,17 @@
import asyncio
import json
import os
import sys
import time
from datetime import datetime, timedelta, timezone
import pytest
from fastapi.testclient import TestClient
from datetime import datetime, timezone
from zoneinfo import ZoneInfo
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
import litellm
from litellm.proxy.common_utils.timezone_utils import (
get_budget_reset_time,
get_budget_reset_timezone,
)
def test_get_budget_reset_time():
@ -33,3 +32,71 @@ def test_get_budget_reset_time():
# Verify budget_reset_at is set to first of next month
assert get_budget_reset_time(budget_duration="1mo") == expected_reset_at
def test_get_budget_reset_timezone_reads_litellm_attr():
"""
Test that get_budget_reset_timezone reads from litellm.timezone attribute.
"""
original = getattr(litellm, "timezone", None)
try:
litellm.timezone = "Asia/Tokyo"
assert get_budget_reset_timezone() == "Asia/Tokyo"
finally:
if original is None:
if hasattr(litellm, "timezone"):
delattr(litellm, "timezone")
else:
litellm.timezone = original
def test_get_budget_reset_timezone_fallback_utc():
"""
Test that get_budget_reset_timezone falls back to UTC when litellm.timezone is not set.
"""
original = getattr(litellm, "timezone", None)
try:
if hasattr(litellm, "timezone"):
delattr(litellm, "timezone")
assert get_budget_reset_timezone() == "UTC"
finally:
if original is not None:
litellm.timezone = original
def test_get_budget_reset_timezone_fallback_on_none():
"""
Test that get_budget_reset_timezone falls back to UTC when litellm.timezone is None.
"""
original = getattr(litellm, "timezone", None)
try:
litellm.timezone = None
assert get_budget_reset_timezone() == "UTC"
finally:
if original is None:
if hasattr(litellm, "timezone"):
delattr(litellm, "timezone")
else:
litellm.timezone = original
def test_get_budget_reset_time_respects_timezone():
"""
Test that get_budget_reset_time uses the configured timezone for reset calculation.
A daily reset should align to midnight in the configured timezone.
"""
original = getattr(litellm, "timezone", None)
try:
litellm.timezone = "Asia/Tokyo"
reset_at = get_budget_reset_time(budget_duration="1d")
# The reset time should be midnight in Asia/Tokyo
tokyo_reset = reset_at.astimezone(ZoneInfo("Asia/Tokyo"))
assert tokyo_reset.hour == 0
assert tokyo_reset.minute == 0
assert tokyo_reset.second == 0
finally:
if original is None:
if hasattr(litellm, "timezone"):
delattr(litellm, "timezone")
else:
litellm.timezone = original