scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

fix(proxy): lock each team in sorted order before the member spend upsert

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-16 03:35:34 +00:00 • committed by ryan-crabbe-berri
parent 60077e90aa
commit 499c334fce
2 changed files with 17 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/db/db_spend_update_writer.py
View file

@ -164,12 +164,9 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
# The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL),
# taken in sorted order so the roster check below cannot interleave with their writes. A row lock would
# deadlock with the access-group endpoints, which lock a team row after an access-group lock.
_TEAM_ADVISORY_LOCKS_SQL: Final = """
SELECT pg_advisory_xact_lock(hashtext(teams.team_id))
FROM (SELECT DISTINCT team_id FROM unnest($1::text[]) AS team_id ORDER BY team_id) AS teams
"""
# so the roster check below cannot interleave with their writes. A row lock would deadlock with the
# access-group endpoints, which lock a team row after an access-group lock.
_TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
# One statement adds every member's cost to their membership row. A missing row is created only
# while the user is still on the team's roster, so a spend flush landing after a removal never
@ -195,7 +192,8 @@ async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_memb
# keeping lock order consistent across pods to prevent deadlocks
keys: Final = tuple(sorted(spend_by_member_key))
team_ids: Final = [key.split("::")[1] for key in keys]
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCKS_SQL, team_ids)
for team_id in dict.fromkeys(team_ids):
_ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id)
_ = await transaction.execute_raw(
_TEAM_MEMBER_SPEND_SQL,
[key.split("::")[3] for key in keys],

21
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -17,7 +17,7 @@ from redis.exceptions import DataError
import litellm
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_spend_update_writer import (
_TEAM_ADVISORY_LOCKS_SQL,
_TEAM_ADVISORY_LOCK_SQL,
_TEAM_MEMBER_SPEND_SQL,
DBSpendUpdateWriter,
)
@ -973,11 +973,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster
)
lock_call, spend_call = mock_transaction.execute_raw.await_args_list
lock_statement, locked_team_ids = lock_call.args
assert lock_statement is _TEAM_ADVISORY_LOCKS_SQL
assert locked_team_ids == [team_id]
assert "pg_advisory_xact_lock(hashtext(teams.team_id))" in lock_statement
assert "ORDER BY team_id" in lock_statement
lock_statement, locked_team_id = lock_call.args
assert lock_statement is _TEAM_ADVISORY_LOCK_SQL
assert locked_team_id == team_id
assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement
statement, user_ids, team_ids, costs = spend_call.args
assert statement is _TEAM_MEMBER_SPEND_SQL
assert (user_ids, team_ids, costs) == ([user_id], [team_id], [response_cost])
@ -993,7 +992,7 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u
"""
The member spend statement touches rows in the order of its input arrays, so the batch
is handed over sorted by (team_id, user_id), with each cost kept next to its member, and
the advisory lock statement receives the same team ids, so concurrent pods lock in the
each distinct team is locked once, in that same order, so concurrent pods lock in the
same order and cannot deadlock.
"""
db_writer = DBSpendUpdateWriter()
@ -1016,9 +1015,13 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u
),
)
lock_call, spend_call = mock_transaction.execute_raw.await_args_list
*lock_calls, spend_call = mock_transaction.execute_raw.await_args_list
_statement, user_ids, team_ids, costs = spend_call.args
assert lock_call.args == (_TEAM_ADVISORY_LOCKS_SQL, ["team_a", "team_a", "team_b", "team_c"])
assert [lock_call.args for lock_call in lock_calls] == [
(_TEAM_ADVISORY_LOCK_SQL, "team_a"),
(_TEAM_ADVISORY_LOCK_SQL, "team_b"),
(_TEAM_ADVISORY_LOCK_SQL, "team_c"),
]
assert list(zip(team_ids, user_ids, costs)) == [
("team_a", "user_x", 0.3),
("team_a", "user_y", 0.2),