scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

Merge pull request #41488 from BerriAI/litellm_bound_enduser_reset_invalidation

Browse source 
fix(budgets): page end-user cache invalidation after a budget reset
This commit is contained in:
ryan-crabbe-berri 2026-09-17 13:21:23 -07:00 • committed by GitHub
commit fbbddb922e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 419 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/caching/dual_cache.py
View file

@ -521,6 +521,18 @@ class DualCache(BaseCache):
if self.redis_cache is not None:
await self.redis_cache.async_delete_cache(key)
async def async_delete_cache_keys(self, keys: Sequence[str]) -> None:
"""Batch twin of ``async_delete_cache``, chunked because Redis takes the
whole list as one DELETE command."""
if not keys:
return
for key in keys:
self.in_memory_cache.delete_cache(key)
if self.redis_cache is None:
return
for start in range(0, len(keys), DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE):
await self.redis_cache.delete_cache_keys(keys[start : start + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE])
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in in-memory cache or redis

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

@ -26,7 +26,6 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import (
DB_RETRY_SAFE_ERROR_TYPES,
LiteLLM_BudgetTableFull,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_TeamTable,
LiteLLM_UserTable,
@ -193,13 +192,6 @@ def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]:
return (end_user_cache_key(row.user_id),)
def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float:
if not caps:
return 0.0
effective_budget_id: Final[str | None] = row.budget_id or litellm.max_end_user_budget_id
return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None)
def _budget_link_where(
budget_ids: Sequence[str],
extra: Mapping[str, object] = MappingProxyType({}),
@ -207,6 +199,19 @@ def _budget_link_where(
return {"budget_id": {"in": list(budget_ids)}, **extra}
def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
"""Customers whose cached spend a committed reset of these tiers invalidated.
Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which
post-commit would match nobody.
"""
linked: Final = _budget_link_where(budget_ids)
default_budget_id: Final = litellm.max_end_user_budget_id
if default_budget_id is None or default_budget_id not in budget_ids:
return linked
return {"OR": [linked, {"budget_id": None}]} # mutable-ok: prisma where filter must be a dict
def _queue_budget_linked_resets(
writes: LinkedSpendResetWrites,
cascade: "_BudgetCascade",
@ -265,16 +270,29 @@ class _BudgetCascade:
budgets: tuple[LiteLLM_BudgetTableFull, ...] = ()
budget_ids: tuple[str, ...] = ()
budget_resets: tuple[tuple[str, datetime], ...] = ()
endusers: tuple[_EndUserRow, ...] = ()
counter_resets: tuple[tuple[str, float], ...] = ()
cache_keys: tuple[str, ...] = ()
rollover_caps: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({}))
@dataclass(frozen=True, slots=True)
class _EndUserWalk:
"""Where the customer walk stands. ``cursor`` is None once it is done, and
``truncated`` says a failed page read cut it short of the tail."""
cursor: str | None = ""
invalidated: int = 0
truncated: bool = False
_ENDUSER_WALK_DONE: Final = _EndUserWalk(cursor=None)
@dataclass(frozen=True, slots=True)
class _BudgetCascadeCommitted:
cascade: _BudgetCascade
advanced: int
endusers: _EndUserWalk
@dataclass(frozen=True, slots=True)
@ -285,6 +303,8 @@ class _BudgetCascadeFailed:
_EMPTY_CASCADE: Final = _BudgetCascade()
_InvalidatedCache = Literal["spend counter", "user_api_key_cache"]
@dataclass(frozen=True, slots=True)
class _ChunkOutcome:
@ -416,10 +436,12 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
)
def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]:
def _budget_cascade_event_metadata(
cascade: _BudgetCascade, endusers: _EndUserWalk = _ENDUSER_WALK_DONE
) -> dict[str, object]:
return {
"num_budgets_found": len(cascade.budgets),
"num_endusers_found": len(cascade.endusers),
"num_endusers_found": endusers.invalidated,
}
@ -593,6 +615,38 @@ class ResetBudgetJob:
e,
)
@staticmethod
async def _invalidate_caches(counter_keys: Sequence[str], cache_keys: Sequence[str]) -> None:
"""Batch twin of ``_invalidate_spend_counter`` and
``_invalidate_user_api_key_cache_entry``, after the commit like both:
one round trip per chunk where a tier's dependents are unbounded."""
await ResetBudgetJob._invalidate_cache("spend counter", counter_keys)
await ResetBudgetJob._invalidate_cache("user_api_key_cache", cache_keys)
@staticmethod
async def _invalidate_cache(cache: _InvalidatedCache, keys: Sequence[str]) -> None:
"""One cache's share of a batch, awaited separately so either failing
still leaves the other invalidated."""
if not keys:
return
try:
from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache
match cache:
case "spend counter":
await spend_counter_cache.async_delete_cache_keys(keys)
case "user_api_key_cache":
await user_api_key_cache.async_delete_cache_keys(keys)
case _:
assert_never(cache)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to invalidate %d %s entries: %s. Budgets may be over-enforced until they expire.",
len(keys),
cache,
e,
)
async def _fetch_linked_rows(
self,
table: SpendLinkedTable[_RowT],
@ -612,18 +666,57 @@ class ResetBudgetJob:
verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e)
return ()
async def _collect_endusers_to_reset(self, budget_ids: Sequence[str]) -> tuple[_EndUserRow, ...]:
linked: Final[Sequence[_EndUserRow] | None] = await self._with_db_retry(
lambda: self.prisma_client.get_data(
table_name="enduser",
query_type="find_all",
budget_id_list=list(budget_ids),
),
reason="reset_budget_read_endusers_failure",
async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserWalk:
"""Drop the cached spend of every customer the committed tier reset zeroed.
Paged like ``_reset_windows_for``, and capless for its reason too: the
customers on one tier are unbounded, and a cap cannot keep its position
across pod elections, so it would restart at the first customer forever.
"""
if not budget_ids:
return _ENDUSER_WALK_DONE
where: Final = _enduser_invalidation_where(budget_ids)
walk = _EndUserWalk()
while walk.cursor is not None:
walk = await self._invalidate_enduser_page(where=where, cursor=walk.cursor, reached=walk.invalidated)
return walk
async def _invalidate_enduser_page(self, where: Mapping[str, object], cursor: str, reached: int) -> _EndUserWalk:
"""Invalidate one page of customers and say where the walk goes next."""
try:
rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor)
except Exception as e:
verbose_proxy_logger.warning(
"Failed to fetch end users for cache invalidation after %s customers (cursor %r): %s. "
"The customers past that page keep their cached spend until it expires.",
reached,
cursor,
e,
)
return _EndUserWalk(cursor=None, invalidated=reached, truncated=True)
if not rows:
return _EndUserWalk(cursor=None, invalidated=reached)
await self._invalidate_caches(
counter_keys=tuple(_enduser_counter_key(row) for row in rows),
cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)),
)
walked: Final = reached + len(rows)
if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE:
return _EndUserWalk(cursor=None, invalidated=walked)
return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked)
async def _fetch_enduser_page(self, where: Mapping[str, object], cursor: str) -> tuple[_EndUserRow, ...]:
"""One keyset page of customers, ordered by primary key so the cursor never repeats a row."""
return tuple(
await self._with_db_retry(
lambda: EndUserRepository(self.prisma_client).table.find_many(
where={**where, "user_id": {"gt": cursor}}, # mutable-ok: prisma where filter must be a dict
order={"user_id": "asc"}, # mutable-ok: prisma order filter must be a dict
take=RESET_BUDGET_JOB_BATCH_SIZE,
),
reason="reset_budget_read_endusers_failure",
)
)
if litellm.max_end_user_budget_id is None or litellm.max_end_user_budget_id not in budget_ids:
return tuple(linked or ())
return (*(linked or ()), *await self._get_endusers_with_no_budget_id())
async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade:
"""Resolve every row the expiring budget tiers gate, before any write.
@ -670,7 +763,6 @@ class ResetBudgetJob:
if _rollover_enabled()
else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType
)
endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids)
return _BudgetCascade(
budgets=tuple(budgets_to_reset),
budget_ids=budget_ids,
@ -682,7 +774,6 @@ class ResetBudgetJob:
for b in budgets_to_reset
if b.budget_id is not None and b.budget_duration is not None
),
endusers=endusers,
counter_resets=(
*(
(_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps))
@ -695,7 +786,6 @@ class ResetBudgetJob:
(_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
for row in model_access_groups
),
*((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers),
),
rollover_caps=rollover_caps,
cache_keys=(
@ -704,7 +794,6 @@ class ResetBudgetJob:
*(key for row in orgs for key in _org_cache_keys(row)),
*(key for row in tags for key in _tag_cache_keys(row)),
*(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
*(key for row in endusers for key in _enduser_cache_keys(row)),
),
)
@ -736,10 +825,10 @@ class ResetBudgetJob:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)
async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None:
for counter_key, _ in cascade.counter_resets:
await self._invalidate_spend_counter(counter_key)
for cache_key in cascade.cache_keys:
await self._invalidate_user_api_key_cache_entry(cache_key)
await self._invalidate_caches(
counter_keys=tuple(counter_key for counter_key, _ in cascade.counter_resets),
cache_keys=cascade.cache_keys,
)
async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed:
now: Final = datetime.now(timezone.utc)
@ -769,6 +858,7 @@ class ResetBudgetJob:
(reset_at for _, reset_at in cascade.budget_resets),
cutoff=datetime.now(timezone.utc),
),
endusers=await self._invalidate_enduser_caches(cascade.budget_ids),
)
async def reset_budget_for_litellm_budget_table(self) -> None:
@ -788,7 +878,7 @@ class ResetBudgetJob:
end_time: Final = time.time()
match outcome:
case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced):
case _BudgetCascadeCommitted() as committed:
asyncio.create_task(
self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
service=ServiceTypes.RESET_BUDGET_JOB,
@ -797,13 +887,14 @@ class ResetBudgetJob:
start_time=start_time,
end_time=end_time,
event_metadata={
**_budget_cascade_event_metadata(cascade),
"num_endusers_updated": len(cascade.endusers),
**_budget_cascade_event_metadata(committed.cascade, committed.endusers),
"num_endusers_updated": committed.endusers.invalidated,
"num_endusers_failed": 0,
"enduser_invalidation_truncated": committed.endusers.truncated,
},
)
)
return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced)
return _ChunkOutcome(fetched=len(committed.cascade.budgets), advanced=committed.advanced)
case _BudgetCascadeFailed(cascade=cascade, error=error):
verbose_proxy_logger.exception(
"Failed to reset the budget table cascade (team member, enduser, org, tag and model access "
@ -827,27 +918,6 @@ class ResetBudgetJob:
case _:
assert_never(outcome)
async def _get_endusers_with_no_budget_id(
self,
) -> list[LiteLLM_EndUserTable]:
"""
Fetch end users that have no explicit budget_id set (NULL) and have
accumulated spend > 0. These are implicitly-created end users that
rely on the default budget (litellm.max_end_user_budget_id) applied
in-memory during auth checks.
"""
table: Final = EndUserRepository(self.prisma_client).table
rows: Final = await self._with_db_retry(
lambda: table.find_many(
where={
"budget_id": None,
"spend": {"gt": 0},
},
),
reason="reset_budget_read_endusers_without_budget_id_failure",
)
return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows]
async def _write_key_reset_updates(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None:
"""
Write per-row {spend, budget_reset_at} updates for keys.

