mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(budgets): compare budget_limits reset_at as an instant, not naive wall-clock
This commit is contained in:
parent
c274cf321c
commit
1dfabcbc7f
2 changed files with 49 additions and 2 deletions
|
|
@ -695,7 +695,9 @@ class ResetBudgetJob:
|
|||
reset_at_str = window.get("reset_at")
|
||||
if not reset_at_str:
|
||||
return False
|
||||
reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
|
||||
reset_at = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))
|
||||
if reset_at.tzinfo is None:
|
||||
reset_at = reset_at.replace(tzinfo=timezone.utc)
|
||||
if reset_at > now:
|
||||
return False
|
||||
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0)
|
||||
|
|
@ -717,7 +719,7 @@ class ResetBudgetJob:
|
|||
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
now = datetime.utcnow()
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# Note on raw SQL: prisma-client-python does not support null-filtering
|
||||
# on `Json?` columns (no DbNull/JsonNull sentinel — see
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from datetime import datetime, timedelta, timezone
|
|||
from datetime import time as dt_time
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -1171,6 +1172,50 @@ def test_reset_budget_windows_skips_unexpired_key_window(monkeypatch):
|
|||
prisma_client.db.litellm_verificationtoken.update.assert_not_awaited()
|
||||
|
||||
|
||||
def test_reset_budget_windows_resets_window_with_positive_utc_offset(monkeypatch):
|
||||
"""A `reset_at` written in a non-UTC timezone must be compared as an instant, not as a
|
||||
naive wall-clock value. With `timezone: Asia/Tokyo` the stored offset is +09:00, so a
|
||||
window that expired a minute ago has a wall-clock time 9 hours ahead of UTC; comparing
|
||||
it naively kept keys blocked for the length of the offset (issue #34896).
|
||||
"""
|
||||
tokyo = ZoneInfo("Asia/Tokyo")
|
||||
expired = (datetime.now(timezone.utc) - timedelta(minutes=1)).astimezone(tokyo).isoformat()
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-tokyo",
|
||||
"budget_limits": [{"budget_duration": "1mo", "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-tokyo:window:1mo", value=0.0)
|
||||
|
||||
|
||||
def test_reset_budget_windows_skips_unexpired_window_with_negative_utc_offset(monkeypatch):
|
||||
"""The mirror case: a negative-offset `reset_at` whose wall-clock time already passed in
|
||||
UTC terms is still in the future as an instant, so the window must not reset early."""
|
||||
honolulu = ZoneInfo("Pacific/Honolulu")
|
||||
future = (datetime.now(timezone.utc) + timedelta(minutes=5)).astimezone(honolulu).isoformat()
|
||||
|
||||
key_rows = [
|
||||
{
|
||||
"token": "sk-honolulu",
|
||||
"budget_limits": [{"budget_duration": "1mo", "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()
|
||||
|
||||
|
||||
def test_reset_budget_windows_resets_expired_team_window(monkeypatch):
|
||||
"""Same as the key test, but for teams."""
|
||||
now = datetime.utcnow()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue