scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

refactor(proxy): extract redis tag spend drain and commit into a helper

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-13 03:02:54 +00:00
parent f3c64029e1
commit d3fae8a260
1 changed files with 35 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1107,20 +1107,12 @@ class DBSpendUpdateWriter:
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
):
verbose_proxy_logger.debug("acquired lock for daily tag spend updates")
daily_tag_spend_update_transactions: dict[str, DailyTagSpendTransaction] | None = None
committed = False
try:
daily_tag_spend_update_transactions: Final = (
await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
await self._drain_and_commit_daily_tag_spend_from_redis(
prisma_client=prisma_client,
n_retry_times=n_retry_times,
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,
)
committed = True
except Exception as e:
spend_log_error(
"Spend tracking - failed to commit daily tag spend updates from Redis to DB. "
@ -1129,14 +1121,41 @@ class DBSpendUpdateWriter:
exc=e,
)
finally:
if not committed and daily_tag_spend_update_transactions:
await self.redis_update_buffer.restore_transactions_to_redis(
daily_tag_spend_update_transactions=daily_tag_spend_update_transactions,
)
await self.pod_lock_manager.release_lock(
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
)
async def _drain_and_commit_daily_tag_spend_from_redis(
self,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
) -> None:
"""
Drain the Redis tag spend buffer and commit it, restoring the drained transactions if the commit fails.
The drain is destructive, so a failed commit must push the transactions back for the next tick
or their spend is lost permanently.
"""
daily_tag_spend_update_transactions: Final = (
await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
)
if not daily_tag_spend_update_transactions:
return
try:
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,
)
except Exception:
await self.redis_update_buffer.restore_transactions_to_redis(
daily_tag_spend_update_transactions=daily_tag_spend_update_transactions,
)
raise
async def _flush_tool_discovery_queue(
self,
prisma_client: PrismaClient,