From 21e3b764f570bdf23ebca9afda4875f83503f643 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 16:31:23 -0700 Subject: [PATCH 01/40] use DBSpendUpdateWriter class for managing DB spend updates --- litellm/proxy/db/db_spend_update_writer.py | 233 ++++++++++++++++++ .../proxy/hooks/proxy_track_cost_callback.py | 8 +- litellm/proxy/proxy_server.py | 221 +---------------- .../hooks/test_proxy_track_cost_callback.py | 3 +- .../test_unit_test_proxy_hooks.py | 13 +- 5 files changed, 255 insertions(+), 223 deletions(-) create mode 100644 litellm/proxy/db/db_spend_update_writer.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py new file mode 100644 index 0000000000..d5f2ef5228 --- /dev/null +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -0,0 +1,233 @@ +import asyncio +import os +import traceback +from datetime import datetime +from typing import Any, Optional, Union + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload +from litellm.proxy.proxy_server import hash_token +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend + + +class DBSpendUpdateWriter: + + @staticmethod + async def update_database( # noqa: PLR0915 + # LiteLLM management object fields + token: Optional[str], + user_id: Optional[str], + end_user_id: Optional[str], + team_id: Optional[str], + org_id: Optional[str], + # Completion object fields + kwargs: Optional[dict], + completion_response: Optional[Union[litellm.ModelResponse, Any, Exception]], + start_time: Optional[datetime], + end_time: Optional[datetime], + response_cost: Optional[float], + ): + from litellm.proxy.proxy_server import ( + disable_spend_logs, + litellm_proxy_budget_name, + prisma_client, + user_api_key_cache, + ) + + try: + verbose_proxy_logger.debug( + f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}" + ) + if ProxyUpdateSpend.disable_spend_updates() is True: + return + if token is not None and isinstance(token, str) and token.startswith("sk-"): + hashed_token = hash_token(token=token) + else: + hashed_token = token + + ### UPDATE USER SPEND ### + async def _update_user_db(): + """ + - Update that user's row + - Update litellm-proxy-budget row (global proxy spend) + """ + ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db + existing_user_obj = await user_api_key_cache.async_get_cache( + key=user_id + ) + if existing_user_obj is not None and isinstance( + existing_user_obj, dict + ): + existing_user_obj = LiteLLM_UserTable(**existing_user_obj) + try: + if prisma_client is not None: # update + user_ids = [user_id] + if ( + litellm.max_budget > 0 + ): # track global proxy budget, if user set max budget + user_ids.append(litellm_proxy_budget_name) + ### KEY CHANGE ### + for _id in user_ids: + if _id is not None: + prisma_client.user_list_transactons[_id] = ( + response_cost + + prisma_client.user_list_transactons.get(_id, 0) + ) + if end_user_id is not None: + prisma_client.end_user_list_transactons[end_user_id] = ( + response_cost + + prisma_client.end_user_list_transactons.get( + end_user_id, 0 + ) + ) + except Exception as e: + verbose_proxy_logger.info( + "\033[91m" + + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" + ) + + ### UPDATE KEY SPEND ### + async def _update_key_db(): + try: + verbose_proxy_logger.debug( + f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." + ) + if hashed_token is None: + return + if prisma_client is not None: + prisma_client.key_list_transactons[hashed_token] = ( + response_cost + + prisma_client.key_list_transactons.get(hashed_token, 0) + ) + except Exception as e: + verbose_proxy_logger.exception( + f"Update Key DB Call failed to execute - {str(e)}" + ) + raise e + + ### UPDATE SPEND LOGS ### + async def _insert_spend_log_to_db(): + try: + if prisma_client: + payload = get_logging_payload( + kwargs=kwargs, + response_obj=completion_response, + start_time=start_time, + end_time=end_time, + ) + payload["spend"] = response_cost or 0.0 + DBSpendUpdateWriter._set_spend_logs_payload( + payload=payload, + spend_logs_url=os.getenv("SPEND_LOGS_URL"), + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + + ### UPDATE TEAM SPEND ### + async def _update_team_db(): + try: + verbose_proxy_logger.debug( + f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}." + ) + if team_id is None: + verbose_proxy_logger.debug( + "track_cost_callback: team_id is None. Not tracking spend for team" + ) + return + if prisma_client is not None: + prisma_client.team_list_transactons[team_id] = ( + response_cost + + prisma_client.team_list_transactons.get(team_id, 0) + ) + + try: + # Track spend of the team member within this team + # key is "team_id::::user_id::" + team_member_key = f"team_id::{team_id}::user_id::{user_id}" + prisma_client.team_member_list_transactons[ + team_member_key + ] = ( + response_cost + + prisma_client.team_member_list_transactons.get( + team_member_key, 0 + ) + ) + except Exception: + pass + except Exception as e: + verbose_proxy_logger.info( + f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + + ### UPDATE ORG SPEND ### + async def _update_org_db(): + try: + verbose_proxy_logger.debug( + "adding spend to org db. Response cost: {}. org_id: {}.".format( + response_cost, org_id + ) + ) + if org_id is None: + verbose_proxy_logger.debug( + "track_cost_callback: org_id is None. Not tracking spend for org" + ) + return + if prisma_client is not None: + prisma_client.org_list_transactons[org_id] = ( + response_cost + + prisma_client.org_list_transactons.get(org_id, 0) + ) + except Exception as e: + verbose_proxy_logger.info( + f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + + asyncio.create_task(_update_user_db()) + asyncio.create_task(_update_key_db()) + asyncio.create_task(_update_team_db()) + asyncio.create_task(_update_org_db()) + if disable_spend_logs is False: + await _insert_spend_log_to_db() + else: + verbose_proxy_logger.info( + "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." + ) + + verbose_proxy_logger.debug("Runs spend update on all tables") + except Exception: + verbose_proxy_logger.debug( + f"Error updating Prisma database: {traceback.format_exc()}" + ) + + @staticmethod + def _set_spend_logs_payload( + payload: Union[dict, SpendLogsPayload], + prisma_client: PrismaClient, + spend_logs_url: Optional[str] = None, + ) -> PrismaClient: + verbose_proxy_logger.info( + "Writing spend log to db - request_id: {}, spend: {}".format( + payload.get("request_id"), payload.get("spend") + ) + ) + if prisma_client is not None and spend_logs_url is not None: + if isinstance(payload["startTime"], datetime): + payload["startTime"] = payload["startTime"].isoformat() + if isinstance(payload["endTime"], datetime): + payload["endTime"] = payload["endTime"].isoformat() + prisma_client.spend_log_transactions.append(payload) + 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() + ) + return prisma_client diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index e8a947329d..f205b0146f 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( StandardLoggingPayload, @@ -33,8 +34,6 @@ class _ProxyDBLogger(CustomLogger): original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, ): - from litellm.proxy.proxy_server import update_database - if _ProxyDBLogger._should_track_errors_in_db() is False: return @@ -67,7 +66,7 @@ class _ProxyDBLogger(CustomLogger): request_data.get("proxy_server_request") or {} ) request_data["litellm_params"]["metadata"] = existing_metadata - await update_database( + await DBSpendUpdateWriter.update_database( token=user_api_key_dict.api_key, response_cost=0.0, user_id=user_api_key_dict.user_id, @@ -94,7 +93,6 @@ class _ProxyDBLogger(CustomLogger): prisma_client, proxy_logging_obj, update_cache, - update_database, ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") @@ -138,7 +136,7 @@ class _ProxyDBLogger(CustomLogger): end_user_id=end_user_id, ): ## UPDATE DATABASE - await update_database( + await DBSpendUpdateWriter.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d7e92fd73..6a2da7d83b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -897,211 +897,6 @@ def cost_tracking(): litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger()) -def _set_spend_logs_payload( - payload: Union[dict, SpendLogsPayload], - prisma_client: PrismaClient, - spend_logs_url: Optional[str] = None, -): - verbose_proxy_logger.info( - "Writing spend log to db - request_id: {}, spend: {}".format( - payload.get("request_id"), payload.get("spend") - ) - ) - if prisma_client is not None and spend_logs_url is not None: - if isinstance(payload["startTime"], datetime): - payload["startTime"] = payload["startTime"].isoformat() - if isinstance(payload["endTime"], datetime): - payload["endTime"] = payload["endTime"].isoformat() - prisma_client.spend_log_transactions.append(payload) - 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()) - return prisma_client - - -async def update_database( # noqa: PLR0915 - token, - response_cost, - user_id=None, - end_user_id=None, - team_id=None, - kwargs=None, - completion_response=None, - start_time=None, - end_time=None, - org_id=None, -): - try: - global prisma_client - verbose_proxy_logger.debug( - f"Enters prisma db call, response_cost: {response_cost}, token: {token}; user_id: {user_id}; team_id: {team_id}" - ) - if ProxyUpdateSpend.disable_spend_updates() is True: - return - if token is not None and isinstance(token, str) and token.startswith("sk-"): - hashed_token = hash_token(token=token) - else: - hashed_token = token - - ### UPDATE USER SPEND ### - async def _update_user_db(): - """ - - Update that user's row - - Update litellm-proxy-budget row (global proxy spend) - """ - ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db - existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id) - if existing_user_obj is not None and isinstance(existing_user_obj, dict): - existing_user_obj = LiteLLM_UserTable(**existing_user_obj) - try: - if prisma_client is not None: # update - user_ids = [user_id] - if ( - litellm.max_budget > 0 - ): # track global proxy budget, if user set max budget - user_ids.append(litellm_proxy_budget_name) - ### KEY CHANGE ### - for _id in user_ids: - if _id is not None: - prisma_client.user_list_transactons[_id] = ( - response_cost - + prisma_client.user_list_transactons.get(_id, 0) - ) - if end_user_id is not None: - prisma_client.end_user_list_transactons[end_user_id] = ( - response_cost - + prisma_client.end_user_list_transactons.get( - end_user_id, 0 - ) - ) - except Exception as e: - verbose_proxy_logger.info( - "\033[91m" - + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" - ) - - ### UPDATE KEY SPEND ### - async def _update_key_db(): - try: - verbose_proxy_logger.debug( - f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." - ) - if hashed_token is None: - return - if prisma_client is not None: - prisma_client.key_list_transactons[hashed_token] = ( - response_cost - + prisma_client.key_list_transactons.get(hashed_token, 0) - ) - except Exception as e: - verbose_proxy_logger.exception( - f"Update Key DB Call failed to execute - {str(e)}" - ) - raise e - - ### UPDATE SPEND LOGS ### - async def _insert_spend_log_to_db(): - try: - global prisma_client - if prisma_client is not None: - # Helper to generate payload to log - payload = get_logging_payload( - kwargs=kwargs, - response_obj=completion_response, - start_time=start_time, - end_time=end_time, - ) - payload["spend"] = response_cost - prisma_client = _set_spend_logs_payload( - payload=payload, - spend_logs_url=os.getenv("SPEND_LOGS_URL"), - prisma_client=prisma_client, - ) - except Exception as e: - verbose_proxy_logger.debug( - f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - ### UPDATE TEAM SPEND ### - async def _update_team_db(): - try: - verbose_proxy_logger.debug( - f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}." - ) - if team_id is None: - verbose_proxy_logger.debug( - "track_cost_callback: team_id is None. Not tracking spend for team" - ) - return - if prisma_client is not None: - prisma_client.team_list_transactons[team_id] = ( - response_cost - + prisma_client.team_list_transactons.get(team_id, 0) - ) - - try: - # Track spend of the team member within this team - # key is "team_id::::user_id::" - team_member_key = f"team_id::{team_id}::user_id::{user_id}" - prisma_client.team_member_list_transactons[team_member_key] = ( - response_cost - + prisma_client.team_member_list_transactons.get( - team_member_key, 0 - ) - ) - except Exception: - pass - except Exception as e: - verbose_proxy_logger.info( - f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - ### UPDATE ORG SPEND ### - async def _update_org_db(): - try: - verbose_proxy_logger.debug( - "adding spend to org db. Response cost: {}. org_id: {}.".format( - response_cost, org_id - ) - ) - if org_id is None: - verbose_proxy_logger.debug( - "track_cost_callback: org_id is None. Not tracking spend for org" - ) - return - if prisma_client is not None: - prisma_client.org_list_transactons[org_id] = ( - response_cost - + prisma_client.org_list_transactons.get(org_id, 0) - ) - except Exception as e: - verbose_proxy_logger.info( - f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - asyncio.create_task(_update_user_db()) - asyncio.create_task(_update_key_db()) - asyncio.create_task(_update_team_db()) - asyncio.create_task(_update_org_db()) - # asyncio.create_task(_insert_spend_log_to_db()) - if disable_spend_logs is False: - await _insert_spend_log_to_db() - else: - verbose_proxy_logger.info( - "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." - ) - - verbose_proxy_logger.debug("Runs spend update on all tables") - except Exception: - verbose_proxy_logger.debug( - f"Error updating Prisma database: {traceback.format_exc()}" - ) - - async def update_cache( # noqa: PLR0915 token: Optional[str], user_id: Optional[str], @@ -3294,14 +3089,14 @@ class ProxyStartupEvent: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=10, - args=[prisma_client], - ) - await proxy_config.get_credentials(prisma_client=prisma_client) + ### GET STORED CREDENTIALS ### + scheduler.add_job( + proxy_config.get_credentials, + "interval", + seconds=1, + args=[prisma_client], + ) + await proxy_config.get_credentials(prisma_client=prisma_client) if ( proxy_logging_obj is not None and proxy_logging_obj.slack_alerting_instance.alerting is not None diff --git a/tests/litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/litellm/proxy/hooks/test_proxy_track_cost_callback.py index 1e3b22ae2d..8850436329 100644 --- a/tests/litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -47,7 +47,8 @@ async def test_async_post_call_failure_hook(): # Mock update_database function with patch( - "litellm.proxy.proxy_server.update_database", new_callable=AsyncMock + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, ) as mock_update_database: # Call the method await logger.async_post_call_failure_hook( diff --git a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py index 535f5bf019..129be6d754 100644 --- a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py +++ b/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py @@ -5,6 +5,7 @@ from unittest.mock import Mock, patch, AsyncMock import pytest from fastapi import Request from litellm.proxy.utils import _get_redoc_url, _get_docs_url +from datetime import datetime sys.path.insert(0, os.path.abspath("../..")) import litellm @@ -22,16 +23,20 @@ async def test_disable_spend_logs(): with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): - from litellm.proxy.proxy_server import update_database + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter # Call update_database with disable_spend_logs=True - await update_database( + await DBSpendUpdateWriter.update_database( token="fake-token", response_cost=0.1, user_id="user123", completion_response=None, - start_time="2024-01-01", - end_time="2024-01-01", + start_time=datetime.now(), + end_time=datetime.now(), + end_user_id="end_user_id", + team_id="team_id", + org_id="org_id", + kwargs={}, ) # Verify no spend logs were added assert len(mock_prisma_client.spend_log_transactions) == 0 From 072be44a54a5cd37162f0cd2d0fb166d1d670cad Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 16:34:06 -0700 Subject: [PATCH 02/40] fix get_credentials job --- litellm/proxy/proxy_server.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6a2da7d83b..d7e62f98d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3089,14 +3089,14 @@ class ProxyStartupEvent: prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=1, - args=[prisma_client], - ) - await proxy_config.get_credentials(prisma_client=prisma_client) + ### GET STORED CREDENTIALS ### + scheduler.add_job( + proxy_config.get_credentials, + "interval", + seconds=10, + args=[prisma_client], + ) + await proxy_config.get_credentials(prisma_client=prisma_client) if ( proxy_logging_obj is not None and proxy_logging_obj.slack_alerting_instance.alerting is not None From a0fd508de405040e22b00cbb043e8bde0afd5e75 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 16:43:18 -0700 Subject: [PATCH 03/40] DBSpendUpdateWriter --- litellm/proxy/db/db_spend_update_writer.py | 3 +-- .../spend_tracking/test_spend_management_endpoints.py | 9 ++++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d5f2ef5228..af1d3294c5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -7,9 +7,8 @@ from typing import Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload -from litellm.proxy.proxy_server import hash_token from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend +from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token class DBSpendUpdateWriter: diff --git a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py index a5ee9ddf70..e08bce432d 100644 --- a/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -416,7 +416,8 @@ class TestSpendLogsPayload: # litellm._turn_on_debug() with patch.object( - litellm.proxy.proxy_server, "_set_spend_logs_payload" + litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter, + "_set_spend_logs_payload", ) as mock_client, patch.object(litellm.proxy.proxy_server, "prisma_client"): response = await litellm.acompletion( model="gpt-4o", @@ -509,7 +510,8 @@ class TestSpendLogsPayload: client = AsyncHTTPHandler() with patch.object( - litellm.proxy.proxy_server, "_set_spend_logs_payload" + litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter, + "_set_spend_logs_payload", ) as mock_client, patch.object( litellm.proxy.proxy_server, "prisma_client" ), patch.object( @@ -604,7 +606,8 @@ class TestSpendLogsPayload: ) with patch.object( - litellm.proxy.proxy_server, "_set_spend_logs_payload" + litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter, + "_set_spend_logs_payload", ) as mock_client, patch.object( litellm.proxy.proxy_server, "prisma_client" ), patch.object( From 7995fd7c98bf6fe3a058edfb60c64b05fc330622 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 16:44:46 -0700 Subject: [PATCH 04/40] fix DBSpendUpdateWriter --- litellm/proxy/db/db_spend_update_writer.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index af1d3294c5..ecafdcc3df 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1,3 +1,10 @@ +""" +Module responsible for + +1. Writing spend increments to either in memory list of transactions or to redis +2. Reading increments from redis or in memory list of transactions and committing them to db +""" + import asyncio import os import traceback @@ -12,6 +19,12 @@ from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token class DBSpendUpdateWriter: + """ + Module responsible for + + 1. Writing spend increments to either in memory list of transactions or to redis + 2. Reading increments from redis or in memory list of transactions and committing them to db + """ @staticmethod async def update_database( # noqa: PLR0915 From 403f2ef68dae3735fc61b5ada98b336ee5ab044c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 16:57:12 -0700 Subject: [PATCH 05/40] use simple static methods for updating spend --- litellm/proxy/db/db_spend_update_writer.py | 349 ++++++++++++--------- 1 file changed, 201 insertions(+), 148 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ecafdcc3df..37d709c0f1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,6 +13,7 @@ from typing import Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token @@ -27,7 +28,7 @@ class DBSpendUpdateWriter: """ @staticmethod - async def update_database( # noqa: PLR0915 + async def update_database( # LiteLLM management object fields token: Optional[str], user_id: Optional[str], @@ -59,155 +60,47 @@ class DBSpendUpdateWriter: else: hashed_token = token - ### UPDATE USER SPEND ### - async def _update_user_db(): - """ - - Update that user's row - - Update litellm-proxy-budget row (global proxy spend) - """ - ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db - existing_user_obj = await user_api_key_cache.async_get_cache( - key=user_id + asyncio.create_task( + DBSpendUpdateWriter._update_user_db( + response_cost=response_cost, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + end_user_id=end_user_id, ) - if existing_user_obj is not None and isinstance( - existing_user_obj, dict - ): - existing_user_obj = LiteLLM_UserTable(**existing_user_obj) - try: - if prisma_client is not None: # update - user_ids = [user_id] - if ( - litellm.max_budget > 0 - ): # track global proxy budget, if user set max budget - user_ids.append(litellm_proxy_budget_name) - ### KEY CHANGE ### - for _id in user_ids: - if _id is not None: - prisma_client.user_list_transactons[_id] = ( - response_cost - + prisma_client.user_list_transactons.get(_id, 0) - ) - if end_user_id is not None: - prisma_client.end_user_list_transactons[end_user_id] = ( - response_cost - + prisma_client.end_user_list_transactons.get( - end_user_id, 0 - ) - ) - except Exception as e: - verbose_proxy_logger.info( - "\033[91m" - + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" - ) - - ### UPDATE KEY SPEND ### - async def _update_key_db(): - try: - verbose_proxy_logger.debug( - f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." - ) - if hashed_token is None: - return - if prisma_client is not None: - prisma_client.key_list_transactons[hashed_token] = ( - response_cost - + prisma_client.key_list_transactons.get(hashed_token, 0) - ) - except Exception as e: - verbose_proxy_logger.exception( - f"Update Key DB Call failed to execute - {str(e)}" - ) - raise e - - ### UPDATE SPEND LOGS ### - async def _insert_spend_log_to_db(): - try: - if prisma_client: - payload = get_logging_payload( - kwargs=kwargs, - response_obj=completion_response, - start_time=start_time, - end_time=end_time, - ) - payload["spend"] = response_cost or 0.0 - DBSpendUpdateWriter._set_spend_logs_payload( - payload=payload, - spend_logs_url=os.getenv("SPEND_LOGS_URL"), - prisma_client=prisma_client, - ) - except Exception as e: - verbose_proxy_logger.debug( - f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - ### UPDATE TEAM SPEND ### - async def _update_team_db(): - try: - verbose_proxy_logger.debug( - f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}." - ) - if team_id is None: - verbose_proxy_logger.debug( - "track_cost_callback: team_id is None. Not tracking spend for team" - ) - return - if prisma_client is not None: - prisma_client.team_list_transactons[team_id] = ( - response_cost - + prisma_client.team_list_transactons.get(team_id, 0) - ) - - try: - # Track spend of the team member within this team - # key is "team_id::::user_id::" - team_member_key = f"team_id::{team_id}::user_id::{user_id}" - prisma_client.team_member_list_transactons[ - team_member_key - ] = ( - response_cost - + prisma_client.team_member_list_transactons.get( - team_member_key, 0 - ) - ) - except Exception: - pass - except Exception as e: - verbose_proxy_logger.info( - f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - ### UPDATE ORG SPEND ### - async def _update_org_db(): - try: - verbose_proxy_logger.debug( - "adding spend to org db. Response cost: {}. org_id: {}.".format( - response_cost, org_id - ) - ) - if org_id is None: - verbose_proxy_logger.debug( - "track_cost_callback: org_id is None. Not tracking spend for org" - ) - return - if prisma_client is not None: - prisma_client.org_list_transactons[org_id] = ( - response_cost - + prisma_client.org_list_transactons.get(org_id, 0) - ) - except Exception as e: - verbose_proxy_logger.info( - f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" - ) - raise e - - asyncio.create_task(_update_user_db()) - asyncio.create_task(_update_key_db()) - asyncio.create_task(_update_team_db()) - asyncio.create_task(_update_org_db()) + ) + asyncio.create_task( + DBSpendUpdateWriter._update_key_db( + response_cost=response_cost, + hashed_token=hashed_token, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + DBSpendUpdateWriter._update_team_db( + response_cost=response_cost, + team_id=team_id, + user_id=user_id, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + DBSpendUpdateWriter._update_org_db( + response_cost=response_cost, + org_id=org_id, + prisma_client=prisma_client, + ) + ) if disable_spend_logs is False: - await _insert_spend_log_to_db() + await DBSpendUpdateWriter._insert_spend_log_to_db( + kwargs=kwargs, + completion_response=completion_response, + start_time=start_time, + end_time=end_time, + response_cost=response_cost, + prisma_client=prisma_client, + ) else: verbose_proxy_logger.info( "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." @@ -219,6 +112,166 @@ class DBSpendUpdateWriter: f"Error updating Prisma database: {traceback.format_exc()}" ) + @staticmethod + async def _update_key_db( + response_cost: Optional[float], + hashed_token: Optional[str], + prisma_client: Optional[PrismaClient], + ): + try: + verbose_proxy_logger.debug( + f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." + ) + if hashed_token is None: + return + if prisma_client is not None: + prisma_client.key_list_transactons[hashed_token] = ( + response_cost + + prisma_client.key_list_transactons.get(hashed_token, 0) + ) + except Exception as e: + verbose_proxy_logger.exception( + f"Update Key DB Call failed to execute - {str(e)}" + ) + raise e + + @staticmethod + async def _update_user_db( + response_cost: Optional[float], + user_id: Optional[str], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + litellm_proxy_budget_name: Optional[str], + end_user_id: Optional[str] = None, + ): + """ + - Update that user's row + - Update litellm-proxy-budget row (global proxy spend) + """ + ## if an end-user is passed in, do an upsert - we can't guarantee they already exist in db + existing_user_obj = await user_api_key_cache.async_get_cache(key=user_id) + if existing_user_obj is not None and isinstance(existing_user_obj, dict): + existing_user_obj = LiteLLM_UserTable(**existing_user_obj) + try: + if prisma_client is not None: # update + user_ids = [user_id] + if ( + litellm.max_budget > 0 + ): # track global proxy budget, if user set max budget + user_ids.append(litellm_proxy_budget_name) + ### KEY CHANGE ### + for _id in user_ids: + if _id is not None: + prisma_client.user_list_transactons[_id] = ( + response_cost + + prisma_client.user_list_transactons.get(_id, 0) + ) + if end_user_id is not None: + prisma_client.end_user_list_transactons[end_user_id] = ( + response_cost + + prisma_client.end_user_list_transactons.get(end_user_id, 0) + ) + except Exception as e: + verbose_proxy_logger.info( + "\033[91m" + + f"Update User DB call failed to execute {str(e)}\n{traceback.format_exc()}" + ) + + @staticmethod + async def _update_team_db( + response_cost: Optional[float], + team_id: Optional[str], + user_id: Optional[str], + prisma_client: Optional[PrismaClient], + ): + try: + verbose_proxy_logger.debug( + f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}." + ) + if team_id is None: + verbose_proxy_logger.debug( + "track_cost_callback: team_id is None. Not tracking spend for team" + ) + return + if prisma_client is not None: + prisma_client.team_list_transactons[team_id] = ( + response_cost + prisma_client.team_list_transactons.get(team_id, 0) + ) + + try: + # Track spend of the team member within this team + # key is "team_id::::user_id::" + team_member_key = f"team_id::{team_id}::user_id::{user_id}" + prisma_client.team_member_list_transactons[team_member_key] = ( + response_cost + + prisma_client.team_member_list_transactons.get( + team_member_key, 0 + ) + ) + except Exception: + pass + except Exception as e: + verbose_proxy_logger.info( + f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + + @staticmethod + async def _update_org_db( + response_cost: Optional[float], + org_id: Optional[str], + prisma_client: Optional[PrismaClient], + ): + try: + verbose_proxy_logger.debug( + "adding spend to org db. Response cost: {}. org_id: {}.".format( + response_cost, org_id + ) + ) + if org_id is None: + verbose_proxy_logger.debug( + "track_cost_callback: org_id is None. Not tracking spend for org" + ) + return + if prisma_client is not None: + prisma_client.org_list_transactons[org_id] = ( + response_cost + prisma_client.org_list_transactons.get(org_id, 0) + ) + except Exception as e: + verbose_proxy_logger.info( + f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + + @staticmethod + async def _insert_spend_log_to_db( + kwargs: Optional[dict], + completion_response: Optional[Union[litellm.ModelResponse, Any, Exception]], + start_time: Optional[datetime], + end_time: Optional[datetime], + response_cost: Optional[float], + prisma_client: Optional[PrismaClient], + ): + try: + if prisma_client: + payload = get_logging_payload( + kwargs=kwargs, + response_obj=completion_response, + start_time=start_time, + end_time=end_time, + ) + payload["spend"] = response_cost or 0.0 + DBSpendUpdateWriter._set_spend_logs_payload( + payload=payload, + spend_logs_url=os.getenv("SPEND_LOGS_URL"), + prisma_client=prisma_client, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Update Spend Logs DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + @staticmethod def _set_spend_logs_payload( payload: Union[dict, SpendLogsPayload], From b721b2b4acfb84e4b5b73e488a33562b12759c75 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 17:04:50 -0700 Subject: [PATCH 06/40] use DBSpendUpdateWriter common function for --- litellm/proxy/_types.py | 15 +++ litellm/proxy/db/db_spend_update_writer.py | 142 ++++++++++++++------- 2 files changed, 111 insertions(+), 46 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 6e242ddacb..17b4acd138 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -144,6 +144,21 @@ class LitellmTableNames(str, enum.Enum): PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" +class Litellm_EntityType(enum.Enum): + """ + Enum for types of entities on litellm + + This enum allows specifying the type of entity that is being tracked in the database. + """ + + KEY = "key" + USER = "user" + END_USER = "end_user" + TEAM = "team" + TEAM_MEMBER = "team_member" + ORGANIZATION = "organization" + + def hash_token(token: str): import hashlib diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 37d709c0f1..34e8eae173 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,7 +14,7 @@ from typing import Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache -from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload +from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token @@ -112,6 +112,52 @@ class DBSpendUpdateWriter: f"Error updating Prisma database: {traceback.format_exc()}" ) + @staticmethod + async def _update_transaction_list( + response_cost: Optional[float], + entity_id: Optional[str], + transaction_list: dict, + entity_type: Litellm_EntityType, + debug_msg: Optional[str] = None, + ) -> bool: + """ + Common helper method to update a transaction list for an entity + + Args: + response_cost: The cost to add + entity_id: The ID of the entity to update + transaction_list: The transaction list dictionary to update + entity_type: The type of entity (from EntityType enum) + debug_msg: Optional custom debug message + + Returns: + bool: True if update happened, False otherwise + """ + try: + if debug_msg: + verbose_proxy_logger.debug(debug_msg) + else: + verbose_proxy_logger.debug( + f"adding spend to {entity_type.value} db. Response cost: {response_cost}. {entity_type.value}_id: {entity_id}." + ) + + if entity_id is None: + verbose_proxy_logger.debug( + f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}" + ) + return False + + transaction_list[entity_id] = response_cost + transaction_list.get( + entity_id, 0 + ) + return True + + except Exception as e: + verbose_proxy_logger.info( + f"Update {entity_type.value.capitalize()} DB failed to execute - {str(e)}\n{traceback.format_exc()}" + ) + raise e + @staticmethod async def _update_key_db( response_cost: Optional[float], @@ -119,16 +165,16 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient], ): try: - verbose_proxy_logger.debug( - f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}." - ) - if hashed_token is None: + if hashed_token is None or prisma_client is None: return - if prisma_client is not None: - prisma_client.key_list_transactons[hashed_token] = ( - response_cost - + prisma_client.key_list_transactons.get(hashed_token, 0) - ) + + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=hashed_token, + transaction_list=prisma_client.key_list_transactons, + entity_type=Litellm_EntityType.KEY, + debug_msg=f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}.", + ) except Exception as e: verbose_proxy_logger.exception( f"Update Key DB Call failed to execute - {str(e)}" @@ -159,17 +205,22 @@ class DBSpendUpdateWriter: litellm.max_budget > 0 ): # track global proxy budget, if user set max budget user_ids.append(litellm_proxy_budget_name) - ### KEY CHANGE ### + for _id in user_ids: if _id is not None: - prisma_client.user_list_transactons[_id] = ( - response_cost - + prisma_client.user_list_transactons.get(_id, 0) + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=_id, + transaction_list=prisma_client.user_list_transactons, + entity_type=Litellm_EntityType.USER, ) + if end_user_id is not None: - prisma_client.end_user_list_transactons[end_user_id] = ( - response_cost - + prisma_client.end_user_list_transactons.get(end_user_id, 0) + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=end_user_id, + transaction_list=prisma_client.end_user_list_transactons, + entity_type=Litellm_EntityType.END_USER, ) except Exception as e: verbose_proxy_logger.info( @@ -185,31 +236,32 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient], ): try: - verbose_proxy_logger.debug( - f"adding spend to team db. Response cost: {response_cost}. team_id: {team_id}." - ) - if team_id is None: + if team_id is None or prisma_client is None: verbose_proxy_logger.debug( - "track_cost_callback: team_id is None. Not tracking spend for team" + "track_cost_callback: team_id is None or prisma_client is None. Not tracking spend for team" ) return - if prisma_client is not None: - prisma_client.team_list_transactons[team_id] = ( - response_cost + prisma_client.team_list_transactons.get(team_id, 0) - ) - try: - # Track spend of the team member within this team + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=team_id, + transaction_list=prisma_client.team_list_transactons, + entity_type=Litellm_EntityType.TEAM, + ) + + try: + # Track spend of the team member within this team + if user_id is not None: # key is "team_id::::user_id::" team_member_key = f"team_id::{team_id}::user_id::{user_id}" - prisma_client.team_member_list_transactons[team_member_key] = ( - response_cost - + prisma_client.team_member_list_transactons.get( - team_member_key, 0 - ) + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=team_member_key, + transaction_list=prisma_client.team_member_list_transactons, + entity_type=Litellm_EntityType.TEAM_MEMBER, ) - except Exception: - pass + except Exception: + pass except Exception as e: verbose_proxy_logger.info( f"Update Team DB failed to execute - {str(e)}\n{traceback.format_exc()}" @@ -223,20 +275,18 @@ class DBSpendUpdateWriter: prisma_client: Optional[PrismaClient], ): try: - verbose_proxy_logger.debug( - "adding spend to org db. Response cost: {}. org_id: {}.".format( - response_cost, org_id - ) - ) - if org_id is None: + if org_id is None or prisma_client is None: verbose_proxy_logger.debug( - "track_cost_callback: org_id is None. Not tracking spend for org" + "track_cost_callback: org_id is None or prisma_client is None. Not tracking spend for org" ) return - if prisma_client is not None: - prisma_client.org_list_transactons[org_id] = ( - response_cost + prisma_client.org_list_transactons.get(org_id, 0) - ) + + await DBSpendUpdateWriter._update_transaction_list( + response_cost=response_cost, + entity_id=org_id, + transaction_list=prisma_client.org_list_transactons, + entity_type=Litellm_EntityType.ORGANIZATION, + ) except Exception as e: verbose_proxy_logger.info( f"Update Org DB failed to execute - {str(e)}\n{traceback.format_exc()}" From 894306141e4a34393a86fa5a45e8fccd45a44792 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 18:07:23 -0700 Subject: [PATCH 07/40] refactor, use commit_update_transactions_to_db --- litellm/proxy/db/db_spend_update_writer.py | 250 ++++++++++++++++++++- litellm/proxy/utils.py | 200 +---------------- 2 files changed, 252 insertions(+), 198 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 34e8eae173..cc45b6b96a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -7,16 +7,28 @@ Module responsible for import asyncio import os +import time import traceback -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache -from litellm.proxy._types import Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload +from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, + Litellm_EntityType, + LiteLLM_UserTable, + SpendLogsPayload, +) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token +from litellm.proxy.utils import ( + PrismaClient, + ProxyLogging, + ProxyUpdateSpend, + _raise_failed_update_spend_exception, + hash_token, +) class DBSpendUpdateWriter: @@ -346,3 +358,235 @@ class DBSpendUpdateWriter: payload.copy() ) return prisma_client + + @staticmethod + async def commit_update_transactions_to_db( # noqa: PLR0915 + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Handles commiting update spend transactions to db + + UPDATES can lead to deadlocks, hence we handle them separately + + Args: + prisma_client: PrismaClient object + n_retry_times: int, number of retry times + proxy_logging_obj: ProxyLogging object + """ + ### UPDATE USER TABLE ### + if len(prisma_client.user_list_transactons.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with prisma_client.db.tx( + timeout=timedelta(seconds=60) + ) as transaction: + async with transaction.batch_() as batcher: + for ( + user_id, + response_cost, + ) in prisma_client.user_list_transactons.items(): + batcher.litellm_usertable.update_many( + where={"user_id": user_id}, + data={"spend": {"increment": response_cost}}, + ) + prisma_client.user_list_transactons = ( + {} + ) # Clear the remaining transactions after processing all batches in the loop. + break + except DB_CONNECTION_ERROR_TYPES as e: + if ( + i >= n_retry_times + ): # If we've reached the maximum number of retries + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + # Optionally, sleep for a bit before retrying + await asyncio.sleep(2**i) # Exponential backoff + except Exception as e: + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) + + ### UPDATE END-USER TABLE ### + verbose_proxy_logger.debug( + "End-User Spend transactions: {}".format( + len(prisma_client.end_user_list_transactons.keys()) + ) + ) + if len(prisma_client.end_user_list_transactons.keys()) > 0: + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ### UPDATE KEY TABLE ### + verbose_proxy_logger.debug( + "KEY Spend transactions: {}".format( + len(prisma_client.key_list_transactons.keys()) + ) + ) + if len(prisma_client.key_list_transactons.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with prisma_client.db.tx( + timeout=timedelta(seconds=60) + ) as transaction: + async with transaction.batch_() as batcher: + for ( + token, + response_cost, + ) in prisma_client.key_list_transactons.items(): + batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists + where={"token": token}, + data={"spend": {"increment": response_cost}}, + ) + prisma_client.key_list_transactons = ( + {} + ) # Clear the remaining transactions after processing all batches in the loop. + break + except DB_CONNECTION_ERROR_TYPES as e: + if ( + i >= n_retry_times + ): # If we've reached the maximum number of retries + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + # Optionally, sleep for a bit before retrying + await asyncio.sleep(2**i) # Exponential backoff + except Exception as e: + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) + + ### UPDATE TEAM TABLE ### + verbose_proxy_logger.debug( + "Team Spend transactions: {}".format( + len(prisma_client.team_list_transactons.keys()) + ) + ) + if len(prisma_client.team_list_transactons.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with prisma_client.db.tx( + timeout=timedelta(seconds=60) + ) as transaction: + async with transaction.batch_() as batcher: + for ( + team_id, + response_cost, + ) in prisma_client.team_list_transactons.items(): + verbose_proxy_logger.debug( + "Updating spend for team id={} by {}".format( + team_id, response_cost + ) + ) + batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"team_id": team_id}, + data={"spend": {"increment": response_cost}}, + ) + prisma_client.team_list_transactons = ( + {} + ) # Clear the remaining transactions after processing all batches in the loop. + break + except DB_CONNECTION_ERROR_TYPES as e: + if ( + i >= n_retry_times + ): # If we've reached the maximum number of retries + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + # Optionally, sleep for a bit before retrying + await asyncio.sleep(2**i) # Exponential backoff + except Exception as e: + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) + + ### UPDATE TEAM Membership TABLE with spend ### + if len(prisma_client.team_member_list_transactons.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with prisma_client.db.tx( + timeout=timedelta(seconds=60) + ) as transaction: + async with transaction.batch_() as batcher: + for ( + key, + response_cost, + ) in prisma_client.team_member_list_transactons.items(): + # key is "team_id::::user_id::" + team_id = key.split("::")[1] + user_id = key.split("::")[3] + + batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists + where={"team_id": team_id, "user_id": user_id}, + data={"spend": {"increment": response_cost}}, + ) + prisma_client.team_member_list_transactons = ( + {} + ) # Clear the remaining transactions after processing all batches in the loop. + break + except DB_CONNECTION_ERROR_TYPES as e: + if ( + i >= n_retry_times + ): # If we've reached the maximum number of retries + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + # Optionally, sleep for a bit before retrying + await asyncio.sleep(2**i) # Exponential backoff + except Exception as e: + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) + + ### UPDATE ORG TABLE ### + if len(prisma_client.org_list_transactons.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with prisma_client.db.tx( + timeout=timedelta(seconds=60) + ) as transaction: + async with transaction.batch_() as batcher: + for ( + org_id, + response_cost, + ) in prisma_client.org_list_transactons.items(): + batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"organization_id": org_id}, + data={"spend": {"increment": response_cost}}, + ) + prisma_client.org_list_transactons = ( + {} + ) # Clear the remaining transactions after processing all batches in the loop. + break + except DB_CONNECTION_ERROR_TYPES as e: + if ( + i >= n_retry_times + ): # If we've reached the maximum number of retries + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + # Optionally, sleep for a bit before retrying + await asyncio.sleep(2**i) # Exponential backoff + except Exception as e: + _raise_failed_update_spend_exception( + e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj + ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7f1ac814a8..013ecf97c4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -62,6 +62,7 @@ from litellm.proxy.db.create_views import ( create_missing_views, should_create_missing_views, ) +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck @@ -2674,202 +2675,11 @@ async def update_spend( # noqa: PLR0915 spend_logs: list, """ n_retry_times = 3 - i = None - ### UPDATE USER TABLE ### - if len(prisma_client.user_list_transactons.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: - async with transaction.batch_() as batcher: - for ( - user_id, - response_cost, - ) in prisma_client.user_list_transactons.items(): - batcher.litellm_usertable.update_many( - where={"user_id": user_id}, - data={"spend": {"increment": response_cost}}, - ) - prisma_client.user_list_transactons = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. - break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff - except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - - ### UPDATE END-USER TABLE ### - verbose_proxy_logger.debug( - "End-User Spend transactions: {}".format( - len(prisma_client.end_user_list_transactons.keys()) - ) + await DBSpendUpdateWriter.commit_update_transactions_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, ) - if len(prisma_client.end_user_list_transactons.keys()) > 0: - await ProxyUpdateSpend.update_end_user_spend( - n_retry_times=n_retry_times, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - ) - ### UPDATE KEY TABLE ### - verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format( - len(prisma_client.key_list_transactons.keys()) - ) - ) - if len(prisma_client.key_list_transactons.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: - async with transaction.batch_() as batcher: - for ( - token, - response_cost, - ) in prisma_client.key_list_transactons.items(): - batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists - where={"token": token}, - data={"spend": {"increment": response_cost}}, - ) - prisma_client.key_list_transactons = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. - break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff - except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - - ### UPDATE TEAM TABLE ### - verbose_proxy_logger.debug( - "Team Spend transactions: {}".format( - len(prisma_client.team_list_transactons.keys()) - ) - ) - if len(prisma_client.team_list_transactons.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: - async with transaction.batch_() as batcher: - for ( - team_id, - response_cost, - ) in prisma_client.team_list_transactons.items(): - verbose_proxy_logger.debug( - "Updating spend for team id={} by {}".format( - team_id, response_cost - ) - ) - batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"team_id": team_id}, - data={"spend": {"increment": response_cost}}, - ) - prisma_client.team_list_transactons = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. - break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff - except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - - ### UPDATE TEAM Membership TABLE with spend ### - if len(prisma_client.team_member_list_transactons.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: - async with transaction.batch_() as batcher: - for ( - key, - response_cost, - ) in prisma_client.team_member_list_transactons.items(): - # key is "team_id::::user_id::" - team_id = key.split("::")[1] - user_id = key.split("::")[3] - - batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists - where={"team_id": team_id, "user_id": user_id}, - data={"spend": {"increment": response_cost}}, - ) - prisma_client.team_member_list_transactons = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. - break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff - except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - - ### UPDATE ORG TABLE ### - if len(prisma_client.org_list_transactons.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with prisma_client.db.tx( - timeout=timedelta(seconds=60) - ) as transaction: - async with transaction.batch_() as batcher: - for ( - org_id, - response_cost, - ) in prisma_client.org_list_transactons.items(): - batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"organization_id": org_id}, - data={"spend": {"increment": response_cost}}, - ) - prisma_client.org_list_transactons = ( - {} - ) # Clear the remaining transactions after processing all batches in the loop. - break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff - except Exception as e: - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) ### UPDATE SPEND LOGS ### verbose_proxy_logger.debug( From ad720781676e441b24eb7c772c1d853c23e9cd0a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 18:21:33 -0700 Subject: [PATCH 08/40] basic structure for commit update txs to redis --- litellm/proxy/db/db_spend_update_writer.py | 70 +++++++++++++++++++++- litellm/proxy/utils.py | 2 +- 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index cc45b6b96a..d85e69c472 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -29,6 +29,7 @@ from litellm.proxy.utils import ( _raise_failed_update_spend_exception, hash_token, ) +from litellm.secret_managers.main import str_to_bool class DBSpendUpdateWriter: @@ -360,7 +361,7 @@ class DBSpendUpdateWriter: return prisma_client @staticmethod - async def commit_update_transactions_to_db( # noqa: PLR0915 + async def db_spend_transaction_handler( prisma_client: PrismaClient, n_retry_times: int, proxy_logging_obj: ProxyLogging, @@ -368,12 +369,77 @@ class DBSpendUpdateWriter: """ Handles commiting update spend transactions to db - UPDATES can lead to deadlocks, hence we handle them separately + `UPDATES` can lead to deadlocks, hence we handle them separately Args: prisma_client: PrismaClient object n_retry_times: int, number of retry times proxy_logging_obj: ProxyLogging object + + How this works: + - Check `general_settings.use_redis_transaction_buffer` + - If enabled, write in-memory transactions to Redis + - Check if this Pod should read from the DB + else: + - Regular flow of this method + """ + if DBSpendUpdateWriter._should_commit_spend_updates_to_redis(): + pass + + if DBSpendUpdateWriter._should_commit_spend_updates_to_db(): + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) + + pass + + @staticmethod + def _should_commit_spend_updates_to_redis() -> bool: + """ + Checks if the Pod should commit spend updates to Redis + + This setting enables buffering database transactions in Redis + to improve reliability and reduce database contention + """ + from litellm.proxy.proxy_server import general_settings + + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) + if isinstance(_use_redis_transaction_buffer, str): + _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) + if _use_redis_transaction_buffer is None: + return False + return _use_redis_transaction_buffer + + @staticmethod + async def _commit_spend_updates_to_redis( + prisma_client: PrismaClient, + ): + """ + Commits all the spend updates to Redis for each entity type + + once committed, the transactions are cleared from the in-memory variables + """ + pass + + @staticmethod + def _should_commit_spend_updates_to_db() -> bool: + """ + Checks if the Pod should commit spend updates to the Database + """ + return False + + @staticmethod + async def _commit_spend_updates_to_db( # noqa: PLR0915 + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Commits all the spend updates to the Database """ ### UPDATE USER TABLE ### if len(prisma_client.user_list_transactons.keys()) > 0: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 013ecf97c4..435bf38c0e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2675,7 +2675,7 @@ async def update_spend( # noqa: PLR0915 spend_logs: list, """ n_retry_times = 3 - await DBSpendUpdateWriter.commit_update_transactions_to_db( + await DBSpendUpdateWriter._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, From 963791bbb58f8d36c8551c6b46c3572d89a213f3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 19:12:51 -0700 Subject: [PATCH 09/40] use redis update buffer class --- litellm/proxy/db/db_spend_update_writer.py | 68 +++++++------ litellm/proxy/db/redis_update_buffer.py | 110 +++++++++++++++++++++ litellm/proxy/utils.py | 5 +- 3 files changed, 147 insertions(+), 36 deletions(-) create mode 100644 litellm/proxy/db/redis_update_buffer.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d85e69c472..00f60a76f2 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -10,26 +10,24 @@ import os import time import traceback from datetime import datetime, timedelta -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache +from litellm.caching import DualCache, RedisCache, RedisClusterCache from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload, ) -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload -from litellm.proxy.utils import ( - PrismaClient, - ProxyLogging, - ProxyUpdateSpend, - _raise_failed_update_spend_exception, - hash_token, -) -from litellm.secret_managers.main import str_to_bool +from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient, ProxyLogging +else: + PrismaClient = Any + ProxyLogging = Any class DBSpendUpdateWriter: @@ -40,6 +38,12 @@ class DBSpendUpdateWriter: 2. Reading increments from redis or in memory list of transactions and committing them to db """ + def __init__( + self, redis_cache: Optional[Union[RedisCache, RedisClusterCache]] = None + ): + self.redis_cache = redis_cache + self.redis_update_buffer = RedisUpdateBuffer(redis_cache=redis_cache) + @staticmethod async def update_database( # LiteLLM management object fields @@ -61,6 +65,7 @@ class DBSpendUpdateWriter: prisma_client, user_api_key_cache, ) + from litellm.proxy.utils import ProxyUpdateSpend, hash_token try: verbose_proxy_logger.debug( @@ -315,6 +320,10 @@ class DBSpendUpdateWriter: response_cost: Optional[float], prisma_client: Optional[PrismaClient], ): + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_logging_payload, + ) + try: if prisma_client: payload = get_logging_payload( @@ -360,8 +369,8 @@ class DBSpendUpdateWriter: ) return prisma_client - @staticmethod - async def db_spend_transaction_handler( + async def db_update_spend_transaction_handler( + self, prisma_client: PrismaClient, n_retry_times: int, proxy_logging_obj: ProxyLogging, @@ -383,8 +392,10 @@ class DBSpendUpdateWriter: else: - Regular flow of this method """ - if DBSpendUpdateWriter._should_commit_spend_updates_to_redis(): - pass + if RedisUpdateBuffer._should_commit_spend_updates_to_redis(): + await self.redis_update_buffer.store_in_memory_spend_updates_in_redis( + prisma_client=prisma_client, + ) if DBSpendUpdateWriter._should_commit_spend_updates_to_db(): await DBSpendUpdateWriter._commit_spend_updates_to_db( @@ -395,25 +406,6 @@ class DBSpendUpdateWriter: pass - @staticmethod - def _should_commit_spend_updates_to_redis() -> bool: - """ - Checks if the Pod should commit spend updates to Redis - - This setting enables buffering database transactions in Redis - to improve reliability and reduce database contention - """ - from litellm.proxy.proxy_server import general_settings - - _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( - general_settings.get("use_redis_transaction_buffer", False) - ) - if isinstance(_use_redis_transaction_buffer, str): - _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) - if _use_redis_transaction_buffer is None: - return False - return _use_redis_transaction_buffer - @staticmethod async def _commit_spend_updates_to_redis( prisma_client: PrismaClient, @@ -439,8 +431,14 @@ class DBSpendUpdateWriter: proxy_logging_obj: ProxyLogging, ): """ - Commits all the spend updates to the Database + Commits all the spend `UPDATE` transactions to the Database + """ + from litellm.proxy.utils import ( + ProxyUpdateSpend, + _raise_failed_update_spend_exception, + ) + ### UPDATE USER TABLE ### if len(prisma_client.user_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py new file mode 100644 index 0000000000..843ec445a3 --- /dev/null +++ b/litellm/proxy/db/redis_update_buffer.py @@ -0,0 +1,110 @@ +""" +Handles buffering database `UPDATE` transactions in Redis before committing them to the database + +This is to prevent deadlocks and improve reliability +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union, cast + +from litellm.caching import RedisCache, RedisClusterCache +from litellm.secret_managers.main import str_to_bool + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient +else: + PrismaClient = Any + + +class DBSpendUpdateTransactions(TypedDict): + user_list_transactons: Dict[str, float] + end_user_list_transactons: Dict[str, float] + key_list_transactons: Dict[str, float] + team_list_transactons: Dict[str, float] + team_member_list_transactons: Dict[str, float] + org_list_transactons: Dict[str, float] + + +class RedisUpdateBuffer: + """ + Handles buffering database `UPDATE` transactions in Redis before committing them to the database + + This is to prevent deadlocks and improve reliability + """ + + def __init__( + self, redis_cache: Optional[Union[RedisCache, RedisClusterCache]] = None + ): + self.redis_cache = redis_cache + + @staticmethod + def _should_commit_spend_updates_to_redis() -> bool: + """ + Checks if the Pod should commit spend updates to Redis + + This setting enables buffering database transactions in Redis + to improve reliability and reduce database contention + """ + from litellm.proxy.proxy_server import general_settings + + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) + if isinstance(_use_redis_transaction_buffer, str): + _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) + if _use_redis_transaction_buffer is None: + return False + return _use_redis_transaction_buffer + + async def store_in_memory_spend_updates_in_redis( + self, + prisma_client: PrismaClient, + ): + """ + 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 + + ``` + { + "0929880201": 10, + "0929880202": 20, + "0929880203": 30, + } + ``` + """ + IN_MEMORY_UPDATE_TRANSACTIONS: DBSpendUpdateTransactions = ( + DBSpendUpdateTransactions( + user_list_transactons=prisma_client.user_list_transactons, + end_user_list_transactons=prisma_client.end_user_list_transactons, + key_list_transactons=prisma_client.key_list_transactons, + team_list_transactons=prisma_client.team_list_transactons, + team_member_list_transactons=prisma_client.team_member_list_transactons, + org_list_transactons=prisma_client.org_list_transactons, + ) + ) + for key, _transactions in IN_MEMORY_UPDATE_TRANSACTIONS.items(): + await self.increment_all_transaction_objects_in_redis( + key=key, + transactions=cast(Dict, _transactions), + ) + + async def increment_all_transaction_objects_in_redis( + self, + key: str, + transactions: Dict, + ): + """ + Increments all transaction objects in Redis + """ + if self.redis_cache is None: + return + for transaction_id, transaction_amount in transactions.items(): + await self.redis_cache.async_increment( + key=f"{key}:{transaction_id}", + value=transaction_amount, + ) + + async def get_all_update_transactions_from_redis(self): + pass diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 435bf38c0e..c208f60529 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -265,6 +265,9 @@ class ProxyLogging: ) self.premium_user = premium_user self.service_logging_obj = ServiceLogging() + self.db_spend_update_writer = DBSpendUpdateWriter( + redis_cache=self.internal_usage_cache.dual_cache.redis_cache + ) def startup_event( self, @@ -2675,7 +2678,7 @@ async def update_spend( # noqa: PLR0915 spend_logs: list, """ n_retry_times = 3 - await DBSpendUpdateWriter._commit_spend_updates_to_db( + await proxy_logging_obj.db_spend_update_writer.db_update_spend_transaction_handler( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, From e65b708fa2d2ff425143da10970be96c8bfc277a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 19:15:30 -0700 Subject: [PATCH 10/40] get_all_update_transactions_from_redis --- litellm/proxy/db/redis_update_buffer.py | 29 +++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 843ec445a3..d3591032ff 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -106,5 +106,30 @@ class RedisUpdateBuffer: value=transaction_amount, ) - async def get_all_update_transactions_from_redis(self): - pass + async def get_all_update_transactions_from_redis( + self, + ) -> Optional[DBSpendUpdateTransactions]: + """ + Gets all the update transactions from Redis + """ + if self.redis_cache is None: + return None + expected_keys = [ + "user_list_transactons", + "end_user_list_transactons", + "key_list_transactons", + "team_list_transactons", + "team_member_list_transactons", + "org_list_transactons", + ] + result = await self.redis_cache.async_batch_get_cache(expected_keys) + if result is None: + return None + return DBSpendUpdateTransactions( + user_list_transactons=result.get("user_list_transactons", {}), + end_user_list_transactons=result.get("end_user_list_transactons", {}), + key_list_transactons=result.get("key_list_transactons", {}), + team_list_transactons=result.get("team_list_transactons", {}), + team_member_list_transactons=result.get("team_member_list_transactons", {}), + org_list_transactons=result.get("org_list_transactons", {}), + ) From ea93a09b7a4632f6b5fcb47afdc4aa66df3ab944 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 21:10:26 -0700 Subject: [PATCH 11/40] add model CronJob --- litellm/proxy/schema.prisma | 15 +++++++++++++++ schema.prisma | 15 +++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9269e89014..0df675fd84 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -336,3 +336,18 @@ model LiteLLM_DailyUserSpend { @@index([api_key]) @@index([model]) } + + +// Track the status of cron jobs running. Only allow one pod to run the job at a time +model CronJob { + id String @id @default(cuid()) // Unique ID for the record + podId String // Unique identifier for the pod acting as the leader + status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) + lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record + ttl DateTime // Time when the leader's lease expires +} + +enum JobStatus { + ACTIVE + INACTIVE +} diff --git a/schema.prisma b/schema.prisma index 3312b26354..5d1535d2ff 100644 --- a/schema.prisma +++ b/schema.prisma @@ -335,3 +335,18 @@ model LiteLLM_DailyUserSpend { @@index([api_key]) @@index([model]) } + + +// Track the status of cron jobs running. Only allow one pod to run the job at a time +model CronJob { + id String @id @default(cuid()) // Unique ID for the record + podId String // Unique identifier for the pod acting as the leader + status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) + lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record + ttl DateTime // Time when the leader's lease expires +} + +enum JobStatus { + ACTIVE + INACTIVE +} From dd00429adecf49aeeca5f8ba7466a4bbb85ffbef Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 21:36:22 -0700 Subject: [PATCH 12/40] update schema --- litellm/constants.py | 4 ++++ litellm/proxy/schema.prisma | 1 + 2 files changed, 5 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index da66f897c9..0d5a9dcddf 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -441,3 +441,7 @@ HEALTH_CHECK_TIMEOUT_SECONDS = 60 # 60 seconds UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" + +########################### DB CRON JOB NAMES ########################### +DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" +DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = 600 # 5 minutes diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 0df675fd84..6e4937390a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -351,3 +351,4 @@ enum JobStatus { ACTIVE INACTIVE } + From aa6c34cb9757159e6e92a7155baf46157ddf99a5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 21:43:02 -0700 Subject: [PATCH 13/40] use lock manager for update spend job --- litellm/proxy/db/db_spend_update_writer.py | 37 +++---- litellm/proxy/db/pod_leader_manager.py | 111 +++++++++++++++++++++ 2 files changed, 127 insertions(+), 21 deletions(-) create mode 100644 litellm/proxy/db/pod_leader_manager.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 00f60a76f2..21717b8cbd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -15,12 +15,14 @@ from typing import TYPE_CHECKING, Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache, RedisCache, RedisClusterCache +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload, ) +from litellm.proxy.db.pod_leader_manager import PodLockManager from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer if TYPE_CHECKING: @@ -41,8 +43,14 @@ class DBSpendUpdateWriter: def __init__( self, redis_cache: Optional[Union[RedisCache, RedisClusterCache]] = None ): + from litellm.proxy.proxy_server import prisma_client + self.redis_cache = redis_cache self.redis_update_buffer = RedisUpdateBuffer(redis_cache=redis_cache) + self.pod_leader_manager = PodLockManager( + cronjob_id=DB_SPEND_UPDATE_JOB_NAME, + prisma_client=prisma_client, + ) @staticmethod async def update_database( @@ -397,33 +405,20 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, ) - if DBSpendUpdateWriter._should_commit_spend_updates_to_db(): + if await self.pod_leader_manager.acquire_lock(): + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) + await self.pod_leader_manager.release_lock() + else: await DBSpendUpdateWriter._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, ) - pass - - @staticmethod - async def _commit_spend_updates_to_redis( - prisma_client: PrismaClient, - ): - """ - Commits all the spend updates to Redis for each entity type - - once committed, the transactions are cleared from the in-memory variables - """ - pass - - @staticmethod - def _should_commit_spend_updates_to_db() -> bool: - """ - Checks if the Pod should commit spend updates to the Database - """ - return False - @staticmethod async def _commit_spend_updates_to_db( # noqa: PLR0915 prisma_client: PrismaClient, diff --git a/litellm/proxy/db/pod_leader_manager.py b/litellm/proxy/db/pod_leader_manager.py new file mode 100644 index 0000000000..73ecd68b1b --- /dev/null +++ b/litellm/proxy/db/pod_leader_manager.py @@ -0,0 +1,111 @@ +import uuid +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient, ProxyLogging +else: + PrismaClient = Any + ProxyLogging = Any + + +class PodLockManager: + """ + Manager for acquiring and releasing locks for cron jobs. + + Ensures that only one pod can run a cron job at a time. + """ + + def __init__(self, prisma_client: Optional[PrismaClient], cronjob_id: str): + self.pod_id = str(uuid.uuid4()) + self.prisma = prisma_client + self.cronjob_id = cronjob_id + + async def acquire_lock(self) -> bool: + """ + Attempt to acquire the lock for a specific cron job. + """ + if not self.prisma: + return False + try: + current_time = datetime.now(timezone.utc) + # Lease expiry time + ttl_expiry = current_time + timedelta( + seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + ) + + # Attempt to acquire the lock by upserting the record in the `cronjob_locks` table + cronjob_lock = await self.prisma.db.cronJob.upsert( + where={"cronjob_id": self.cronjob_id}, + create={ + "cronjob_id": self.cronjob_id, + "pod_id": self.pod_id, + "status": "ACTIVE", + "last_updated": current_time, + "ttl": ttl_expiry, + }, + update={ + "status": "ACTIVE", + "last_updated": current_time, + "ttl": ttl_expiry, + }, + ) + + if cronjob_lock.status == "ACTIVE" and cronjob_lock.pod_id == self.pod_id: + verbose_proxy_logger.debug( + f"Pod {self.pod_id} has acquired the lock for {self.cronjob_id}." + ) + return True # Lock successfully acquired + return False + except Exception as e: + verbose_proxy_logger.error( + f"Error acquiring the lock for {self.cronjob_id}: {e}" + ) + return False + + async def renew_lock(self): + """ + Renew the lock (update the TTL) for the pod holding the lock. + """ + if not self.prisma: + return False + try: + current_time = datetime.now(timezone.utc) + # Extend the TTL for another DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + ttl_expiry = current_time + timedelta( + seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + ) + + await self.prisma.db.cronJob.update( + where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, + data={"ttl": ttl_expiry, "last_updated": current_time}, + ) + verbose_proxy_logger.info( + f"Renewed the lock for Pod {self.pod_id} for {self.cronjob_id}" + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error renewing the lock for {self.cronjob_id}: {e}" + ) + + async def release_lock(self): + """ + Release the lock and mark the pod as inactive. + """ + if not self.prisma: + return False + try: + await self.prisma.db.cronJob.update( + where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, + data={"status": "INACTIVE"}, + ) + verbose_proxy_logger.info( + f"Pod {self.pod_id} has released the lock for {self.cronjob_id}." + ) + except Exception as e: + verbose_proxy_logger.error( + f"Error releasing the lock for {self.cronjob_id}: {e}" + ) From 91392305316200a0075e8d2c5a538f0bc8842fc9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 21:54:12 -0700 Subject: [PATCH 14/40] update the correct set of txs --- litellm/proxy/db/db_spend_update_writer.py | 64 +++++++++++++++------- 1 file changed, 45 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 21717b8cbd..c49d5aa989 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -23,7 +23,10 @@ from litellm.proxy._types import ( SpendLogsPayload, ) from litellm.proxy.db.pod_leader_manager import PodLockManager -from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.redis_update_buffer import ( + DBSpendUpdateTransactions, + RedisUpdateBuffer, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -406,17 +409,31 @@ class DBSpendUpdateWriter: ) if await self.pod_leader_manager.acquire_lock(): - await DBSpendUpdateWriter._commit_spend_updates_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions = ( + await self.redis_update_buffer.get_all_update_transactions_from_redis() ) + if db_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, + ) await self.pod_leader_manager.release_lock() else: + db_spend_update_transactions = DBSpendUpdateTransactions( + user_list_transactons=prisma_client.user_list_transactons, + end_user_list_transactons=prisma_client.end_user_list_transactons, + key_list_transactons=prisma_client.key_list_transactons, + team_list_transactons=prisma_client.team_list_transactons, + team_member_list_transactons=prisma_client.team_member_list_transactons, + org_list_transactons=prisma_client.org_list_transactons, + ) await DBSpendUpdateWriter._commit_spend_updates_to_db( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, ) @staticmethod @@ -424,6 +441,7 @@ class DBSpendUpdateWriter: prisma_client: PrismaClient, n_retry_times: int, proxy_logging_obj: ProxyLogging, + db_spend_update_transactions: DBSpendUpdateTransactions, ): """ Commits all the spend `UPDATE` transactions to the Database @@ -435,7 +453,8 @@ class DBSpendUpdateWriter: ) ### UPDATE USER TABLE ### - if len(prisma_client.user_list_transactons.keys()) > 0: + user_list_transactons = db_spend_update_transactions["user_list_transactons"] + if len(user_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -446,7 +465,7 @@ class DBSpendUpdateWriter: for ( user_id, response_cost, - ) in prisma_client.user_list_transactons.items(): + ) in user_list_transactons.items(): batcher.litellm_usertable.update_many( where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, @@ -477,19 +496,21 @@ class DBSpendUpdateWriter: len(prisma_client.end_user_list_transactons.keys()) ) ) - if len(prisma_client.end_user_list_transactons.keys()) > 0: + end_user_list_transactons = db_spend_update_transactions[ + "end_user_list_transactons" + ] + if len(end_user_list_transactons.keys()) > 0: await ProxyUpdateSpend.update_end_user_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) ### UPDATE KEY TABLE ### + key_list_transactons = db_spend_update_transactions["key_list_transactons"] verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format( - len(prisma_client.key_list_transactons.keys()) - ) + "KEY Spend transactions: {}".format(len(key_list_transactons.keys())) ) - if len(prisma_client.key_list_transactons.keys()) > 0: + if len(key_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -500,7 +521,7 @@ class DBSpendUpdateWriter: for ( token, response_cost, - ) in prisma_client.key_list_transactons.items(): + ) in key_list_transactons.items(): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={"spend": {"increment": response_cost}}, @@ -531,7 +552,8 @@ class DBSpendUpdateWriter: len(prisma_client.team_list_transactons.keys()) ) ) - if len(prisma_client.team_list_transactons.keys()) > 0: + team_list_transactons = db_spend_update_transactions["team_list_transactons"] + if len(team_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -542,7 +564,7 @@ class DBSpendUpdateWriter: for ( team_id, response_cost, - ) in prisma_client.team_list_transactons.items(): + ) in team_list_transactons.items(): verbose_proxy_logger.debug( "Updating spend for team id={} by {}".format( team_id, response_cost @@ -573,7 +595,10 @@ class DBSpendUpdateWriter: ) ### UPDATE TEAM Membership TABLE with spend ### - if len(prisma_client.team_member_list_transactons.keys()) > 0: + team_member_list_transactons = db_spend_update_transactions[ + "team_member_list_transactons" + ] + if len(team_member_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -584,7 +609,7 @@ class DBSpendUpdateWriter: for ( key, response_cost, - ) in prisma_client.team_member_list_transactons.items(): + ) in team_member_list_transactons.items(): # key is "team_id::::user_id::" team_id = key.split("::")[1] user_id = key.split("::")[3] @@ -614,7 +639,8 @@ class DBSpendUpdateWriter: ) ### UPDATE ORG TABLE ### - if len(prisma_client.org_list_transactons.keys()) > 0: + org_list_transactons = db_spend_update_transactions["org_list_transactons"] + if len(org_list_transactons.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -625,7 +651,7 @@ class DBSpendUpdateWriter: for ( org_id, response_cost, - ) in prisma_client.org_list_transactons.items(): + ) in org_list_transactons.items(): batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"organization_id": org_id}, data={"spend": {"increment": response_cost}}, From 1bfffadd052a87236ed4cd3cc3fdd4812c00c469 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 22:31:38 -0700 Subject: [PATCH 15/40] fix typo --- litellm/proxy/proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7e62f98d0..91d700b843 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1532,7 +1532,7 @@ class ProxyConfig: cache_params = {} if "cache_params" in litellm_settings: cache_params_in_config = litellm_settings["cache_params"] - # overwrie cache_params with cache_params_in_config + # overwrite cache_params with cache_params_in_config cache_params.update(cache_params_in_config) cache_type = cache_params.get("type", "redis") From fc46f6b86126b76d292d3e512ae079a1ebe8f032 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 22:34:15 -0700 Subject: [PATCH 16/40] fix db spend update buffer --- litellm/proxy/db/db_spend_update_writer.py | 8 +++++--- litellm/proxy/db/redis_update_buffer.py | 9 +++++++-- litellm/proxy/proxy_config.yaml | 7 ++++++- litellm/proxy/utils.py | 5 ++--- 4 files changed, 20 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c49d5aa989..0183fa433b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching import DualCache, RedisCache, RedisClusterCache +from litellm.caching import DualCache, RedisCache from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, @@ -44,12 +44,13 @@ class DBSpendUpdateWriter: """ def __init__( - self, redis_cache: Optional[Union[RedisCache, RedisClusterCache]] = None + self, + redis_cache: Optional[RedisCache] = None, ): from litellm.proxy.proxy_server import prisma_client self.redis_cache = redis_cache - self.redis_update_buffer = RedisUpdateBuffer(redis_cache=redis_cache) + self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_leader_manager = PodLockManager( cronjob_id=DB_SPEND_UPDATE_JOB_NAME, prisma_client=prisma_client, @@ -408,6 +409,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, ) + # Only commit from redis to db if this pod is the leader if await self.pod_leader_manager.acquire_lock(): db_spend_update_transactions = ( await self.redis_update_buffer.get_all_update_transactions_from_redis() diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index d3591032ff..ea7f14d321 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -6,7 +6,8 @@ This is to prevent deadlocks and improve reliability from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union, cast -from litellm.caching import RedisCache, RedisClusterCache +from litellm._logging import verbose_proxy_logger +from litellm.caching import RedisCache from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -32,7 +33,8 @@ class RedisUpdateBuffer: """ def __init__( - self, redis_cache: Optional[Union[RedisCache, RedisClusterCache]] = None + self, + redis_cache: Optional[RedisCache] = None, ): self.redis_cache = redis_cache @@ -99,6 +101,9 @@ class RedisUpdateBuffer: Increments all transaction objects in Redis """ if self.redis_cache is None: + verbose_proxy_logger.debug( + "redis_cache is None, skipping increment_all_transaction_objects_in_redis" + ) return for transaction_id, transaction_amount in transactions.items(): await self.redis_cache.async_increment( diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 4912a35f89..a1be54421b 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,4 +6,9 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ general_settings: - allow_requests_on_db_unavailable: True \ No newline at end of file + use_redis_transaction_buffer: True + +litellm_settings: + cache: true + cache_params: + type: redis diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c208f60529..040c6c14ef 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -265,9 +265,7 @@ class ProxyLogging: ) self.premium_user = premium_user self.service_logging_obj = ServiceLogging() - self.db_spend_update_writer = DBSpendUpdateWriter( - redis_cache=self.internal_usage_cache.dual_cache.redis_cache - ) + self.db_spend_update_writer = DBSpendUpdateWriter() def startup_event( self, @@ -340,6 +338,7 @@ class ProxyLogging: if redis_cache is not None: self.internal_usage_cache.dual_cache.redis_cache = redis_cache + self.db_spend_update_writer.redis_update_buffer.redis_cache = redis_cache def _init_litellm_callbacks(self, llm_router: Optional[Router] = None): litellm.logging_callback_manager.add_litellm_callback(self.max_parallel_request_limiter) # type: ignore From 758182fc7f66a96403013ffe0ba970bec635ea2d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 22:36:00 -0700 Subject: [PATCH 17/40] fix typo on codebase --- litellm/proxy/db/db_spend_update_writer.py | 78 ++++++++++----------- litellm/proxy/db/redis_update_buffer.py | 50 ++++++------- litellm/proxy/utils.py | 16 ++--- tests/local_testing/test_update_spend.py | 2 +- tests/proxy_unit_tests/test_proxy_utils.py | 4 +- tests/proxy_unit_tests/test_update_spend.py | 12 ++-- 6 files changed, 82 insertions(+), 80 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 0183fa433b..7184262ee5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -201,7 +201,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=hashed_token, - transaction_list=prisma_client.key_list_transactons, + transaction_list=prisma_client.key_list_transactions, entity_type=Litellm_EntityType.KEY, debug_msg=f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}.", ) @@ -241,7 +241,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=_id, - transaction_list=prisma_client.user_list_transactons, + transaction_list=prisma_client.user_list_transactions, entity_type=Litellm_EntityType.USER, ) @@ -249,7 +249,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=end_user_id, - transaction_list=prisma_client.end_user_list_transactons, + transaction_list=prisma_client.end_user_list_transactions, entity_type=Litellm_EntityType.END_USER, ) except Exception as e: @@ -275,7 +275,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=team_id, - transaction_list=prisma_client.team_list_transactons, + transaction_list=prisma_client.team_list_transactions, entity_type=Litellm_EntityType.TEAM, ) @@ -287,7 +287,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=team_member_key, - transaction_list=prisma_client.team_member_list_transactons, + transaction_list=prisma_client.team_member_list_transactions, entity_type=Litellm_EntityType.TEAM_MEMBER, ) except Exception: @@ -314,7 +314,7 @@ class DBSpendUpdateWriter: await DBSpendUpdateWriter._update_transaction_list( response_cost=response_cost, entity_id=org_id, - transaction_list=prisma_client.org_list_transactons, + transaction_list=prisma_client.org_list_transactions, entity_type=Litellm_EntityType.ORGANIZATION, ) except Exception as e: @@ -424,12 +424,12 @@ class DBSpendUpdateWriter: await self.pod_leader_manager.release_lock() else: db_spend_update_transactions = DBSpendUpdateTransactions( - user_list_transactons=prisma_client.user_list_transactons, - end_user_list_transactons=prisma_client.end_user_list_transactons, - key_list_transactons=prisma_client.key_list_transactons, - team_list_transactons=prisma_client.team_list_transactons, - team_member_list_transactons=prisma_client.team_member_list_transactons, - org_list_transactons=prisma_client.org_list_transactons, + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, ) await DBSpendUpdateWriter._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -455,8 +455,8 @@ class DBSpendUpdateWriter: ) ### UPDATE USER TABLE ### - user_list_transactons = db_spend_update_transactions["user_list_transactons"] - if len(user_list_transactons.keys()) > 0: + user_list_transactions = db_spend_update_transactions["user_list_transactions"] + if len(user_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -467,12 +467,12 @@ class DBSpendUpdateWriter: for ( user_id, response_cost, - ) in user_list_transactons.items(): + ) in user_list_transactions.items(): batcher.litellm_usertable.update_many( where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.user_list_transactons = ( + prisma_client.user_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -495,24 +495,24 @@ class DBSpendUpdateWriter: ### UPDATE END-USER TABLE ### verbose_proxy_logger.debug( "End-User Spend transactions: {}".format( - len(prisma_client.end_user_list_transactons.keys()) + len(prisma_client.end_user_list_transactions.keys()) ) ) - end_user_list_transactons = db_spend_update_transactions[ - "end_user_list_transactons" + end_user_list_transactions = db_spend_update_transactions[ + "end_user_list_transactions" ] - if len(end_user_list_transactons.keys()) > 0: + if len(end_user_list_transactions.keys()) > 0: await ProxyUpdateSpend.update_end_user_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, ) ### UPDATE KEY TABLE ### - key_list_transactons = db_spend_update_transactions["key_list_transactons"] + key_list_transactions = db_spend_update_transactions["key_list_transactions"] verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format(len(key_list_transactons.keys())) + "KEY Spend transactions: {}".format(len(key_list_transactions.keys())) ) - if len(key_list_transactons.keys()) > 0: + if len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -523,12 +523,12 @@ class DBSpendUpdateWriter: for ( token, response_cost, - ) in key_list_transactons.items(): + ) in key_list_transactions.items(): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, data={"spend": {"increment": response_cost}}, ) - prisma_client.key_list_transactons = ( + prisma_client.key_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -551,11 +551,11 @@ class DBSpendUpdateWriter: ### UPDATE TEAM TABLE ### verbose_proxy_logger.debug( "Team Spend transactions: {}".format( - len(prisma_client.team_list_transactons.keys()) + len(prisma_client.team_list_transactions.keys()) ) ) - team_list_transactons = db_spend_update_transactions["team_list_transactons"] - if len(team_list_transactons.keys()) > 0: + team_list_transactions = db_spend_update_transactions["team_list_transactions"] + if len(team_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -566,7 +566,7 @@ class DBSpendUpdateWriter: for ( team_id, response_cost, - ) in team_list_transactons.items(): + ) in team_list_transactions.items(): verbose_proxy_logger.debug( "Updating spend for team id={} by {}".format( team_id, response_cost @@ -576,7 +576,7 @@ class DBSpendUpdateWriter: where={"team_id": team_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_list_transactons = ( + prisma_client.team_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -597,10 +597,10 @@ class DBSpendUpdateWriter: ) ### UPDATE TEAM Membership TABLE with spend ### - team_member_list_transactons = db_spend_update_transactions[ - "team_member_list_transactons" + team_member_list_transactions = db_spend_update_transactions[ + "team_member_list_transactions" ] - if len(team_member_list_transactons.keys()) > 0: + if len(team_member_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -611,7 +611,7 @@ class DBSpendUpdateWriter: for ( key, response_cost, - ) in team_member_list_transactons.items(): + ) in team_member_list_transactions.items(): # key is "team_id::::user_id::" team_id = key.split("::")[1] user_id = key.split("::")[3] @@ -620,7 +620,7 @@ class DBSpendUpdateWriter: where={"team_id": team_id, "user_id": user_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.team_member_list_transactons = ( + prisma_client.team_member_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break @@ -641,8 +641,8 @@ class DBSpendUpdateWriter: ) ### UPDATE ORG TABLE ### - org_list_transactons = db_spend_update_transactions["org_list_transactons"] - if len(org_list_transactons.keys()) > 0: + org_list_transactions = db_spend_update_transactions["org_list_transactions"] + if len(org_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -653,12 +653,12 @@ class DBSpendUpdateWriter: for ( org_id, response_cost, - ) in org_list_transactons.items(): + ) in org_list_transactions.items(): batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists where={"organization_id": org_id}, data={"spend": {"increment": response_cost}}, ) - prisma_client.org_list_transactons = ( + prisma_client.org_list_transactions = ( {} ) # Clear the remaining transactions after processing all batches in the loop. break diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index ea7f14d321..22afc56483 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -17,12 +17,12 @@ else: class DBSpendUpdateTransactions(TypedDict): - user_list_transactons: Dict[str, float] - end_user_list_transactons: Dict[str, float] - key_list_transactons: Dict[str, float] - team_list_transactons: Dict[str, float] - team_member_list_transactons: Dict[str, float] - org_list_transactons: Dict[str, float] + user_list_transactions: Dict[str, float] + end_user_list_transactions: Dict[str, float] + key_list_transactions: Dict[str, float] + team_list_transactions: Dict[str, float] + team_member_list_transactions: Dict[str, float] + org_list_transactions: Dict[str, float] class RedisUpdateBuffer: @@ -78,12 +78,12 @@ class RedisUpdateBuffer: """ IN_MEMORY_UPDATE_TRANSACTIONS: DBSpendUpdateTransactions = ( DBSpendUpdateTransactions( - user_list_transactons=prisma_client.user_list_transactons, - end_user_list_transactons=prisma_client.end_user_list_transactons, - key_list_transactons=prisma_client.key_list_transactons, - team_list_transactons=prisma_client.team_list_transactons, - team_member_list_transactons=prisma_client.team_member_list_transactons, - org_list_transactons=prisma_client.org_list_transactons, + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, ) ) for key, _transactions in IN_MEMORY_UPDATE_TRANSACTIONS.items(): @@ -120,21 +120,23 @@ class RedisUpdateBuffer: if self.redis_cache is None: return None expected_keys = [ - "user_list_transactons", - "end_user_list_transactons", - "key_list_transactons", - "team_list_transactons", - "team_member_list_transactons", - "org_list_transactons", + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", ] result = await self.redis_cache.async_batch_get_cache(expected_keys) if result is None: return None return DBSpendUpdateTransactions( - user_list_transactons=result.get("user_list_transactons", {}), - end_user_list_transactons=result.get("end_user_list_transactons", {}), - key_list_transactons=result.get("key_list_transactons", {}), - team_list_transactons=result.get("team_list_transactons", {}), - team_member_list_transactons=result.get("team_member_list_transactons", {}), - org_list_transactons=result.get("org_list_transactons", {}), + user_list_transactions=result.get("user_list_transactions", {}), + end_user_list_transactions=result.get("end_user_list_transactions", {}), + key_list_transactions=result.get("key_list_transactions", {}), + team_list_transactions=result.get("team_list_transactions", {}), + team_member_list_transactions=result.get( + "team_member_list_transactions", {} + ), + org_list_transactions=result.get("org_list_transactions", {}), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 040c6c14ef..d0d4ea9b4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1101,12 +1101,12 @@ def jsonify_object(data: dict) -> dict: class PrismaClient: - user_list_transactons: dict = {} - end_user_list_transactons: dict = {} - key_list_transactons: dict = {} - team_list_transactons: dict = {} - team_member_list_transactons: dict = {} # key is ["team_id" + "user_id"] - org_list_transactons: dict = {} + user_list_transactions: dict = {} + end_user_list_transactions: dict = {} + key_list_transactions: dict = {} + team_list_transactions: dict = {} + team_member_list_transactions: dict = {} # key is ["team_id" + "user_id"] + org_list_transactions: dict = {} spend_log_transactions: List = [] daily_user_spend_transactions: Dict[str, DailyUserSpendTransaction] = {} @@ -2433,7 +2433,7 @@ class ProxyUpdateSpend: for ( end_user_id, response_cost, - ) in prisma_client.end_user_list_transactons.items(): + ) in prisma_client.end_user_list_transactions.items(): if litellm.max_end_user_budget is not None: pass batcher.litellm_endusertable.upsert( @@ -2461,7 +2461,7 @@ class ProxyUpdateSpend: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) finally: - prisma_client.end_user_list_transactons = ( + prisma_client.end_user_list_transactions = ( {} ) # reset the end user list transactions - prevent bad data from causing issues diff --git a/tests/local_testing/test_update_spend.py b/tests/local_testing/test_update_spend.py index 6aeae851ab..fffa3062d7 100644 --- a/tests/local_testing/test_update_spend.py +++ b/tests/local_testing/test_update_spend.py @@ -93,7 +93,7 @@ def prisma_client(): @pytest.mark.asyncio async def test_batch_update_spend(prisma_client): - prisma_client.user_list_transactons["test-litellm-user-5"] = 23 + prisma_client.user_list_transactions["test-litellm-user-5"] = 23 setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") await litellm.proxy.proxy_server.prisma_client.connect() diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index dccf0d1842..d613118fc8 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1509,7 +1509,7 @@ from litellm.proxy.utils import ProxyUpdateSpend async def test_end_user_transactions_reset(): # Setup mock_client = MagicMock() - mock_client.end_user_list_transactons = {"1": 10.0} # Bad log + mock_client.end_user_list_transactions = {"1": 10.0} # Bad log mock_client.db.tx = AsyncMock(side_effect=Exception("DB Error")) # Call function - should raise error @@ -1520,7 +1520,7 @@ async def test_end_user_transactions_reset(): # Verify cleanup happened assert ( - mock_client.end_user_list_transactons == {} + mock_client.end_user_list_transactions == {} ), "Transactions list should be empty after error" diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 6efc68a077..355cb8a403 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -22,12 +22,12 @@ class MockPrismaClient: def __init__(self): self.db = MagicMock() self.spend_log_transactions = [] - self.user_list_transactons = {} - self.end_user_list_transactons = {} - self.key_list_transactons = {} - self.team_list_transactons = {} - self.team_member_list_transactons = {} - self.org_list_transactons = {} + self.user_list_transactions = {} + self.end_user_list_transactions = {} + self.key_list_transactions = {} + self.team_list_transactions = {} + self.team_member_list_transactions = {} + self.org_list_transactions = {} def jsonify_object(self, obj): return obj From 0cb6cb6c6656b53d943644d21c0ee24acf832aec Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 22:54:46 -0700 Subject: [PATCH 18/40] fix schema to work with lock acquisition --- litellm/proxy/db/pod_leader_manager.py | 54 +++++++++++++++++--------- litellm/proxy/schema.prisma | 7 ++-- schema.prisma | 6 +-- 3 files changed, 42 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/db/pod_leader_manager.py b/litellm/proxy/db/pod_leader_manager.py index 73ecd68b1b..c762eabf66 100644 --- a/litellm/proxy/db/pod_leader_manager.py +++ b/litellm/proxy/db/pod_leader_manager.py @@ -19,16 +19,19 @@ class PodLockManager: Ensures that only one pod can run a cron job at a time. """ - def __init__(self, prisma_client: Optional[PrismaClient], cronjob_id: str): + def __init__(self, cronjob_id: str): self.pod_id = str(uuid.uuid4()) - self.prisma = prisma_client self.cronjob_id = cronjob_id async def acquire_lock(self) -> bool: """ Attempt to acquire the lock for a specific cron job. """ - if not self.prisma: + from litellm.proxy.proxy_server import prisma_client + + verbose_proxy_logger.debug("acquiring lock for cronjob_id=%s", self.cronjob_id) + if not prisma_client: + verbose_proxy_logger.debug("prisma is None, returning False") return False try: current_time = datetime.now(timezone.utc) @@ -38,21 +41,24 @@ class PodLockManager: ) # Attempt to acquire the lock by upserting the record in the `cronjob_locks` table - cronjob_lock = await self.prisma.db.cronJob.upsert( + cronjob_lock = await prisma_client.db.cronjob.upsert( where={"cronjob_id": self.cronjob_id}, - create={ - "cronjob_id": self.cronjob_id, - "pod_id": self.pod_id, - "status": "ACTIVE", - "last_updated": current_time, - "ttl": ttl_expiry, - }, - update={ - "status": "ACTIVE", - "last_updated": current_time, - "ttl": ttl_expiry, + data={ + "create": { + "cronjob_id": self.cronjob_id, + "pod_id": self.pod_id, + "status": "ACTIVE", + "last_updated": current_time, + "ttl": ttl_expiry, + }, + "update": { + "status": "ACTIVE", + "last_updated": current_time, + "ttl": ttl_expiry, + }, }, ) + verbose_proxy_logger.debug("cronjob_lock=%s", cronjob_lock) if cronjob_lock.status == "ACTIVE" and cronjob_lock.pod_id == self.pod_id: verbose_proxy_logger.debug( @@ -70,16 +76,21 @@ class PodLockManager: """ Renew the lock (update the TTL) for the pod holding the lock. """ - if not self.prisma: + from litellm.proxy.proxy_server import prisma_client + + if not prisma_client: return False try: + verbose_proxy_logger.debug( + "renewing lock for cronjob_id=%s", self.cronjob_id + ) current_time = datetime.now(timezone.utc) # Extend the TTL for another DEFAULT_CRON_JOB_LOCK_TTL_SECONDS ttl_expiry = current_time + timedelta( seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS ) - await self.prisma.db.cronJob.update( + await prisma_client.db.cronjob.update( where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, data={"ttl": ttl_expiry, "last_updated": current_time}, ) @@ -95,10 +106,15 @@ class PodLockManager: """ Release the lock and mark the pod as inactive. """ - if not self.prisma: + from litellm.proxy.proxy_server import prisma_client + + if not prisma_client: return False try: - await self.prisma.db.cronJob.update( + verbose_proxy_logger.debug( + "releasing lock for cronjob_id=%s", self.cronjob_id + ) + await prisma_client.db.cronjob.update( where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, data={"status": "INACTIVE"}, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6e4937390a..ccc22f8af1 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -340,10 +340,10 @@ model LiteLLM_DailyUserSpend { // Track the status of cron jobs running. Only allow one pod to run the job at a time model CronJob { - id String @id @default(cuid()) // Unique ID for the record - podId String // Unique identifier for the pod acting as the leader + cronjob_id String @id @default(cuid()) // Unique ID for the record + pod_id String // Unique identifier for the pod acting as the leader status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) - lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record + last_updated DateTime @default(now()) // Timestamp for the last update of the cron job record ttl DateTime // Time when the leader's lease expires } @@ -352,3 +352,4 @@ enum JobStatus { INACTIVE } + diff --git a/schema.prisma b/schema.prisma index 5d1535d2ff..d0de6e2e27 100644 --- a/schema.prisma +++ b/schema.prisma @@ -339,10 +339,10 @@ model LiteLLM_DailyUserSpend { // Track the status of cron jobs running. Only allow one pod to run the job at a time model CronJob { - id String @id @default(cuid()) // Unique ID for the record - podId String // Unique identifier for the pod acting as the leader + cronjob_id String @id @default(cuid()) // Unique ID for the record + pod_id String // Unique identifier for the pod acting as the leader status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) - lastUpdated DateTime @default(now()) // Timestamp for the last update of the cron job record + last_updated DateTime @default(now()) // Timestamp for the last update of the cron job record ttl DateTime // Time when the leader's lease expires } From 666cd8d4e336d4a625bcb4fd68e0c6ebe7f07bd5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 22:55:36 -0700 Subject: [PATCH 19/40] handle commit updates to the DB --- litellm/proxy/db/db_spend_update_writer.py | 59 +++++++++++++--------- litellm/proxy/db/redis_update_buffer.py | 12 ++--- 2 files changed, 42 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 7184262ee5..eee598f79a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -47,14 +47,9 @@ class DBSpendUpdateWriter: self, redis_cache: Optional[RedisCache] = None, ): - from litellm.proxy.proxy_server import prisma_client - self.redis_cache = redis_cache self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) - self.pod_leader_manager = PodLockManager( - cronjob_id=DB_SPEND_UPDATE_JOB_NAME, - prisma_client=prisma_client, - ) + self.pod_leader_manager = PodLockManager(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) @staticmethod async def update_database( @@ -411,17 +406,23 @@ class DBSpendUpdateWriter: # Only commit from redis to db if this pod is the leader if await self.pod_leader_manager.acquire_lock(): - db_spend_update_transactions = ( - await self.redis_update_buffer.get_all_update_transactions_from_redis() - ) - if db_spend_update_transactions is not None: - await DBSpendUpdateWriter._commit_spend_updates_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, - db_spend_update_transactions=db_spend_update_transactions, + verbose_proxy_logger.debug("acquired lock for spend updates") + + try: + db_spend_update_transactions = ( + await self.redis_update_buffer.get_all_update_transactions_from_redis() ) - await self.pod_leader_manager.release_lock() + if db_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, + ) + except Exception as e: + verbose_proxy_logger.error(f"Error committing spend updates: {e}") + finally: + await self.pod_leader_manager.release_lock() else: db_spend_update_transactions = DBSpendUpdateTransactions( user_list_transactions=prisma_client.user_list_transactions, @@ -456,7 +457,10 @@ class DBSpendUpdateWriter: ### UPDATE USER TABLE ### user_list_transactions = db_spend_update_transactions["user_list_transactions"] - if len(user_list_transactions.keys()) > 0: + if ( + user_list_transactions is not None + and len(user_list_transactions.keys()) > 0 + ): for i in range(n_retry_times + 1): start_time = time.time() try: @@ -501,7 +505,10 @@ class DBSpendUpdateWriter: end_user_list_transactions = db_spend_update_transactions[ "end_user_list_transactions" ] - if len(end_user_list_transactions.keys()) > 0: + if ( + end_user_list_transactions is not None + and len(end_user_list_transactions.keys()) > 0 + ): await ProxyUpdateSpend.update_end_user_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, @@ -510,9 +517,9 @@ class DBSpendUpdateWriter: ### UPDATE KEY TABLE ### key_list_transactions = db_spend_update_transactions["key_list_transactions"] verbose_proxy_logger.debug( - "KEY Spend transactions: {}".format(len(key_list_transactions.keys())) + "KEY Spend transactions: {}".format(key_list_transactions) ) - if len(key_list_transactions.keys()) > 0: + if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: @@ -555,7 +562,10 @@ class DBSpendUpdateWriter: ) ) team_list_transactions = db_spend_update_transactions["team_list_transactions"] - if len(team_list_transactions.keys()) > 0: + if ( + team_list_transactions is not None + and len(team_list_transactions.keys()) > 0 + ): for i in range(n_retry_times + 1): start_time = time.time() try: @@ -600,7 +610,10 @@ class DBSpendUpdateWriter: team_member_list_transactions = db_spend_update_transactions[ "team_member_list_transactions" ] - if len(team_member_list_transactions.keys()) > 0: + if ( + team_member_list_transactions is not None + and len(team_member_list_transactions.keys()) > 0 + ): for i in range(n_retry_times + 1): start_time = time.time() try: @@ -642,7 +655,7 @@ class DBSpendUpdateWriter: ### UPDATE ORG TABLE ### org_list_transactions = db_spend_update_transactions["org_list_transactions"] - if len(org_list_transactions.keys()) > 0: + if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 22afc56483..94cd1e47c7 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -17,12 +17,12 @@ else: class DBSpendUpdateTransactions(TypedDict): - user_list_transactions: Dict[str, float] - end_user_list_transactions: Dict[str, float] - key_list_transactions: Dict[str, float] - team_list_transactions: Dict[str, float] - team_member_list_transactions: Dict[str, float] - org_list_transactions: Dict[str, float] + user_list_transactions: Optional[Dict[str, float]] + end_user_list_transactions: Optional[Dict[str, float]] + key_list_transactions: Optional[Dict[str, float]] + team_list_transactions: Optional[Dict[str, float]] + team_member_list_transactions: Optional[Dict[str, float]] + org_list_transactions: Optional[Dict[str, float]] class RedisUpdateBuffer: From 1ba73f491a2b07eb0a92be6d97848c1c2f723bca Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 23:13:01 -0700 Subject: [PATCH 20/40] working cron job updates --- litellm/constants.py | 2 +- litellm/proxy/db/pod_leader_manager.py | 6 +- litellm/proxy/db/redis_update_buffer.py | 90 +++++++++++++++++++------ litellm/proxy/schema.prisma | 2 +- schema.prisma | 2 +- 5 files changed, 77 insertions(+), 25 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 0d5a9dcddf..1276c08fae 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -444,4 +444,4 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id" ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" -DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = 600 # 5 minutes +DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = 60 # 1 minute diff --git a/litellm/proxy/db/pod_leader_manager.py b/litellm/proxy/db/pod_leader_manager.py index c762eabf66..9c055e96f6 100644 --- a/litellm/proxy/db/pod_leader_manager.py +++ b/litellm/proxy/db/pod_leader_manager.py @@ -41,7 +41,7 @@ class PodLockManager: ) # Attempt to acquire the lock by upserting the record in the `cronjob_locks` table - cronjob_lock = await prisma_client.db.cronjob.upsert( + cronjob_lock = await prisma_client.db.litellm_cronjob.upsert( where={"cronjob_id": self.cronjob_id}, data={ "create": { @@ -90,7 +90,7 @@ class PodLockManager: seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS ) - await prisma_client.db.cronjob.update( + await prisma_client.db.litellm_cronjob.update( where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, data={"ttl": ttl_expiry, "last_updated": current_time}, ) @@ -114,7 +114,7 @@ class PodLockManager: verbose_proxy_logger.debug( "releasing lock for cronjob_id=%s", self.cronjob_id ) - await prisma_client.db.cronjob.update( + await prisma_client.db.litellm_cronjob.update( where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, data={"status": "INACTIVE"}, ) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 94cd1e47c7..2207911024 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -111,6 +111,13 @@ class RedisUpdateBuffer: value=transaction_amount, ) + @staticmethod + def _remove_prefix_from_keys(data: Dict[str, Any], prefix: str) -> Dict[str, Any]: + """ + Removes the specified prefix from the keys of a dictionary. + """ + return {key.replace(prefix, "", 1): value for key, value in data.items()} + async def get_all_update_transactions_from_redis( self, ) -> Optional[DBSpendUpdateTransactions]: @@ -119,24 +126,69 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - expected_keys = [ - "user_list_transactions", - "end_user_list_transactions", - "key_list_transactions", - "team_list_transactions", - "team_member_list_transactions", - "org_list_transactions", - ] - result = await self.redis_cache.async_batch_get_cache(expected_keys) - if result is None: - return None + user_transaction_keys = await self.redis_cache.async_scan_iter( + "user_list_transactions:*" + ) + end_user_transaction_keys = await self.redis_cache.async_scan_iter( + "end_user_list_transactions:*" + ) + key_transaction_keys = await self.redis_cache.async_scan_iter( + "key_list_transactions:*" + ) + team_transaction_keys = await self.redis_cache.async_scan_iter( + "team_list_transactions:*" + ) + team_member_transaction_keys = await self.redis_cache.async_scan_iter( + "team_member_list_transactions:*" + ) + org_transaction_keys = await self.redis_cache.async_scan_iter( + "org_list_transactions:*" + ) + + user_list_transactions = await self.redis_cache.async_batch_get_cache( + user_transaction_keys + ) + end_user_list_transactions = await self.redis_cache.async_batch_get_cache( + end_user_transaction_keys + ) + key_list_transactions = await self.redis_cache.async_batch_get_cache( + key_transaction_keys + ) + team_list_transactions = await self.redis_cache.async_batch_get_cache( + team_transaction_keys + ) + team_member_list_transactions = await self.redis_cache.async_batch_get_cache( + team_member_transaction_keys + ) + org_list_transactions = await self.redis_cache.async_batch_get_cache( + org_transaction_keys + ) + + # filter out the "prefix" from the keys using the helper method + user_list_transactions = self._remove_prefix_from_keys( + user_list_transactions, "user_list_transactions:" + ) + end_user_list_transactions = self._remove_prefix_from_keys( + end_user_list_transactions, "end_user_list_transactions:" + ) + key_list_transactions = self._remove_prefix_from_keys( + key_list_transactions, "key_list_transactions:" + ) + team_list_transactions = self._remove_prefix_from_keys( + team_list_transactions, "team_list_transactions:" + ) + team_member_list_transactions = self._remove_prefix_from_keys( + team_member_list_transactions, "team_member_list_transactions:" + ) + org_list_transactions = self._remove_prefix_from_keys( + org_list_transactions, "org_list_transactions:" + ) + return DBSpendUpdateTransactions( - user_list_transactions=result.get("user_list_transactions", {}), - end_user_list_transactions=result.get("end_user_list_transactions", {}), - key_list_transactions=result.get("key_list_transactions", {}), - team_list_transactions=result.get("team_list_transactions", {}), - team_member_list_transactions=result.get( - "team_member_list_transactions", {} - ), - org_list_transactions=result.get("org_list_transactions", {}), + user_list_transactions=user_list_transactions, + end_user_list_transactions=end_user_list_transactions, + key_list_transactions=key_list_transactions, + team_list_transactions=team_list_transactions, + team_member_list_transactions=team_member_list_transactions, + org_list_transactions=org_list_transactions, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index ccc22f8af1..ae7560ead3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -339,7 +339,7 @@ model LiteLLM_DailyUserSpend { // Track the status of cron jobs running. Only allow one pod to run the job at a time -model CronJob { +model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record pod_id String // Unique identifier for the pod acting as the leader status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) diff --git a/schema.prisma b/schema.prisma index d0de6e2e27..ad46f96f6d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -338,7 +338,7 @@ model LiteLLM_DailyUserSpend { // Track the status of cron jobs running. Only allow one pod to run the job at a time -model CronJob { +model LiteLLM_CronJob { cronjob_id String @id @default(cuid()) // Unique ID for the record pod_id String // Unique identifier for the pod acting as the leader status JobStatus @default(INACTIVE) // Status of the cron job (active or inactive) From 0edd4aa8a7f471a6e57e134932b4db717e3a42f8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 27 Mar 2025 23:29:48 -0700 Subject: [PATCH 21/40] fix PodLockManager --- litellm/proxy/db/pod_leader_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/db/pod_leader_manager.py b/litellm/proxy/db/pod_leader_manager.py index 9c055e96f6..9b739b8695 100644 --- a/litellm/proxy/db/pod_leader_manager.py +++ b/litellm/proxy/db/pod_leader_manager.py @@ -1,6 +1,6 @@ import uuid from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS From c53d172b062bc1cbbb5b4097b167d925fb73a19a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 12:57:00 -0700 Subject: [PATCH 22/40] rename pod lock manager --- litellm/proxy/db/db_spend_update_writer.py | 8 +- ..._leader_manager.py => pod_lock_manager.py} | 0 .../litellm/proxy/db/test_pod_lock_manager.py | 127 ++++++++++++++++++ 3 files changed, 131 insertions(+), 4 deletions(-) rename litellm/proxy/db/{pod_leader_manager.py => pod_lock_manager.py} (100%) create mode 100644 tests/litellm/proxy/db/test_pod_lock_manager.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index eee598f79a..9047bd4c38 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -22,7 +22,7 @@ from litellm.proxy._types import ( LiteLLM_UserTable, SpendLogsPayload, ) -from litellm.proxy.db.pod_leader_manager import PodLockManager +from litellm.proxy.db.pod_lock_manager import PodLockManager from litellm.proxy.db.redis_update_buffer import ( DBSpendUpdateTransactions, RedisUpdateBuffer, @@ -49,7 +49,7 @@ class DBSpendUpdateWriter: ): self.redis_cache = redis_cache self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) - self.pod_leader_manager = PodLockManager(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) + self.pod_lock_manager = PodLockManager(cronjob_id=DB_SPEND_UPDATE_JOB_NAME) @staticmethod async def update_database( @@ -405,7 +405,7 @@ class DBSpendUpdateWriter: ) # Only commit from redis to db if this pod is the leader - if await self.pod_leader_manager.acquire_lock(): + if await self.pod_lock_manager.acquire_lock(): verbose_proxy_logger.debug("acquired lock for spend updates") try: @@ -422,7 +422,7 @@ class DBSpendUpdateWriter: except Exception as e: verbose_proxy_logger.error(f"Error committing spend updates: {e}") finally: - await self.pod_leader_manager.release_lock() + await self.pod_lock_manager.release_lock() else: db_spend_update_transactions = DBSpendUpdateTransactions( user_list_transactions=prisma_client.user_list_transactions, diff --git a/litellm/proxy/db/pod_leader_manager.py b/litellm/proxy/db/pod_lock_manager.py similarity index 100% rename from litellm/proxy/db/pod_leader_manager.py rename to litellm/proxy/db/pod_lock_manager.py diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/test_pod_lock_manager.py new file mode 100644 index 0000000000..8894ea6716 --- /dev/null +++ b/tests/litellm/proxy/db/test_pod_lock_manager.py @@ -0,0 +1,127 @@ +import json +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # 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 + + +# Mock Prisma client class +class MockPrismaClient: + def __init__(self): + self.db = MagicMock() + self.db.litellm_cronjob = AsyncMock() + + +@pytest.fixture +def mock_prisma(monkeypatch): + mock_client = MockPrismaClient() + + # Mock the prisma_client import in proxy_server + def mock_get_prisma(): + return mock_client + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client) + return mock_client + + +@pytest.fixture +def pod_lock_manager(): + return PodLockManager(cronjob_id="test_job") + + +@pytest.mark.asyncio +async def test_acquire_lock_success(pod_lock_manager, mock_prisma): + # Mock successful lock acquisition + mock_response = AsyncMock() + mock_response.status = "ACTIVE" + mock_response.pod_id = pod_lock_manager.pod_id + mock_prisma.db.litellm_cronjob.upsert.return_value = mock_response + + result = await pod_lock_manager.acquire_lock() + assert result == True + + # Verify upsert was called with correct parameters + mock_prisma.db.litellm_cronjob.upsert.assert_called_once() + call_args = mock_prisma.db.litellm_cronjob.upsert.call_args[1] + assert call_args["where"]["cronjob_id"] == "test_job" + assert "create" in call_args["data"] + assert "update" in call_args["data"] + + +@pytest.mark.asyncio +async def test_acquire_lock_failure(pod_lock_manager, mock_prisma): + """ + Test that the lock is not acquired if the lock is held by a different pod + """ + # Mock failed lock acquisition (different pod holds the lock) + mock_response = AsyncMock() + mock_response.status = "ACTIVE" + mock_response.pod_id = "different_pod_id" + mock_prisma.db.litellm_cronjob.upsert.return_value = mock_response + + result = await pod_lock_manager.acquire_lock() + assert result == False + + +@pytest.mark.asyncio +async def test_renew_lock(pod_lock_manager, mock_prisma): + # Mock successful lock renewal + mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() + + await pod_lock_manager.renew_lock() + + # Verify update was called with correct parameters + mock_prisma.db.litellm_cronjob.update.assert_called_once() + call_args = mock_prisma.db.litellm_cronjob.update.call_args[1] + assert call_args["where"]["cronjob_id"] == "test_job" + assert call_args["where"]["pod_id"] == pod_lock_manager.pod_id + assert "ttl" in call_args["data"] + assert "last_updated" in call_args["data"] + + +@pytest.mark.asyncio +async def test_release_lock(pod_lock_manager, mock_prisma): + # Mock successful lock release + mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() + + await pod_lock_manager.release_lock() + + # Verify update was called with correct parameters + mock_prisma.db.litellm_cronjob.update.assert_called_once() + call_args = mock_prisma.db.litellm_cronjob.update.call_args[1] + assert call_args["where"]["cronjob_id"] == "test_job" + assert call_args["where"]["pod_id"] == pod_lock_manager.pod_id + assert call_args["data"]["status"] == "INACTIVE" + + +@pytest.mark.asyncio +async def test_prisma_client_none(pod_lock_manager, monkeypatch): + # Mock prisma_client as None + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + # Test all methods with None client + assert await pod_lock_manager.acquire_lock() == False + assert await pod_lock_manager.renew_lock() == False + assert await pod_lock_manager.release_lock() == False + + +@pytest.mark.asyncio +async def test_database_error_handling(pod_lock_manager, mock_prisma): + # Mock database errors + mock_prisma.db.litellm_cronjob.upsert.side_effect = Exception("Database error") + mock_prisma.db.litellm_cronjob.update.side_effect = Exception("Database error") + + # Test error handling in all methods + assert await pod_lock_manager.acquire_lock() == False + await pod_lock_manager.renew_lock() # Should not raise exception + await pod_lock_manager.release_lock() # Should not raise exception From 021eedaf693ff34bfe8d824bce48f4f76bbbdee6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 12:59:16 -0700 Subject: [PATCH 23/40] test pod lock manager --- tests/litellm/proxy/db/test_pod_lock_manager.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/test_pod_lock_manager.py index 8894ea6716..e685172c1e 100644 --- a/tests/litellm/proxy/db/test_pod_lock_manager.py +++ b/tests/litellm/proxy/db/test_pod_lock_manager.py @@ -41,7 +41,9 @@ def pod_lock_manager(): @pytest.mark.asyncio async def test_acquire_lock_success(pod_lock_manager, mock_prisma): - # Mock successful lock acquisition + """ + Test that the lock is acquired successfully if the DB response is successful + """ mock_response = AsyncMock() mock_response.status = "ACTIVE" mock_response.pod_id = pod_lock_manager.pod_id @@ -75,7 +77,9 @@ async def test_acquire_lock_failure(pod_lock_manager, mock_prisma): @pytest.mark.asyncio async def test_renew_lock(pod_lock_manager, mock_prisma): - # Mock successful lock renewal + """ + Test that the renew lock calls the DB update method with the correct parameters + """ mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() await pod_lock_manager.renew_lock() @@ -91,7 +95,11 @@ async def test_renew_lock(pod_lock_manager, mock_prisma): @pytest.mark.asyncio async def test_release_lock(pod_lock_manager, mock_prisma): - # Mock successful lock release + """ + Test that the release lock calls the DB update method with the correct parameters + + specifically, the status should be set to INACTIVE + """ mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() await pod_lock_manager.release_lock() From 1eaf847f8a7801fe93a0b2fef3b5f03b5c0c5326 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 13:31:45 -0700 Subject: [PATCH 24/40] test pod lock manager --- litellm/proxy/db/pod_lock_manager.py | 50 ++- .../test_e2e_pod_lock_manager.py | 354 ++++++++++++++++++ 2 files changed, 384 insertions(+), 20 deletions(-) create mode 100644 tests/proxy_unit_tests/test_e2e_pod_lock_manager.py diff --git a/litellm/proxy/db/pod_lock_manager.py b/litellm/proxy/db/pod_lock_manager.py index 9b739b8695..84c92c9daa 100644 --- a/litellm/proxy/db/pod_lock_manager.py +++ b/litellm/proxy/db/pod_lock_manager.py @@ -25,47 +25,57 @@ class PodLockManager: async def acquire_lock(self) -> bool: """ - Attempt to acquire the lock for a specific cron job. + Attempt to acquire the lock for a specific cron job using database locking. """ from litellm.proxy.proxy_server import prisma_client - verbose_proxy_logger.debug("acquiring lock for cronjob_id=%s", self.cronjob_id) + verbose_proxy_logger.debug( + "Pod %s acquiring lock for cronjob_id=%s", self.pod_id, self.cronjob_id + ) if not prisma_client: verbose_proxy_logger.debug("prisma is None, returning False") return False + try: current_time = datetime.now(timezone.utc) - # Lease expiry time ttl_expiry = current_time + timedelta( seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS ) - # Attempt to acquire the lock by upserting the record in the `cronjob_locks` table - cronjob_lock = await prisma_client.db.litellm_cronjob.upsert( + # Use Prisma's findUnique with FOR UPDATE lock to prevent race conditions + lock_record = await prisma_client.db.litellm_cronjob.find_unique( where={"cronjob_id": self.cronjob_id}, - data={ - "create": { - "cronjob_id": self.cronjob_id, + ) + + if lock_record: + # If record exists, only update if it's inactive or expired + if lock_record.status == "ACTIVE" and lock_record.ttl > current_time: + return lock_record.pod_id == self.pod_id + + # Update existing record + updated_lock = await prisma_client.db.litellm_cronjob.update( + where={"cronjob_id": self.cronjob_id}, + data={ "pod_id": self.pod_id, "status": "ACTIVE", "last_updated": current_time, "ttl": ttl_expiry, }, - "update": { + ) + else: + # Create new record if none exists + updated_lock = await prisma_client.db.litellm_cronjob.create( + data={ + "cronjob_id": self.cronjob_id, + "pod_id": self.pod_id, "status": "ACTIVE", "last_updated": current_time, "ttl": ttl_expiry, - }, - }, - ) - verbose_proxy_logger.debug("cronjob_lock=%s", cronjob_lock) - - if cronjob_lock.status == "ACTIVE" and cronjob_lock.pod_id == self.pod_id: - verbose_proxy_logger.debug( - f"Pod {self.pod_id} has acquired the lock for {self.cronjob_id}." + } ) - return True # Lock successfully acquired - return False + + return updated_lock.pod_id == self.pod_id + except Exception as e: verbose_proxy_logger.error( f"Error acquiring the lock for {self.cronjob_id}: {e}" @@ -112,7 +122,7 @@ class PodLockManager: return False try: verbose_proxy_logger.debug( - "releasing lock for cronjob_id=%s", self.cronjob_id + "Pod %s releasing lock for cronjob_id=%s", self.pod_id, self.cronjob_id ) await prisma_client.db.litellm_cronjob.update( where={"cronjob_id": self.cronjob_id, "pod_id": self.pod_id}, diff --git a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py new file mode 100644 index 0000000000..7d36bb4791 --- /dev/null +++ b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py @@ -0,0 +1,354 @@ +import os +import sys +import traceback +import uuid +from datetime import datetime, timezone, timedelta + +from dotenv import load_dotenv +from fastapi import Request +from fastapi.routing import APIRoute +import httpx + +load_dotenv() +import io +import os +import time + +# this file is to test litellm/proxy + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import logging + +import pytest +from litellm.proxy.db.pod_lock_manager import PodLockManager +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + new_user, + user_info, + user_update, +) +from litellm.proxy.auth.auth_checks import get_key_object +from litellm.proxy.management_endpoints.key_management_endpoints import ( + delete_key_fn, + generate_key_fn, + generate_key_helper_fn, + info_key_fn, + list_keys, + regenerate_key_fn, + update_key_fn, +) +from litellm.proxy.management_endpoints.team_endpoints import ( + new_team, + team_info, + update_team, +) +from litellm.proxy.proxy_server import ( + LitellmUserRoles, + audio_transcriptions, + chat_completion, + completion, + embeddings, + image_generation, + model_list, + moderations, + user_api_key_auth, +) +from litellm.proxy.management_endpoints.customer_endpoints import ( + new_end_user, +) +from litellm.proxy.spend_tracking.spend_management_endpoints import ( + global_spend, + spend_key_fn, + spend_user_fn, + view_spend_logs, +) +from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend + +verbose_proxy_logger.setLevel(level=logging.DEBUG) + +from starlette.datastructures import URL + +from litellm.caching.caching import DualCache +from litellm.proxy._types import ( + DynamoDBArgs, + GenerateKeyRequest, + KeyRequest, + LiteLLM_UpperboundKeyGenerateParams, + NewCustomerRequest, + NewTeamRequest, + NewUserRequest, + ProxyErrorTypes, + ProxyException, + UpdateKeyRequest, + UpdateTeamRequest, + UpdateUserRequest, + UserAPIKeyAuth, +) + +proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + + +request_data = { + "model": "azure-gpt-3.5", + "messages": [ + {"role": "user", "content": "this is my new test. respond in 50 lines"} + ], +} + + +@pytest.fixture +def prisma_client(): + from litellm.proxy.proxy_cli import append_query_params + + ### add connection pool + pool timeout args + params = {"connection_limit": 100, "pool_timeout": 60} + database_url = os.getenv("DATABASE_URL") + modified_url = append_query_params(database_url, params) + os.environ["DATABASE_URL"] = modified_url + + # Assuming PrismaClient is a class that needs to be instantiated + prisma_client = PrismaClient( + database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj + ) + + # Reset litellm.proxy.proxy_server.prisma_client to None + litellm.proxy.proxy_server.litellm_proxy_budget_name = ( + f"litellm-proxy-budget-{time.time()}" + ) + litellm.proxy.proxy_server.user_custom_key_generate = None + + return prisma_client + + +async def setup_db_connection(prisma_client): + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + + +@pytest.mark.asyncio +async def test_pod_lock_acquisition_when_no_active_lock(prisma_client): + """Test if a pod can acquire a lock when no lock is active""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + lock_manager = PodLockManager(cronjob_id=cronjob_id) + + # Attempt to acquire lock + result = await lock_manager.acquire_lock() + + assert result == True, "Pod should be able to acquire lock when no lock exists" + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id == lock_manager.pod_id + + +@pytest.mark.asyncio +async def test_pod_lock_acquisition_after_completion(prisma_client): + """Test if a new pod can acquire lock after previous pod completes""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + # First pod acquires and releases lock + first_lock_manager = PodLockManager(cronjob_id=cronjob_id) + await first_lock_manager.acquire_lock() + await first_lock_manager.release_lock() + + # Second pod attempts to acquire lock + second_lock_manager = PodLockManager(cronjob_id=cronjob_id) + result = await second_lock_manager.acquire_lock() + + assert result == True, "Second pod should acquire lock after first pod releases it" + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id == second_lock_manager.pod_id + + +@pytest.mark.asyncio +async def test_pod_lock_acquisition_after_expiry(prisma_client): + """Test if a new pod can acquire lock after previous pod's lock expires""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + # First pod acquires lock + first_lock_manager = PodLockManager(cronjob_id=cronjob_id) + await first_lock_manager.acquire_lock() + + # release the lock from the first pod + await first_lock_manager.release_lock() + + # Second pod attempts to acquire lock + second_lock_manager = PodLockManager(cronjob_id=cronjob_id) + result = await second_lock_manager.acquire_lock() + + assert ( + result == True + ), "Second pod should acquire lock after first pod's lock expires" + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id == second_lock_manager.pod_id + + +@pytest.mark.asyncio +async def test_pod_lock_release(prisma_client): + """Test if a pod can successfully release its lock""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + lock_manager = PodLockManager(cronjob_id=cronjob_id) + + # Acquire and then release lock + await lock_manager.acquire_lock() + await lock_manager.release_lock() + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "INACTIVE" + + +@pytest.mark.asyncio +async def test_concurrent_lock_acquisition(prisma_client): + """Test that only one pod can acquire the lock when multiple pods try simultaneously""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + # Create multiple lock managers simulating different pods + lock_manager1 = PodLockManager(cronjob_id=cronjob_id) + lock_manager2 = PodLockManager(cronjob_id=cronjob_id) + lock_manager3 = PodLockManager(cronjob_id=cronjob_id) + + # Try to acquire locks concurrently + results = await asyncio.gather( + lock_manager1.acquire_lock(), + lock_manager2.acquire_lock(), + lock_manager3.acquire_lock(), + ) + + # Only one should succeed + print("all results=", results) + assert sum(results) == 1, "Only one pod should acquire the lock" + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id in [ + lock_manager1.pod_id, + lock_manager2.pod_id, + lock_manager3.pod_id, + ] + + +@pytest.mark.asyncio +async def test_lock_renewal(prisma_client): + """Test that a pod can successfully renew its lock""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + lock_manager = PodLockManager(cronjob_id=cronjob_id) + + # Acquire initial lock + await lock_manager.acquire_lock() + + # Get initial TTL + initial_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + initial_ttl = initial_record.ttl + + # Wait a short time + await asyncio.sleep(1) + + # Renew the lock + await lock_manager.renew_lock() + + # Get updated record + renewed_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + + assert renewed_record.ttl > initial_ttl, "Lock TTL should be extended after renewal" + assert renewed_record.status == "ACTIVE" + assert renewed_record.pod_id == lock_manager.pod_id + + +@pytest.mark.asyncio +async def test_lock_acquisition_with_expired_ttl(prisma_client): + """Test that a pod can acquire a lock when existing lock has expired TTL""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + first_lock_manager = PodLockManager(cronjob_id=cronjob_id) + + # First pod acquires lock + await first_lock_manager.acquire_lock() + + # Manually expire the TTL + expired_time = datetime.now(timezone.utc) - timedelta(seconds=10) + await prisma_client.db.litellm_cronjob.update( + where={"cronjob_id": cronjob_id}, data={"ttl": expired_time} + ) + + # Second pod tries to acquire without explicit release + second_lock_manager = PodLockManager(cronjob_id=cronjob_id) + result = await second_lock_manager.acquire_lock() + + assert result == True, "Should acquire lock when existing lock has expired TTL" + + # Verify in database + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id == second_lock_manager.pod_id + + +@pytest.mark.asyncio +async def test_release_expired_lock(prisma_client): + """Test that a pod cannot release a lock that has been taken over by another pod""" + await setup_db_connection(prisma_client) + + cronjob_id = str(uuid.uuid4()) + first_lock_manager = PodLockManager(cronjob_id=cronjob_id) + + # First pod acquires lock + await first_lock_manager.acquire_lock() + + # Manually expire the TTL + expired_time = datetime.now(timezone.utc) - timedelta(seconds=10) + await prisma_client.db.litellm_cronjob.update( + where={"cronjob_id": cronjob_id}, data={"ttl": expired_time} + ) + + # Second pod acquires the lock + second_lock_manager = PodLockManager(cronjob_id=cronjob_id) + await second_lock_manager.acquire_lock() + + # First pod attempts to release its lock + await first_lock_manager.release_lock() + + # Verify that second pod's lock is still active + lock_record = await prisma_client.db.litellm_cronjob.find_first( + where={"cronjob_id": cronjob_id} + ) + assert lock_record.status == "ACTIVE" + assert lock_record.pod_id == second_lock_manager.pod_id From 193052ed7054292e7d01794cf3b3fb76bbb83503 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:05:17 -0700 Subject: [PATCH 25/40] test pod lock manager --- .../litellm/proxy/db/test_pod_lock_manager.py | 215 ++++++++++++++++-- 1 file changed, 200 insertions(+), 15 deletions(-) diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/test_pod_lock_manager.py index e685172c1e..bce7b66409 100644 --- a/tests/litellm/proxy/db/test_pod_lock_manager.py +++ b/tests/litellm/proxy/db/test_pod_lock_manager.py @@ -42,38 +42,75 @@ def pod_lock_manager(): @pytest.mark.asyncio async def test_acquire_lock_success(pod_lock_manager, mock_prisma): """ - Test that the lock is acquired successfully if the DB response is successful + Test that the lock is acquired successfully when no existing lock exists """ + # Mock find_unique to return None (no existing lock) + mock_prisma.db.litellm_cronjob.find_unique.return_value = None + + # Mock successful creation of new lock mock_response = AsyncMock() mock_response.status = "ACTIVE" mock_response.pod_id = pod_lock_manager.pod_id - mock_prisma.db.litellm_cronjob.upsert.return_value = mock_response + mock_prisma.db.litellm_cronjob.create.return_value = mock_response result = await pod_lock_manager.acquire_lock() assert result == True - # Verify upsert was called with correct parameters - mock_prisma.db.litellm_cronjob.upsert.assert_called_once() - call_args = mock_prisma.db.litellm_cronjob.upsert.call_args[1] - assert call_args["where"]["cronjob_id"] == "test_job" - assert "create" in call_args["data"] - assert "update" in call_args["data"] + # Verify find_unique was called + mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() + # Verify create was called with correct parameters + mock_prisma.db.litellm_cronjob.create.assert_called_once() + call_args = mock_prisma.db.litellm_cronjob.create.call_args[1] + assert call_args["data"]["cronjob_id"] == "test_job" + assert call_args["data"]["pod_id"] == pod_lock_manager.pod_id + assert call_args["data"]["status"] == "ACTIVE" @pytest.mark.asyncio -async def test_acquire_lock_failure(pod_lock_manager, mock_prisma): +async def test_acquire_lock_existing_active(pod_lock_manager, mock_prisma): """ - Test that the lock is not acquired if the lock is held by a different pod + Test that the lock is not acquired if there's an active lock by different pod """ - # Mock failed lock acquisition (different pod holds the lock) - mock_response = AsyncMock() - mock_response.status = "ACTIVE" - mock_response.pod_id = "different_pod_id" - mock_prisma.db.litellm_cronjob.upsert.return_value = mock_response + # Mock existing active lock + mock_existing = AsyncMock() + mock_existing.status = "ACTIVE" + mock_existing.pod_id = "different_pod_id" + mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) # Future TTL + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing result = await pod_lock_manager.acquire_lock() assert result == False + # Verify find_unique was called but update/create were not + mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() + mock_prisma.db.litellm_cronjob.update.assert_not_called() + mock_prisma.db.litellm_cronjob.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_acquire_lock_expired(pod_lock_manager, mock_prisma): + """ + Test that the lock can be acquired if existing lock is expired + """ + # Mock existing expired lock + mock_existing = AsyncMock() + mock_existing.status = "ACTIVE" + mock_existing.pod_id = "different_pod_id" + mock_existing.ttl = datetime.now(timezone.utc) - timedelta(seconds=30) # Past TTL + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + + # Mock successful update + mock_updated = AsyncMock() + mock_updated.pod_id = pod_lock_manager.pod_id + mock_prisma.db.litellm_cronjob.update.return_value = mock_updated + + result = await pod_lock_manager.acquire_lock() + assert result == True + + # Verify both find_unique and update were called + mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() + mock_prisma.db.litellm_cronjob.update.assert_called_once() + @pytest.mark.asyncio async def test_renew_lock(pod_lock_manager, mock_prisma): @@ -133,3 +170,151 @@ async def test_database_error_handling(pod_lock_manager, mock_prisma): assert await pod_lock_manager.acquire_lock() == False await pod_lock_manager.renew_lock() # Should not raise exception await pod_lock_manager.release_lock() # Should not raise exception + + +@pytest.mark.asyncio +async def test_acquire_lock_inactive_status(pod_lock_manager, mock_prisma): + """ + Test that the lock can be acquired if existing lock is INACTIVE + """ + # Mock existing inactive lock + mock_existing = AsyncMock() + mock_existing.status = "INACTIVE" + mock_existing.pod_id = "different_pod_id" + mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + + # Mock successful update + mock_updated = AsyncMock() + mock_updated.pod_id = pod_lock_manager.pod_id + mock_prisma.db.litellm_cronjob.update.return_value = mock_updated + + result = await pod_lock_manager.acquire_lock() + assert result == True + + mock_prisma.db.litellm_cronjob.update.assert_called_once() + + +@pytest.mark.asyncio +async def test_acquire_lock_same_pod(pod_lock_manager, mock_prisma): + """ + Test that the lock returns True if the same pod already holds the lock + """ + # Mock existing active lock held by same pod + mock_existing = AsyncMock() + mock_existing.status = "ACTIVE" + mock_existing.pod_id = pod_lock_manager.pod_id + mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + + result = await pod_lock_manager.acquire_lock() + assert result == True + + # Verify no update was needed + mock_prisma.db.litellm_cronjob.update.assert_not_called() + mock_prisma.db.litellm_cronjob.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_acquire_lock_race_condition(pod_lock_manager, mock_prisma): + """ + Test handling of potential race conditions during lock acquisition + """ + # First find_unique returns None + mock_prisma.db.litellm_cronjob.find_unique.return_value = None + + # But create raises unique constraint violation + mock_prisma.db.litellm_cronjob.create.side_effect = Exception( + "Unique constraint violation" + ) + + result = await pod_lock_manager.acquire_lock() + assert result == False + + +@pytest.mark.asyncio +async def test_ttl_calculation(pod_lock_manager, mock_prisma): + """ + Test that TTL is calculated correctly when acquiring lock + """ + mock_prisma.db.litellm_cronjob.find_unique.return_value = None + mock_prisma.db.litellm_cronjob.create.return_value = AsyncMock() + + await pod_lock_manager.acquire_lock() + + call_args = mock_prisma.db.litellm_cronjob.create.call_args[1] + ttl = call_args["data"]["ttl"] + + # Verify TTL is in the future by DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + expected_ttl = datetime.now(timezone.utc) + timedelta( + seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + ) + assert abs((ttl - expected_ttl).total_seconds()) < 1 # Allow 1 second difference + + +@pytest.mark.asyncio +async def test_concurrent_lock_acquisition_simulation(mock_prisma): + """ + Simulate multiple pods trying to acquire the lock simultaneously + """ + pod1 = PodLockManager(cronjob_id="test_job") + pod2 = PodLockManager(cronjob_id="test_job") + pod3 = PodLockManager(cronjob_id="test_job") + + # Simulate first pod getting the lock + mock_prisma.db.litellm_cronjob.find_unique.return_value = None + mock_response = AsyncMock() + mock_response.pod_id = pod1.pod_id + mock_response.status = "ACTIVE" + mock_prisma.db.litellm_cronjob.create.return_value = mock_response + + # First pod should get the lock + result1 = await pod1.acquire_lock() + assert result1 == True + + # Simulate other pods trying to acquire same lock immediately after + mock_existing = AsyncMock() + mock_existing.status = "ACTIVE" + mock_existing.pod_id = pod1.pod_id + mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + + # Other pods should fail to acquire + result2 = await pod2.acquire_lock() + result3 = await pod3.acquire_lock() + assert result2 == False + assert result3 == False + + +@pytest.mark.asyncio +async def test_lock_takeover_race_condition(mock_prisma): + """ + Test scenario where multiple pods try to take over an expired lock + """ + pod1 = PodLockManager(cronjob_id="test_job") + pod2 = PodLockManager(cronjob_id="test_job") + + # Simulate expired lock + mock_existing = AsyncMock() + mock_existing.status = "ACTIVE" + mock_existing.pod_id = "old_pod" + mock_existing.ttl = datetime.now(timezone.utc) - timedelta(seconds=30) + mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + + # Simulate pod1's update succeeding + mock_update1 = AsyncMock() + mock_update1.pod_id = pod1.pod_id + mock_prisma.db.litellm_cronjob.update.return_value = mock_update1 + + # First pod should successfully take over + result1 = await pod1.acquire_lock() + assert result1 == True + + # Simulate pod2's update failing due to race condition + mock_prisma.db.litellm_cronjob.update.side_effect = Exception( + "Row was updated by another transaction" + ) + + # Second pod should fail to take over + result2 = await pod2.acquire_lock() + assert result2 == False From 3df839d8783f4c1ca2c4d22e4d018601cec18780 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:12:19 -0700 Subject: [PATCH 26/40] use helper methods _commit_spend_updates_to_db_without_redis_buffer --- litellm/proxy/db/db_spend_update_writer.py | 109 ++++++++++++++------- 1 file changed, 75 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 9047bd4c38..17247aff74 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -400,45 +400,86 @@ class DBSpendUpdateWriter: - Regular flow of this method """ if RedisUpdateBuffer._should_commit_spend_updates_to_redis(): - await self.redis_update_buffer.store_in_memory_spend_updates_in_redis( - prisma_client=prisma_client, - ) - - # Only commit from redis to db if this pod is the leader - if await self.pod_lock_manager.acquire_lock(): - verbose_proxy_logger.debug("acquired lock for spend updates") - - try: - db_spend_update_transactions = ( - await self.redis_update_buffer.get_all_update_transactions_from_redis() - ) - if db_spend_update_transactions is not None: - await DBSpendUpdateWriter._commit_spend_updates_to_db( - prisma_client=prisma_client, - n_retry_times=n_retry_times, - proxy_logging_obj=proxy_logging_obj, - db_spend_update_transactions=db_spend_update_transactions, - ) - except Exception as e: - verbose_proxy_logger.error(f"Error committing spend updates: {e}") - finally: - await self.pod_lock_manager.release_lock() - else: - db_spend_update_transactions = DBSpendUpdateTransactions( - user_list_transactions=prisma_client.user_list_transactions, - end_user_list_transactions=prisma_client.end_user_list_transactions, - key_list_transactions=prisma_client.key_list_transactions, - team_list_transactions=prisma_client.team_list_transactions, - team_member_list_transactions=prisma_client.team_member_list_transactions, - org_list_transactions=prisma_client.org_list_transactions, - ) - await DBSpendUpdateWriter._commit_spend_updates_to_db( + await self._commit_spend_updates_to_db_with_redis( prisma_client=prisma_client, n_retry_times=n_retry_times, proxy_logging_obj=proxy_logging_obj, - db_spend_update_transactions=db_spend_update_transactions, ) + else: + await self._commit_spend_updates_to_db_without_redis_buffer( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + ) + + async def _commit_spend_updates_to_db_with_redis( + self, + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Handler to commit spend updates to Redis and attempt to acquire lock to commit to db + + This is a v2 scalable approach to first commit spend updates to redis, then commit to db + + This minimizes DB Deadlocks since + - All pods only need to write their spend updates to redis + - Only 1 pod will commit to db at a time (based on if it can acquire the lock over writing to DB) + """ + await self.redis_update_buffer.store_in_memory_spend_updates_in_redis( + prisma_client=prisma_client, + ) + + # Only commit from redis to db if this pod is the leader + if await self.pod_lock_manager.acquire_lock(): + verbose_proxy_logger.debug("acquired lock for spend updates") + + try: + db_spend_update_transactions = ( + await self.redis_update_buffer.get_all_update_transactions_from_redis() + ) + if db_spend_update_transactions is not None: + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, + ) + except Exception as e: + verbose_proxy_logger.error(f"Error committing spend updates: {e}") + finally: + await self.pod_lock_manager.release_lock() + + async def _commit_spend_updates_to_db_without_redis_buffer( + self, + prisma_client: PrismaClient, + n_retry_times: int, + proxy_logging_obj: ProxyLogging, + ): + """ + Commits all the spend `UPDATE` transactions to the Database + + This is the regular flow of committing to db without using a redis buffer + + Note: This flow causes Deadlocks in production (1K RPS+). Use self._commit_spend_updates_to_db_with_redis() instead if you expect 1K+ RPS. + """ + db_spend_update_transactions = DBSpendUpdateTransactions( + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, + ) + await DBSpendUpdateWriter._commit_spend_updates_to_db( + prisma_client=prisma_client, + n_retry_times=n_retry_times, + proxy_logging_obj=proxy_logging_obj, + db_spend_update_transactions=db_spend_update_transactions, + ) + @staticmethod async def _commit_spend_updates_to_db( # noqa: PLR0915 prisma_client: PrismaClient, From 7c73987066dd1a003f5960c26847506237e2f323 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:17:01 -0700 Subject: [PATCH 27/40] fix loc of DBSpendUpdateTransactions --- litellm/proxy/_types.py | 9 +++++++++ litellm/proxy/db/db_spend_update_writer.py | 6 ++---- litellm/proxy/db/redis_update_buffer.py | 12 ++---------- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 17b4acd138..1c3251c12a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2733,3 +2733,12 @@ class DailyUserSpendTransaction(TypedDict): prompt_tokens: int completion_tokens: int spend: float + + +class DBSpendUpdateTransactions(TypedDict): + user_list_transactions: Optional[Dict[str, float]] + end_user_list_transactions: Optional[Dict[str, float]] + key_list_transactions: Optional[Dict[str, float]] + team_list_transactions: Optional[Dict[str, float]] + team_member_list_transactions: Optional[Dict[str, float]] + org_list_transactions: Optional[Dict[str, float]] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 17247aff74..8f175fc905 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -18,15 +18,13 @@ from litellm.caching import DualCache, RedisCache from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DBSpendUpdateTransactions, Litellm_EntityType, LiteLLM_UserTable, SpendLogsPayload, ) from litellm.proxy.db.pod_lock_manager import PodLockManager -from litellm.proxy.db.redis_update_buffer import ( - DBSpendUpdateTransactions, - RedisUpdateBuffer, -) +from litellm.proxy.db.redis_update_buffer import RedisUpdateBuffer if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 2207911024..6f334e90ed 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -4,10 +4,11 @@ Handles buffering database `UPDATE` transactions in Redis before committing them This is to prevent deadlocks and improve reliability """ -from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache +from litellm.proxy._types import DBSpendUpdateTransactions from litellm.secret_managers.main import str_to_bool if TYPE_CHECKING: @@ -16,15 +17,6 @@ else: PrismaClient = Any -class DBSpendUpdateTransactions(TypedDict): - user_list_transactions: Optional[Dict[str, float]] - end_user_list_transactions: Optional[Dict[str, float]] - key_list_transactions: Optional[Dict[str, float]] - team_list_transactions: Optional[Dict[str, float]] - team_member_list_transactions: Optional[Dict[str, float]] - org_list_transactions: Optional[Dict[str, float]] - - class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database From df136c5689af818ed44c1b7c28d8eb3cdc66f769 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:20:28 -0700 Subject: [PATCH 28/40] add docstring for DBSpendUpdateTransactions --- litellm/proxy/_types.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1c3251c12a..24beef6b63 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2736,6 +2736,10 @@ class DailyUserSpendTransaction(TypedDict): class DBSpendUpdateTransactions(TypedDict): + """ + Internal Data Structure for buffering spend updates in Redis or in memory before committing them to the database + """ + user_list_transactions: Optional[Dict[str, float]] end_user_list_transactions: Optional[Dict[str, float]] key_list_transactions: Optional[Dict[str, float]] From cec5280e033e6ff57ed1a1da0016622ad34ae1e9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 16:16:44 -0700 Subject: [PATCH 29/40] debugging for Redis TX management --- litellm/proxy/db/db_spend_update_writer.py | 27 ++++++++++++++-------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8f175fc905..fea40e58e7 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -496,6 +496,9 @@ class DBSpendUpdateWriter: ### UPDATE USER TABLE ### user_list_transactions = db_spend_update_transactions["user_list_transactions"] + verbose_proxy_logger.debug( + "User Spend transactions: {}".format(user_list_transactions) + ) if ( user_list_transactions is not None and len(user_list_transactions.keys()) > 0 @@ -536,14 +539,12 @@ class DBSpendUpdateWriter: ) ### UPDATE END-USER TABLE ### - verbose_proxy_logger.debug( - "End-User Spend transactions: {}".format( - len(prisma_client.end_user_list_transactions.keys()) - ) - ) end_user_list_transactions = db_spend_update_transactions[ "end_user_list_transactions" ] + verbose_proxy_logger.debug( + "End-User Spend transactions: {}".format(end_user_list_transactions) + ) if ( end_user_list_transactions is not None and len(end_user_list_transactions.keys()) > 0 @@ -595,12 +596,10 @@ class DBSpendUpdateWriter: ) ### UPDATE TEAM TABLE ### - verbose_proxy_logger.debug( - "Team Spend transactions: {}".format( - len(prisma_client.team_list_transactions.keys()) - ) - ) team_list_transactions = db_spend_update_transactions["team_list_transactions"] + verbose_proxy_logger.debug( + "Team Spend transactions: {}".format(team_list_transactions) + ) if ( team_list_transactions is not None and len(team_list_transactions.keys()) > 0 @@ -649,6 +648,11 @@ class DBSpendUpdateWriter: team_member_list_transactions = db_spend_update_transactions[ "team_member_list_transactions" ] + verbose_proxy_logger.debug( + "Team Membership Spend transactions: {}".format( + team_member_list_transactions + ) + ) if ( team_member_list_transactions is not None and len(team_member_list_transactions.keys()) > 0 @@ -694,6 +698,9 @@ class DBSpendUpdateWriter: ### UPDATE ORG TABLE ### org_list_transactions = db_spend_update_transactions["org_list_transactions"] + verbose_proxy_logger.debug( + "Org Spend transactions: {}".format(org_list_transactions) + ) if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() From e5e0ffa28a6b182dec898e986e839029a9cb0f62 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 17:23:47 -0700 Subject: [PATCH 30/40] add async_rpush --- litellm/caching/redis_cache.py | 135 +++++++++++++++++++++++++++++++++ 1 file changed, 135 insertions(+) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 0571ac9f15..620ba5d73a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1045,3 +1045,138 @@ class RedisCache(BaseCache): except Exception as e: verbose_logger.debug(f"Redis TTL Error: {e}") return None + + async def _generic_redis_command( + self, + command: str, + key: str, + *args, + parent_otel_span: Optional[Span] = None, + **kwargs, + ) -> Any: + """ + Generic method to execute any Redis command + + Args: + command: The Redis command to execute (e.g. 'rpush', 'lpop') + key: The Redis key to operate on + args: Additional arguments to pass to the Redis command + parent_otel_span: Optional parent OpenTelemetry span + kwargs: Additional keyword arguments + + Returns: + Any: The result of the Redis command + """ + _redis_client = self.init_async_client() + key = self.check_and_fix_namespace(key=key) + start_time = time.time() + + try: + # Get the method from the Redis client + redis_method = getattr(_redis_client, command) + if not redis_method: + raise AttributeError(f"Redis client has no method named '{command}'") + + # Execute the command + result = await redis_method(key, *args) + + # Log success + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type=f"async_{command}", + start_time=start_time, + end_time=end_time, + parent_otel_span=parent_otel_span, + event_metadata={"key": key}, + ) + ) + return result + except Exception as e: + # Log failure + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type=f"async_{command}", + start_time=start_time, + end_time=end_time, + parent_otel_span=parent_otel_span, + event_metadata={"key": key}, + ) + ) + verbose_logger.error( + f"LiteLLM Redis Caching: async {command}() - Got exception from REDIS: {str(e)}" + ) + raise e + + async def async_rpush( + self, + key: str, + values: List[Any], + parent_otel_span: Optional[Span] = None, + **kwargs, + ) -> int: + """ + Append one or multiple values to a list stored at key + + Args: + key: The Redis key of the list + values: One or more values to append to the list + parent_otel_span: Optional parent OpenTelemetry span + + Returns: + int: The length of the list after the push operation + """ + print_verbose(f"RPUSH to Redis list: key: {key}, values: {values}") + return await self._generic_redis_command( + "rpush", key, *values, parent_otel_span=parent_otel_span, **kwargs + ) + + async def async_lpop( + self, + key: str, + count: Optional[int] = None, + parent_otel_span: Optional[Span] = None, + **kwargs, + ) -> Union[Any, List[Any]]: + """ + Remove and return the first element(s) of a list stored at key + + Args: + key: The Redis key of the list + count: Number of elements to pop (if None, pops a single element) + parent_otel_span: Optional parent OpenTelemetry span + + Returns: + Union[Any, List[Any]]: The popped value(s) or None if list is empty + """ + print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") + args = [count] if count is not None else [] + + # Fix: Pass command as a positional argument, not as a keyword argument + result = await self._generic_redis_command( + "lpop", key, *args, parent_otel_span=parent_otel_span, **kwargs + ) + + # Handle result parsing if needed + if isinstance(result, bytes): + try: + return result.decode("utf-8") + except Exception: + return result + elif isinstance(result, list) and all( + isinstance(item, bytes) for item in result + ): + try: + return [item.decode("utf-8") for item in result] + except Exception: + return result + + return result From 88458a6568f933d00ddefe9d9e57622b19e53964 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 17:24:54 -0700 Subject: [PATCH 31/40] redis update buffer queue --- litellm/proxy/db/db_spend_update_writer.py | 2 +- litellm/proxy/db/redis_update_buffer.py | 240 +++++++++++++-------- 2 files changed, 152 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index fea40e58e7..55b95d16bf 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -436,7 +436,7 @@ class DBSpendUpdateWriter: try: db_spend_update_transactions = ( - await self.redis_update_buffer.get_all_update_transactions_from_redis() + await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() ) if db_spend_update_transactions is not None: await DBSpendUpdateWriter._commit_spend_updates_to_db( diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 6f334e90ed..ab43dc575d 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -4,10 +4,12 @@ Handles buffering database `UPDATE` transactions in Redis before committing them This is to prevent deadlocks and improve reliability """ -from typing import TYPE_CHECKING, Any, Dict, Optional, Union, cast +import json +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DBSpendUpdateTransactions from litellm.secret_managers.main import str_to_bool @@ -16,6 +18,9 @@ if TYPE_CHECKING: else: PrismaClient = Any +REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" +MAX_REDIS_BUFFER_DEQUEUE_COUNT = 100 + class RedisUpdateBuffer: """ @@ -61,14 +66,21 @@ class RedisUpdateBuffer: - value is the spend amount ``` - { - "0929880201": 10, - "0929880202": 20, - "0929880203": 30, - } + Redis List: + key_list_transactions: + [ + "0929880201": 1.2, + "0929880202": 0.01, + "0929880203": 0.001, + ] ``` """ - IN_MEMORY_UPDATE_TRANSACTIONS: DBSpendUpdateTransactions = ( + if self.redis_cache is None: + verbose_proxy_logger.debug( + "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" + ) + return + db_spend_update_transactions: DBSpendUpdateTransactions = ( DBSpendUpdateTransactions( user_list_transactions=prisma_client.user_list_transactions, end_user_list_transactions=prisma_client.end_user_list_transactions, @@ -78,30 +90,47 @@ class RedisUpdateBuffer: org_list_transactions=prisma_client.org_list_transactions, ) ) - for key, _transactions in IN_MEMORY_UPDATE_TRANSACTIONS.items(): - await self.increment_all_transaction_objects_in_redis( - key=key, - transactions=cast(Dict, _transactions), - ) - async def increment_all_transaction_objects_in_redis( - self, - key: str, - transactions: Dict, + # only store in redis if there are any updates to commit + if ( + self._number_of_transactions_to_store_in_redis(db_spend_update_transactions) + == 0 + ): + return + + list_of_transactions = [safe_dumps(db_spend_update_transactions)] + await self.redis_cache.async_rpush( + key=REDIS_UPDATE_BUFFER_KEY, + values=list_of_transactions, + ) + self._clear_all_in_memory_spend_updates(prisma_client) + + @staticmethod + def _number_of_transactions_to_store_in_redis( + db_spend_update_transactions: DBSpendUpdateTransactions, + ) -> int: + """ + Gets the number of transactions to store in Redis + """ + num_transactions = 0 + for v in db_spend_update_transactions.values(): + if isinstance(v, dict): + num_transactions += len(v) + return num_transactions + + @staticmethod + def _clear_all_in_memory_spend_updates( + prisma_client: PrismaClient, ): """ - Increments all transaction objects in Redis + Clears all in-memory spend updates """ - if self.redis_cache is None: - verbose_proxy_logger.debug( - "redis_cache is None, skipping increment_all_transaction_objects_in_redis" - ) - return - for transaction_id, transaction_amount in transactions.items(): - await self.redis_cache.async_increment( - key=f"{key}:{transaction_id}", - value=transaction_amount, - ) + prisma_client.user_list_transactions = {} + prisma_client.end_user_list_transactions = {} + prisma_client.key_list_transactions = {} + prisma_client.team_list_transactions = {} + prisma_client.team_member_list_transactions = {} + prisma_client.org_list_transactions = {} @staticmethod def _remove_prefix_from_keys(data: Dict[str, Any], prefix: str) -> Dict[str, Any]: @@ -110,77 +139,110 @@ class RedisUpdateBuffer: """ return {key.replace(prefix, "", 1): value for key, value in data.items()} - async def get_all_update_transactions_from_redis( + async def get_all_update_transactions_from_redis_buffer( self, ) -> Optional[DBSpendUpdateTransactions]: """ Gets all the update transactions from Redis + + On Redis we store a list of transactions as a JSON string + + eg. + [ + DBSpendUpdateTransactions( + user_list_transactions={ + "user_id_1": 1.2, + "user_id_2": 0.01, + }, + end_user_list_transactions={}, + key_list_transactions={ + "0929880201": 1.2, + "0929880202": 0.01, + }, + team_list_transactions={}, + team_member_list_transactions={}, + org_list_transactions={}, + ), + DBSpendUpdateTransactions( + user_list_transactions={ + "user_id_3": 1.2, + "user_id_4": 0.01, + }, + end_user_list_transactions={}, + key_list_transactions={ + "key_id_1": 1.2, + "key_id_2": 0.01, + }, + team_list_transactions={}, + team_member_list_transactions={}, + org_list_transactions={}, + ] """ if self.redis_cache is None: return None - user_transaction_keys = await self.redis_cache.async_scan_iter( - "user_list_transactions:*" + list_of_transactions = await self.redis_cache.async_lpop( + key=REDIS_UPDATE_BUFFER_KEY, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) - end_user_transaction_keys = await self.redis_cache.async_scan_iter( - "end_user_list_transactions:*" - ) - key_transaction_keys = await self.redis_cache.async_scan_iter( - "key_list_transactions:*" - ) - team_transaction_keys = await self.redis_cache.async_scan_iter( - "team_list_transactions:*" - ) - team_member_transaction_keys = await self.redis_cache.async_scan_iter( - "team_member_list_transactions:*" - ) - org_transaction_keys = await self.redis_cache.async_scan_iter( - "org_list_transactions:*" + if list_of_transactions is None: + return None + + # Parse the list of transactions from JSON strings + parsed_transactions = self._parse_list_of_transactions(list_of_transactions) + + # If there are no transactions, return None + if len(parsed_transactions) == 0: + return None + + # Combine all transactions into a single transaction + combined_transaction = self._combine_list_of_transactions(parsed_transactions) + + return combined_transaction + + @staticmethod + def _parse_list_of_transactions( + list_of_transactions: List[str], + ) -> List[DBSpendUpdateTransactions]: + """ + Parses the list of transactions from Redis + """ + return [json.loads(transaction) for transaction in list_of_transactions] + + @staticmethod + def _combine_list_of_transactions( + list_of_transactions: List[DBSpendUpdateTransactions], + ) -> DBSpendUpdateTransactions: + """ + Combines the list of transactions into a single DBSpendUpdateTransactions object + """ + # Initialize a new combined transaction object with empty dictionaries + combined_transaction = DBSpendUpdateTransactions( + user_list_transactions={}, + end_user_list_transactions={}, + key_list_transactions={}, + team_list_transactions={}, + team_member_list_transactions={}, + org_list_transactions={}, ) - user_list_transactions = await self.redis_cache.async_batch_get_cache( - user_transaction_keys - ) - end_user_list_transactions = await self.redis_cache.async_batch_get_cache( - end_user_transaction_keys - ) - key_list_transactions = await self.redis_cache.async_batch_get_cache( - key_transaction_keys - ) - team_list_transactions = await self.redis_cache.async_batch_get_cache( - team_transaction_keys - ) - team_member_list_transactions = await self.redis_cache.async_batch_get_cache( - team_member_transaction_keys - ) - org_list_transactions = await self.redis_cache.async_batch_get_cache( - org_transaction_keys - ) + # Define the transaction fields to process + transaction_fields = [ + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", + ] - # filter out the "prefix" from the keys using the helper method - user_list_transactions = self._remove_prefix_from_keys( - user_list_transactions, "user_list_transactions:" - ) - end_user_list_transactions = self._remove_prefix_from_keys( - end_user_list_transactions, "end_user_list_transactions:" - ) - key_list_transactions = self._remove_prefix_from_keys( - key_list_transactions, "key_list_transactions:" - ) - team_list_transactions = self._remove_prefix_from_keys( - team_list_transactions, "team_list_transactions:" - ) - team_member_list_transactions = self._remove_prefix_from_keys( - team_member_list_transactions, "team_member_list_transactions:" - ) - org_list_transactions = self._remove_prefix_from_keys( - org_list_transactions, "org_list_transactions:" - ) + # Loop through each transaction and combine the values + for transaction in list_of_transactions: + # Process each field type + for field in transaction_fields: + if transaction.get(field): + for entity_id, amount in transaction[field].items(): + combined_transaction[field][entity_id] = ( + combined_transaction[field].get(entity_id, 0) + amount + ) - return DBSpendUpdateTransactions( - user_list_transactions=user_list_transactions, - end_user_list_transactions=end_user_list_transactions, - key_list_transactions=key_list_transactions, - team_list_transactions=team_list_transactions, - team_member_list_transactions=team_member_list_transactions, - org_list_transactions=org_list_transactions, - ) + return combined_transaction From eae2adcee0a4ff98544eeabd4b4850f0794ecdc3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 19:49:38 -0700 Subject: [PATCH 32/40] redis cache add async push and pop methods --- litellm/caching/redis_cache.py | 183 ++++++++++++++------------------- 1 file changed, 77 insertions(+), 106 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 620ba5d73a..7378ed878a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1046,76 +1046,6 @@ class RedisCache(BaseCache): verbose_logger.debug(f"Redis TTL Error: {e}") return None - async def _generic_redis_command( - self, - command: str, - key: str, - *args, - parent_otel_span: Optional[Span] = None, - **kwargs, - ) -> Any: - """ - Generic method to execute any Redis command - - Args: - command: The Redis command to execute (e.g. 'rpush', 'lpop') - key: The Redis key to operate on - args: Additional arguments to pass to the Redis command - parent_otel_span: Optional parent OpenTelemetry span - kwargs: Additional keyword arguments - - Returns: - Any: The result of the Redis command - """ - _redis_client = self.init_async_client() - key = self.check_and_fix_namespace(key=key) - start_time = time.time() - - try: - # Get the method from the Redis client - redis_method = getattr(_redis_client, command) - if not redis_method: - raise AttributeError(f"Redis client has no method named '{command}'") - - # Execute the command - result = await redis_method(key, *args) - - # Log success - end_time = time.time() - _duration = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.REDIS, - duration=_duration, - call_type=f"async_{command}", - start_time=start_time, - end_time=end_time, - parent_otel_span=parent_otel_span, - event_metadata={"key": key}, - ) - ) - return result - except Exception as e: - # Log failure - end_time = time.time() - _duration = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_failure_hook( - service=ServiceTypes.REDIS, - duration=_duration, - error=e, - call_type=f"async_{command}", - start_time=start_time, - end_time=end_time, - parent_otel_span=parent_otel_span, - event_metadata={"key": key}, - ) - ) - verbose_logger.error( - f"LiteLLM Redis Caching: async {command}() - Got exception from REDIS: {str(e)}" - ) - raise e - async def async_rpush( self, key: str, @@ -1134,10 +1064,38 @@ class RedisCache(BaseCache): Returns: int: The length of the list after the push operation """ - print_verbose(f"RPUSH to Redis list: key: {key}, values: {values}") - return await self._generic_redis_command( - "rpush", key, *values, parent_otel_span=parent_otel_span, **kwargs - ) + _redis_client: Any = self.init_async_client() + start_time = time.time() + try: + response = await _redis_client.rpush(key, *values) + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type="async_rpush", + ) + ) + return response + except Exception as e: + # NON blocking - notify users Redis is throwing an exception + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type="async_rpush", + ) + ) + verbose_logger.error( + f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}" + ) + raise e async def async_lpop( self, @@ -1146,37 +1104,50 @@ class RedisCache(BaseCache): parent_otel_span: Optional[Span] = None, **kwargs, ) -> Union[Any, List[Any]]: - """ - Remove and return the first element(s) of a list stored at key - - Args: - key: The Redis key of the list - count: Number of elements to pop (if None, pops a single element) - parent_otel_span: Optional parent OpenTelemetry span - - Returns: - Union[Any, List[Any]]: The popped value(s) or None if list is empty - """ + _redis_client: Any = self.init_async_client() + start_time = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") - args = [count] if count is not None else [] + try: + result = await _redis_client.lpop(key, count) + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type="async_lpop", + ) + ) - # Fix: Pass command as a positional argument, not as a keyword argument - result = await self._generic_redis_command( - "lpop", key, *args, parent_otel_span=parent_otel_span, **kwargs - ) - - # Handle result parsing if needed - if isinstance(result, bytes): - try: - return result.decode("utf-8") - except Exception: - return result - elif isinstance(result, list) and all( - isinstance(item, bytes) for item in result - ): - try: - return [item.decode("utf-8") for item in result] - except Exception: - return result - - return result + # Handle result parsing if needed + if isinstance(result, bytes): + try: + return result.decode("utf-8") + except Exception: + return result + elif isinstance(result, list) and all( + isinstance(item, bytes) for item in result + ): + try: + return [item.decode("utf-8") for item in result] + except Exception: + return result + return result + except Exception as e: + # NON blocking - notify users Redis is throwing an exception + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type="async_lpop", + ) + ) + verbose_logger.error( + f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}" + ) + raise e From 21bf15263cf23d2e2a5031c3ce4985c83fe468f9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 19:51:15 -0700 Subject: [PATCH 33/40] use asyncio lock for updating PrismaClient txs --- litellm/proxy/db/redis_update_buffer.py | 54 ++++++++++++++----------- litellm/proxy/utils.py | 1 + 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index ab43dc575d..3d27e41d15 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -80,30 +80,35 @@ class RedisUpdateBuffer: "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" ) return - db_spend_update_transactions: DBSpendUpdateTransactions = ( - DBSpendUpdateTransactions( - user_list_transactions=prisma_client.user_list_transactions, - end_user_list_transactions=prisma_client.end_user_list_transactions, - key_list_transactions=prisma_client.key_list_transactions, - team_list_transactions=prisma_client.team_list_transactions, - team_member_list_transactions=prisma_client.team_member_list_transactions, - org_list_transactions=prisma_client.org_list_transactions, + async with prisma_client.in_memory_transaction_lock: + db_spend_update_transactions: DBSpendUpdateTransactions = ( + DBSpendUpdateTransactions( + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, + ) ) - ) - # only store in redis if there are any updates to commit - if ( - self._number_of_transactions_to_store_in_redis(db_spend_update_transactions) - == 0 - ): - return + # only store in redis if there are any updates to commit + if ( + self._number_of_transactions_to_store_in_redis( + db_spend_update_transactions + ) + == 0 + ): + return - list_of_transactions = [safe_dumps(db_spend_update_transactions)] - await self.redis_cache.async_rpush( - key=REDIS_UPDATE_BUFFER_KEY, - values=list_of_transactions, - ) - self._clear_all_in_memory_spend_updates(prisma_client) + list_of_transactions = [safe_dumps(db_spend_update_transactions)] + await self.redis_cache.async_rpush( + key=REDIS_UPDATE_BUFFER_KEY, + values=list_of_transactions, + ) + + # clear the in-memory spend updates + RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client) @staticmethod def _number_of_transactions_to_store_in_redis( @@ -201,12 +206,15 @@ class RedisUpdateBuffer: @staticmethod def _parse_list_of_transactions( - list_of_transactions: List[str], + list_of_transactions: Union[Any, List[Any]], ) -> List[DBSpendUpdateTransactions]: """ Parses the list of transactions from Redis """ - return [json.loads(transaction) for transaction in list_of_transactions] + if isinstance(list_of_transactions, list): + return [json.loads(transaction) for transaction in list_of_transactions] + else: + return [json.loads(list_of_transactions)] @staticmethod def _combine_list_of_transactions( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d0d4ea9b4e..66dd224dff 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1121,6 +1121,7 @@ class PrismaClient: self.iam_token_db_auth: Optional[bool] = str_to_bool( os.getenv("IAM_TOKEN_DB_AUTH") ) + self.in_memory_transaction_lock = asyncio.Lock() verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma # type: ignore From 183c33bb026ba17ec91122611e8baedd2a937fe7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 19:52:07 -0700 Subject: [PATCH 34/40] prisma client in_memory_transaction_lock --- litellm/proxy/db/db_spend_update_writer.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 55b95d16bf..72692339c8 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -142,6 +142,7 @@ class DBSpendUpdateWriter: transaction_list: dict, entity_type: Litellm_EntityType, debug_msg: Optional[str] = None, + prisma_client: Optional[PrismaClient] = None, ) -> bool: """ Common helper method to update a transaction list for an entity @@ -163,16 +164,18 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug( f"adding spend to {entity_type.value} db. Response cost: {response_cost}. {entity_type.value}_id: {entity_id}." ) + if prisma_client is None: + return False if entity_id is None: verbose_proxy_logger.debug( f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}" ) return False - - transaction_list[entity_id] = response_cost + transaction_list.get( - entity_id, 0 - ) + async with prisma_client.in_memory_transaction_lock: + transaction_list[entity_id] = response_cost + transaction_list.get( + entity_id, 0 + ) return True except Exception as e: @@ -197,6 +200,7 @@ class DBSpendUpdateWriter: transaction_list=prisma_client.key_list_transactions, entity_type=Litellm_EntityType.KEY, debug_msg=f"adding spend to key db. Response cost: {response_cost}. Token: {hashed_token}.", + prisma_client=prisma_client, ) except Exception as e: verbose_proxy_logger.exception( @@ -236,6 +240,7 @@ class DBSpendUpdateWriter: entity_id=_id, transaction_list=prisma_client.user_list_transactions, entity_type=Litellm_EntityType.USER, + prisma_client=prisma_client, ) if end_user_id is not None: @@ -244,6 +249,7 @@ class DBSpendUpdateWriter: entity_id=end_user_id, transaction_list=prisma_client.end_user_list_transactions, entity_type=Litellm_EntityType.END_USER, + prisma_client=prisma_client, ) except Exception as e: verbose_proxy_logger.info( @@ -270,6 +276,7 @@ class DBSpendUpdateWriter: entity_id=team_id, transaction_list=prisma_client.team_list_transactions, entity_type=Litellm_EntityType.TEAM, + prisma_client=prisma_client, ) try: @@ -282,6 +289,7 @@ class DBSpendUpdateWriter: entity_id=team_member_key, transaction_list=prisma_client.team_member_list_transactions, entity_type=Litellm_EntityType.TEAM_MEMBER, + prisma_client=prisma_client, ) except Exception: pass @@ -309,6 +317,7 @@ class DBSpendUpdateWriter: entity_id=org_id, transaction_list=prisma_client.org_list_transactions, entity_type=Litellm_EntityType.ORGANIZATION, + prisma_client=prisma_client, ) except Exception as e: verbose_proxy_logger.info( From ba550e214705abd20411e4b2446a27a25c567be2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 19:52:39 -0700 Subject: [PATCH 35/40] test local spend accuracy --- litellm/proxy/proxy_config.yaml | 1 + .../local_test_spend_accuracy_tests.py | 279 ++++++++++++++++++ 2 files changed, 280 insertions(+) create mode 100644 tests/otel_tests/local_test_spend_accuracy_tests.py diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index a1be54421b..9ca63f18ea 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -12,3 +12,4 @@ litellm_settings: cache: true cache_params: type: redis + supported_call_types: [] diff --git a/tests/otel_tests/local_test_spend_accuracy_tests.py b/tests/otel_tests/local_test_spend_accuracy_tests.py new file mode 100644 index 0000000000..6d756219c7 --- /dev/null +++ b/tests/otel_tests/local_test_spend_accuracy_tests.py @@ -0,0 +1,279 @@ +import pytest +import asyncio +import aiohttp +import json +from httpx import AsyncClient +from typing import Any, Optional +import uuid + +""" +Tests to run + +Basic Tests: +1. Basic Spend Accuracy Test: + - 1 Request costs $0.037 + - Make 12 requests + - Expect the spend for each of the following to be 12 * $0.037 + Key: $0.444 (call /info endpoint for each object to validate) + Team: $0.444 + User: $0.444 + Org: $0.444 + End User: $0.444 + +2. Long term spend accuracy test (with 2 bursts of requests) + - 1 Request costs $0.037 + - Burst 1: 12 requests + - Burst 2: 22 requests + + - Expect the spend for each of the following to be (12 + 22) * $0.037 + Key: $1.296 + Team: $1.296 + User: $1.296 + Org: $1.296 + End User: $1.296 + +Additional Test Scenarios: + +3. Concurrent Request Accuracy Test: + - Make 20 concurrent requests + - Verify total spend is 20 * $0.037 + - Check for race conditions in spend tracking + +4. Error Case Test: + - Make 10 successful requests ($0.037 each) + - Make 5 failed requests + - Verify spend is only counted for successful requests (10 * $0.037) + +5. Mixed Request Type Test: + - Make different types of requests with varying costs + - Verify accurate total spend calculation +""" + + +async def create_organization(session, organization_alias: str): + """Helper function to create a new organization""" + url = "http://0.0.0.0:4002/organization/new" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = {"organization_alias": organization_alias} + async with session.post(url, headers=headers, json=data) as response: + return await response.json() + + +async def create_team(session, org_id: str): + """Helper function to create a new team under an organization""" + url = "http://0.0.0.0:4002/team/new" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = {"organization_id": org_id, "team_alias": f"test-team-{uuid.uuid4()}"} + async with session.post(url, headers=headers, json=data) as response: + return await response.json() + + +async def create_user(session, org_id: str): + """Helper function to create a new user""" + url = "http://0.0.0.0:4002/user/new" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = {"user_name": f"test-user-{uuid.uuid4()}"} + async with session.post(url, headers=headers, json=data) as response: + return await response.json() + + +async def generate_key(session, user_id: str, team_id: str): + """Helper function to generate a key for a specific user and team""" + url = "http://0.0.0.0:4002/key/generate" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = {"user_id": user_id, "team_id": team_id} + async with session.post(url, headers=headers, json=data) as response: + return await response.json() + + +async def chat_completion(session, key: str): + """Make a chat completion request""" + from openai import AsyncOpenAI + import uuid + + client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4002/v1") + + response = await client.chat.completions.create( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": f"Test message {uuid.uuid4()}"}], + ) + return response + + +async def get_spend_info(session, entity_type: str, entity_id: str): + """Helper function to get spend information for an entity""" + url = f"http://0.0.0.0:4002/{entity_type}/info" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + if entity_type == "key": + data = {"key": entity_id} + else: + data = {f"{entity_type}_id": entity_id} + + async with session.get(url, headers=headers, params=data) as response: + return await response.json() + + +@pytest.mark.asyncio +async def test_basic_spend_accuracy(): + """ + Test basic spend accuracy across different entities: + 1. Create org, team, user, and key + 2. Make 12 requests at $0.037 each + 3. Verify spend accuracy for key, team, user, org, and end user + """ + SPEND_PER_REQUEST = 3.75 * 10**-5 + NUM_LLM_REQUESTS = 20 + expected_spend = NUM_LLM_REQUESTS * SPEND_PER_REQUEST # 12 requests at $0.037 each + + # Add tolerance constant at the top of the test + TOLERANCE = 1e-10 # Small number to account for floating-point precision + + async with aiohttp.ClientSession() as session: + # Create organization + org_response = await create_organization( + session=session, organization_alias=f"test-org-{uuid.uuid4()}" + ) + print("org_response: ", org_response) + org_id = org_response["organization_id"] + + # Create team under organization + team_response = await create_team(session, org_id) + print("team_response: ", team_response) + team_id = team_response["team_id"] + + # Create user + user_response = await create_user(session, org_id) + print("user_response: ", user_response) + user_id = user_response["user_id"] + + # Generate key + key_response = await generate_key(session, user_id, team_id) + print("key_response: ", key_response) + key = key_response["key"] + + # Make 12 requests + for _ in range(NUM_LLM_REQUESTS): + response = await chat_completion(session, key) + print("response: ", response) + + # wait 10 seconds for spend to be updated + await asyncio.sleep(15) + + # Get spend information for each entity + key_info = await get_spend_info(session, "key", key) + print("key_info: ", key_info) + team_info = await get_spend_info(session, "team", team_id) + print("team_info: ", team_info) + user_info = await get_spend_info(session, "user", user_id) + print("user_info: ", user_info) + org_info = await get_spend_info(session, "organization", org_id) + print("org_info: ", org_info) + + # Verify spend for each entity + assert ( + abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE + ), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE + ), f"User spend {user_info['info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE + ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(org_info["spend"] - expected_spend) < TOLERANCE + ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" + + +@pytest.mark.asyncio +async def test_long_term_spend_accuracy_with_bursts(): + """ + Test long-term spend accuracy with multiple bursts of requests: + 1. Create org, team, user, and key + 2. Burst 1: Make 12 requests + 3. Burst 2: Make 22 more requests + 4. Verify the total spend (34 requests) is tracked accurately across all entities + """ + SPEND_PER_REQUEST = 3.75 * 10**-5 # Cost per request + BURST_1_REQUESTS = 22 # Number of requests in first burst + BURST_2_REQUESTS = 12 # Number of requests in second burst + TOTAL_REQUESTS = BURST_1_REQUESTS + BURST_2_REQUESTS + expected_spend = TOTAL_REQUESTS * SPEND_PER_REQUEST + + # Tolerance for floating-point comparison + TOLERANCE = 1e-10 + + async with aiohttp.ClientSession() as session: + # Create organization + org_response = await create_organization( + session=session, organization_alias=f"test-org-{uuid.uuid4()}" + ) + print("org_response: ", org_response) + org_id = org_response["organization_id"] + + # Create team under organization + team_response = await create_team(session, org_id) + print("team_response: ", team_response) + team_id = team_response["team_id"] + + # Create user + user_response = await create_user(session, org_id) + print("user_response: ", user_response) + user_id = user_response["user_id"] + + # Generate key + key_response = await generate_key(session, user_id, team_id) + print("key_response: ", key_response) + key = key_response["key"] + + # First burst: 12 requests + print(f"Starting first burst of {BURST_1_REQUESTS} requests...") + for i in range(BURST_1_REQUESTS): + response = await chat_completion(session, key) + print(f"Burst 1 - Request {i+1}/{BURST_1_REQUESTS} completed") + + # Wait for spend to be updated + await asyncio.sleep(8) + + # Check intermediate spend + intermediate_key_info = await get_spend_info(session, "key", key) + print(f"After Burst 1 - Key spend: {intermediate_key_info['info']['spend']}") + + # Second burst: 22 requests + print(f"Starting second burst of {BURST_2_REQUESTS} requests...") + for i in range(BURST_2_REQUESTS): + response = await chat_completion(session, key) + print(f"Burst 2 - Request {i+1}/{BURST_2_REQUESTS} completed") + + # Wait for spend to be updated + await asyncio.sleep(8) + + # Get final spend information for each entity + key_info = await get_spend_info(session, "key", key) + team_info = await get_spend_info(session, "team", team_id) + user_info = await get_spend_info(session, "user", user_id) + org_info = await get_spend_info(session, "organization", org_id) + + print(f"Final key spend: {key_info['info']['spend']}") + print(f"Final team spend: {team_info['team_info']['spend']}") + print(f"Final user spend: {user_info['user_info']['spend']}") + print(f"Final org spend: {org_info['spend']}") + + # Verify total spend for each entity + assert ( + abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE + ), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE + ), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE + ), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}" + + assert ( + abs(org_info["spend"] - expected_spend) < TOLERANCE + ), f"Organization spend {org_info['spend']} does not match expected {expected_spend}" From 69d573468575478811cfe3680782141753dab189 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 20:10:58 -0700 Subject: [PATCH 36/40] fix - locking in memory leads to failing tests --- litellm/proxy/db/db_spend_update_writer.py | 7 ++-- litellm/proxy/db/redis_update_buffer.py | 47 ++++++++++------------ litellm/proxy/proxy_config.yaml | 8 ---- litellm/proxy/utils.py | 1 - 4 files changed, 25 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 72692339c8..f46b03b57a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -172,10 +172,9 @@ class DBSpendUpdateWriter: f"track_cost_callback: {entity_type.value}_id is None. Not tracking spend for {entity_type.value}" ) return False - async with prisma_client.in_memory_transaction_lock: - transaction_list[entity_id] = response_cost + transaction_list.get( - entity_id, 0 - ) + transaction_list[entity_id] = response_cost + transaction_list.get( + entity_id, 0 + ) return True except Exception as e: diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 3d27e41d15..f842fc85d2 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -80,35 +80,32 @@ class RedisUpdateBuffer: "redis_cache is None, skipping store_in_memory_spend_updates_in_redis" ) return - async with prisma_client.in_memory_transaction_lock: - db_spend_update_transactions: DBSpendUpdateTransactions = ( - DBSpendUpdateTransactions( - user_list_transactions=prisma_client.user_list_transactions, - end_user_list_transactions=prisma_client.end_user_list_transactions, - key_list_transactions=prisma_client.key_list_transactions, - team_list_transactions=prisma_client.team_list_transactions, - team_member_list_transactions=prisma_client.team_member_list_transactions, - org_list_transactions=prisma_client.org_list_transactions, - ) + db_spend_update_transactions: DBSpendUpdateTransactions = ( + DBSpendUpdateTransactions( + user_list_transactions=prisma_client.user_list_transactions, + end_user_list_transactions=prisma_client.end_user_list_transactions, + key_list_transactions=prisma_client.key_list_transactions, + team_list_transactions=prisma_client.team_list_transactions, + team_member_list_transactions=prisma_client.team_member_list_transactions, + org_list_transactions=prisma_client.org_list_transactions, ) + ) - # only store in redis if there are any updates to commit - if ( - self._number_of_transactions_to_store_in_redis( - db_spend_update_transactions - ) - == 0 - ): - return + # only store in redis if there are any updates to commit + if ( + self._number_of_transactions_to_store_in_redis(db_spend_update_transactions) + == 0 + ): + return - list_of_transactions = [safe_dumps(db_spend_update_transactions)] - await self.redis_cache.async_rpush( - key=REDIS_UPDATE_BUFFER_KEY, - values=list_of_transactions, - ) + list_of_transactions = [safe_dumps(db_spend_update_transactions)] + await self.redis_cache.async_rpush( + key=REDIS_UPDATE_BUFFER_KEY, + values=list_of_transactions, + ) - # clear the in-memory spend updates - RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client) + # clear the in-memory spend updates + RedisUpdateBuffer._clear_all_in_memory_spend_updates(prisma_client) @staticmethod def _number_of_transactions_to_store_in_redis( diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 9ca63f18ea..106003f996 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -5,11 +5,3 @@ model_list: api_key: fake-key api_base: https://exampleopenaiendpoint-production.up.railway.app/ -general_settings: - use_redis_transaction_buffer: True - -litellm_settings: - cache: true - cache_params: - type: redis - supported_call_types: [] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 66dd224dff..d0d4ea9b4e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1121,7 +1121,6 @@ class PrismaClient: self.iam_token_db_auth: Optional[bool] = str_to_bool( os.getenv("IAM_TOKEN_DB_AUTH") ) - self.in_memory_transaction_lock = asyncio.Lock() verbose_proxy_logger.debug("Creating Prisma Client..") try: from prisma import Prisma # type: ignore From 4fb9d27a35459ede5e22056c99a513138f27713c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 20:55:43 -0700 Subject: [PATCH 37/40] use constants for redis buffer in DB --- litellm/constants.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index ea8bc76ac4..e224b3d33e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -18,6 +18,8 @@ DEFAULT_IMAGE_WIDTH = 300 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" +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. #### Networking settings #### From 29d6968f58d546244b1744ea1a42fa5d9ab4a798 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 20:57:28 -0700 Subject: [PATCH 38/40] fix linting --- litellm/proxy/db/redis_update_buffer.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index f842fc85d2..0b090db56a 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -5,10 +5,11 @@ This is to prevent deadlocks and improve reliability """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, 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.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DBSpendUpdateTransactions from litellm.secret_managers.main import str_to_bool @@ -18,9 +19,6 @@ if TYPE_CHECKING: else: PrismaClient = Any -REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" -MAX_REDIS_BUFFER_DEQUEUE_COUNT = 100 - class RedisUpdateBuffer: """ @@ -245,9 +243,9 @@ class RedisUpdateBuffer: # Process each field type for field in transaction_fields: if transaction.get(field): - for entity_id, amount in transaction[field].items(): - combined_transaction[field][entity_id] = ( - combined_transaction[field].get(entity_id, 0) + amount + for entity_id, amount in transaction[field].items(): # type: ignore + combined_transaction[field][entity_id] = ( # type: ignore + combined_transaction[field].get(entity_id, 0) + amount # type: ignore ) return combined_transaction From c95fb6c692d1ee27f3c6e64eabefa83ab3722b45 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 20:59:28 -0700 Subject: [PATCH 39/40] MAX_REDIS_BUFFER_DEQUEUE_COUNT --- litellm/proxy/db/redis_update_buffer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/db/redis_update_buffer.py b/litellm/proxy/db/redis_update_buffer.py index 0b090db56a..f77c839aaf 100644 --- a/litellm/proxy/db/redis_update_buffer.py +++ b/litellm/proxy/db/redis_update_buffer.py @@ -5,7 +5,7 @@ This is to prevent deadlocks and improve reliability """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache From 7c93b19b001e4d7b6c701c92965071c4333395d8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 21:47:26 -0700 Subject: [PATCH 40/40] add migration.sql --- .../migration.sql | 407 ++++++++++++++++++ 1 file changed, 407 insertions(+) create mode 100644 migrations/20250329044652_add_cron_job_table/migration.sql diff --git a/migrations/20250329044652_add_cron_job_table/migration.sql b/migrations/20250329044652_add_cron_job_table/migration.sql new file mode 100644 index 0000000000..edfa99abd5 --- /dev/null +++ b/migrations/20250329044652_add_cron_job_table/migration.sql @@ -0,0 +1,407 @@ +-- CreateEnum +CREATE TYPE "JobStatus" AS ENUM ('ACTIVE', 'INACTIVE'); + +-- CreateTable +CREATE TABLE "LiteLLM_BudgetTable" ( + "budget_id" TEXT NOT NULL, + "max_budget" DOUBLE PRECISION, + "soft_budget" DOUBLE PRECISION, + "max_parallel_requests" INTEGER, + "tpm_limit" BIGINT, + "rpm_limit" BIGINT, + "model_max_budget" JSONB, + "budget_duration" TEXT, + "budget_reset_at" TIMESTAMP(3), + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_BudgetTable_pkey" PRIMARY KEY ("budget_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_CredentialsTable" ( + "credential_id" TEXT NOT NULL, + "credential_name" TEXT NOT NULL, + "credential_values" JSONB NOT NULL, + "credential_info" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_CredentialsTable_pkey" PRIMARY KEY ("credential_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_ProxyModelTable" ( + "model_id" TEXT NOT NULL, + "model_name" TEXT NOT NULL, + "litellm_params" JSONB NOT NULL, + "model_info" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_ProxyModelTable_pkey" PRIMARY KEY ("model_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_OrganizationTable" ( + "organization_id" TEXT NOT NULL, + "organization_alias" TEXT NOT NULL, + "budget_id" TEXT NOT NULL, + "metadata" JSONB NOT NULL DEFAULT '{}', + "models" TEXT[], + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "model_spend" JSONB NOT NULL DEFAULT '{}', + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_OrganizationTable_pkey" PRIMARY KEY ("organization_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_ModelTable" ( + "id" SERIAL NOT NULL, + "aliases" JSONB, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_ModelTable_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_TeamTable" ( + "team_id" TEXT NOT NULL, + "team_alias" TEXT, + "organization_id" TEXT, + "admins" TEXT[], + "members" TEXT[], + "members_with_roles" JSONB NOT NULL DEFAULT '{}', + "metadata" JSONB NOT NULL DEFAULT '{}', + "max_budget" DOUBLE PRECISION, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "models" TEXT[], + "max_parallel_requests" INTEGER, + "tpm_limit" BIGINT, + "rpm_limit" BIGINT, + "budget_duration" TEXT, + "budget_reset_at" TIMESTAMP(3), + "blocked" BOOLEAN NOT NULL DEFAULT false, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "model_spend" JSONB NOT NULL DEFAULT '{}', + "model_max_budget" JSONB NOT NULL DEFAULT '{}', + "model_id" INTEGER, + + CONSTRAINT "LiteLLM_TeamTable_pkey" PRIMARY KEY ("team_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_UserTable" ( + "user_id" TEXT NOT NULL, + "user_alias" TEXT, + "team_id" TEXT, + "sso_user_id" TEXT, + "organization_id" TEXT, + "password" TEXT, + "teams" TEXT[] DEFAULT ARRAY[]::TEXT[], + "user_role" TEXT, + "max_budget" DOUBLE PRECISION, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "user_email" TEXT, + "models" TEXT[], + "metadata" JSONB NOT NULL DEFAULT '{}', + "max_parallel_requests" INTEGER, + "tpm_limit" BIGINT, + "rpm_limit" BIGINT, + "budget_duration" TEXT, + "budget_reset_at" TIMESTAMP(3), + "allowed_cache_controls" TEXT[] DEFAULT ARRAY[]::TEXT[], + "model_spend" JSONB NOT NULL DEFAULT '{}', + "model_max_budget" JSONB NOT NULL DEFAULT '{}', + "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_UserTable_pkey" PRIMARY KEY ("user_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_VerificationToken" ( + "token" TEXT NOT NULL, + "key_name" TEXT, + "key_alias" TEXT, + "soft_budget_cooldown" BOOLEAN NOT NULL DEFAULT false, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "expires" TIMESTAMP(3), + "models" TEXT[], + "aliases" JSONB NOT NULL DEFAULT '{}', + "config" JSONB NOT NULL DEFAULT '{}', + "user_id" TEXT, + "team_id" TEXT, + "permissions" JSONB NOT NULL DEFAULT '{}', + "max_parallel_requests" INTEGER, + "metadata" JSONB NOT NULL DEFAULT '{}', + "blocked" BOOLEAN, + "tpm_limit" BIGINT, + "rpm_limit" BIGINT, + "max_budget" DOUBLE PRECISION, + "budget_duration" TEXT, + "budget_reset_at" TIMESTAMP(3), + "allowed_cache_controls" TEXT[] DEFAULT ARRAY[]::TEXT[], + "model_spend" JSONB NOT NULL DEFAULT '{}', + "model_max_budget" JSONB NOT NULL DEFAULT '{}', + "budget_id" TEXT, + "organization_id" TEXT, + "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "created_by" TEXT, + "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "updated_by" TEXT, + + CONSTRAINT "LiteLLM_VerificationToken_pkey" PRIMARY KEY ("token") +); + +-- CreateTable +CREATE TABLE "LiteLLM_EndUserTable" ( + "user_id" TEXT NOT NULL, + "alias" TEXT, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "allowed_model_region" TEXT, + "default_model" TEXT, + "budget_id" TEXT, + "blocked" BOOLEAN NOT NULL DEFAULT false, + + CONSTRAINT "LiteLLM_EndUserTable_pkey" PRIMARY KEY ("user_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_Config" ( + "param_name" TEXT NOT NULL, + "param_value" JSONB, + + CONSTRAINT "LiteLLM_Config_pkey" PRIMARY KEY ("param_name") +); + +-- CreateTable +CREATE TABLE "LiteLLM_SpendLogs" ( + "request_id" TEXT NOT NULL, + "call_type" TEXT NOT NULL, + "api_key" TEXT NOT NULL DEFAULT '', + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "total_tokens" INTEGER NOT NULL DEFAULT 0, + "prompt_tokens" INTEGER NOT NULL DEFAULT 0, + "completion_tokens" INTEGER NOT NULL DEFAULT 0, + "startTime" TIMESTAMP(3) NOT NULL, + "endTime" TIMESTAMP(3) NOT NULL, + "completionStartTime" TIMESTAMP(3), + "model" TEXT NOT NULL DEFAULT '', + "model_id" TEXT DEFAULT '', + "model_group" TEXT DEFAULT '', + "custom_llm_provider" TEXT DEFAULT '', + "api_base" TEXT DEFAULT '', + "user" TEXT DEFAULT '', + "metadata" JSONB DEFAULT '{}', + "cache_hit" TEXT DEFAULT '', + "cache_key" TEXT DEFAULT '', + "request_tags" JSONB DEFAULT '[]', + "team_id" TEXT, + "end_user" TEXT, + "requester_ip_address" TEXT, + "messages" JSONB DEFAULT '{}', + "response" JSONB DEFAULT '{}', + + CONSTRAINT "LiteLLM_SpendLogs_pkey" PRIMARY KEY ("request_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_ErrorLogs" ( + "request_id" TEXT NOT NULL, + "startTime" TIMESTAMP(3) NOT NULL, + "endTime" TIMESTAMP(3) NOT NULL, + "api_base" TEXT NOT NULL DEFAULT '', + "model_group" TEXT NOT NULL DEFAULT '', + "litellm_model_name" TEXT NOT NULL DEFAULT '', + "model_id" TEXT NOT NULL DEFAULT '', + "request_kwargs" JSONB NOT NULL DEFAULT '{}', + "exception_type" TEXT NOT NULL DEFAULT '', + "exception_string" TEXT NOT NULL DEFAULT '', + "status_code" TEXT NOT NULL DEFAULT '', + + CONSTRAINT "LiteLLM_ErrorLogs_pkey" PRIMARY KEY ("request_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_UserNotifications" ( + "request_id" TEXT NOT NULL, + "user_id" TEXT NOT NULL, + "models" TEXT[], + "justification" TEXT NOT NULL, + "status" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_UserNotifications_pkey" PRIMARY KEY ("request_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_TeamMembership" ( + "user_id" TEXT NOT NULL, + "team_id" TEXT NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "budget_id" TEXT, + + CONSTRAINT "LiteLLM_TeamMembership_pkey" PRIMARY KEY ("user_id","team_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_OrganizationMembership" ( + "user_id" TEXT NOT NULL, + "organization_id" TEXT NOT NULL, + "user_role" TEXT, + "spend" DOUBLE PRECISION DEFAULT 0.0, + "budget_id" TEXT, + "created_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_OrganizationMembership_pkey" PRIMARY KEY ("user_id","organization_id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_InvitationLink" ( + "id" TEXT NOT NULL, + "user_id" TEXT NOT NULL, + "is_accepted" BOOLEAN NOT NULL DEFAULT false, + "accepted_at" TIMESTAMP(3), + "expires_at" TIMESTAMP(3) NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL, + "created_by" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL, + "updated_by" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_InvitationLink_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_AuditLog" ( + "id" TEXT NOT NULL, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "changed_by" TEXT NOT NULL DEFAULT '', + "changed_by_api_key" TEXT NOT NULL DEFAULT '', + "action" TEXT NOT NULL, + "table_name" TEXT NOT NULL, + "object_id" TEXT NOT NULL, + "before_value" JSONB, + "updated_values" JSONB, + + CONSTRAINT "LiteLLM_AuditLog_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_DailyUserSpend" ( + "id" TEXT NOT NULL, + "user_id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "model" TEXT NOT NULL, + "model_group" TEXT, + "custom_llm_provider" TEXT, + "prompt_tokens" INTEGER NOT NULL DEFAULT 0, + "completion_tokens" INTEGER NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "api_requests" INTEGER NOT NULL DEFAULT 0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_DailyUserSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateTable +CREATE TABLE "LiteLLM_CronJob" ( + "cronjob_id" TEXT NOT NULL, + "pod_id" TEXT NOT NULL, + "status" "JobStatus" NOT NULL DEFAULT 'INACTIVE', + "last_updated" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "ttl" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_CronJob_pkey" PRIMARY KEY ("cronjob_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_CredentialsTable_credential_name_key" ON "LiteLLM_CredentialsTable"("credential_name"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_TeamTable_model_id_key" ON "LiteLLM_TeamTable"("model_id"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_UserTable_sso_user_id_key" ON "LiteLLM_UserTable"("sso_user_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_SpendLogs_startTime_idx" ON "LiteLLM_SpendLogs"("startTime"); + +-- CreateIndex +CREATE INDEX "LiteLLM_SpendLogs_end_user_idx" ON "LiteLLM_SpendLogs"("end_user"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_OrganizationMembership_user_id_organization_id_key" ON "LiteLLM_OrganizationMembership"("user_id", "organization_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyUserSpend_date_idx" ON "LiteLLM_DailyUserSpend"("date"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyUserSpend_user_id_idx" ON "LiteLLM_DailyUserSpend"("user_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyUserSpend_api_key_idx" ON "LiteLLM_DailyUserSpend"("api_key"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyUserSpend_model_idx" ON "LiteLLM_DailyUserSpend"("model"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key" ON "LiteLLM_DailyUserSpend"("user_id", "date", "api_key", "model", "custom_llm_provider"); + +-- AddForeignKey +ALTER TABLE "LiteLLM_OrganizationTable" ADD CONSTRAINT "LiteLLM_OrganizationTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_TeamTable" ADD CONSTRAINT "LiteLLM_TeamTable_model_id_fkey" FOREIGN KEY ("model_id") REFERENCES "LiteLLM_ModelTable"("id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_UserTable" ADD CONSTRAINT "LiteLLM_UserTable_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_VerificationToken" ADD CONSTRAINT "LiteLLM_VerificationToken_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_EndUserTable" ADD CONSTRAINT "LiteLLM_EndUserTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_TeamMembership" ADD CONSTRAINT "LiteLLM_TeamMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_organization_id_fkey" FOREIGN KEY ("organization_id") REFERENCES "LiteLLM_OrganizationTable"("organization_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_OrganizationMembership" ADD CONSTRAINT "LiteLLM_OrganizationMembership_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_created_by_fkey" FOREIGN KEY ("created_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE; + +-- AddForeignKey +ALTER TABLE "LiteLLM_InvitationLink" ADD CONSTRAINT "LiteLLM_InvitationLink_updated_by_fkey" FOREIGN KEY ("updated_by") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE RESTRICT ON UPDATE CASCADE;