attempt to avoid/minimize deadlocks

Author: CAFxX

litellm/proxy/db/db_spend_update_writer.py

@@ -10,6 +10,7 @@
 import os
 import time
 import traceback
+import random
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Union, cast, overload
 
@@ -849,8 +850,23 @@ async def _update_daily_spend(
         try:
             for i in range(n_retry_times + 1):
                 try:
+                    # Sort the transactions to ensure consistent processing order and minimize the probability of deadlocks
                     transactions_to_process = dict(
-                        list(daily_spend_transactions.items())[:BATCH_SIZE]
+                        sorted(
+                            daily_spend_transactions.items(),
+                            # Normally to avoid deadlocks we would sort by the index, but since we have sprinkled indexes
+                            # on our schema like we're discount Salt Bae, we just sort by all fields that have an index,
+                            # in an ad-hoc (but hopefully sensible) order of indexes.
+                            # If _update_daily_spend ever gets the ability to write to multiple tables at once, the sorting
+                            # should sort by the table first.
+                            key=lambda x: (
+                                x[1]["date"],
+                                x[1][entity_id_field],
+                                x[1]["api_key"],
+                                x[1]["model"],
+                                x[1]["custom_llm_provider"],
+                            ),
+                        )[:BATCH_SIZE]
                     )
 
                     if len(transactions_to_process) == 0:
@@ -893,7 +909,8 @@ async def _update_daily_spend(
                                 "model_group": transaction.get("model_group"),
                                 "mcp_namespaced_tool_name": transaction.get(
                                     "mcp_namespaced_tool_name"
-                                ) or "",
+                                )
+                                or "",
                                 "custom_llm_provider": transaction.get(
                                     "custom_llm_provider"
                                 ),
@@ -909,13 +926,13 @@ async def _update_daily_spend(
 
                             # Add cache-related fields if they exist
                             if "cache_read_input_tokens" in transaction:
-                                common_data["cache_read_input_tokens"] = (
-                                    transaction.get("cache_read_input_tokens", 0)
-                                )
+                                common_data[
+                                    "cache_read_input_tokens"
+                                ] = transaction.get("cache_read_input_tokens", 0)
                             if "cache_creation_input_tokens" in transaction:
-                                common_data["cache_creation_input_tokens"] = (
-                                    transaction.get("cache_creation_input_tokens", 0)
-                                )
+                                common_data[
+                                    "cache_creation_input_tokens"
+                                ] = transaction.get("cache_creation_input_tokens", 0)
 
                             # Create update data structure
                             update_data = {
@@ -976,7 +993,9 @@ async def _update_daily_spend(
                             start_time=start_time,
                             proxy_logging_obj=proxy_logging_obj,
                         )
-                    await asyncio.sleep(2**i)
+                    await asyncio.sleep(
+                        random.uniform(2**i, 2 ** (i + 1))
+                    )  # Sleep a random amount to avoid retrying and deadlocking again
 
         except Exception as e:
             if "transactions_to_process" in locals():