ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Reorder counter invalidation to run after DB write

Browse Source
This commit is contained in:
Michael Riad Zaky 2026-04-29 18:32:48 -07:00
parent ff2a938847
commit 4e26835098
2 changed files with 52 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
litellm/proxy/common_utils/reset_budget_job.py
View File

@ -54,15 +54,22 @@ class ResetBudgetJob:
@staticmethod
async def _invalidate_spend_counter(counter_key: str) -> None:
"""Zero a spend counter so a DB-row reset takes effect immediately."""
"""Zero a spend counter so a DB-row reset takes effect immediately.
Call AFTER the DB write commits. Clearing Redis before the DB
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key, value=0.0, ttl=60
)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key, value=0.0
key=counter_key, value=0.0, ttl=60
)
except Exception as redis_err:
verbose_proxy_logger.warning(
@ -92,22 +99,26 @@ class ResetBudgetJob:
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to fetch team memberships for counter invalidation: %s", e
)
return await self.prisma_client.db.litellm_teammembership.update_many(
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
)
for m in memberships:
await self._invalidate_spend_counter(
f"spend:team_member:{m.user_id}:{m.team_id}"
)
return update_result
async def reset_budget_for_keys_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
@ -140,20 +151,26 @@ class ResetBudgetJob:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
return await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
return update_result
async def reset_budget_for_litellm_budget_table(self):
"""
Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired
@ -377,15 +394,15 @@ class ResetBudgetJob:
)
if updated_keys:
for k in updated_keys:
token = getattr(k, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
await self.prisma_client.update_data(
query_type="update_many",
data_list=updated_keys,
table_name="key",
)
for k in updated_keys:
token = getattr(k, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
end_time = time.time()
if len(failed_keys) > 0: # If any keys failed to reset
@ -466,17 +483,17 @@ class ResetBudgetJob:
"Updated users %s", json.dumps(updated_users, indent=4, default=str)
)
if updated_users:
await self.prisma_client.update_data(
query_type="update_many",
data_list=updated_users,
table_name="user",
)
for u in updated_users:
user_id = getattr(u, "user_id", None)
if user_id:
await self._invalidate_spend_counter(
f"spend:user:{user_id}"
)
await self.prisma_client.update_data(
query_type="update_many",
data_list=updated_users,
table_name="user",
)
end_time = time.time()
if len(failed_users) > 0: # If any users failed to reset
@ -563,17 +580,17 @@ class ResetBudgetJob:
"Updated teams %s", json.dumps(updated_teams, indent=4, default=str)
)
if updated_teams:
await self.prisma_client.update_data(
query_type="update_many",
data_list=updated_teams,
table_name="team",
)
for t in updated_teams:
team_id = getattr(t, "team_id", None)
if team_id:
await self._invalidate_spend_counter(
f"spend:team:{team_id}"
)
await self.prisma_client.update_data(
query_type="update_many",
data_list=updated_teams,
table_name="team",
)
end_time = time.time()
if len(failed_teams) > 0: # If any teams failed to reset

12
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View File

@ -1093,10 +1093,10 @@ def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch):
asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team_member:alice:team-x", value=0.0
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:team_member:alice:team-x", value=0.0
key="spend:team_member:alice:team-x", value=0.0, ttl=60
)
@ -1124,7 +1124,7 @@ def test_reset_budget_for_keys_invalidates_redis_counter(
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-abc", value=0.0
key="spend:key:sk-abc", value=0.0, ttl=60
)
@ -1152,7 +1152,7 @@ def test_reset_budget_for_users_invalidates_redis_counter(
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:user:alice", value=0.0
key="spend:user:alice", value=0.0, ttl=60
)
@ -1180,7 +1180,7 @@ def test_reset_budget_for_teams_invalidates_redis_counter(
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:team:team-x", value=0.0
key="spend:team:team-x", value=0.0, ttl=60
)
@ -1203,5 +1203,5 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monke
asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-linked", value=0.0
key="spend:key:sk-linked", value=0.0, ttl=60
)