From 4c14550721b91cdaba3d5886bd22085dad837696 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 17:45:19 -0700 Subject: [PATCH 1/8] refactor daily spend updates to use new Queue DS --- litellm/constants.py | 1 + litellm/proxy/db/db_spend_update_writer.py | 231 ++++++++++++++++++++- litellm/proxy/db/redis_update_buffer.py | 97 +++++++-- litellm/proxy/db/spend_update_queue.py | 61 +++++- litellm/proxy/utils.py | 212 ------------------- 5 files changed, 365 insertions(+), 237 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index d5e0215ebf0..cace674f2f0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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. diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 5bf255feae2..20c33fb2403 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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,6 +18,7 @@ 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, @@ -26,7 +27,7 @@ from litellm.proxy._types import ( ) 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.spend_update_queue import DailySpendUpdateQueue, SpendUpdateQueue if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -51,6 +52,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 +121,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 +280,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 +302,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 +313,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 +333,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 +393,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 @@ -400,6 +404,9 @@ class DBSpendUpdateWriter: db_spend_update_transactions = ( await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() ) + daily_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer() + ) if db_spend_update_transactions is not None: await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -407,6 +414,13 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, db_spend_update_transactions=db_spend_update_transactions, ) + 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: @@ -428,6 +442,9 @@ class DBSpendUpdateWriter: db_spend_update_transactions = ( await self.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() ) + daily_spend_update_transactions = ( + await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, @@ -435,6 +452,13 @@ class DBSpendUpdateWriter: db_spend_update_transactions=db_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 +703,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 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 = 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 diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 1a3fd3d42d1..6ffada4deeb 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -9,10 +9,14 @@ 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.spend_update_queue import DailySpendUpdateQueue, SpendUpdateQueue from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -56,23 +60,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 +118,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 +138,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 +226,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]], diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 28e05246fad..47b4b9724ab 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -1,8 +1,9 @@ import asyncio -from typing import TYPE_CHECKING, Any, List +from typing import TYPE_CHECKING, Any, Dict, List from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( + DailyUserSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem, @@ -130,3 +131,61 @@ class SpendUpdateQueue: transactions_dict[entity_id] += response_cost or 0 return db_spend_update_transactions + + +class DailySpendUpdateQueue: + def __init__( + self, + ): + self.update_queue: asyncio.Queue[ + Dict[str, DailyUserSpendTransaction] + ] = asyncio.Queue() + + async def add_update(self, update: Dict[str, DailyUserSpendTransaction]) -> 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[Dict[str, DailyUserSpendTransaction]]: + """Get all updates from the queue.""" + updates: List[Dict[str, DailyUserSpendTransaction]] = [] + while not self.update_queue.empty(): + updates.append(await self.update_queue.get()) + return updates + + 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() + return DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + updates + ) + + @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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 67d1882a11d..0b87444628d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 From 827ad38e7f97e7d1d3e59e1208005907378b85e5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 17:49:28 -0700 Subject: [PATCH 2/8] fix spend update queue --- litellm/proxy/db/db_spend_update_writer.py | 21 ++++++++++++++------- litellm/proxy/db/spend_update_queue.py | 13 ++++++++++++- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 20c33fb2403..56f7664c738 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -404,9 +404,6 @@ class DBSpendUpdateWriter: db_spend_update_transactions = ( await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() ) - daily_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer() - ) if db_spend_update_transactions is not None: await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -414,6 +411,10 @@ 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, @@ -439,12 +440,12 @@ 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() ) - daily_spend_update_transactions = ( - await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, @@ -452,6 +453,12 @@ 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, @@ -843,7 +850,7 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient, ): """ - Add a spend log transaction to the daily user transaction list + Add a spend log transaction to the `daily_spend_update_queue` Key = @@unique([user_id, date, api_key, model, custom_llm_provider]) ) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index 47b4b9724ab..c17c1cddf9f 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -142,7 +142,18 @@ class DailySpendUpdateQueue: ] = asyncio.Queue() async def add_update(self, update: Dict[str, DailyUserSpendTransaction]) -> None: - """Enqueue an update. Each update might be a dict like {'entity_type': 'user', 'entity_id': '123', 'amount': 1.2}.""" + """Enqueue an update. Each update might be a dict like + { + "user_date_api_key_model_custom_llm_provider": { + "spend": 1.2, + "prompt_tokens": 1000, + "completion_tokens": 1000, + "api_requests": 1000, + "successful_requests": 1000, + "failed_requests": 1000, + } + } + .""" verbose_proxy_logger.debug("Adding update to queue: %s", update) await self.update_queue.put(update) From 07fc5a72f10713a903258c42fb42168518c0be21 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 17:54:52 -0700 Subject: [PATCH 3/8] add debug statement --- litellm/proxy/db/spend_update_queue.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/spend_update_queue.py index c17c1cddf9f..e77f9450943 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/spend_update_queue.py @@ -171,9 +171,16 @@ class DailySpendUpdateQueue: ) -> 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() - return DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( - updates + 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( From 290e8375158cb44766e2849eec65e353fa1bd8a9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:15:01 -0700 Subject: [PATCH 4/8] test_update_logs_with_spend_logs_url --- tests/proxy_unit_tests/test_key_generate_prisma.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index ea837b717bb..0400a71ceab 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -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 ) From 8dc792139e3ae211e9f87761848e9eb5cffca14c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:30:48 -0700 Subject: [PATCH 5/8] refactor file structure --- litellm/proxy/db/db_spend_update_writer.py | 9 +- .../db_transaction_queue/base_update_queue.py | 22 ++++ .../daily_spend_update_queue.py | 95 ++++++++++++++++ .../pod_lock_manager.py | 0 .../redis_update_buffer.py | 5 +- .../spend_update_queue.py | 106 +----------------- .../litellm/proxy/db/test_pod_lock_manager.py | 2 +- .../proxy/db/test_spend_update_queue.py | 2 +- .../test_e2e_pod_lock_manager.py | 2 +- 9 files changed, 135 insertions(+), 108 deletions(-) create mode 100644 litellm/proxy/db/db_transaction_queue/base_update_queue.py create mode 100644 litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py rename litellm/proxy/db/{ => db_transaction_queue}/pod_lock_manager.py (100%) rename litellm/proxy/db/{ => db_transaction_queue}/redis_update_buffer.py (98%) rename litellm/proxy/db/{ => db_transaction_queue}/spend_update_queue.py (52%) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 56f7664c738..f4f045b2a20 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -25,9 +25,12 @@ from litellm.proxy._types import ( 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 DailySpendUpdateQueue, 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 diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py new file mode 100644 index 00000000000..b74ed439df9 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -0,0 +1,22 @@ +import asyncio + +from litellm._logging import verbose_proxy_logger + + +class BaseUpdateQueue: + """Base class for spend update queues with common functionality""" + + 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 diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py new file mode 100644 index 00000000000..dedb8c8f8fa --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -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 diff --git a/litellm/proxy/db/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py similarity index 100% rename from litellm/proxy/db/pod_lock_manager.py rename to litellm/proxy/db/db_transaction_queue/pod_lock_manager.py diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py similarity index 98% rename from litellm/proxy/db/redis_update_buffer.py rename to litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6ffada4deeb..ea1356159a0 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -16,7 +16,10 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyUserSpendTransaction, DBSpendUpdateTransactions -from litellm.proxy.db.spend_update_queue import DailySpendUpdateQueue, SpendUpdateQueue +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: diff --git a/litellm/proxy/db/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py similarity index 52% rename from litellm/proxy/db/spend_update_queue.py rename to litellm/proxy/db/db_transaction_queue/spend_update_queue.py index e77f9450943..ce181d14784 100644 --- a/litellm/proxy/db/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -1,44 +1,24 @@ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List +from typing import List from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( - DailyUserSpendTransaction, DBSpendUpdateTransactions, 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: @@ -131,79 +111,3 @@ class SpendUpdateQueue: transactions_dict[entity_id] += response_cost or 0 return db_spend_update_transactions - - -class DailySpendUpdateQueue: - def __init__( - self, - ): - self.update_queue: asyncio.Queue[ - Dict[str, DailyUserSpendTransaction] - ] = asyncio.Queue() - - async def add_update(self, update: Dict[str, DailyUserSpendTransaction]) -> None: - """Enqueue an update. Each update might be a dict like - { - "user_date_api_key_model_custom_llm_provider": { - "spend": 1.2, - "prompt_tokens": 1000, - "completion_tokens": 1000, - "api_requests": 1000, - "successful_requests": 1000, - "failed_requests": 1000, - } - } - .""" - 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[Dict[str, DailyUserSpendTransaction]]: - """Get all updates from the queue.""" - updates: List[Dict[str, DailyUserSpendTransaction]] = [] - while not self.update_queue.empty(): - updates.append(await self.update_queue.get()) - return updates - - 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 diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/test_pod_lock_manager.py index bce7b66409e..cde43158379 100644 --- a/tests/litellm/proxy/db/test_pod_lock_manager.py +++ b/tests/litellm/proxy/db/test_pod_lock_manager.py @@ -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 diff --git a/tests/litellm/proxy/db/test_spend_update_queue.py b/tests/litellm/proxy/db/test_spend_update_queue.py index 89d494a0701..98d3b4e4c73 100644 --- a/tests/litellm/proxy/db/test_spend_update_queue.py +++ b/tests/litellm/proxy/db/test_spend_update_queue.py @@ -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("../../..") diff --git a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py index 7d36bb4791f..3522c8e1e26 100644 --- a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py +++ b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py @@ -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 ( From 44bc8974c8c0d3bc8726b2544aeb20d5c7ee1e22 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:31:54 -0700 Subject: [PATCH 6/8] BaseUpdateQueue --- litellm/proxy/db/db_transaction_queue/base_update_queue.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index b74ed439df9..b3c3c26c847 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -1,10 +1,13 @@ +""" +Base class for in memory buffer for database transactions +""" import asyncio from litellm._logging import verbose_proxy_logger class BaseUpdateQueue: - """Base class for spend update queues with common functionality""" + """Base class for in memory buffer for database transactions""" def __init__(self): self.update_queue = asyncio.Queue() From 4a091a34b0053525f112161c9f8a64396ff8d901 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:33:33 -0700 Subject: [PATCH 7/8] move test loc --- .../proxy/db/{ => db_transaction_queue}/test_pod_lock_manager.py | 0 .../db/{ => db_transaction_queue}/test_spend_update_queue.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/litellm/proxy/db/{ => db_transaction_queue}/test_pod_lock_manager.py (100%) rename tests/litellm/proxy/db/{ => db_transaction_queue}/test_spend_update_queue.py (100%) diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py similarity index 100% rename from tests/litellm/proxy/db/test_pod_lock_manager.py rename to tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py diff --git a/tests/litellm/proxy/db/test_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py similarity index 100% rename from tests/litellm/proxy/db/test_spend_update_queue.py rename to tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py From feba274a89eb5b5924352a9a994268cc1313eddc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 1 Apr 2025 18:39:23 -0700 Subject: [PATCH 8/8] test DailySpendUpdateQueue --- .../test_daily_spend_update_queue.py | 264 ++++++++++++++++++ 1 file changed, 264 insertions(+) create mode 100644 tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py new file mode 100644 index 00000000000..228dfe64b1e --- /dev/null +++ b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py @@ -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