From 7e661ca176d01e8cf45bf5d50b96456d328910b2 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 01:13:16 -0800 Subject: [PATCH] fix(proxy): retry a model rename when a team starts listing the model mid-lock The team advisory locks come from a snapshot, so a team that gained the old name before its row lock was locked without its advisory lock. That attempt now commits without writing and the next one locks the fresh set of teams. Co-Authored-By: Claude Opus 5.5 --- .../model_allowlist_rename_sync.py | 42 ++++++++++++------ .../test_model_management_endpoints.py | 43 +++++++++++++++++++ 2 files changed, 72 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index 3d12cdd99f0..0e65e98c557 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -6,7 +6,7 @@ Those allowlists store public model names, not ids, so a deployment rename that alone denies the new name while the old entry grants a name nothing serves any more. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final @@ -98,7 +98,7 @@ _LOCK_ORGS_LISTING_SQL: Final = ( 'ORDER BY organization_id COLLATE "C" FOR UPDATE' ) _LOCK_TEAMS_LISTING_SQL: Final = ( - 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C" FOR UPDATE' + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C" FOR UPDATE' ) _LOCK_USERS_LISTING_SQL: Final = ( 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE $1 = ANY("models") ORDER BY user_id COLLATE "C" FOR UPDATE' @@ -115,6 +115,30 @@ _REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1 _APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))') +def _team_ids(rows: Sequence[object]) -> frozenset[str]: + return frozenset(_ListingTeam.model_validate(row).team_id for row in rows) + + +async def _rewrite_allowlists( + prisma_client: object, old_name: str, new_name: str, *, append: bool, attempts: int +) -> Sequence[object]: + """Team advisory locks come from a snapshot, so a team that starts listing the name before its row lock would + be locked without one. That attempt commits having written nothing, and the next one locks the fresh set.""" + async with raw_transactions(prisma_client).tx() as transaction: + listed: Final = _team_ids(await transaction.query_raw(_TEAMS_LISTING_SQL, old_name)) + for team_id in sorted(listed): + _ = await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + _ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name) + locked: Final = _team_ids(await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name)) + retry: Final = attempts > 1 and not locked <= listed + if not retry: + _ = await transaction.query_raw(_LOCK_USERS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_KEYS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_PROJECTS_LISTING_SQL, old_name) + return await transaction.query_raw(_APPEND_SQL if append else _REPLACE_SQL, old_name, new_name) + return await _rewrite_allowlists(prisma_client, old_name, new_name, append=append, attempts=attempts - 1) + + async def sync_model_allowlists_for_renamed_model( prisma_client: object, *, @@ -128,17 +152,9 @@ async def sync_model_allowlists_for_renamed_model( return executor: Final = raw_executor(prisma_client) old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id) - async with raw_transactions(prisma_client).tx() as transaction: - for team in await transaction.query_raw(_TEAMS_LISTING_SQL, old_name): - _ = await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, _ListingTeam.model_validate(team).team_id) - _ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name) - _ = await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name) - _ = await transaction.query_raw(_LOCK_USERS_LISTING_SQL, old_name) - _ = await transaction.query_raw(_LOCK_KEYS_LISTING_SQL, old_name) - _ = await transaction.query_raw(_LOCK_PROJECTS_LISTING_SQL, old_name) - touched_rows: Final = await transaction.query_raw( - _APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name - ) + touched_rows: Final = await _rewrite_allowlists( + prisma_client, old_name, new_name, append=old_name_still_backed, attempts=3 + ) touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows) await evict_and_broadcast( tuple(cache_key for row in touched for cache_key in _CACHE_KEYS_BY_KIND[row.kind](row)), diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 4bd463d8e82..9665b15bd60 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -11427,3 +11427,46 @@ async def test_add_new_model_reports_a_database_outage_during_the_insert_as_503( ) assert (exc_info.value.code, exc_info.value.type) == ("503", "no_db_connection") + + +@pytest.mark.asyncio +async def test_rename_rewrites_only_in_a_transaction_holding_every_locked_teams_advisory_lock(): + """A team that starts listing the old name between the advisory snapshot and the row lock must not be + rewritten without its advisory lock, or a roster change on that team can still deadlock with the rename.""" + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model + + snapshots = iter(([{"team_id": "team-1"}], [{"team_id": "team-1"}, {"team_id": "team-2"}])) + transactions: list[list[tuple[str, tuple[str, ...]]]] = [] + + async def query_raw(sql: str, *params: str) -> list[dict[str, object]]: + if transactions: + transactions[-1].append((sql, params)) + if sql.startswith("SELECT COUNT(*)"): + return [{"deployment_count": 0}] + if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"') and not sql.endswith("FOR UPDATE"): + return next(snapshots) + if sql.startswith('SELECT team_id FROM "LiteLLM_TeamTable"'): + return [{"team_id": "team-1"}, {"team_id": "team-2"}] + return [] + + @contextlib.asynccontextmanager + async def transaction(): + transactions.append([]) + yield db + + db = SimpleNamespace(query_raw=query_raw, tx=transaction) + with patch("litellm.proxy.management_helpers.model_allowlist_rename_sync.evict_and_broadcast", AsyncMock()): + await sync_model_allowlists_for_renamed_model( + SimpleNamespace(db=db), + model_id="deployment", + old_name="old-model", + new_name="new-model", + llm_router=None, + user_api_key_cache=UserApiKeyCache(), + ) + + rewrites = [calls for calls in transactions if any(sql.startswith("WITH ") for sql, _ in calls)] + assert len(rewrites) == 1 + advisory = [params[0] for sql, params in rewrites[0] if "pg_advisory_xact_lock" in sql] + assert advisory == ["team-1", "team-2"]