Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
pnookala-godaddy wants to merge 1 commit into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 27 additions & 21 deletions litellm/proxy/db/db_spend_update_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1133,10 +1133,12 @@ async def _commit_spend_updates_to_db( # noqa: PLR0915
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}},
Expand Down Expand Up @@ -1188,10 +1190,10 @@ async def _commit_spend_updates_to_db( # noqa: PLR0915
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={
Expand Down Expand Up @@ -1232,10 +1234,10 @@ async def _commit_spend_updates_to_db( # noqa: PLR0915
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
Expand Down Expand Up @@ -1290,10 +1292,11 @@ async def _commit_spend_updates_to_db( # noqa: PLR0915
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]
Expand Down Expand Up @@ -1350,10 +1353,10 @@ async def _commit_spend_updates_to_db( # noqa: PLR0915
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}},
Expand Down Expand Up @@ -1441,7 +1444,10 @@ async def _update_entity_spend_in_db(
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}"
)
Expand Down
8 changes: 4 additions & 4 deletions litellm/proxy/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4874,10 +4874,10 @@ async def update_end_user_spend(
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(
Expand Down
143 changes: 143 additions & 0 deletions tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading