From f52f26d9c636038c6e3609cc7d17cf625d434c3e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 11 Nov 2026 12:37:28 -0800 Subject: [PATCH] fix(proxy): let /team/delete transactions wait for a pooled connection (#51099) The team-model and team-key transactions in /team/delete used Prisma's 2s default max_wait, so a busy writer pool failed the delete with P2028 while every plain query in the same request waited up to the 60s pool timeout. Both now get the 30s team-delete budget for max_wait, and the key transaction gets it for timeout too. --- .../model_management_endpoints.py | 7 +- litellm/proxy/utils.py | 4 +- .../test_team_delete_large_membership.py | 92 +++++++++++++++++++ 3 files changed, 100 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 7846922bae2..0b4f266f838 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -295,7 +295,9 @@ class _TxModelTables(Protocol): @runtime_checkable class _TransactionFactory(Protocol): - def __call__(self, *, timeout: datetime.timedelta = ...) -> AbstractAsyncContextManager[_TxModelTables]: ... + def __call__( + self, *, timeout: datetime.timedelta = ..., max_wait: datetime.timedelta = ... + ) -> AbstractAsyncContextManager[_TxModelTables]: ... class _ModelTransactionClient(LiteLLMBaseModel): @@ -2449,7 +2451,8 @@ async def delete_team_models( """ deleted_model_ids: Final[list[str]] = [] transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) - async with transaction_client.tx(timeout=datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS)) as tx: + tx_budget: Final = datetime.timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS) + async with transaction_client.tx(max_wait=tx_budget, timeout=tx_budget) as tx: for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bf00bdfb5fb..51f1c8a614d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -73,6 +73,7 @@ from litellm.constants import ( SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, + TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS, VLLM_CHOICE_FIELDS, ) from litellm.cost_calculator import BaseTokenUsageProcessor @@ -5361,6 +5362,7 @@ _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(ma _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) _PRISMA_DEFAULT_TX_MAX_WAIT: Final = timedelta(seconds=2) +_TEAM_DELETE_TX_BUDGET: Final = timedelta(seconds=TEAM_DELETE_TRANSACTION_TIMEOUT_SECONDS) _SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS) _STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5) @@ -6983,7 +6985,7 @@ class PrismaClient: return {"deleted_teams": team_id_list} elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - async with self.tx() as transaction: + async with self.tx(max_wait=_TEAM_DELETE_TX_BUDGET, timeout=_TEAM_DELETE_TX_BUDGET) as transaction: _ = await transaction.query_raw(_LOCK_KEYS_BY_TEAM_SQL, team_id_list) _ = await transaction.litellm_verificationtoken.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: diff --git a/tests/integration/management/test_team_delete_large_membership.py b/tests/integration/management/test_team_delete_large_membership.py index 2a8508ff3c9..361cbf037a3 100644 --- a/tests/integration/management/test_team_delete_large_membership.py +++ b/tests/integration/management/test_team_delete_large_membership.py @@ -21,6 +21,7 @@ import httpx import psycopg import pytest import yaml +from psycopg.rows import TupleRow from pydantic import JsonValue from redis import Redis @@ -55,7 +56,11 @@ FROM generate_series(1, %s::int) AS n """ TAKE_TEAM_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext(%s))" LOCK_MODEL_ROW_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s FOR UPDATE' +LOCK_TEAM_KEY_ROWS_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE team_id = %s FOR UPDATE' +LOCK_BUDGET_ROW_SQL: Final = 'SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s FOR UPDATE' LOCK_HELD_PAST_DEFAULT_TRANSACTION_SECONDS: Final = 7 +POOL_HELD_PAST_DELETE_RETRIES_SECONDS: Final = 12 +SATURATED_POOL_LIMIT: Final = 2 OLDEST_TRANSACTION_BLOCKED_BY_SQL: Final = """ SELECT coalesce(max(extract(epoch FROM clock_timestamp() - xact_start)), 0)::float AS age FROM pg_stat_activity @@ -146,6 +151,16 @@ def _waiters_on_lock_held_by(backend_pid: int) -> int: return waiting +def _sessions_blocked_by(backend_pid: int) -> int: + rows: Final = read_rows( + "SELECT count(*)::int AS blocked FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))", + (str(backend_pid),), + ) + blocked: Final = rows[0]["blocked"] + assert isinstance(blocked, int) + return blocked + + def _oldest_transaction_blocked_by(backend_pid: int) -> float: rows: Final = read_rows(OLDEST_TRANSACTION_BLOCKED_BY_SQL, (str(backend_pid),)) age: Final = rows[0]["age"] @@ -468,6 +483,83 @@ def test_delete_completes_after_waiting_on_the_team_lock_past_the_default_transa assert len(_tombstones(team)) == 1, _tombstones(team) +def _park_on_team_key_rows(holder: psycopg.Connection[TupleRow], team_id: str) -> None: + _ = holder.execute(LOCK_TEAM_KEY_ROWS_SQL, (team_id,)) + + +def _park_on_deleted_key_records(holder: psycopg.Connection[TupleRow], _team_id: str) -> None: + _ = holder.execute('LOCK TABLE "LiteLLM_DeletedVerificationToken" IN SHARE MODE') + + +def _requests_in_flight_including_this_one(proxy: Gateway) -> JsonValue: + return proxy.get("/health/backlog")["in_flight_requests"] + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize( + "park", + [ + pytest.param(_park_on_team_key_rows, id="next-transaction-is-the-team-model-delete"), + pytest.param(_park_on_deleted_key_records, id="next-transaction-is-the-team-key-delete"), + ], +) +def test_delete_waits_for_a_pooled_connection_when_every_one_is_busy( + gateway: Gateway, + tmp_path: Path, + park: Callable[[psycopg.Connection[TupleRow], str], None], +) -> None: + database_url: Final = os.environ["DATABASE_URL"] + with ( + gateway.scenario() as scenario, + owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=_config_with_pool_limit(tmp_path, SATURATED_POOL_LIMIT), + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned, + ThreadPoolExecutor(max_workers=3) as pool, + ): + team: Final = string_value( + owned.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"] + ) + scenario.cleanups.callback(_remove_team_by_sql, team) + scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_VerificationToken" WHERE team_id = %s', (team,)) + keys: Final = _seeded_team_keys(team, 3) + budget: Final = f"integration-{uuid.uuid4().hex}" + owned.gateway.post("/budget/new", {"budget_id": budget, "max_budget": 100}) + scenario.cleanups.callback(write_rows, 'DELETE FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget,)) + update_budget: Final = {"budget_id": budget, "max_budget": 1} + + with psycopg.connect(database_url) as budget_holder: + budget_holder.execute(LOCK_BUDGET_ROW_SQL, (budget,)) + first_connection_taker: Final = _post_in_background(pool, owned.gateway, "/budget/update", update_budget) + eventually(lambda: _sessions_blocked_by(budget_holder.info.backend_pid), lambda count: count == 1) + with psycopg.connect(database_url) as parking: + park(parking, team) + pending_delete: Final = _post_in_background(pool, owned.gateway, "/team/delete", {"team_ids": [team]}) + eventually(lambda: _sessions_blocked_by(parking.info.backend_pid), lambda count: count == 1) + second_connection_taker: Final = _post_in_background( + pool, owned.gateway, "/budget/update", update_budget + ) + eventually(lambda: _requests_in_flight_including_this_one(owned.gateway), lambda count: count == 4) + eventually( + pending_delete.done, + lambda done: done, + seconds=POOL_HELD_PAST_DELETE_RETRIES_SECONDS, + return_last_on_timeout=True, + ) + + response: Final = pending_delete.result(timeout=60) + assert response.status_code == 200, response.text + assert response.json() == {"deleted_teams": [team]} + takers: Final = (first_connection_taker, second_connection_taker) + assert [taker.result(timeout=60).status_code for taker in takers] == [200, 200] + assert _team_rows(team) == [] + assert [_live_token(_hashed(key)) for key in keys] == [[], [], []] + assert len(_tombstones(team)) == 1, _tombstones(team) + + def test_deleting_two_teams_sharing_a_member_in_one_call_clears_both_from_the_member(gateway: Gateway) -> None: with gateway.scenario() as scenario: user: Final = scenario.user()