diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c06e1850d9..418fc2d2f0 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1133,10 +1133,12 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - user_id, - response_cost, - ) in user_list_transactions.items(): + # Sort by ID for consistent lock ordering across pods to prevent deadlocks. + # batch_() issues statements sequentially within the tx, so iteration + # order = lock acquisition order. + for user_id, response_cost in sorted( + user_list_transactions.items() + ): batcher.litellm_usertable.update_many( where={"user_id": user_id}, data={"spend": {"increment": response_cost}}, @@ -1188,10 +1190,10 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - token, - response_cost, - ) in key_list_transactions.items(): + # Sort by token 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}, data={ @@ -1232,10 +1234,10 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - team_id, - response_cost, - ) in team_list_transactions.items(): + # Sort by team_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( team_id, response_cost @@ -1290,10 +1292,11 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - key, - response_cost, - ) in team_member_list_transactions.items(): + # Sort by composite key for consistent lock ordering across pods to prevent deadlocks. + # Key format "team_id::::user_id::" makes the string sort equivalent to sorting by (team_id, user_id). + for key, response_cost in sorted( + team_member_list_transactions.items() + ): # key is "team_id::::user_id::" team_id = key.split("::")[1] user_id = key.split("::")[3] @@ -1350,10 +1353,10 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - org_id, - response_cost, - ) in org_list_transactions.items(): + # Sort by org_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}, data={"spend": {"increment": response_cost}}, @@ -1441,7 +1444,10 @@ class DBSpendUpdateWriter: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for entity_id, response_cost in transactions.items(): + # Sort by entity_id for consistent lock ordering across pods to prevent deadlocks. + for entity_id, response_cost in sorted( + transactions.items() + ): verbose_proxy_logger.debug( f"Updating spend for {entity_name} {where_field}={entity_id} by {response_cost}" ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d2dfa17751..800a4be37a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4874,10 +4874,10 @@ class ProxyUpdateSpend: timeout=timedelta(seconds=60) ) as transaction: async with transaction.batch_() as batcher: - for ( - end_user_id, - response_cost, - ) in end_user_list_transactions.items(): + # Sort by end_user_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 batcher.litellm_endusertable.upsert( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4d58434934..9d4d3c4a56 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1508,3 +1508,146 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called() mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called() mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() + + +@pytest.mark.parametrize( + "bucket_name,input_dict,table_attr,method_name,where_key,expected_order", + [ + pytest.param( + "user_list_transactions", + {"user_c": 0.1, "user_a": 0.2, "user_b": 0.3}, + "litellm_usertable", + "update_many", + "user_id", + ["user_a", "user_b", "user_c"], + id="user", + ), + pytest.param( + "key_list_transactions", + {"tok_c": 0.1, "tok_a": 0.2, "tok_b": 0.3}, + "litellm_verificationtoken", + "update_many", + "token", + ["tok_a", "tok_b", "tok_c"], + id="key", + ), + pytest.param( + "team_list_transactions", + {"team_c": 0.1, "team_a": 0.2, "team_b": 0.3}, + "litellm_teamtable", + "update_many", + "team_id", + ["team_a", "team_b", "team_c"], + id="team", + ), + pytest.param( + "team_member_list_transactions", + { + "team_id::team_c::user_id::user_x": 0.1, + "team_id::team_a::user_id::user_x": 0.2, + "team_id::team_b::user_id::user_x": 0.3, + }, + "litellm_teammembership", + "update_many", + "team_id", + ["team_a", "team_b", "team_c"], + id="team_member", + ), + pytest.param( + "org_list_transactions", + {"org_c": 0.1, "org_a": 0.2, "org_b": 0.3}, + "litellm_organizationtable", + "update_many", + "organization_id", + ["org_a", "org_b", "org_c"], + id="org", + ), + pytest.param( + "end_user_list_transactions", + {"eu_c": 0.1, "eu_a": 0.2, "eu_b": 0.3}, + "litellm_endusertable", + "upsert", + "user_id", + ["eu_a", "eu_b", "eu_c"], + id="end_user", + ), + pytest.param( + "tag_list_transactions", + {"prod": 0.1, "customer-x": 0.2, "test": 0.3}, + "litellm_tagtable", + "update_many", + "tag_name", + ["customer-x", "prod", "test"], + id="tag", + ), + pytest.param( + "agent_list_transactions", + {"agent_c": 0.1, "agent_a": 0.2, "agent_b": 0.3}, + "litellm_agentstable", + "update_many", + "agent_id", + ["agent_a", "agent_b", "agent_c"], + id="agent", + ), + ], +) +@pytest.mark.asyncio +async def test_commit_spend_updates_iterates_in_sorted_order( + bucket_name, input_dict, table_attr, method_name, where_key, expected_order +): + """ + Every spend-bucket code path in _commit_spend_updates_to_db must iterate + in sorted order so concurrent pods acquire row locks in the same order + and avoid PostgreSQL deadlocks. Covers the 5 direct loops (user/key/team/ + team_member/org), the end_user path in ProxyUpdateSpend.update_end_user_spend, + and the shared _update_entity_spend_in_db helper (tag, agent). + """ + db_writer = DBSpendUpdateWriter() + + captured_where_values = [] + + def capture(*, where, data): + captured_where_values.append(where[where_key]) + + mock_batcher = MagicMock() + table_mock = MagicMock() + setattr(table_mock, method_name, MagicMock(side_effect=capture)) + setattr(mock_batcher, table_attr, table_mock) + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + mock_proxy_logging = MagicMock() + mock_proxy_logging.call_details = {} + + buckets = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + buckets[bucket_name] = input_dict + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=3, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=buckets, + ) + + assert captured_where_values == expected_order