scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(reset_budget_job): use spend_counter_cache from proxy_server in reset_budget_windows

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-11 11:38:47 -07:00
parent 9de2ae6544
commit 3df367bdb8
No known key found for this signature in database
1 changed files with 12 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -565,9 +565,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")).replace(
tzinfo=None
)
if reset_at > now:
return False
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0)
@ -591,6 +591,8 @@ class ResetBudgetJob:
reset_at <= now. Only the expired windows are reset; other windows are untouched.
"""
from litellm.proxy.proxy_server import spend_counter_cache
now = datetime.utcnow()
# --- Keys ---
@ -605,9 +607,11 @@ class ResetBudgetJob:
windows: list = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = f"spend:key:{key.token}:window:{window['budget_duration']}"
counter_key = (
f"spend:key:{key.token}:window:{window['budget_duration']}"
)
if await ResetBudgetJob._reset_expired_window(
window, counter_key, ProxyLogging.spend_counter_cache, now
window, counter_key, spend_counter_cache, now
):
changed = True
if changed:
@ -632,7 +636,9 @@ class ResetBudgetJob:
windows = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = f"spend:team:{team.team_id}:window:{window['budget_duration']}"
counter_key = (
f"spend:team:{team.team_id}:window:{window['budget_duration']}"
)
if await ResetBudgetJob._reset_expired_window(
window, counter_key, spend_counter_cache, now
):