From e30de5d7fdd0f2399183d6d8923d03e984d9f271 Mon Sep 17 00:00:00 2001 From: pnookala-godaddy Date: Tue, 17 Mar 2026 16:07:50 -0700 Subject: [PATCH] refactor: consolidate deadlock helpers, add sorted() comments, revert _update_entity_spend_in_db - Move PRISMA_DEADLOCK_CODE and _is_deadlock_error to _types.py (single source of truth) - Add comments explaining why sorted() is used (consistent lock ordering) - Revert _update_entity_spend_in_db to original (low-contention tables don't need deadlock handling) --- litellm/proxy/_types.py | 12 +++++++++++ litellm/proxy/db/db_spend_update_writer.py | 24 ++++++++-------------- litellm/proxy/utils.py | 11 ++-------- 3 files changed, 22 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b7ac4212cbd..27f652f2241 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4,6 +4,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union import httpx +import prisma.errors from pydantic import ( BaseModel, ConfigDict, @@ -3572,6 +3573,17 @@ DB_CONNECTION_ERROR_TYPES = ( httpx.ReadTimeout, ) +PRISMA_DEADLOCK_CODE = "P2034" +"""Prisma error code for PostgreSQL deadlock / write conflict.""" + + +def _is_deadlock_error(e: Exception) -> bool: + """Check if a Prisma error is a PostgreSQL deadlock / write conflict (P2034).""" + return ( + isinstance(e, prisma.errors.DataError) + and getattr(e, "code", None) == PRISMA_DEADLOCK_CODE + ) + class SSOUserDefinedValues(TypedDict): models: List[str] diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2775f4865a9..38d12f0f80a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -25,8 +25,6 @@ from typing import ( overload, ) -import prisma.errors - import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache, RedisCache @@ -48,6 +46,7 @@ from litellm.proxy._types import ( SpendLogsPayload, SpendUpdateQueueItem, ToolDiscoveryQueueItem, + _is_deadlock_error, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, @@ -66,14 +65,6 @@ else: PrismaClient = Any ProxyLogging = Any -PRISMA_DEADLOCK_CODE = "P2034" - - -def _is_deadlock_error(e: Exception) -> bool: - """Check if a Prisma error is a PostgreSQL deadlock / write conflict (P2034).""" - return isinstance(e, prisma.errors.DataError) and getattr(e, "code", None) == PRISMA_DEADLOCK_CODE - - class DBSpendUpdateWriter: """ Module responsible for @@ -1093,6 +1084,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for user_id, response_cost in sorted(user_list_transactions.items()): batcher.litellm_usertable.update_many( where={"user_id": user_id}, @@ -1148,6 +1140,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for token, response_cost in sorted(key_list_transactions.items()): batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists where={"token": token}, @@ -1192,6 +1185,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for team_id, response_cost in sorted(team_list_transactions.items()): verbose_proxy_logger.debug( "Updating spend for team id={} by {}".format( @@ -1250,6 +1244,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for key, response_cost in sorted(team_member_list_transactions.items()): # key is "team_id::::user_id::" team_id = key.split("::")[1] @@ -1307,6 +1302,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for org_id, response_cost in sorted(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}, @@ -1395,7 +1391,7 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for entity_id, response_cost in sorted(transactions.items()): + for entity_id, response_cost in transactions.items(): verbose_proxy_logger.debug( f"Updating spend for {entity_name} {where_field}={entity_id} by {response_cost}" ) @@ -1411,12 +1407,8 @@ class DBSpendUpdateWriter: start_time=start_time, proxy_logging_obj=proxy_logging_obj, ) - # Randomized backoff to reduce repeated collisions across pods - await asyncio.sleep(random.uniform(2**i, 2**(i+1))) + await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - if _is_deadlock_error(e) and i < n_retry_times: - await asyncio.sleep(random.uniform(2**i, 2**(i+1))) - continue _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 d324daefdee..33a08c403f7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -25,8 +25,6 @@ from typing import ( overload, ) -import prisma.errors - from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME, MAX_TEAM_LIST_LIMIT from litellm.proxy._types import ( @@ -36,18 +34,12 @@ from litellm.proxy._types import ( ProxyException, SpendLogsMetadata, SpendLogsPayload, + _is_deadlock_error, ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes, CallTypesLiteral -PRISMA_DEADLOCK_CODE = "P2034" - - -def _is_deadlock_error(e: Exception) -> bool: - """Check if a Prisma error is a PostgreSQL deadlock / write conflict (P2034).""" - return isinstance(e, prisma.errors.DataError) and getattr(e, "code", None) == PRISMA_DEADLOCK_CODE - try: from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( BaseEmailLogger, @@ -4549,6 +4541,7 @@ class ProxyUpdateSpend: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: + # Sort by ID for consistent lock ordering across pods to prevent deadlocks for end_user_id, response_cost in sorted(end_user_list_transactions.items()): if litellm.max_end_user_budget is not None: pass