scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

Merge pull request #9690 from BerriAI/litellm_fixes_for_deadlock_round_2

Browse source 
[Reliability Fixes] - Ensure no deadlocks occur when updating `DailyUserSpendTransaction`
This commit is contained in:
Ishaan Jaff 2025-04-01 20:59:41 -07:00 • committed by GitHub
commit ac3399238e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 717 additions and 269 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -19,6 +19,7 @@ DEFAULT_IMAGE_HEIGHT = 300
MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = 1024 # 1MB = 1024KB
SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = 1000 # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic.
REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer"
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer"
MAX_REDIS_BUFFER_DEQUEUE_COUNT = 100
#### RELIABILITY ####
REPEATED_STREAMING_CHUNK_LIMIT = 100 # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives.

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

@ -10,7 +10,7 @@ import os
import time
import traceback
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
@ -18,15 +18,19 @@ from litellm.caching import DualCache, RedisCache
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
LiteLLM_UserTable,
SpendLogsPayload,
SpendUpdateQueueItem,
)
from litellm.proxy.db.pod_lock_manager import PodLockManager
from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.db.spend_update_queue import SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -51,6 +55,7 @@ class DBSpendUpdateWriter:
self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache)
self.pod_lock_manager = PodLockManager(cronjob_id=DB_SPEND_UPDATE_JOB_NAME)
self.spend_update_queue = SpendUpdateQueue()
self.daily_spend_update_queue = DailySpendUpdateQueue()
async def update_database(
# LiteLLM management object fields
@ -119,7 +124,7 @@ class DBSpendUpdateWriter:
)
)
if disable_spend_logs is False:
await DBSpendUpdateWriter._insert_spend_log_to_db(
await self._insert_spend_log_to_db(
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
@ -278,8 +283,8 @@ class DBSpendUpdateWriter:
)
raise e
@staticmethod
async def _insert_spend_log_to_db(
self,
kwargs: Optional[dict],
completion_response: Optional[Union[litellm.ModelResponse, Any, Exception]],
start_time: Optional[datetime],
@ -300,7 +305,7 @@ class DBSpendUpdateWriter:
end_time=end_time,
)
payload["spend"] = response_cost or 0.0
DBSpendUpdateWriter._set_spend_logs_payload(
await self._set_spend_logs_payload(
payload=payload,
spend_logs_url=os.getenv("SPEND_LOGS_URL"),
prisma_client=prisma_client,
@ -311,8 +316,8 @@ class DBSpendUpdateWriter:
)
raise e
@staticmethod
def _set_spend_logs_payload(
async def _set_spend_logs_payload(
self,
payload: Union[dict, SpendLogsPayload],
prisma_client: PrismaClient,
spend_logs_url: Optional[str] = None,
@ -331,8 +336,9 @@ class DBSpendUpdateWriter:
elif prisma_client is not None:
prisma_client.spend_log_transactions.append(payload)
prisma_client.add_spend_log_transaction_to_daily_user_transaction(
payload.copy()
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload.copy(),
prisma_client=prisma_client,
)
return prisma_client
@ -390,6 +396,7 @@ class DBSpendUpdateWriter:
"""
await self.redis_update_buffer.store_in_memory_spend_updates_in_redis(
spend_update_queue=self.spend_update_queue,
daily_spend_update_queue=self.daily_spend_update_queue,
)
# Only commit from redis to db if this pod is the leader
@ -407,6 +414,17 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
db_spend_update_transactions=db_spend_update_transactions,
)
daily_spend_update_transactions = (
await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer()
)
if daily_spend_update_transactions is not None:
await DBSpendUpdateWriter.update_daily_user_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_update_transactions,
)
except Exception as e:
verbose_proxy_logger.error(f"Error committing spend updates: {e}")
finally:
@ -425,6 +443,9 @@ class DBSpendUpdateWriter:
Note: This flow causes Deadlocks in production (1K RPS+). Use self._commit_spend_updates_to_db_with_redis() instead if you expect 1K+ RPS.
"""
# Aggregate all in memory spend updates (key, user, end_user, team, team_member, org) and commit to db
################## Spend Update Transactions ##################
db_spend_update_transactions = (
await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
)
@ -435,6 +456,19 @@ class DBSpendUpdateWriter:
db_spend_update_transactions=db_spend_update_transactions,
)
################## Daily Spend Update Transactions ##################
# Aggregate all in memory daily spend transactions and commit to db
daily_spend_update_transactions = (
await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
await DBSpendUpdateWriter.update_daily_user_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_update_transactions,
)
async def _commit_spend_updates_to_db( # noqa: PLR0915
self,
prisma_client: PrismaClient,
@ -679,3 +713,192 @@ class DBSpendUpdateWriter:
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
@staticmethod
async def update_daily_user_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: Dict[str, DailyUserSpendTransaction],
):
"""
Batch job to update LiteLLM_DailyUserSpend table using in-memory daily_spend_transactions
"""
from litellm.proxy.utils import _raise_failed_update_spend_exception
### UPDATE DAILY USER SPEND ###
verbose_proxy_logger.debug(
"Daily User Spend transactions: {}".format(len(daily_spend_transactions))
)
BATCH_SIZE = (
100 # Number of aggregated records to update in each database operation
)
start_time = time.time()
try:
for i in range(n_retry_times + 1):
try:
# Get transactions to process
transactions_to_process = dict(
list(daily_spend_transactions.items())[:BATCH_SIZE]
)
if len(transactions_to_process) == 0:
verbose_proxy_logger.debug(
"No new transactions to process for daily spend update"
)
break
# Update DailyUserSpend table in batches
async with prisma_client.db.batch_() as batcher:
for _, transaction in transactions_to_process.items():
user_id = transaction.get("user_id")
if not user_id: # Skip if no user_id
continue
batcher.litellm_dailyuserspend.upsert(
where={
"user_id_date_api_key_model_custom_llm_provider": {
"user_id": user_id,
"date": transaction["date"],
"api_key": transaction["api_key"],
"model": transaction["model"],
"custom_llm_provider": transaction.get(
"custom_llm_provider"
),
}
},
data={
"create": {
"user_id": user_id,
"date": transaction["date"],
"api_key": transaction["api_key"],
"model": transaction["model"],
"model_group": transaction.get("model_group"),
"custom_llm_provider": transaction.get(
"custom_llm_provider"
),
"prompt_tokens": transaction["prompt_tokens"],
"completion_tokens": transaction[
"completion_tokens"
],
"spend": transaction["spend"],
"api_requests": transaction["api_requests"],
"successful_requests": transaction[
"successful_requests"
],
"failed_requests": transaction[
"failed_requests"
],
},
"update": {
"prompt_tokens": {
"increment": transaction["prompt_tokens"]
},
"completion_tokens": {
"increment": transaction[
"completion_tokens"
]
},
"spend": {"increment": transaction["spend"]},
"api_requests": {
"increment": transaction["api_requests"]
},
"successful_requests": {
"increment": transaction[
"successful_requests"
]
},
"failed_requests": {
"increment": transaction["failed_requests"]
},
},
},
)
verbose_proxy_logger.info(
f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s"
)
# Remove processed transactions
for key in transactions_to_process.keys():
daily_spend_transactions.pop(key, None)
verbose_proxy_logger.debug(
f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s"
)
break
except DB_CONNECTION_ERROR_TYPES as e:
if i >= n_retry_times:
_raise_failed_update_spend_exception(
e=e,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(2**i) # Exponential backoff
except Exception as e:
# Remove processed transactions even if there was an error
if "transactions_to_process" in locals():
for key in transactions_to_process.keys(): # type: ignore
daily_spend_transactions.pop(key, None)
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
async def add_spend_log_transaction_to_daily_user_transaction(
self,
payload: Union[dict, SpendLogsPayload],
prisma_client: PrismaClient,
):
"""
Add a spend log transaction to the `daily_spend_update_queue`
Key = @@unique([user_id, date, api_key, model, custom_llm_provider]) )
If key exists, update the transaction with the new spend and usage
"""
expected_keys = ["user", "startTime", "api_key", "model", "custom_llm_provider"]
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
f"Missing expected keys: {expected_keys}, in payload, skipping from daily_user_spend_transactions"
)
return
request_status = prisma_client.get_request_status(payload)
verbose_proxy_logger.info(f"Logged request status: {request_status}")
if isinstance(payload["startTime"], datetime):
start_time = payload["startTime"].isoformat()
date = start_time.split("T")[0]
elif isinstance(payload["startTime"], str):
date = payload["startTime"].split("T")[0]
else:
verbose_proxy_logger.debug(
f"Invalid start time: {payload['startTime']}, skipping from daily_user_spend_transactions"
)
return
try:
daily_transaction_key = f"{payload['user']}_{date}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
daily_transaction = DailyUserSpendTransaction(
user_id=payload["user"],
date=date,
api_key=payload["api_key"],
model=payload["model"],
model_group=payload["model_group"],
custom_llm_provider=payload["custom_llm_provider"],
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
spend=payload["spend"],
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
)
await self.daily_spend_update_queue.add_update(
update={daily_transaction_key: daily_transaction}
)
except Exception as e:
raise e

25
litellm/proxy/db/db_transaction_queue/base_update_queue.py Normal file
View file

@ -0,0 +1,25 @@
"""
Base class for in memory buffer for database transactions
"""
import asyncio
from litellm._logging import verbose_proxy_logger
class BaseUpdateQueue:
"""Base class for in memory buffer for database transactions"""
def __init__(self):
self.update_queue = asyncio.Queue()
async def add_update(self, update):
"""Enqueue an update."""
verbose_proxy_logger.debug("Adding update to queue: %s", update)
await self.update_queue.put(update)
async def flush_all_updates_from_in_memory_queue(self):
"""Get all updates from the queue."""
updates = []
while not self.update_queue.empty():
updates.append(await self.update_queue.get())
return updates

95
litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py Normal file
View file

@ -0,0 +1,95 @@
import asyncio
from typing import Dict, List
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import DailyUserSpendTransaction
from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue
class DailySpendUpdateQueue(BaseUpdateQueue):
"""
In memory buffer for daily spend updates that should be committed to the database
To add a new daily spend update transaction, use the following format:
daily_spend_update_queue.add_update({
"user1_date_api_key_model_custom_llm_provider": {
"spend": 10,
"prompt_tokens": 100,
"completion_tokens": 100,
}
})
Queue contains a list of daily spend update transactions
eg
queue = [
{
"user1_date_api_key_model_custom_llm_provider": {
"spend": 10,
"prompt_tokens": 100,
"completion_tokens": 100,
"api_requests": 100,
"successful_requests": 100,
"failed_requests": 100,
}
},
{
"user2_date_api_key_model_custom_llm_provider": {
"spend": 10,
"prompt_tokens": 100,
"completion_tokens": 100,
"api_requests": 100,
"successful_requests": 100,
"failed_requests": 100,
}
}
]
"""
def __init__(self):
super().__init__()
self.update_queue: asyncio.Queue[
Dict[str, DailyUserSpendTransaction]
] = asyncio.Queue()
async def flush_and_get_aggregated_daily_spend_update_transactions(
self,
) -> Dict[str, DailyUserSpendTransaction]:
"""Get all updates from the queue and return all updates aggregated by daily_transaction_key."""
updates = await self.flush_all_updates_from_in_memory_queue()
aggregated_daily_spend_update_transactions = (
DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
updates
)
)
verbose_proxy_logger.debug(
"Aggregated daily spend update transactions: %s",
aggregated_daily_spend_update_transactions,
)
return aggregated_daily_spend_update_transactions
@staticmethod
def get_aggregated_daily_spend_update_transactions(
updates: List[Dict[str, DailyUserSpendTransaction]]
) -> Dict[str, DailyUserSpendTransaction]:
"""Aggregate updates by daily_transaction_key."""
aggregated_daily_spend_update_transactions: Dict[
str, DailyUserSpendTransaction
] = {}
for _update in updates:
for _key, payload in _update.items():
if _key in aggregated_daily_spend_update_transactions:
daily_transaction = aggregated_daily_spend_update_transactions[_key]
daily_transaction["spend"] += payload["spend"]
daily_transaction["prompt_tokens"] += payload["prompt_tokens"]
daily_transaction["completion_tokens"] += payload[
"completion_tokens"
]
daily_transaction["api_requests"] += payload["api_requests"]
daily_transaction["successful_requests"] += payload[
"successful_requests"
]
daily_transaction["failed_requests"] += payload["failed_requests"]
else:
aggregated_daily_spend_update_transactions[_key] = payload
return aggregated_daily_spend_update_transactions

0
litellm/proxy/db/pod_lock_manager.py → litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

100
litellm/proxy/db/redis_update_buffer.py → litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -9,10 +9,17 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import MAX_REDIS_BUFFER_DEQUEUE_COUNT, REDIS_UPDATE_BUFFER_KEY
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import DBSpendUpdateTransactions
from litellm.proxy.db.spend_update_queue import SpendUpdateQueue
from litellm.proxy._types import DailyUserSpendTransaction, DBSpendUpdateTransactions
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.secret_managers.main import str_to_bool
if TYPE_CHECKING:
@ -56,23 +63,51 @@ class RedisUpdateBuffer:
async def store_in_memory_spend_updates_in_redis(
self,
spend_update_queue: SpendUpdateQueue,
daily_spend_update_queue: DailySpendUpdateQueue,
):
"""
Stores the in-memory spend updates to Redis
Each transaction is a dict stored as following:
- key is the entity id
- value is the spend amount
Stores the following in memory data structures in Redis:
- SpendUpdateQueue - Key, User, Team, TeamMember, Org, EndUser Spend updates
- DailySpendUpdateQueue - Daily Spend updates Aggregate view
```
Redis List:
key_list_transactions:
[
"0929880201": 1.2,
"0929880202": 0.01,
"0929880203": 0.001,
]
```
For SpendUpdateQueue:
Each transaction is a dict stored as following:
- key is the entity id
- value is the spend amount
```
Redis List:
key_list_transactions:
[
"0929880201": 1.2,
"0929880202": 0.01,
"0929880203": 0.001,
]
```
For DailySpendUpdateQueue:
Each transaction is a Dict[str, DailyUserSpendTransaction] stored as following:
- key is the daily_transaction_key
- value is the DailyUserSpendTransaction
```
Redis List:
daily_spend_update_transactions:
[
{
"user_keyhash_1_model_1": {
"spend": 1.2,
"prompt_tokens": 1000,
"completion_tokens": 1000,
"api_requests": 1000,
"successful_requests": 1000,
},
}
]
```
"""
if self.redis_cache is None:
verbose_proxy_logger.debug(
@ -86,6 +121,12 @@ class RedisUpdateBuffer:
verbose_proxy_logger.debug(
"ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions
)
daily_spend_update_transactions = (
await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
verbose_proxy_logger.debug(
"ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions
)
# only store in redis if there are any updates to commit
if (
@ -100,6 +141,14 @@ class RedisUpdateBuffer:
values=list_of_transactions,
)
list_of_daily_spend_update_transactions = [
safe_dumps(daily_spend_update_transactions)
]
await self.redis_cache.async_rpush(
key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
values=list_of_daily_spend_update_transactions,
)
@staticmethod
def _number_of_transactions_to_store_in_redis(
db_spend_update_transactions: DBSpendUpdateTransactions,
@ -180,6 +229,27 @@ class RedisUpdateBuffer:
return combined_transaction
async def get_all_daily_spend_update_transactions_from_redis_buffer(
self,
) -> Optional[Dict[str, DailyUserSpendTransaction]]:
"""
Gets all the daily spend update transactions from Redis
"""
if self.redis_cache is None:
return None
list_of_transactions = await self.redis_cache.async_lpop(
key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
)
if list_of_transactions is None:
return None
list_of_daily_spend_update_transactions = [
json.loads(transaction) for transaction in list_of_transactions
]
return DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
list_of_daily_spend_update_transactions
)
@staticmethod
def _parse_list_of_transactions(
list_of_transactions: Union[Any, List[Any]],

29
litellm/proxy/db/spend_update_queue.py → litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -1,5 +1,5 @@
import asyncio
from typing import TYPE_CHECKING, Any, List
from typing import List
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
@ -7,37 +7,18 @@ from litellm.proxy._types import (
Litellm_EntityType,
SpendUpdateQueueItem,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
else:
PrismaClient = Any
from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue
class SpendUpdateQueue:
class SpendUpdateQueue(BaseUpdateQueue):
"""
In memory buffer for spend updates that should be committed to the database
"""
def __init__(
self,
):
def __init__(self):
super().__init__()
self.update_queue: asyncio.Queue[SpendUpdateQueueItem] = asyncio.Queue()
async def add_update(self, update: SpendUpdateQueueItem) -> None:
"""Enqueue an update. Each update might be a dict like {'entity_type': 'user', 'entity_id': '123', 'amount': 1.2}."""
verbose_proxy_logger.debug("Adding update to queue: %s", update)
await self.update_queue.put(update)
async def flush_all_updates_from_in_memory_queue(
self,
) -> List[SpendUpdateQueueItem]:
"""Get all updates from the queue."""
updates: List[SpendUpdateQueueItem] = []
while not self.update_queue.empty():
updates.append(await self.update_queue.get())
return updates
async def flush_and_get_aggregated_db_spend_update_transactions(
self,
) -> DBSpendUpdateTransactions:

212
litellm/proxy/utils.py
View file

@ -25,7 +25,6 @@ from typing import (
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
CommonProxyErrors,
DailyUserSpendTransaction,
ProxyErrorTypes,
ProxyException,
SpendLogsMetadata,
@ -1112,7 +1111,6 @@ def jsonify_object(data: dict) -> dict:
class PrismaClient:
spend_log_transactions: List = []
daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {}
def __init__(
self,
@ -1185,74 +1183,6 @@ class PrismaClient:
# Default to success if metadata parsing fails
return "success"
def add_spend_log_transaction_to_daily_user_transaction(
self, payload: Union[dict, SpendLogsPayload]
):
"""
Add a spend log transaction to the daily user transaction list
Key = @@unique([user_id, date, api_key, model, custom_llm_provider]) )
If key exists, update the transaction with the new spend and usage
"""
expected_keys = ["user", "startTime", "api_key", "model", "custom_llm_provider"]
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
f"Missing expected keys: {expected_keys}, in payload, skipping from daily_user_spend_transactions"
)
return
request_status = self.get_request_status(payload)
verbose_proxy_logger.info(f"Logged request status: {request_status}")
if isinstance(payload["startTime"], datetime):
start_time = payload["startTime"].isoformat()
date = start_time.split("T")[0]
elif isinstance(payload["startTime"], str):
date = payload["startTime"].split("T")[0]
else:
verbose_proxy_logger.debug(
f"Invalid start time: {payload['startTime']}, skipping from daily_user_spend_transactions"
)
return
try:
daily_transaction_key = f"{payload['user']}_{date}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}"
if daily_transaction_key in self.daily_user_spend_transactions:
daily_transaction = self.daily_user_spend_transactions[
daily_transaction_key
]
daily_transaction["spend"] += payload["spend"]
daily_transaction["prompt_tokens"] += payload["prompt_tokens"]
daily_transaction["completion_tokens"] += payload["completion_tokens"]
daily_transaction["api_requests"] += 1
if request_status == "success":
daily_transaction["successful_requests"] += 1
else:
daily_transaction["failed_requests"] += 1
else:
daily_transaction = DailyUserSpendTransaction(
user_id=payload["user"],
date=date,
api_key=payload["api_key"],
model=payload["model"],
model_group=payload["model_group"],
custom_llm_provider=payload["custom_llm_provider"],
prompt_tokens=payload["prompt_tokens"],
completion_tokens=payload["completion_tokens"],
spend=payload["spend"],
api_requests=1,
successful_requests=1 if request_status == "success" else 0,
failed_requests=1 if request_status != "success" else 0,
)
self.daily_user_spend_transactions[
daily_transaction_key
] = daily_transaction
except Exception as e:
raise e
def hash_token(self, token: str):
# Hash the string using SHA-256
hashed_token = hashlib.sha256(token.encode()).hexdigest()
@ -2588,134 +2518,6 @@ class ProxyUpdateSpend:
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
@staticmethod
async def update_daily_user_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
):
"""
Batch job to update LiteLLM_DailyUserSpend table using in-memory daily_spend_transactions
"""
BATCH_SIZE = (
100 # Number of aggregated records to update in each database operation
)
start_time = time.time()
try:
for i in range(n_retry_times + 1):
try:
# Get transactions to process
transactions_to_process = dict(
list(prisma_client.daily_user_spend_transactions.items())[
:BATCH_SIZE
]
)
if len(transactions_to_process) == 0:
verbose_proxy_logger.debug(
"No new transactions to process for daily spend update"
)
break
# Update DailyUserSpend table in batches
async with prisma_client.db.batch_() as batcher:
for _, transaction in transactions_to_process.items():
user_id = transaction.get("user_id")
if not user_id: # Skip if no user_id
continue
batcher.litellm_dailyuserspend.upsert(
where={
"user_id_date_api_key_model_custom_llm_provider": {
"user_id": user_id,
"date": transaction["date"],
"api_key": transaction["api_key"],
"model": transaction["model"],
"custom_llm_provider": transaction.get(
"custom_llm_provider"
),
}
},
data={
"create": {
"user_id": user_id,
"date": transaction["date"],
"api_key": transaction["api_key"],
"model": transaction["model"],
"model_group": transaction.get("model_group"),
"custom_llm_provider": transaction.get(
"custom_llm_provider"
),
"prompt_tokens": transaction["prompt_tokens"],
"completion_tokens": transaction[
"completion_tokens"
],
"spend": transaction["spend"],
"api_requests": transaction["api_requests"],
"successful_requests": transaction[
"successful_requests"
],
"failed_requests": transaction[
"failed_requests"
],
},
"update": {
"prompt_tokens": {
"increment": transaction["prompt_tokens"]
},
"completion_tokens": {
"increment": transaction[
"completion_tokens"
]
},
"spend": {"increment": transaction["spend"]},
"api_requests": {
"increment": transaction["api_requests"]
},
"successful_requests": {
"increment": transaction[
"successful_requests"
]
},
"failed_requests": {
"increment": transaction["failed_requests"]
},
},
},
)
verbose_proxy_logger.info(
f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s"
)
# Remove processed transactions
for key in transactions_to_process.keys():
prisma_client.daily_user_spend_transactions.pop(key, None)
verbose_proxy_logger.debug(
f"Processed {len(transactions_to_process)} daily spend transactions in {time.time() - start_time:.2f}s"
)
break
except DB_CONNECTION_ERROR_TYPES as e:
if i >= n_retry_times:
_raise_failed_update_spend_exception(
e=e,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(2**i) # Exponential backoff
except Exception as e:
# Remove processed transactions even if there was an error
if "transactions_to_process" in locals():
for key in transactions_to_process.keys(): # type: ignore
prisma_client.daily_user_spend_transactions.pop(key, None)
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
@staticmethod
def disable_spend_updates() -> bool:
"""
@ -2765,20 +2567,6 @@ async def update_spend( # noqa: PLR0915
db_writer_client=db_writer_client,
)
### UPDATE DAILY USER SPEND ###
verbose_proxy_logger.debug(
"Daily User Spend transactions: {}".format(
len(prisma_client.daily_user_spend_transactions)
)
)
if len(prisma_client.daily_user_spend_transactions) > 0:
await ProxyUpdateSpend.update_daily_user_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
def _raise_failed_update_spend_exception(
e: Exception, start_time: float, proxy_logging_obj: ProxyLogging

264
tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py Normal file
View file

@ -0,0 +1,264 @@
import asyncio
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
from litellm.proxy._types import (
DailyUserSpendTransaction,
Litellm_EntityType,
SpendUpdateQueueItem,
)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
@pytest.fixture
def daily_spend_update_queue():
return DailySpendUpdateQueue()
@pytest.mark.asyncio
async def test_empty_queue_flush(daily_spend_update_queue):
"""Test flushing an empty queue returns an empty list"""
result = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue()
assert result == []
@pytest.mark.asyncio
async def test_add_single_update(daily_spend_update_queue):
"""Test adding a single update to the queue"""
test_key = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
# Add update to queue
await daily_spend_update_queue.add_update({test_key: test_transaction})
# Flush and check
updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue()
assert len(updates) == 1
assert test_key in updates[0]
assert updates[0][test_key] == test_transaction
@pytest.mark.asyncio
async def test_add_multiple_updates(daily_spend_update_queue):
"""Test adding multiple updates to the queue"""
test_key1 = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction1 = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
test_key2 = "user2_2023-01-01_key456_gpt-3.5-turbo_openai"
test_transaction2 = {
"spend": 5.0,
"prompt_tokens": 200,
"completion_tokens": 30,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
# Add updates to queue
await daily_spend_update_queue.add_update({test_key1: test_transaction1})
await daily_spend_update_queue.add_update({test_key2: test_transaction2})
# Flush and check
updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue()
assert len(updates) == 2
# Find each transaction in the list of updates
found_transaction1 = False
found_transaction2 = False
for update in updates:
if test_key1 in update:
assert update[test_key1] == test_transaction1
found_transaction1 = True
if test_key2 in update:
assert update[test_key2] == test_transaction2
found_transaction2 = True
assert found_transaction1
assert found_transaction2
@pytest.mark.asyncio
async def test_aggregated_daily_spend_update_empty(daily_spend_update_queue):
"""Test aggregating updates from an empty queue"""
result = (
await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
assert result == {}
@pytest.mark.asyncio
async def test_get_aggregated_daily_spend_update_transactions_single_key():
"""Test static method for aggregating a single key"""
test_key = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
updates = [{test_key: test_transaction}]
# Test aggregation
result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
updates
)
assert len(result) == 1
assert test_key in result
assert result[test_key] == test_transaction
@pytest.mark.asyncio
async def test_get_aggregated_daily_spend_update_transactions_multiple_keys():
"""Test static method for aggregating multiple different keys"""
test_key1 = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction1 = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
test_key2 = "user2_2023-01-01_key456_gpt-3.5-turbo_openai"
test_transaction2 = {
"spend": 5.0,
"prompt_tokens": 200,
"completion_tokens": 30,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
updates = [{test_key1: test_transaction1}, {test_key2: test_transaction2}]
# Test aggregation
result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
updates
)
assert len(result) == 2
assert test_key1 in result
assert test_key2 in result
assert result[test_key1] == test_transaction1
assert result[test_key2] == test_transaction2
@pytest.mark.asyncio
async def test_get_aggregated_daily_spend_update_transactions_same_key():
"""Test static method for aggregating updates with the same key"""
test_key = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction1 = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
test_transaction2 = {
"spend": 5.0,
"prompt_tokens": 200,
"completion_tokens": 30,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
expected_transaction = {
"spend": 15.0, # 10 + 5
"prompt_tokens": 300, # 100 + 200
"completion_tokens": 80, # 50 + 30
"api_requests": 2, # 1 + 1
"successful_requests": 2, # 1 + 1
"failed_requests": 0, # 0 + 0
}
updates = [{test_key: test_transaction1}, {test_key: test_transaction2}]
# Test aggregation
result = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
updates
)
assert len(result) == 1
assert test_key in result
assert result[test_key] == expected_transaction
@pytest.mark.asyncio
async def test_flush_and_get_aggregated_daily_spend_update_transactions(
daily_spend_update_queue,
):
"""Test the full workflow of adding, flushing, and aggregating updates"""
test_key = "user1_2023-01-01_key123_gpt-4_openai"
test_transaction1 = {
"spend": 10.0,
"prompt_tokens": 100,
"completion_tokens": 50,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
test_transaction2 = {
"spend": 5.0,
"prompt_tokens": 200,
"completion_tokens": 30,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
}
expected_transaction = {
"spend": 15.0, # 10 + 5
"prompt_tokens": 300, # 100 + 200
"completion_tokens": 80, # 50 + 30
"api_requests": 2, # 1 + 1
"successful_requests": 2, # 1 + 1
"failed_requests": 0, # 0 + 0
}
# Add updates to queue
await daily_spend_update_queue.add_update({test_key: test_transaction1})
await daily_spend_update_queue.add_update({test_key: test_transaction2})
# Test full workflow
result = (
await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
assert len(result) == 1
assert test_key in result
assert result[test_key] == expected_transaction

2
tests/litellm/proxy/db/test_pod_lock_manager.py → tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py
View file

@ -12,7 +12,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS
from litellm.proxy.db.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
# Mock Prisma client class

2
tests/litellm/proxy/db/test_spend_update_queue.py → tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
View file

@ -7,7 +7,7 @@ import pytest
from fastapi.testclient import TestClient
from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
from litellm.proxy.db.spend_update_queue import SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
sys.path.insert(
0, os.path.abspath("../../..")

2
tests/proxy_unit_tests/test_e2e_pod_lock_manager.py
View file

@ -23,7 +23,7 @@ import asyncio
import logging
import pytest
from litellm.proxy.db.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy.management_endpoints.internal_user_endpoints import (

9
tests/proxy_unit_tests/test_key_generate_prisma.py
View file

@ -2280,15 +2280,16 @@ def test_get_bearer_token():
result = _get_bearer_token(api_key)
assert result == "sk-1234", f"Expected 'valid_token', got '{result}'"
def test_update_logs_with_spend_logs_url(prisma_client):
@pytest.mark.asyncio
async def test_update_logs_with_spend_logs_url(prisma_client):
"""
Unit test for making sure spend logs list is still updated when url passed in
"""
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
db_spend_update_writer = DBSpendUpdateWriter()
payload = {"startTime": datetime.now(), "endTime": datetime.now()}
DBSpendUpdateWriter._set_spend_logs_payload(payload=payload, prisma_client=prisma_client)
await db_spend_update_writer._set_spend_logs_payload(payload=payload, prisma_client=prisma_client)
assert len(prisma_client.spend_log_transactions) > 0
@ -2296,7 +2297,7 @@ def test_update_logs_with_spend_logs_url(prisma_client):
spend_logs_url = ""
payload = {"startTime": datetime.now(), "endTime": datetime.now()}
DBSpendUpdateWriter._set_spend_logs_payload(
await db_spend_update_writer._set_spend_logs_payload(
payload=payload, spend_logs_url=spend_logs_url, prisma_client=prisma_client
)