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"]