19
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import re
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload
@ -221,6 +222,24 @@ class UserApiKeyCache(DualCache):
return
await super().async_delete_cache(key)
async def async_delete_cache_keys(self, keys: Sequence[str]) -> None:
"""Batch twin of ``async_delete_cache``, partitioned like
``async_set_cache_pipeline``.
Both partitions are cleared even when one raises, because a caller
batching these has already committed the rows they cache.
"""
key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key))
other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key))
outcomes: Final = await asyncio.gather(
self.key_object_cache.async_delete_cache_keys(key_object_keys),
super().async_delete_cache_keys(other_keys),
return_exceptions=True,
)
failed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, BaseException))
if failed:
raise failed[0]
def flush_cache(self) -> None:
super().flush_cache()
self.key_object_cache.in_memory_cache.flush_cache()

22
tests/litellm_utils_tests/test_proxy_budget_reset.py
View file

@ -102,21 +102,24 @@ def _wire_batcher_for_test(prisma_client, fail_commit=False):
return batch_calls
def _wire_cascade_reads_for_test(prisma_client):
def _wire_cascade_reads_for_test(prisma_client, endusers=()):
"""
The budget tier's cascade reads the rows it is about to zero, so their
spend counters can be invalidated after the commit. Give each of those
tables an awaitable find_many so the reads resolve instead of falling into
the job's warn-and-continue path.
End users are read by the post-commit invalidation walk rather than by
``get_data``, so callers that care about customers pass them here.
"""
for table in (
"litellm_teammembership",
"litellm_verificationtoken",
"litellm_organizationtable",
"litellm_tagtable",
"litellm_endusertable",
):
getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=list(endusers))
@pytest.mark.asyncio
@ -556,7 +559,7 @@ async def test_reset_budget_continues_other_categories_on_failure():
**{u["user_id"]: u["spend"] for u in [user2]},
**{t["team_id"]: t["spend"] for t in [team1, team2]},
}
_wire_cascade_reads_for_test(prisma_client)
_wire_cascade_reads_for_test(prisma_client, endusers=[enduser1])
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@ -607,7 +610,10 @@ async def test_reset_budget_continues_other_categories_on_failure():
called_tables = {
call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list
}
assert called_tables == {"key", "user", "team", "budget", "enduser"}
assert called_tables == {"key", "user", "team", "budget"}
# Customers are not part of that set: the cascade zeroes them by budget link
# and reads them only afterwards, to invalidate their cached spend.
prisma_client.db.litellm_endusertable.find_many.assert_awaited()
# Every category writes through the batch path now, so update_data is unused.
prisma_client.update_data.assert_not_awaited()
@ -1029,7 +1035,7 @@ async def test_service_logger_endusers_success():
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
prisma_client.update_data = AsyncMock()
batch_calls = _wire_batcher_for_test(prisma_client)
_wire_cascade_reads_for_test(prisma_client)
_wire_cascade_reads_for_test(prisma_client, endusers=endusers)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@ -1094,7 +1100,7 @@ async def test_service_logger_endusers_failure():
prisma_client.get_data = AsyncMock(side_effect=fake_get_data)
prisma_client.update_data = AsyncMock()
_wire_batcher_for_test(prisma_client, fail_commit=True)
_wire_cascade_reads_for_test(prisma_client)
_wire_cascade_reads_for_test(prisma_client, endusers=endusers)
proxy_logging_obj = MagicMock()
proxy_logging_obj.service_logging_obj = MagicMock()
@ -1121,7 +1127,9 @@ async def test_service_logger_endusers_failure():
) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args
event_metadata = kwargs.get("event_metadata", {})
assert event_metadata.get("num_budgets_found") == len(budgets)
assert event_metadata.get("num_endusers_found") == len(endusers)
# Customers are read by the post-commit invalidation walk, which a failed
# commit never reaches, so a failure reports none touched.
assert event_metadata.get("num_endusers_found") == 0
assert "endusers_found" not in event_metadata
assert "budgets_found" not in event_metadata
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()

