diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ecf107ef4be..f617302d8ad 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -163,10 +163,17 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx +# 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 +""" + # 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; FOR SHARE on the team row makes that check and -# the insert atomic against /team/member_delete and /team/delete, which update or delete that row -# before removing memberships, so a spend flush landing after a removal never recreates the member. +# while the user is still on the team's roster, so a spend flush landing after a removal never +# recreates the member. _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) SELECT p.user_id, p.team_id, p.cost, p.cost @@ -175,7 +182,6 @@ WHERE EXISTS ( SELECT 1 FROM "LiteLLM_TeamTable" t WHERE t.team_id = p.team_id AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id)) - FOR SHARE ) OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id) ON CONFLICT (user_id, team_id) DO UPDATE @@ -188,10 +194,12 @@ async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_memb # key is "team_id::::user_id::"; 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] + _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCKS_SQL, team_ids) _ = await transaction.execute_raw( _TEAM_MEMBER_SPEND_SQL, [key.split("::")[3] for key in keys], - [key.split("::")[1] for key in keys], + team_ids, [spend_by_member_key[key] for key in keys], ) 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 05769eea343..9f3bc5acd92 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 @@ -16,7 +16,11 @@ from redis.exceptions import DataError import litellm from litellm.proxy._types import Litellm_EntityType -from litellm.proxy.db.db_spend_update_writer import _TEAM_MEMBER_SPEND_SQL, DBSpendUpdateWriter +from litellm.proxy.db.db_spend_update_writer import ( + _TEAM_ADVISORY_LOCKS_SQL, + _TEAM_MEMBER_SPEND_SQL, + DBSpendUpdateWriter, +) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) @@ -945,9 +949,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster Regression (LIT-5502): members added without a budget had no membership row, and the previous update_many matched zero rows, so their spend was silently dropped. - The flush now runs one INSERT ... ON CONFLICT statement for the whole batch that adds - the cost to both spend and total_spend and creates the missing row for a user still on - the team roster, so no per-team read can fail or time out ahead of the writes. + The flush now takes the same per-team advisory lock the team endpoints hold, then runs + one INSERT ... ON CONFLICT statement for the whole batch that adds the cost to both spend + and total_spend and creates the missing row for a user still on the team roster, so no + per-team read can fail or time out ahead of the writes. """ db_writer = DBSpendUpdateWriter() team_id = "team-abc" @@ -967,13 +972,17 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster ), ) - mock_transaction.execute_raw.assert_awaited_once() - statement, user_ids, team_ids, costs = mock_transaction.execute_raw.await_args.args + 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 + 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]) assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement - assert "FOR SHARE" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement @@ -982,9 +991,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster @pytest.mark.asyncio async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_user(): """ - The single member spend statement locks 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, so concurrent pods lock in the same order and cannot deadlock. + 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 + same order and cannot deadlock. """ db_writer = DBSpendUpdateWriter() mock_transaction, mock_prisma_client = _team_member_flush_fixtures() @@ -1006,7 +1016,9 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u ), ) - _statement, user_ids, team_ids, costs = mock_transaction.execute_raw.await_args.args + lock_call, 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 list(zip(team_ids, user_ids, costs)) == [ ("team_a", "user_x", 0.3), ("team_a", "user_y", 0.2),