tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-21 00:24:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(proxy): sort spend updates to prevent DB deadlocks

Browse Source 
Iterate user/key/team/team_member/org/end_user/tag spend dicts in sorted
order inside each Prisma transaction so concurrent pods acquire row locks
in the same order, avoiding PostgreSQL deadlocks under load.
This commit is contained in:
pnookala-godaddy
2026-05-05 13:49:19 -07:00
parent 934ecdca78
commit 209bd0b906
3 changed files with 174 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/db/db_spend_update_writer.py
+27 -21
View File
@@ -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::<v>::user_id::<v>" 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::<value>::user_id::<value>"
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}"
)
litellm/proxy/utils.py
+4 -4
View File
@@ -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(
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
+143
View File
@@ -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