tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 23:07:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: reset org and tag budgets (#27326)

Browse Source 
* reset org budgets

* reset tag budgets

---------

Co-authored-by: Michael Riad Zaky <michaelr@Mac.localdomain>
This commit is contained in:
Michael-RZ-Berri
2026-05-09 16:15:32 -07:00
committed by GitHub
parent b888177ea6
commit 02edaef50c
3 changed files with 363 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/common_utils/reset_budget_job.py
+82 -70
View File
@@ -2,7 +2,7 @@ import asyncio
import json
import time
from datetime import datetime, timezone
from typing import Any, List, Literal, Optional, Union
from typing import Any, Callable, List, Literal, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
@@ -83,93 +83,97 @@ class ResetBudgetJob:
"Failed to reset spend counter %s: %s", counter_key, e
)
async def _cascade_reset_spend_for_budget_link(
self,
budgets_to_reset: List[LiteLLM_BudgetTableFull],
table: Any,
counter_key_fn: Callable[[Any], str],
log_subject: str,
extra_where: Optional[dict] = None,
):
"""
Generic cascade: zero spend on rows whose budget_id is in the reset set.
"""
budget_ids = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
if not budget_ids:
return
where: dict = {"budget_id": {"in": budget_ids}}
if extra_where:
where.update(extra_where)
try:
rows = await table.find_many(where=where)
except Exception as e:
rows = []
verbose_proxy_logger.warning(
"Failed to fetch %s for counter invalidation: %s", log_subject, e
)
update_result = await table.update_many(where=where, data={"spend": 0})
for row in rows:
await self._invalidate_spend_counter(counter_key_fn(row))
return update_result
async def reset_budget_for_litellm_team_members(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the budget for all LiteLLM Team Members if their budget has expired
"""
budget_ids = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
try:
memberships = await self.prisma_client.db.litellm_teammembership.find_many(
where={"budget_id": {"in": budget_ids}}
)
except Exception as e:
memberships = []
verbose_proxy_logger.warning(
"Failed to fetch team memberships for counter invalidation: %s", e
)
update_result = await self.prisma_client.db.litellm_teammembership.update_many(
where={"budget_id": {"in": budget_ids}},
data={
"spend": 0,
},
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_teammembership,
counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}",
log_subject="team memberships",
)
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]
):
"""
Resets the spend for keys linked to budget tiers that are being reset.
This handles keys that have budget_id but no budget_duration set on the key
itself. Keys with budget_id rely on their linked budget tier's reset schedule
rather than having their own budget_duration.
Keys that have their own budget_duration are already handled by
reset_budget_for_litellm_keys() and are excluded here to avoid
double-resetting.
Excludes keys with their own budget_duration; those are reset by
reset_budget_for_litellm_keys() to avoid double-resetting.
"""
budget_ids = [
budget.budget_id
for budget in budgets_to_reset
if budget.budget_id is not None
]
if not budget_ids:
return
where_clause: dict = {
"budget_id": {"in": budget_ids},
"budget_duration": None, # only keys without their own reset schedule
"spend": {"gt": 0}, # only reset keys that have accumulated spend
}
try:
keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
where=where_clause
)
except Exception as e:
keys = []
verbose_proxy_logger.warning(
"Failed to fetch keys for counter invalidation: %s", e
)
update_result = (
await self.prisma_client.db.litellm_verificationtoken.update_many(
where=where_clause,
data={
"spend": 0,
},
)
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_verificationtoken,
counter_key_fn=lambda k: f"spend:key:{k.token}",
log_subject="keys",
extra_where={"budget_duration": None, "spend": {"gt": 0}},
)
for k in keys:
await self._invalidate_spend_counter(f"spend:key:{k.token}")
async def reset_budget_for_orgs_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the spend for orgs linked to budget tiers that are being reset.
"""
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_organizationtable,
counter_key_fn=lambda o: f"spend:org:{o.organization_id}",
log_subject="orgs",
extra_where={"spend": {"gt": 0}},
)
return update_result
async def reset_budget_for_tags_linked_to_budgets(
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
"""
Resets the spend for tags linked to budget tiers that are being reset.
"""
return await self._cascade_reset_spend_for_budget_link(
budgets_to_reset=budgets_to_reset,
table=self.prisma_client.db.litellm_tagtable,
counter_key_fn=lambda t: f"spend:tag:{t.tag_name}",
log_subject="tags",
extra_where={"spend": {"gt": 0}},
)
async def reset_budget_for_litellm_budget_table(self):
"""
@@ -237,6 +241,14 @@ class ResetBudgetJob:
budgets_to_reset=budgets_to_reset
)
await self.reset_budget_for_orgs_linked_to_budgets(
budgets_to_reset=budgets_to_reset
)
await self.reset_budget_for_tags_linked_to_budgets(
budgets_to_reset=budgets_to_reset
)
if endusers_to_reset is not None and len(endusers_to_reset) > 0:
for enduser in endusers_to_reset:
try:
tests/litellm_utils_tests/test_proxy_budget_reset.py
+28
View File
@@ -233,6 +233,12 @@ async def test_reset_budget_endusers_partial_failure():
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -400,6 +406,12 @@ async def test_reset_budget_continues_other_categories_on_failure():
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -884,6 +896,12 @@ async def test_service_logger_endusers_success():
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -966,6 +984,12 @@ async def test_service_logger_endusers_failure():
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 0}
)
# Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@@ -1060,6 +1084,10 @@ async def test_reset_budget_for_litellm_team_members_called():
prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(
return_value={"count": 0}
)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 0}
)
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0})
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
+253
View File
@@ -39,6 +39,46 @@ class MockLiteLLMVerificationToken:
return {"count": 1}
class MockLiteLLMOrganizationTable:
def __init__(self):
self.update_many_calls: List[Dict[str, Any]] = []
self.find_many_calls: List[Dict[str, Any]] = []
self._find_many_results: List[Any] = []
def set_find_many_results(self, results: List[Any]):
self._find_many_results = results
async def find_many(self, where: Dict[str, Any]) -> List[Any]:
self.find_many_calls.append({"where": where})
return self._find_many_results
async def update_many(
self, where: Dict[str, Any], data: Dict[str, Any]
) -> Dict[str, Any]:
self.update_many_calls.append({"where": where, "data": data})
return {"count": 1}
class MockLiteLLMTagTable:
def __init__(self):
self.update_many_calls: List[Dict[str, Any]] = []
self.find_many_calls: List[Dict[str, Any]] = []
self._find_many_results: List[Any] = []
def set_find_many_results(self, results: List[Any]):
self._find_many_results = results
async def find_many(self, where: Dict[str, Any]) -> List[Any]:
self.find_many_calls.append({"where": where})
return self._find_many_results
async def update_many(
self, where: Dict[str, Any], data: Dict[str, Any]
) -> Dict[str, Any]:
self.update_many_calls.append({"where": where, "data": data})
return {"count": 1}
class MockLiteLLMEndUserTable:
def __init__(self):
self.find_many_calls: List[Dict[str, Any]] = []
@@ -57,6 +97,8 @@ class MockDB:
self.litellm_teammembership = MockLiteLLMTeamMembership()
self.litellm_verificationtoken = MockLiteLLMVerificationToken()
self.litellm_endusertable = MockLiteLLMEndUserTable()
self.litellm_organizationtable = MockLiteLLMOrganizationTable()
self.litellm_tagtable = MockLiteLLMTagTable()
class MockPrismaClient:
@@ -459,6 +501,100 @@ def test_reset_budget_for_keys_linked_to_budgets_empty(
assert len(calls) == 0
def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_client):
"""
Test that when a budget tier is reset, orgs linked to that budget
(via budget_id) also get their spend reset.
"""
now = datetime.now(timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 100.0,
"budget_duration": "30d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "30d-org-budget",
"created_at": now - timedelta(days=30),
},
)
asyncio.run(
reset_budget_job.reset_budget_for_orgs_linked_to_budgets(
budgets_to_reset=[test_budget]
)
)
calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
assert len(calls) == 1
call = calls[0]
assert call["where"]["budget_id"] == {"in": ["30d-org-budget"]}
assert call["where"]["spend"] == {"gt": 0}
assert call["data"]["spend"] == 0
def test_reset_budget_for_orgs_linked_to_budgets_empty(
reset_budget_job, mock_prisma_client
):
"""
Test that when there are no budgets to reset, no update is performed
on the organization table.
"""
asyncio.run(
reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[])
)
calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
assert len(calls) == 0
def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_client):
"""
Test that when a budget tier is reset, tags linked to that budget
(via budget_id) also get their spend reset.
"""
now = datetime.now(timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 50.0,
"budget_duration": "30d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "30d-tag-budget",
"created_at": now - timedelta(days=30),
},
)
asyncio.run(
reset_budget_job.reset_budget_for_tags_linked_to_budgets(
budgets_to_reset=[test_budget]
)
)
calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
assert len(calls) == 1
call = calls[0]
assert call["where"]["budget_id"] == {"in": ["30d-tag-budget"]}
assert call["where"]["spend"] == {"gt": 0}
assert call["data"]["spend"] == 0
def test_reset_budget_for_tags_linked_to_budgets_empty(
reset_budget_job, mock_prisma_client
):
"""
Test that when there are no budgets to reset, no update is performed
on the tag table.
"""
asyncio.run(
reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[])
)
calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
assert len(calls) == 0
@pytest.mark.parametrize(
"budget_duration, expected_day, expected_month",
[
@@ -618,6 +754,75 @@ def test_budget_table_reset_also_resets_linked_keys(
assert calls[0]["data"]["spend"] == 0
def test_budget_table_reset_also_resets_linked_orgs(
reset_budget_job, mock_prisma_client
):
"""
Integration-style test: when reset_budget_for_litellm_budget_table runs,
it should also reset spend for orgs linked to the expiring budget tiers
(in addition to end-users, team members, and keys).
"""
now = datetime.now(timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 100.0,
"budget_duration": "30d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "30d-org-budget",
"created_at": now - timedelta(days=30),
},
)
mock_prisma_client.data["budget"] = [test_budget]
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls
assert len(calls) == 1, (
"Expected reset_budget_for_litellm_budget_table to also reset orgs "
f"linked to expiring budgets, but got {len(calls)} update_many calls"
)
assert calls[0]["where"]["budget_id"] == {"in": ["30d-org-budget"]}
assert calls[0]["data"]["spend"] == 0
def test_budget_table_reset_also_resets_linked_tags(
reset_budget_job, mock_prisma_client
):
"""
Integration-style test: when reset_budget_for_litellm_budget_table runs,
it should also reset spend for tags linked to the expiring budget tiers.
"""
now = datetime.now(timezone.utc)
test_budget = type(
"LiteLLM_BudgetTableFull",
(),
{
"max_budget": 50.0,
"budget_duration": "30d",
"budget_reset_at": now - timedelta(hours=1),
"budget_id": "30d-tag-budget",
"created_at": now - timedelta(days=30),
},
)
mock_prisma_client.data["budget"] = [test_budget]
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
calls = mock_prisma_client.db.litellm_tagtable.update_many_calls
assert len(calls) == 1, (
"Expected reset_budget_for_litellm_budget_table to also reset tags "
f"linked to expiring budgets, but got {len(calls)} update_many calls"
)
assert calls[0]["where"]["budget_id"] == {"in": ["30d-tag-budget"]}
assert calls[0]["data"]["spend"] == 0
def test_reset_budget_resets_endusers_with_null_budget_id(
reset_budget_job, mock_prisma_client
):
@@ -1205,3 +1410,51 @@ def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monke
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:key:sk-linked", value=0.0, ttl=60
)
def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monkeypatch):
"""Resetting orgs via budget tier must clear each linked org's counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_org = type("Org", (), {"organization_id": "org-acme"})
prisma_client = MagicMock()
prisma_client.db.litellm_organizationtable.find_many = AsyncMock(
return_value=[linked_org]
)
prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
return_value={"count": 1}
)
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:org:org-acme", value=0.0, ttl=60
)
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:org:org-acme", value=0.0, ttl=60
)
def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monkeypatch):
"""Resetting tags via budget tier must clear each linked tag's counter."""
counter_cache = _make_counter_invalidation_job(monkeypatch)
expired_budget = type("B", (), {"budget_id": "budget-1"})
linked_tag = type("Tag", (), {"tag_name": "tenant-42"})
prisma_client = MagicMock()
prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag])
prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1})
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget]))
counter_cache.in_memory_cache.set_cache.assert_any_call(
key="spend:tag:tenant-42", value=0.0, ttl=60
)
counter_cache.redis_cache.async_set_cache.assert_any_await(
key="spend:tag:tenant-42", value=0.0, ttl=60
)