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 teams in sorted team id order and keep existing member budgets on re-add

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-17 00:37:03 +00:00 • committed by ryan-crabbe-berri
parent 499c334fce
commit 2d61fa66b1
4 changed files with 24 additions and 20 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

@ -188,17 +188,17 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None:
# key is "team_id::<value>::user_id::<value>"; the string sort orders rows by (team_id, user_id),
# 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]
# key is "team_id::<value>::user_id::<value>"; rows are sorted by (team_id, user_id) so the teams are
# locked in the same `sorted(team_ids)` order the team endpoints use, preventing deadlocks
rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items())
team_ids: Final = [team_id for team_id, _user_id, _cost in rows]
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],
[user_id for _team_id, user_id, _cost in rows],
team_ids,
[spend_by_member_key[key] for key in keys],
[cost for _team_id, _user_id, cost in rows],
)

2
litellm/proxy/management_helpers/utils.py
View file

@ -479,7 +479,7 @@ async def add_new_member(
budget_link: Final[Mapping[str, object]] = {"budget_id": _budget_id} if _budget_id is not None else {}
_returned_team_membership: Final = await membership_table.upsert(
where={"user_id_team_id": membership_key},
data={"create": {**membership_key, **budget_link}, "update": {**budget_link}},
data={"create": {**membership_key, **budget_link}, "update": {}},
include={"litellm_budget_table": True},
)

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

@ -992,8 +992,9 @@ 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
each distinct team is locked once, in that same order, so concurrent pods lock in the
same order and cannot deadlock.
each distinct team is locked once, in `sorted(team_ids)` order, the order /team/delete
locks in, so a concurrent flush and delete cannot deadlock. `eng` and `eng2` pin that:
sorting the composite keys instead would lock `eng2` first because `2` < `:`.
"""
db_writer = DBSpendUpdateWriter()
mock_transaction, mock_prisma_client = _team_member_flush_fixtures()
@ -1007,10 +1008,10 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=_team_member_only_transactions(
{
"team_id::team_c::user_id::user_x": 0.1,
"team_id::team_a::user_id::user_y": 0.2,
"team_id::team_a::user_id::user_x": 0.3,
"team_id::team_b::user_id::user_x": 0.4,
"team_id::eng2::user_id::user_x": 0.1,
"team_id::eng::user_id::user_y": 0.2,
"team_id::eng::user_id::user_x": 0.3,
"team_id::eng-b::user_id::user_x": 0.4,
}
),
)
@ -1018,15 +1019,15 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u
*lock_calls, spend_call = mock_transaction.execute_raw.await_args_list
_statement, user_ids, team_ids, costs = spend_call.args
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"),
(_TEAM_ADVISORY_LOCK_SQL, "eng"),
(_TEAM_ADVISORY_LOCK_SQL, "eng-b"),
(_TEAM_ADVISORY_LOCK_SQL, "eng2"),
]
assert list(zip(team_ids, user_ids, costs)) == [
("team_a", "user_x", 0.3),
("team_a", "user_y", 0.2),
("team_b", "user_x", 0.4),
("team_c", "user_x", 0.1),
("eng", "user_x", 0.3),
("eng", "user_y", 0.2),
("eng-b", "user_x", 0.4),
("eng2", "user_x", 0.1),
]

3
tests/test_litellm/proxy/management_helpers/test_management_helpers_utils.py
View file

@ -446,6 +446,8 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided():
1. When max_budget_in_team is provided
2. A new budget is created in the litellm_budgettable
3. The new budget_id is used for the team membership
4. The upsert's update branch stays empty, so a bulk /team/member_add that names a member
already on the team does not replace the budget_id (and the spend) their existing row carries
"""
from litellm.proxy._types import LitellmUserRoles
@ -529,6 +531,7 @@ async def test_add_new_member_creates_new_budget_when_max_budget_provided():
assert team_membership_call_args is not None
create_data = team_membership_call_args.kwargs["data"]["create"]
assert create_data["budget_id"] == test_new_budget_id
assert team_membership_call_args.kwargs["data"]["update"] == {}
@pytest.mark.asyncio