scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #22029 from BerriAI/litellm_spend_tracking_logging

Browse source 
[Infra] Add Spend Tracking Lifecycle Logging
This commit is contained in:
yuneng-jiang 2026-02-24 12:52:32 -08:00 • committed by GitHub
commit 80ebe722d9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 166 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -224,8 +224,17 @@ class DBSpendUpdateWriter:
verbose_proxy_logger.debug("Runs spend update on all tables")
except Exception:
verbose_proxy_logger.debug(
f"Error updating Prisma database: {traceback.format_exc()}"
verbose_proxy_logger.error(
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
"may not have completed for this request. "
"response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, end_user_id=%s - %s",
response_cost,
token,
user_id,
team_id,
org_id,
end_user_id,
traceback.format_exc(),
)
async def _update_key_db(
@ -295,9 +304,14 @@ class DBSpendUpdateWriter:
)
)
except Exception as e:
verbose_proxy_logger.debug(
"\033[91m"
+ f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}"
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue user spend update. "
"user_id=%s, end_user_id=%s, response_cost=%s - %s\n%s",
user_id,
end_user_id,
response_cost,
str(e),
traceback.format_exc(),
)
async def _update_team_db(
@ -334,11 +348,24 @@ class DBSpendUpdateWriter:
response_cost=response_cost,
)
)
except Exception:
pass
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue team member spend update. "
"team_id=%s, user_id=%s, response_cost=%s - %s\n%s",
team_id,
user_id,
response_cost,
str(e),
traceback.format_exc(),
)
except Exception as e:
verbose_proxy_logger.debug(
f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}"
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue team spend update. "
"team_id=%s, response_cost=%s - %s\n%s",
team_id,
response_cost,
str(e),
traceback.format_exc(),
)
raise e
@ -363,8 +390,13 @@ class DBSpendUpdateWriter:
)
)
except Exception as e:
verbose_proxy_logger.debug(
f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}"
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue org spend update. "
"org_id=%s, response_cost=%s - %s\n%s",
org_id,
response_cost,
str(e),
traceback.format_exc(),
)
raise e
@ -411,8 +443,13 @@ class DBSpendUpdateWriter:
)
)
except Exception as e:
verbose_proxy_logger.debug(
f"Update Tag DB failed to execute - {str(e)}\n{traceback.format_exc()}"
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue tag spend update. "
"request_tags=%s, response_cost=%s - %s\n%s",
request_tags,
response_cost,
str(e),
traceback.format_exc(),
)
raise e
@ -513,6 +550,17 @@ class DBSpendUpdateWriter:
await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer()
)
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
len(db_spend_update_transactions.get("org_list_transactions") or {}),
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
)
await self._commit_spend_updates_to_db(
prisma_client=prisma_client,
n_retry_times=n_retry_times,
@ -583,7 +631,12 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_agent_spend_update_transactions,
)
except Exception as e:
verbose_proxy_logger.error(f"Error committing spend updates: {e}")
verbose_proxy_logger.error(
"Spend tracking - failed to commit spend updates from Redis to DB. "
"Data already popped from Redis may be lost. Error: %s\n%s",
str(e),
traceback.format_exc(),
)
finally:
await self.pod_lock_manager.release_lock(
cronjob_id=DB_SPEND_UPDATE_JOB_NAME,

5
litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py
View file

@ -86,6 +86,11 @@ class DailySpendUpdateQueue(BaseUpdateQueue):
) -> Dict[str, BaseDailySpendTransaction]:
"""Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates."""
updates = await self.flush_all_updates_from_in_memory_queue()
if len(updates) > 0:
verbose_proxy_logger.info(
"Spend tracking - flushed %d daily spend update items from in-memory queue",
len(updates),
)
aggregated_daily_spend_update_transactions = (
DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
updates

14
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

@ -80,6 +80,14 @@ class PodLockManager:
)
self._emit_acquired_lock_event(cronjob_id, self.pod_id)
return True
else:
verbose_proxy_logger.info(
"Spend tracking - pod %s could not acquire lock for cronjob_id=%s, "
"held by pod %s. Spend updates in Redis will wait for the leader pod to commit.",
self.pod_id,
cronjob_id,
current_value,
)
return False
except Exception as e:
verbose_proxy_logger.error(
@ -124,10 +132,12 @@ class PodLockManager:
pod_id=self.pod_id,
)
else:
verbose_proxy_logger.debug(
"Pod %s failed to release Redis lock for cronjob_id=%s",
verbose_proxy_logger.warning(
"Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. "
"Lock will expire after TTL=%ds.",
self.pod_id,
cronjob_id,
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS,
)
else:
verbose_proxy_logger.debug(

37
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -96,14 +96,28 @@ class RedisUpdateBuffer:
list_of_transactions = [safe_dumps(transactions)]
if self.redis_cache is None:
return
current_redis_buffer_size = await self.redis_cache.async_rpush(
key=redis_key,
values=list_of_transactions,
)
await self._emit_new_item_added_to_redis_buffer_event(
queue_size=current_redis_buffer_size,
service=service_type,
)
try:
current_redis_buffer_size = await self.redis_cache.async_rpush(
key=redis_key,
values=list_of_transactions,
)
verbose_proxy_logger.debug(
"Spend tracking - pushed spend updates to Redis buffer. "
"redis_key=%s, buffer_size=%s",
redis_key,
current_redis_buffer_size,
)
await self._emit_new_item_added_to_redis_buffer_event(
queue_size=current_redis_buffer_size,
service=service_type,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). "
"Error: %s",
redis_key,
str(e),
)
async def store_in_memory_spend_updates_in_redis(
self,
@ -305,6 +319,13 @@ class RedisUpdateBuffer:
if list_of_transactions is None:
return None
verbose_proxy_logger.info(
"Spend tracking - popped %d spend update batches from Redis buffer (key=%s). "
"These items are now removed from Redis and must be committed to DB.",
len(list_of_transactions) if isinstance(list_of_transactions, list) else 1,
REDIS_UPDATE_BUFFER_KEY,
)
# Parse the list of transactions from JSON strings
parsed_transactions = self._parse_list_of_transactions(list_of_transactions)

5
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -31,6 +31,11 @@ class SpendUpdateQueue(BaseUpdateQueue):
) -> DBSpendUpdateTransactions:
"""Flush all updates from the queue and return all updates aggregated by entity type."""
updates = await self.flush_all_updates_from_in_memory_queue()
if len(updates) > 0:
verbose_proxy_logger.info(
"Spend tracking - flushed %d spend update items from in-memory queue",
len(updates),
)
verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates)
return self.get_aggregated_db_spend_update_transactions(updates)

