From b70659841527198aa08d633f3d4d3a6896934f22 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 15:27:06 -0800 Subject: [PATCH 1/6] (v0) fixes --- litellm/proxy/proxy_server.py | 13 ++++++----- litellm/proxy/utils.py | 43 +++++++++++++++++++++++++++-------- 2 files changed, 40 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 47b8c15359d..4ae3840cf10 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9,6 +9,10 @@ import warnings import importlib import warnings +import logging + +logging.getLogger("prisma").setLevel(logging.DEBUG) + def showwarning(message, category, filename, lineno, file=None, line=None): traceback_info = f"{filename}:{lineno}: {category.__name__}: {message}\n" @@ -1138,6 +1142,7 @@ async def update_database( ) # set cooldown on alert soft_budget_cooldown = True + # track cost per model, for the given key spend_per_model = existing_spend_obj.model_spend or {} current_model = kwargs.get("model") @@ -1153,11 +1158,7 @@ async def update_database( # Update the cost column for the given token await prisma_client.update_data( token=token, - data={ - "spend": new_spend, - "model_spend": spend_per_model, - "soft_budget_cooldown": soft_budget_cooldown, - }, + data={"spend": new_spend, "model_spend": spend_per_model}, ) valid_token = user_api_key_cache.get_cache(key=token) @@ -1211,9 +1212,9 @@ async def update_database( payload["spend"] = response_cost if prisma_client is not None: await prisma_client.insert_data(data=payload, table_name="spend") + elif custom_db_client is not None: await custom_db_client.insert_data(payload, table_name="spend") - except Exception as e: verbose_proxy_logger.info(f"Update Spend Logs DB failed to execute") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fc90c3b7b81..f38b78cfa61 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -64,7 +64,7 @@ class ProxyLogging: litellm.callbacks.append(self.max_parallel_request_limiter) litellm.callbacks.append(self.max_budget_limiter) litellm.callbacks.append(self.cache_control_check) - litellm.callbacks.append(self.response_taking_too_long_callback) + # litellm.callbacks.append(self.response_taking_too_long_callback) for callback in litellm.callbacks: if callback not in litellm.input_callback: litellm.input_callback.append(callback) @@ -362,7 +362,7 @@ class ProxyLogging: else: raise Exception("Missing SENTRY_DSN from environment") - async def failure_handler(self, original_exception): + async def failure_handler(self, original_exception, traceback_str=""): """ Log failed db read/writes @@ -373,6 +373,7 @@ class ProxyLogging: error_message = original_exception.detail else: error_message = str(original_exception) + error_message += traceback_str asyncio.create_task( self.alerting_handler( message=f"DB read/write call failed: {error_message}", @@ -706,8 +707,13 @@ class PrismaClient: ) return response except Exception as e: + import traceback + + tracback_str = traceback.format_exc() asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=tracback_str + ) ) raise e @@ -912,9 +918,10 @@ class PrismaClient: return response elif table_name == "team": if query_type == "find_unique": - response = await self.db.litellm_teamtable.find_unique( - where={"team_id": team_id} # type: ignore - ) + response = None + # response = await self.db.litellm_teamtable.find_unique( + # where={"team_id": team_id} # type: ignore + # ) elif query_type == "find_all" and user_id is not None: response = await self.db.litellm_teamtable.find_many( where={ @@ -971,8 +978,12 @@ class PrismaClient: import traceback traceback.print_exc() + # get tracback + traceback_string = traceback.format_exc() asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=traceback_string + ) ) raise e @@ -1093,8 +1104,12 @@ class PrismaClient: except Exception as e: print_verbose(f"LiteLLM Prisma Client Exception: {e}") + import traceback + + traceback_str = traceback.format_exc() + print_verbose(f"Traceback: {traceback_str}") asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler(original_exception=e, traceback_str=traceback_str) # type: ignore # noqa=traceback_str) ) raise e @@ -1277,8 +1292,12 @@ class PrismaClient: "\033[91m" + f"DB User Table Batch update succeeded" + "\033[0m" ) except Exception as e: + import traceback + asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=traceback.format_exc() + ) ) print_verbose("\033[91m" + f"DB write failed: {e}" + "\033[0m") raise e @@ -1331,8 +1350,12 @@ class PrismaClient: where={"team_id": {"in": team_id_list}} ) except Exception as e: + import traceback + asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=traceback.format_exc() + ) ) raise e From d0953a13499dc69f4ea27c85e2707ce2b5a98eee Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 17:04:25 -0800 Subject: [PATCH 2/6] (fix) chore --- litellm/proxy/proxy_server.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4ae3840cf10..17a213f5138 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9,10 +9,6 @@ import warnings import importlib import warnings -import logging - -logging.getLogger("prisma").setLevel(logging.DEBUG) - def showwarning(message, category, filename, lineno, file=None, line=None): traceback_info = f"{filename}:{lineno}: {category.__name__}: {message}\n" @@ -1142,7 +1138,6 @@ async def update_database( ) # set cooldown on alert soft_budget_cooldown = True - # track cost per model, for the given key spend_per_model = existing_spend_obj.model_spend or {} current_model = kwargs.get("model") @@ -1158,7 +1153,11 @@ async def update_database( # Update the cost column for the given token await prisma_client.update_data( token=token, - data={"spend": new_spend, "model_spend": spend_per_model}, + data={ + "spend": new_spend, + "model_spend": spend_per_model, + "soft_budget_cooldown": soft_budget_cooldown, + }, ) valid_token = user_api_key_cache.get_cache(key=token) @@ -1212,7 +1211,6 @@ async def update_database( payload["spend"] = response_cost if prisma_client is not None: await prisma_client.insert_data(data=payload, table_name="spend") - elif custom_db_client is not None: await custom_db_client.insert_data(payload, table_name="spend") except Exception as e: From 3df618bf048e6ef3e39cc724c780aaa9042af5d6 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 17:05:01 -0800 Subject: [PATCH 3/6] (fix) cleanup --- litellm/proxy/proxy_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 17a213f5138..47b8c15359d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1213,6 +1213,7 @@ async def update_database( await prisma_client.insert_data(data=payload, table_name="spend") elif custom_db_client is not None: await custom_db_client.insert_data(payload, table_name="spend") + except Exception as e: verbose_proxy_logger.info(f"Update Spend Logs DB failed to execute") From 735728c4ddada72b388ff2abc8122f175dfe19a2 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 18:16:59 -0800 Subject: [PATCH 4/6] (feat) add better debugging when DB read/write fails --- litellm/proxy/tests/load_test_completion.py | 6 ++- litellm/proxy/utils.py | 57 +++++++++++++++------ 2 files changed, 44 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/tests/load_test_completion.py b/litellm/proxy/tests/load_test_completion.py index b5d711d8131..d708f303684 100644 --- a/litellm/proxy/tests/load_test_completion.py +++ b/litellm/proxy/tests/load_test_completion.py @@ -4,7 +4,9 @@ import uuid import traceback -litellm_client = AsyncOpenAI(base_url="http://0.0.0.0:8000", api_key="any") +litellm_client = AsyncOpenAI( + base_url="http://0.0.0.0:4000", api_key="sk-iNwH_oOtAQ6syi_2gkEOpQ" +) async def litellm_completion(): @@ -27,7 +29,7 @@ async def litellm_completion(): async def main(): for i in range(150): start = time.time() - n = 150 # Number of concurrent tasks + n = 2000 # Number of concurrent tasks tasks = [litellm_completion() for _ in range(n)] chat_completions = await asyncio.gather(*tasks) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f38b78cfa61..959f2b71462 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -373,7 +373,8 @@ class ProxyLogging: error_message = original_exception.detail else: error_message = str(original_exception) - error_message += traceback_str + if isinstance(traceback_str, str): + error_message += traceback_str[:1000] asyncio.create_task( self.alerting_handler( message=f"DB read/write call failed: {error_message}", @@ -709,10 +710,12 @@ class PrismaClient: except Exception as e: import traceback - tracback_str = traceback.format_exc() + error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( self.proxy_logging_obj.failure_handler( - original_exception=e, traceback_str=tracback_str + original_exception=e, traceback_str=error_traceback ) ) raise e @@ -974,15 +977,14 @@ class PrismaClient: response.expires = response.expires.isoformat() return response except Exception as e: - print_verbose(f"LiteLLM Prisma Client Exception: {e}") import traceback - traceback.print_exc() - # get tracback - traceback_string = traceback.format_exc() + error_msg = f"LiteLLM Prisma Client Exception get_data: {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( self.proxy_logging_obj.failure_handler( - original_exception=e, traceback_str=traceback_string + original_exception=e, traceback_str=error_traceback ) ) raise e @@ -1103,13 +1105,15 @@ class PrismaClient: return new_user_notification_row except Exception as e: - print_verbose(f"LiteLLM Prisma Client Exception: {e}") import traceback - traceback_str = traceback.format_exc() - print_verbose(f"Traceback: {traceback_str}") + error_msg = f"LiteLLM Prisma Client Exception in insert_data: {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e, traceback_str=traceback_str) # type: ignore # noqa=traceback_str) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=error_traceback + ) ) raise e @@ -1294,12 +1298,14 @@ class PrismaClient: except Exception as e: import traceback + error_msg = f"LiteLLM Prisma Client Exception - update_data: {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( self.proxy_logging_obj.failure_handler( - original_exception=e, traceback_str=traceback.format_exc() + original_exception=e, traceback_str=error_traceback ) ) - print_verbose("\033[91m" + f"DB write failed: {e}" + "\033[0m") raise e # Define a retrying strategy with exponential backoff @@ -1352,9 +1358,12 @@ class PrismaClient: except Exception as e: import traceback + error_msg = f"LiteLLM Prisma Client Exception - delete_data: {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( self.proxy_logging_obj.failure_handler( - original_exception=e, traceback_str=traceback.format_exc() + original_exception=e, traceback_str=error_traceback ) ) raise e @@ -1378,8 +1387,15 @@ class PrismaClient: ) await self.db.connect() except Exception as e: + import traceback + + error_msg = f"LiteLLM Prisma Client Exception connect(): {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=error_traceback + ) ) raise e @@ -1395,8 +1411,15 @@ class PrismaClient: try: await self.db.disconnect() except Exception as e: + import traceback + + error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e}" + print_verbose(error_msg) + error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( - self.proxy_logging_obj.failure_handler(original_exception=e) + self.proxy_logging_obj.failure_handler( + original_exception=e, traceback_str=error_traceback + ) ) raise e From 00666566c20b3c5b63e6078acc5cd533e8a1f191 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 18:54:04 -0800 Subject: [PATCH 5/6] (fix) chore --- litellm/proxy/utils.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9e3a374ecac..6946904c75c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -921,10 +921,9 @@ class PrismaClient: return response elif table_name == "team": if query_type == "find_unique": - response = None - # response = await self.db.litellm_teamtable.find_unique( - # where={"team_id": team_id} # type: ignore - # ) + response = await self.db.litellm_teamtable.find_unique( + where={"team_id": team_id} # type: ignore + ) elif query_type == "find_all" and user_id is not None: response = await self.db.litellm_teamtable.find_many( where={ From 52e59e050c02dc234c61946c061d46e23e8a5b71 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Tue, 5 Mar 2024 18:57:26 -0800 Subject: [PATCH 6/6] (fix) use str(e) --- litellm/proxy/utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6946904c75c..527b66f8c0c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -710,7 +710,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {e}" + error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -978,7 +978,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception get_data: {e}" + error_msg = f"LiteLLM Prisma Client Exception get_data: {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -1106,7 +1106,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception in insert_data: {e}" + error_msg = f"LiteLLM Prisma Client Exception in insert_data: {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -1297,7 +1297,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception - update_data: {e}" + error_msg = f"LiteLLM Prisma Client Exception - update_data: {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -1357,7 +1357,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception - delete_data: {e}" + error_msg = f"LiteLLM Prisma Client Exception - delete_data: {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -1388,7 +1388,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception connect(): {e}" + error_msg = f"LiteLLM Prisma Client Exception connect(): {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task( @@ -1412,7 +1412,7 @@ class PrismaClient: except Exception as e: import traceback - error_msg = f"LiteLLM Prisma Client Exception disconnect(): {e}" + error_msg = f"LiteLLM Prisma Client Exception disconnect(): {str(e)}" print_verbose(error_msg) error_traceback = error_msg + "\n" + traceback.format_exc() asyncio.create_task(