32
tests/test_litellm/caching/test_dual_cache.py
View file

@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync
@ -759,3 +760,34 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo
" (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)",
)
]
@pytest.mark.asyncio
async def test_async_delete_cache_keys_drops_memory_and_chunks_redis():
"""Batch delete clears both layers, and chunks Redis so one caller's large
key list cannot become a single oversized DELETE command."""
redis_cache = MagicMock(spec=RedisCache)
redis_cache.delete_cache_keys = AsyncMock()
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
keys = [f"key-{i}" for i in range(DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + 7)]
for key in keys:
dual_cache.in_memory_cache.set_cache(key=key, value=1)
await dual_cache.async_delete_cache_keys(keys)
assert all(dual_cache.in_memory_cache.get_cache(key=key) is None for key in keys)
sent = [call.args[0] for call in redis_cache.delete_cache_keys.await_args_list]
assert [len(chunk) for chunk in sent] == [DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, 7]
assert [key for chunk in sent for key in chunk] == keys
@pytest.mark.asyncio
async def test_async_delete_cache_keys_on_empty_list_touches_no_backend():
"""An empty page must not reach Redis: DELETE with no arguments is an error."""
redis_cache = MagicMock(spec=RedisCache)
redis_cache.delete_cache_keys = AsyncMock()
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
await dual_cache.async_delete_cache_keys([])
redis_cache.delete_cache_keys.assert_not_awaited()

174
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -4,7 +4,7 @@ import sys
import types
from datetime import datetime, timedelta, timezone
from datetime import time as dt_time
from typing import Any, Dict, Final, List
from typing import Any, Dict, Final, List, Optional
from unittest.mock import AsyncMock, MagicMock
import httpx
@ -16,6 +16,7 @@ from litellm.proxy._types import LiteLLM_VerificationToken
from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
from litellm.constants import (
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
RESET_BUDGET_JOB_BATCH_SIZE,
RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
RESET_BUDGET_JOB_NAME,
)
@ -31,13 +32,36 @@ class MockTable:
self.find_many_calls: List[Dict[str, Any]] = []
self.update_many_calls: List[Dict[str, Any]] = []
self._find_many_results: List[Any] = []
self._find_many_error: Optional[tuple[int, Exception]] = None
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
def set_find_many_error(self, after_reads: int, error: Exception):
"""Fail every read past the first ``after_reads``, the way a connection
dropping partway through a paged walk does."""
self._find_many_error = (after_reads, error)
async def find_many(
self,
where: Dict[str, Any],
order: Optional[Dict[str, str]] = None,
take: Optional[int] = None,
) -> List[Any]:
"""Replays canned rows, honouring the keyset cursor + ``take`` a paged
caller relies on: without that a paged walk never advances and the
test would hang instead of failing."""
if self._find_many_error is not None and len(self.find_many_calls) >= self._find_many_error[0]:
raise self._find_many_error[1]
paging = {k: v for k, v in (("order", order), ("take", take)) if v is not None}
self.find_many_calls.append({"where": where, **paging})
rows = list(self._find_many_results)
for field, condition in where.items():
if isinstance(condition, dict) and "gt" in condition and field != "spend":
rows = [row for row in rows if getattr(row, field, "") > condition["gt"]]
for field, direction in (order or {}).items():
rows.sort(key=lambda row: getattr(row, field, ""), reverse=direction == "desc")
return rows[:take] if take is not None else rows
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})
@ -801,10 +825,16 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock
},
]
# Verify find_many was called to fetch NULL-budget-id end users
# The post-commit invalidation walk covers both branches, so implicitly
# created customers on the default tier get their cached spend dropped too,
# and it is paged rather than reading the whole customer population.
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 1
assert find_many_calls[0]["where"] == {"budget_id": None, "spend": {"gt": 0}}
assert find_many_calls[0]["where"]["OR"] == [
{"budget_id": {"in": [default_budget_id]}},
{"budget_id": None},
]
assert find_many_calls[0]["take"] == RESET_BUDGET_JOB_BATCH_SIZE
litellm.max_end_user_budget_id = None
@ -835,9 +865,12 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured(
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
# Should NOT have queried for NULL-budget-id end users
# The invalidation walk must not reach for NULL-budget-id customers: they
# ride a default tier that is not expiring, so their spend stays put.
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 0
assert [call["where"] for call in find_many_calls] == [
{"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}}
]
litellm.max_end_user_budget_id = None
@ -872,9 +905,12 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
# Should NOT have queried for NULL-budget-id end users
# The invalidation walk must not reach for NULL-budget-id customers: they
# ride a default tier that is not expiring, so their spend stays put.
find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert len(find_many_calls) == 0
assert [call["where"] for call in find_many_calls] == [
{"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}}
]
litellm.max_end_user_budget_id = None
@ -1252,6 +1288,21 @@ def _make_counter_invalidation_job(monkeypatch):
user_api_key_cache = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
# Batch deletes fan out to the same per-key calls the real DualCache makes,
# so an assertion reads "this key was invalidated" whether the caller went
# one key at a time or a page at a time.
async def _delete_counter_keys(keys):
for key in keys:
spend_counter_cache.in_memory_cache.delete_cache(key=key)
await spend_counter_cache.redis_cache.async_delete_cache(key=key)
async def _delete_management_keys(keys):
for key in keys:
await user_api_key_cache.async_delete_cache(key=key)
spend_counter_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_counter_keys)
user_api_key_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_management_keys)
fake_module = types.ModuleType("litellm.proxy.proxy_server")
fake_module.spend_counter_cache = spend_counter_cache
fake_module.user_api_key_cache = user_api_key_cache
@ -1586,7 +1637,7 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j
"user_id": "customer-42",
},
)
mock_prisma_client.data["enduser"] = [test_enduser]
mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser])
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
@ -1596,6 +1647,107 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j
assert "end_user_id:customer-42" in deleted
def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma_client, monkeypatch):
"""The post-commit invalidation walk stays bounded in memory and in round trips.
Reading every customer on an expiring tier into one result set puts a
customer-count-sized list in the proxy's heap on every tick, which is an OOM
on a large enough deployment rather than a slow tick. Awaiting one cache call
per customer makes the last customer wait out every customer ahead of it.
Both regress silently, so pin the page size, the strictly advancing cursor,
and one batched call per page.
"""
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
population: Final = RESET_BUDGET_JOB_BATCH_SIZE * 2 + 3
mock_prisma_client.db.litellm_endusertable.set_find_many_results(
[
type("EndUser", (), {"user_id": f"cust-{i:06d}", "spend": 5.0, "budget_id": "budget-1"})
for i in range(population)
]
)
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls
assert [read["take"] for read in reads] == [RESET_BUDGET_JOB_BATCH_SIZE] * 3
assert [read["where"]["user_id"]["gt"] for read in reads] == [
"",
f"cust-{RESET_BUDGET_JOB_BATCH_SIZE - 1:06d}",
f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}",
]
assert counter_cache.async_delete_cache_keys.await_count == 3
assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 3
counter_cache.async_delete_cache.assert_not_called()
invalidated: Final = {
key for call in counter_cache.async_delete_cache_keys.await_args_list for key in call.args[0]
}
assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)}
evicted: Final = {
key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0]
}
assert evicted == {f"end_user_id:cust-{i:06d}" for i in range(population)}
def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_finish(
mock_prisma_client, monkeypatch
):
"""A page that fails to read is not the end of the customer list.
The tier's window is already advanced by the time this walk runs, so no later
tick comes back for the customers past the page that failed: their cached
spend goes on rejecting requests until it expires. Returning the same empty
page normal end-of-data returns hid that behind a report of a clean pass.
"""
_make_counter_invalidation_job(monkeypatch)
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
endusers: Final = mock_prisma_client.db.litellm_endusertable
endusers.set_find_many_results(
[
type("EndUser", (), {"user_id": f"cust-{i:06d}", "spend": 5.0, "budget_id": "budget-1"})
for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3)
]
)
endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers"))
logging_obj: Final = RecordingProxyLogging()
job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client)
_run_and_drain_hooks(job.reset_budget_for_litellm_budget_table)
metadata: Final = logging_obj.service_logging_obj.success_calls[0]["event_metadata"]
assert metadata["enduser_invalidation_truncated"] is True
assert metadata["num_endusers_updated"] == RESET_BUDGET_JOB_BATCH_SIZE
def test_a_failed_counter_batch_still_evicts_the_management_cache(
reset_budget_job, mock_prisma_client, monkeypatch
):
"""The spend counters and the management cache are invalidated independently.
Sharing one handler meant a Redis failure on the counters returned before the
management cache was touched at all. The commit has already zeroed those rows
by then, so the cached objects keep authorizing against their pre-reset spend
until they expire.
"""
counter_cache: Final = _make_counter_invalidation_job(monkeypatch)
counter_cache.async_delete_cache_keys = AsyncMock(side_effect=RuntimeError("redis unavailable"))
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")]
mock_prisma_client.db.litellm_endusertable.set_find_many_results(
[type("EndUser", (), {"user_id": "customer-42", "spend": 5.0, "budget_id": "budget-1"})]
)
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
evicted: Final = {
key
for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list
for key in call.args[0]
}
assert "end_user_id:customer-42" in evicted
def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch):
"""Eviction runs after the commit, so a broken cache cannot undo the write."""

53
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py
View file

@ -82,6 +82,19 @@ class FakeRedisCache(RedisCache):
async def async_delete_cache(self, key: str): # type: ignore[override]
self._store.pop(key, None)
async def delete_cache_keys(self, keys): # type: ignore[override]
for key in keys:
self._store.pop(key, None)
class PartitionFailingRedisCache(FakeRedisCache):
"""Fails the batch delete for the key-object partition and no other."""
async def delete_cache_keys(self, keys): # type: ignore[override]
if any(is_user_key_cache_key(key) for key in keys):
raise ConnectionError("redis unavailable")
await super().delete_cache_keys(keys)
def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth:
# Minimal object (UserAPIKeyAuth inherits token from base view).
@ -331,6 +344,46 @@ class TestUserKeyObjectPartition:
assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None
assert await redis.async_get_cache(HASHED_TOKEN) is None
@pytest.mark.asyncio
async def test_batch_delete_routes_each_key_to_its_partition(self):
"""A batch delete has to clear the same partition the single delete does.
``DualCache``'s batch delete only knows about the main in-memory cache, so
inheriting it unchanged leaves a key object sitting in ``key_object_cache``
with its pre-reset spend, and the next request is authorized against that
stale copy until the local entry expires.
"""
redis = FakeRedisCache()
cache = UserApiKeyCache(redis_cache=redis)
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"})
await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")])
assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None
assert await cache.async_get_cache(end_user_cache_key("u1")) is None
assert await redis.async_get_cache(HASHED_TOKEN) is None
assert await redis.async_get_cache(end_user_cache_key("u1")) is None
@pytest.mark.asyncio
async def test_batch_delete_clears_the_other_partition_when_one_fails(self):
"""One partition failing must not cost the other its deletions.
A caller batching these has already committed the rows they cache, so a
partition that is skipped keeps authorizing against pre-reset spend until
the entry expires. The failure is still raised for the caller to report.
"""
redis = PartitionFailingRedisCache()
cache = UserApiKeyCache(redis_cache=redis)
await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth)
await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"})
with pytest.raises(ConnectionError):
await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")])
assert await cache.async_get_cache(end_user_cache_key("u1")) is None
assert await redis.async_get_cache(end_user_cache_key("u1")) is None
@pytest.mark.asyncio
async def test_pipeline_write_routes_each_entry_to_its_partition(self):
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2))