40
litellm/proxy/proxy_server.py
View file

@ -1768,8 +1768,14 @@ async def update_cache( # noqa: PLR0915
("{}:spend".format(litellm_proxy_admin_name), increment)
)
except Exception as e:
verbose_proxy_logger.debug(
f"An error occurred updating user cache: {str(e)}\n\n{traceback.format_exc()}"
verbose_proxy_logger.warning(
"Spend tracking - failed to update user spend in cache. "
"Budget enforcement may use stale spend values. "
"user_id=%s, response_cost=%s - %s\n%s",
user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE END-USER SPEND ###
@ -1806,8 +1812,14 @@ async def update_cache( # noqa: PLR0915
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj.json()))
except Exception as e:
verbose_proxy_logger.exception(
f"An error occurred updating end user cache: {str(e)}"
verbose_proxy_logger.warning(
"Spend tracking - failed to update end user spend in cache. "
"Budget enforcement may use stale spend values. "
"end_user_id=%s, response_cost=%s - %s\n%s",
end_user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TEAM SPEND ###
@ -1848,8 +1860,14 @@ async def update_cache( # noqa: PLR0915
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append((_id, existing_spend_obj))
except Exception as e:
verbose_proxy_logger.exception(
f"An error occurred updating end user cache: {str(e)}"
verbose_proxy_logger.warning(
"Spend tracking - failed to update team spend in cache. "
"Budget enforcement may use stale spend values. "
"team_id=%s, response_cost=%s - %s\n%s",
team_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TAG SPEND ###
@ -1894,8 +1912,14 @@ async def update_cache( # noqa: PLR0915
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append((cache_key, existing_tag_obj))
except Exception as e:
verbose_proxy_logger.exception(
f"An error occurred updating tag cache: {str(e)}"
verbose_proxy_logger.warning(
"Spend tracking - failed to update tag spend in cache. "
"Budget enforcement may use stale spend values. "
"tags=%s, response_cost=%s - %s\n%s",
tags,
response_cost,
str(e),
traceback.format_exc(),
)
if token is not None and response_cost is not None:

19
litellm/proxy/utils.py
View file

@ -4457,6 +4457,11 @@ class ProxyUpdateSpend:
len(logs_to_process) :
]
popped_batch = True
if len(logs_to_process) > 0:
verbose_proxy_logger.info(
"Spend tracking - processing %d spend logs for DB write",
len(logs_to_process),
)
start_time = time.time()
try:
for i in range(n_retry_times + 1):
@ -4503,9 +4508,17 @@ class ProxyUpdateSpend:
f"{len(logs_to_process)} logs processed. Remaining in queue: {remaining_count}"
)
break
except DB_CONNECTION_ERROR_TYPES:
except DB_CONNECTION_ERROR_TYPES as e:
if i is None:
i = 0
verbose_proxy_logger.warning(
"Spend tracking - DB connection error writing spend logs, "
"retry %d/%d. logs_count=%d, error=%s",
i + 1,
n_retry_times,
len(logs_to_process),
str(e),
)
if i >= n_retry_times:
raise
await asyncio.sleep(2**i)
@ -4620,8 +4633,8 @@ async def update_spend_logs_job(
logs_to_process=logs_to_process,
)
except Exception as guardrail_tracking_err:
verbose_proxy_logger.debug(
"Guardrail usage tracking failed (non-fatal): %s",
verbose_proxy_logger.warning(
"Spend tracking - guardrail usage tracking failed (non-fatal): %s",
guardrail_tracking_err,
)