mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy): compare budget window reset_at in UTC
Dropping the offset from a non-UTC reset_at compared a local wall clock against UTC now, so the window reset fired one offset late. Because the window start is derived as reset_at - budget_duration, the late reset also metered 24h + offset of spend against a 24h budget.
This commit is contained in:
parent
e484a7c89c
commit
05b2dcc095
2 changed files with 71 additions and 2 deletions
|
|
@ -1392,7 +1392,12 @@ class ResetBudgetJob:
|
|||
reset_at_str: Final = window.get("reset_at")
|
||||
if not reset_at_str:
|
||||
return False
|
||||
reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
|
||||
parsed_reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))
|
||||
reset_at: Final = (
|
||||
parsed_reset_at.replace(tzinfo=timezone.utc)
|
||||
if parsed_reset_at.tzinfo is None
|
||||
else parsed_reset_at.astimezone(timezone.utc)
|
||||
)
|
||||
if reset_at > now:
|
||||
return False
|
||||
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache)
|
||||
|
|
@ -1481,7 +1486,7 @@ class ResetBudgetJob:
|
|||
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
now: Final = datetime.utcnow()
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
for source in _WINDOW_SOURCES:
|
||||
try:
|
||||
await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache)
|
||||
|
|
|
|||
|
|
@ -1197,6 +1197,70 @@ def test_reset_budget_windows_resets_expired_team_window(monkeypatch):
|
|||
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-expired:window:30d", value=0.0)
|
||||
|
||||
|
||||
def test_reset_budget_windows_resets_expired_window_with_non_utc_offset(monkeypatch):
|
||||
"""Regression: a `reset_at` serialized with a non-UTC offset (what
|
||||
`litellm_settings.timezone: Asia/Shanghai` produces) must be compared in UTC.
|
||||
Dropping the offset kept an already-expired window blocked for the length of
|
||||
the offset, so the window never reset on time."""
|
||||
shanghai = timezone(timedelta(hours=8))
|
||||
expired = (datetime.now(shanghai) - timedelta(minutes=5)).isoformat()
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-shanghai-expired",
|
||||
"budget_limits": [{"budget_duration": "24h", "reset_at": expired}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, spend_counter_cache = _make_reset_budget_windows_job(
|
||||
monkeypatch, key_rows=key_rows, team_rows=[]
|
||||
)
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.update.assert_awaited_once()
|
||||
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(
|
||||
key="spend:key:sk-shanghai-expired:window:24h", value=0.0
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_windows_skips_unexpired_window_with_non_utc_offset(monkeypatch):
|
||||
"""The mirror of the guard above: a `+08:00` timestamp still in the future must
|
||||
not reset early just because its wall-clock value already trails naive UTC now."""
|
||||
shanghai = timezone(timedelta(hours=8))
|
||||
future = (datetime.now(shanghai) + timedelta(minutes=5)).isoformat()
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-shanghai-future",
|
||||
"budget_limits": [{"budget_duration": "24h", "reset_at": future}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
prisma_client.db.litellm_verificationtoken.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("minutes_from_now", "expected_resets"), [(-5, 1), (5, 0)])
|
||||
def test_reset_budget_windows_reads_offsetless_reset_at_as_utc(monkeypatch, minutes_from_now, expected_resets):
|
||||
"""A `reset_at` stored without any offset is UTC wall clock, so it must keep
|
||||
resetting exactly at that instant now that the comparison is timezone-aware."""
|
||||
reset_at = (datetime.now(timezone.utc) + timedelta(minutes=minutes_from_now)).replace(tzinfo=None).isoformat()
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-offsetless",
|
||||
"budget_limits": [{"budget_duration": "24h", "reset_at": reset_at}],
|
||||
}
|
||||
]
|
||||
job, prisma_client, _ = _make_reset_budget_windows_job(monkeypatch, key_rows=key_rows, team_rows=[])
|
||||
|
||||
asyncio.run(job.reset_budget_windows())
|
||||
|
||||
assert prisma_client.db.litellm_verificationtoken.update.await_count == expected_resets
|
||||
|
||||
|
||||
def test_reset_budget_windows_handles_string_budget_limits(monkeypatch):
|
||||
"""Defensive: if `query_raw` returns `budget_limits` as a JSON-encoded
|
||||
string (driver-dependent), the code still parses and resets it.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue