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

fix(proxy): requeue the daily tag rollup on commit failure without the Redis buffer

Browse source
This commit is contained in:
mateo-berri 2026-09-18 14:58:34 -07:00
parent 72a14246a6
commit 3edbf60e9c
2 changed files with 37 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/proxy/db/db_spend_update_writer.py
View file

@ -1305,7 +1305,7 @@ class DBSpendUpdateWriter:
async def _flush_daily_spend_queue(
self,
queue: DailySpendUpdateQueue,
entity_type: Literal["user", "team", "org", "end_user", "agent"],
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
commit: _DailySpendCommit[_DailySpendTransactionT],
n_retry_times: int,
prisma_client: PrismaClient,
@ -1447,19 +1447,15 @@ class DBSpendUpdateWriter:
Commit only tag spend updates to database.
This is called by a separate scheduler job at a longer interval.
"""
daily_tag_spend_update_transactions: Final = cast(
dict[str, DailyTagSpendTransaction],
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
await self._flush_daily_spend_queue(
queue=self.daily_tag_spend_update_queue,
entity_type="tag",
commit=DBSpendUpdateWriter.update_daily_tag_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
if daily_tag_spend_update_transactions:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
async def _commit_daily_tag_spend_to_db_with_redis(
self,
prisma_client: PrismaClient,

29
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -2864,6 +2864,35 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other
assert db_writer.daily_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_daily_tag_spend_commit_requeues_the_rows():
"""The tag rollup drains on its own scheduler job with the same no-Redis drop:
a failed LiteLLM_DailyTagSpend commit has to put the rows back for the next tick."""
db_writer = DBSpendUpdateWriter()
tag_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"tag": "tag-1"}
await db_writer.daily_tag_spend_update_queue.add_update({"tag-key": tag_txn})
db = _DailySpendFakeDB(failing_table="LiteLLM_DailyTagSpend")
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
await db_writer._commit_daily_tag_spend_to_db(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == []
assert not db_writer.daily_tag_spend_update_queue.update_queue.empty()
db.failing_table = None
await db_writer._commit_daily_tag_spend_to_db(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
(tag_upsert,) = _daily_upserts(db, "LiteLLM_DailyTagSpend")
assert _row_values(tag_upsert, "tag") == ["tag-1"]
assert _row_values(tag_upsert, "spend") == [0.1]
assert db_writer.daily_tag_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
"""The Redis drain is destructive, so a failed window commit has to push