From bbb85ed48e361c094c2623c45bffe2f395f152c9 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 23:30:32 -0800 Subject: [PATCH] fix(accounting): take spend locks in one order so the ledger cannot deadlock A ledger delivery updates many spend tables in one transaction, so it can deadlock with any writer that locks the same rows in another order. It took the team advisory lock after the user and team rows, while team roster writers take it first The delivery now takes every team advisory lock before any row, then writes tables in one canonical order with ids sorted. It also runs under the rollup lock_timeout, applies 500 deliveries per batch by default, and spreads retries with jitter The writers that could still collide with it now follow the same order. The budget reset cascade, model allowlist rename, global spend reset, bulk user update, org delete and the delete-by-list paths lock their rows in key order first, and access group endpoints update teams and keys in one sorted pass Co-Authored-By: Claude Opus 5.5 --- litellm/accounting/application.py | 2 + litellm/accounting/worker.py | 2 +- .../proxy/common_utils/reset_budget_job.py | 5 +- litellm/proxy/db/db_spend_update_writer.py | 258 ++++++++++-------- .../access_group_endpoints.py | 91 +++--- .../internal_user_endpoints.py | 24 +- .../organization_endpoints.py | 29 +- .../access_group_model_sync.py | 10 + .../model_allowlist_rename_sync.py | 31 ++- .../spend_management_endpoints.py | 28 +- litellm/proxy/utils.py | 38 ++- .../accounting_delivery_repository.py | 2 +- litellm/repositories/unit_of_work.py | 42 +++ .../database/test_spendlogs_ledger_parity.py | 2 +- .../proxy/db/test_db_spend_update_writer.py | 34 ++- .../test_internal_user_endpoints.py | 12 +- .../test_model_management_endpoints.py | 7 + .../test_organization_endpoints.py | 7 + tests/unit/repositories/test_unit_of_work.py | 9 +- 19 files changed, 395 insertions(+), 238 deletions(-) diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index 1e5f23e2997..c911f4e35ce 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -7,6 +7,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_b from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert +from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction if TYPE_CHECKING: @@ -25,6 +26,7 @@ async def apply_deliveries( from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row contributions: Final = project_deliveries(deliveries) + await apply_rollup_lock_timeout(transaction) await writer.apply_accounting_entities( prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) ) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index df71f2da31c..0607c1647e1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -54,7 +54,7 @@ async def flush_accounting_deliveries( ) repository: Final = AccountingDeliveryRepository(database) started: Final = time.monotonic() - batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000")) + batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500")) async def run_batch(*, retry_only: bool) -> bool: result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..9c433af487e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -909,13 +909,14 @@ class ResetBudgetJob: db_span("reset_budget_cascade", "LiteLLM_BudgetTable"), budget_cascade_unit_of_work(self._new_batch) as uow, ): + uow.row_locks.queue_row_locks(cascade.budget_ids) + _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) uow.team_memberships.queue_total_spend_repair(cascade.budget_ids) _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) - _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) - _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 80f09df6c86..ba59b971a24 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -328,13 +328,13 @@ _SpendTableName = Literal[ "agent_list_transactions", ] _SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = ( - "user_list_transactions", - "end_user_list_transactions", - "key_list_transactions", - "team_list_transactions", - "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "team_list_transactions", + "user_list_transactions", + "end_user_list_transactions", + "team_member_list_transactions", + "key_list_transactions", "project_list_transactions", "tag_list_transactions", "model_access_group_list_transactions", @@ -563,6 +563,24 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, """ +def _org_members_in_key_order(spend_by_member_key: Mapping[str, float]) -> tuple[tuple[tuple[str, str], float], ...]: + """Sorted by the decoded (organization_id, user_id), the order every other writer locks these rows in, since + the percent-encoded key does not sort the same way.""" + decoded: Final = ( + ((unquote(key.split("::")[1]), unquote(key.split("::")[3])), cost) for key, cost in spend_by_member_key.items() + ) + return tuple(sorted(decoded)) + + +def _charged_team_ids(updates: DBSpendUpdateTransactions) -> tuple[str, ...]: + members: Final = updates.get("team_member_list_transactions") or {} + return tuple( + sorted( + {*(updates.get("team_list_transactions") or {}), *(key.split("::")[1] for key in members)}, + ) + ) + + async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) @@ -2294,6 +2312,10 @@ class DBSpendUpdateWriter: transaction: _SpendTransaction, updates: DBSpendUpdateTransactions, ) -> None: + """Every team advisory lock is taken before the first row lock, the order the team roster writers use, so + applying a batch of charges cannot deadlock with a concurrent roster change.""" + for team_id in _charged_team_ids(updates): + _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) await self._commit_spend_updates_to_db( prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction ) @@ -2313,6 +2335,104 @@ class DBSpendUpdateWriter: """ from litellm.proxy.utils import ProxyUpdateSpend + ### UPDATE ORG TABLE ### + org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] + verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) + if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. + for org_id, response_cost in sorted(org_list_transactions.items()): + batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"organization_id": org_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_list_transactions") + + org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") + verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) + if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with ( + _spend_update_tx( + prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction + ) as transaction, + transaction.batch_() as batcher, + ): + for (org_id, member_id), response_cost in _org_members_in_key_order( + org_member_list_transactions + ): + batcher.litellm_organizationmembership.update_many( + where={"organization_id": org_id, "user_id": member_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_member_list_transactions") + + ### UPDATE TEAM TABLE ### + team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] + verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) + if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. + for team_id, response_cost in sorted(team_list_transactions.items()): + verbose_proxy_logger.debug( + "Updating spend for team id=%s by %s", team_id, response_cost + ) + batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"team_id": team_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("team_list_transactions") + ### UPDATE USER TABLE ### user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"] verbose_proxy_logger.debug("User Spend transactions: %s", user_list_transactions) @@ -2359,75 +2479,6 @@ class DBSpendUpdateWriter: ) if on_table_committed is not None: on_table_committed("end_user_list_transactions") - ### UPDATE KEY TABLE ### - key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] - verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) - if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by token for consistent lock ordering across pods to prevent deadlocks. - for token, response_cost in sorted(key_list_transactions.items()): - spend_increment: _SpendIncrement = {"increment": response_cost} - batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists - where={"token": token}, - data={ - "spend": spend_increment, - "total_spend": spend_increment, - "last_active": datetime.now(timezone.utc), - }, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("key_list_transactions") - - ### UPDATE TEAM TABLE ### - team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] - verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) - if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. - for team_id, response_cost in sorted(team_list_transactions.items()): - verbose_proxy_logger.debug( - "Updating spend for team id=%s by %s", team_id, response_cost - ) - batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"team_id": team_id}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("team_list_transactions") ### UPDATE TEAM Membership TABLE with spend ### team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"] @@ -2478,22 +2529,27 @@ class DBSpendUpdateWriter: elif on_table_committed is not None: on_table_committed("team_member_list_transactions") - ### UPDATE ORG TABLE ### - org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] - verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) - if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + ### UPDATE KEY TABLE ### + key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] + verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) + if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: async with _spend_update_tx( - prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction ) as transaction: async with transaction.batch_() as batcher: - # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. - for org_id, response_cost in sorted(org_list_transactions.items()): - batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"organization_id": org_id}, - data={"spend": {"increment": response_cost}}, + # Sort by token for consistent lock ordering across pods to prevent deadlocks. + for token, response_cost in sorted(key_list_transactions.items()): + spend_increment: _SpendIncrement = {"increment": response_cost} + batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists + where={"token": token}, + data={ + "spend": spend_increment, + "total_spend": spend_increment, + "last_active": datetime.now(timezone.utc), + }, ) break except Exception as e: @@ -2507,39 +2563,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) if on_table_committed is not None: - on_table_committed("org_list_transactions") - - org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") - verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) - if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with ( - _spend_update_tx( - prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction - ) as transaction, - transaction.batch_() as batcher, - ): - for key, response_cost in sorted(org_member_list_transactions.items()): - _, quoted_org_id, _, quoted_user_id = key.split("::") - batcher.litellm_organizationmembership.update_many( - where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("org_member_list_transactions") + on_table_committed("key_list_transactions") ### UPDATE PROJECT TABLE ### project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 6558c82f80c..4cbc2155af0 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -347,47 +347,45 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None: # --------------------------------------------------------------------------- -async def _sync_add_access_group_to_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Add access_group_id to each team's access_group_ids (idempotent).""" - for team_id in team_ids: +async def _sync_access_group_on_teams( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted team_id order, the order every spend writer locks team rows in, so this + cannot deadlock with them.""" + to_add: Final = frozenset(add) + for team_id in sorted(to_add | frozenset(remove)): team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id not in (team.access_group_ids or []): + if team is None: + continue + current = list(team.access_group_ids or []) + if team_id in to_add and access_group_id not in current: + await tx.litellm_teamtable.update( + where={"team_id": team_id}, data={"access_group_ids": [*current, access_group_id]} + ) + elif team_id not in to_add and access_group_id in current: await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + data={"access_group_ids": list(_without_access_group(current, access_group_id))}, ) -async def _sync_remove_access_group_from_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Remove access_group_id from each team's access_group_ids (idempotent).""" - for team_id in team_ids: - team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id in (team.access_group_ids or []): - await tx.litellm_teamtable.update( - where={"team_id": team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - - -async def _sync_add_access_group_to_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Add access_group_id to each key's access_group_ids (idempotent).""" - for token in key_tokens: +async def _sync_access_group_on_keys( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted token order, the order every spend writer locks key rows in.""" + to_add: Final = frozenset(add) + for token in sorted(to_add | frozenset(remove)): key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id not in (key.access_group_ids or []): + if key is None: + continue + current = list(key.access_group_ids or []) + if token in to_add and access_group_id not in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + where={"token": token}, data={"access_group_ids": [*current, access_group_id]} ) - - -async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Remove access_group_id from each key's access_group_ids (idempotent).""" - for token in key_tokens: - key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id in (key.access_group_ids or []): + elif token not in to_add and access_group_id in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, + where={"token": token}, data={"access_group_ids": list(_without_access_group(current, access_group_id))} ) @@ -582,8 +580,8 @@ async def create_access_group( ) # Sync team and key tables to reference the new access group - await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id) - await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id) + await _sync_access_group_on_teams(tx, data.assigned_team_ids or [], (), record.access_group_id) + await _sync_access_group_on_keys(tx, data.assigned_key_ids or [], (), record.access_group_id) except HTTPException: raise except Exception as e: @@ -718,10 +716,8 @@ async def update_access_group( data=update_data, ) - await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) - await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) - await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) - await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + await _sync_access_group_on_teams(tx, teams_to_add, teams_to_remove, access_group_id) + await _sync_access_group_on_keys(tx, keys_to_add, keys_to_remove, access_group_id) except HTTPException: raise except Exception as e: @@ -790,25 +786,8 @@ async def delete_access_group( ) affected_key_tokens = list(all_affected_key_tokens) - # Update teams returned by find_many directly — we already have their data. - for team in teams_with_group: - await tx.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync teams not found by the hasSome query. - out_of_sync_team_ids: Final = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} - await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) - - # Update keys returned by find_many directly — we already have their data. - for key in keys_with_group: - await tx.litellm_verificationtoken.update( - where={"token": key.token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync keys not found by the hasSome query. - out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} - await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) + await _sync_access_group_on_teams(tx, (), affected_team_ids, access_group_id) + await _sync_access_group_on_keys(tx, (), affected_key_tokens, access_group_id) detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..93a429a2fcc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: handle_update_object_permission_common, set_object_permission, ) +from litellm.proxy.management_helpers.team_roster_sync import user_tx_db from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password from litellm.repositories.organization_repository import OrganizationRepository @@ -125,6 +126,11 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +_LOCK_ALL_USERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE TRUE ORDER BY user_id COLLATE "C" FOR UPDATE' +) + + @with_config(ConfigDict(extra="allow", strict=True)) class _AdminTeamAndKeyRows(TypedDict): keys: ReadOnly[Sequence[object] | None] @@ -2011,14 +2017,16 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( - where={}, - data=( - {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} - if "model_max_budget" in non_default_values - else non_default_values - ), - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_USERS_SQL) + _ = await user_tx_db(transaction).update_many( + where={}, + data=( + {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} + if "model_max_budget" in non_default_values + else non_default_values + ), + ) for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..4e511aee24a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1047,12 +1047,12 @@ async def delete_organization( deleted_orgs: Final = [] for organization_id in data.organization_ids: - # delete all teams in the organization - await _table(TeamRepository(prisma_client)).delete_many(where={"organization_id": organization_id}) - # delete all members in the organization - await _table(OrganizationMembershipRepository(prisma_client)).delete_many( - where={"organization_id": organization_id} - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_TEAMS_SQL, organization_id) + _ = await transaction.litellm_teamtable.delete_many(where={"organization_id": organization_id}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_MEMBERS_SQL, organization_id) + _ = await transaction.litellm_organizationmembership.delete_many(where={"organization_id": organization_id}) await _delete_organization_keys( organization_id=organization_id, prisma_client=prisma_client, @@ -1074,6 +1074,19 @@ async def delete_organization( return deleted_orgs +_LOCK_ORG_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE organization_id = $1 ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_MEMBERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationMembership" WHERE organization_id = $1 ' + 'ORDER BY organization_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE organization_id = $1 ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) + + async def _delete_organization_keys( organization_id: str, prisma_client: PrismaClient, @@ -1087,7 +1100,9 @@ async def _delete_organization_keys( hashed_tokens=hashed_tokens_to_delete, prisma_client=prisma_client, ) - await _table(VerificationTokenRepository(prisma_client)).delete_many(where=key_filter) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_KEYS_SQL, organization_id) + _ = await transaction.litellm_verificationtoken.delete_many(where={"organization_id": organization_id}) await delete_cache_key_objects( hashed_tokens=hashed_tokens_to_delete, user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py index 890715bf9c2..20ea2ed0338 100644 --- a/litellm/proxy/management_helpers/access_group_model_sync.py +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -6,6 +6,7 @@ the arrays alone strands every group on a name nothing serves any more. """ from collections.abc import Sequence +from contextlib import AbstractAsyncContextManager from typing import Final, Protocol from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper @@ -53,6 +54,15 @@ _REMOVE_MODEL_NAME_SQL: Final = ( ) +class RawTransactions(Protocol): + def tx(self) -> AbstractAsyncContextManager[RawExecutor]: ... + + +def raw_transactions(prisma_client: object) -> RawTransactions: + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + def raw_executor(prisma_client: object) -> RawExecutor: db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index f1a9faf6b2e..498318cb63d 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -13,7 +13,7 @@ from typing import Final from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed +from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, raw_transactions, still_backed from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel @@ -80,6 +80,23 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str: return f"WITH {ctes} {rows}" +_LOCK_ORGS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE $1 = ANY("models") ' + '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' +) +_LOCK_USERS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE $1 = ANY("models") ORDER BY user_id COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE $1 = ANY("models") ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_PROJECTS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE $1 = ANY("models") ORDER BY project_id COLLATE "C" FOR UPDATE' +) + _REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")') _APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))') @@ -98,9 +115,15 @@ 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) - touched_rows: Final = await executor.query_raw( - _APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name - ) + async with raw_transactions(prisma_client).tx() as transaction: + _ = 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: 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/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcb9eb4a280..0348d8e1f82 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -77,9 +77,6 @@ from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_stat from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.verification_token_repository import ( - VerificationTokenRepository, -) from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureProvider if TYPE_CHECKING: @@ -292,24 +289,22 @@ class _TeamTable(Protocol): async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... -class _VerificationTokenTable(Protocol): - """The subset of the Prisma verification token table API this module uses.""" - - async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... - - def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]: return SpendLogsRepository(prisma_client).table +_LOCK_ALL_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE TRUE ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ALL_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE TRUE ORDER BY token COLLATE "C" FOR UPDATE' +) + + def _team_table(prisma_client: PrismaClient) -> _TeamTable: return TeamRepository(prisma_client).table -def _verification_token_table(prisma_client: PrismaClient) -> _VerificationTokenTable: - return VerificationTokenRepository(prisma_client).table - - def _spend_logs_daily_summary_sql( *, start_date_iso: str, @@ -3626,8 +3621,11 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) - await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_TEAMS_SQL) + _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) + _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 71ecae74058..19f6bff8a8a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5528,6 +5528,20 @@ class _StaleReadEngine: return self.wrapper is current and self.generation == current.engine_generation +_LOCK_KEYS_BY_TOKEN_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE token = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_BY_TEAM_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_TEAMS_BY_ID_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C" FOR UPDATE' +) + + class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() @@ -6828,24 +6842,26 @@ class PrismaClient: else: hashed_token = token hashed_tokens.append(hashed_token) - filter_query: dict[str, object] = {} - if user_id is not None: - filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]} - else: - filter_query = {"token": {"in": hashed_tokens}} - - deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many( - where=filter_query - ) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_KEYS_BY_TOKEN_SQL, hashed_tokens) + deleted_tokens: Final[int] = await transaction.litellm_verificationtoken.delete_many( + where={"token": {"in": hashed_tokens}, "user_id": user_id} + if user_id is not None + else {"token": {"in": hashed_tokens}} + ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) return {"deleted_keys": deleted_tokens} elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_TEAMS_BY_ID_SQL, team_id_list) + _ = await transaction.litellm_teamtable.delete_many(where={"team_id": {"in": team_id_list}}) 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` - await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() 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: import traceback diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 262992a48c5..d5bbe891089 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -84,7 +84,7 @@ SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """ _FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery" SET attempts = attempts + 1, last_error = $3, - next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), + next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)) * (0.5 + random()))), state = $4::"AccountingDeliveryState" WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer RETURNING delivery_id, attempts""" diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index f6aad883479..b78f7e5a9b7 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,46 @@ class TeamMembershipSpendResetWrites: self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids)) +_LOCK_BUDGET_ORGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY organization_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_MEMBERSHIPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamMembership" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_PROJECTS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY project_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_TAGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TagTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY tag_name COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_ACCESS_GROUPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY access_group_name COLLATE "C" FOR UPDATE' +) + + +@dataclass(frozen=True, slots=True) +class BudgetLinkedRowLocks: + batch: PrismaBatch + + def queue_row_locks(self, budget_ids: Sequence[str]) -> None: + """Locks every row a cascade may reset, in the shared row lock order, before the filtered resets run.""" + self.batch.execute_raw(_LOCK_BUDGET_ORGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_MEMBERSHIPS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_KEYS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_PROJECTS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_TAGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_ACCESS_GROUPS_SQL, list(budget_ids)) + + @dataclass(frozen=True, slots=True) class BudgetWindowWrites: table: BatchTable @@ -137,6 +177,7 @@ class BudgetCascadeUnitOfWork: ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ + row_locks: BudgetLinkedRowLocks team_memberships: TeamMembershipSpendResetWrites keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites @@ -163,6 +204,7 @@ async def budget_cascade_unit_of_work( ) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: batch = new_batch() yield BudgetCascadeUnitOfWork( + row_locks=BudgetLinkedRowLocks(batch=batch), team_memberships=TeamMembershipSpendResetWrites( resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch ), diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index e1894aa66bf..78439e4317e 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -63,7 +63,7 @@ _DAILY_TABLES: Final = ( ) -@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) +@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), (None, 501, 500, False), ("200", 251, 200, False), ("10000", 10001, 10000, False), ("3001", 3001, 3001, True)]) @pytest.mark.asyncio async def test_worker_applies_the_configured_backlog_in_one_transaction( parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d90c7f6d21d..e2debe05c03 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -44,6 +44,15 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush from tests.unit.proxy.db.fake_prisma_engine import engine_call + +@pytest.fixture(autouse=True) +def legacy_spend_writers(monkeypatch: pytest.MonkeyPatch) -> None: + """These tests cover the legacy writers, still reachable as the rollback path with use_spend_logs_ledger off.""" + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) + + @pytest.mark.asyncio async def test_daily_spend_tracking_with_disabled_spend_logs(): """ @@ -3714,11 +3723,12 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti proxy_logging_obj=proxy_logging_obj, ) - mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once() - restored = mock_redis_update_buffer.restore_transactions_to_redis.call_args.kwargs["db_spend_update_transactions"] - assert restored["user_list_transactions"] is None - assert restored["team_list_transactions"] == {"team-1": 1.5} - assert restored["key_list_transactions"] == ({"key-1": 1.5} if safe_to_resend else None) + restores = mock_redis_update_buffer.restore_transactions_to_redis.await_args_list + assert [ + {name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows} + for call in restores + ] == ([{"key_list_transactions": {"key-1": 1.5}}] if safe_to_resend else []) + mock_batcher.litellm_teamtable.update_many.assert_called_once() mock_batcher.litellm_usertable.update_many.assert_called_once() @@ -3892,13 +3902,13 @@ async def test_commit_spend_updates_to_db_reports_each_completed_table(): ) assert on_table_committed.call_args_list == [ - call("user_list_transactions"), - call("end_user_list_transactions"), - call("key_list_transactions"), - call("team_list_transactions"), - call("team_member_list_transactions"), call("org_list_transactions"), call("org_member_list_transactions"), + call("team_list_transactions"), + call("user_list_transactions"), + call("end_user_list_transactions"), + call("team_member_list_transactions"), + call("key_list_transactions"), call("project_list_transactions"), call("tag_list_transactions"), call("model_access_group_list_transactions"), @@ -5808,7 +5818,9 @@ async def test_completed_spend_survives_error_log_suppression( prisma.get_request_status.return_value = status monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "disable_spend_logs", False) - monkeypatch.setattr(proxy_server, "general_settings", {"disable_error_logs": disable_error_logs}) + monkeypatch.setattr( + proxy_server, "general_settings", {"disable_error_logs": disable_error_logs, "use_spend_logs_ledger": False} + ) now: Final = datetime(2026, 10, 20, 12, tzinfo=timezone.utc) result: Final = ModelResponse( model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index d8aa0239e2d..cee0c90f80c 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3,7 +3,7 @@ import hashlib import json import logging from collections.abc import AsyncIterator, Mapping, Sequence -from contextlib import asynccontextmanager +from contextlib import AbstractAsyncContextManager, asynccontextmanager from datetime import datetime, timezone from types import SimpleNamespace from typing import TYPE_CHECKING, Final, cast @@ -130,6 +130,11 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): UserWhereCondition = InsensitiveContains | Sequence[Mapping[str, InsensitiveContains]] +@asynccontextmanager +async def _transaction_over(db: object) -> AsyncIterator[SimpleNamespace]: + yield SimpleNamespace(query_raw=AsyncMock(return_value=[]), litellm_usertable=getattr(db, "litellm_usertable")) + + def _matches_user_where(row: LiteLLM_UserTableFiltered, where: Mapping[str, UserWhereCondition]) -> bool: def matches(field: str, condition: UserWhereCondition) -> bool: if not isinstance(condition, Mapping): @@ -2419,6 +2424,7 @@ async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocke prisma_client: Final = mocker.MagicMock() prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.tx = lambda: _transaction_over(prisma_client.db) mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency cache: Final = UserApiKeyCache() await cache.async_set_cache(key=saved_user.user_id, value=saved_user, model_type=LiteLLM_UserTable) @@ -2465,6 +2471,7 @@ async def test_user_max_budget_update_evicts_cached_user_on_every_worker(mocker: prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(return_value=saved_user) prisma_client.db.litellm_usertable.find_many = mocker.AsyncMock(return_value=[saved_user]) prisma_client.db.litellm_usertable.update_many = mocker.AsyncMock(return_value=1) + prisma_client.tx = lambda: _transaction_over(prisma_client.db) prisma_client.get_data = mocker.AsyncMock(return_value=[saved_user]) prisma_client.update_data = mocker.AsyncMock(return_value={"user_id": saved_user.user_id, "data": saved_user}) mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency @@ -2575,6 +2582,9 @@ async def test_user_update_reaches_cached_user_on_every_worker( self.update_data_payload = data return {"user_id": user_id, "data": self.updated_user} + def tx(self) -> AbstractAsyncContextManager[SimpleNamespace]: + return _transaction_over(self.db) + saved_user: Final = LiteLLM_UserTable( user_id="user-spruce", tpm_limit=100000, 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 44c43713438..303b0b5e83f 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -8567,12 +8567,19 @@ class TestAccessGroupModelSync: return [{"deployment_count": deployment_count}] if sql.startswith('UPDATE "LiteLLM_AccessGroupTable"'): return [{"access_group_id": "ag-1"}] + if sql.startswith("SELECT 1 AS locked FROM"): + return [] assert sql.startswith("WITH ") return TestAccessGroupModelSync._ALLOWLIST_ROWS + @contextlib.asynccontextmanager + async def transaction(): + yield mock_prisma.db + mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=query_raw) + mock_prisma.db.tx = MagicMock(side_effect=transaction) mock_prisma.db.litellm_proxymodeltable = AsyncMock() mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=row) mock_prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=row) diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 5e06756f3a9..8d2c6c0054c 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1,5 +1,6 @@ import asyncio import json +from contextlib import asynccontextmanager from types import MappingProxyType, SimpleNamespace from typing import Final, Mapping, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -1560,6 +1561,12 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke prisma_client.db.litellm_jwtkeymapping = jwt_table prisma_client.db.litellm_organizationtable.delete = AsyncMock(return_value=MagicMock()) + @asynccontextmanager + async def transaction(): + yield prisma_client.db + + prisma_client.tx = transaction + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) diff --git a/tests/unit/repositories/test_unit_of_work.py b/tests/unit/repositories/test_unit_of_work.py index 082a29b7001..049e5275db6 100644 --- a/tests/unit/repositories/test_unit_of_work.py +++ b/tests/unit/repositories/test_unit_of_work.py @@ -168,7 +168,8 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( return batches[-1] async with budget_cascade_unit_of_work(_new_batch) as uow: - writes = [getattr(uow, field.name) for field in fields(uow)] + uow.row_locks.queue_row_locks(["budget-1"]) + writes = [getattr(uow, field.name) for field in fields(uow) if field.name != "row_locks"] for write in writes: if isinstance(write, SpendResetWrites): write.queue_spend_zero(where={"budget_id": "budget-1"}) @@ -177,8 +178,10 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch( assert len(batches) == 1, "the cascade must open exactly one transaction" batch = batches[0] - assert len(batch.calls) == len(writes), "a dependent bound to a batch of its own would not land here" - assert len({call[0] for call in batch.calls}) == len(writes), "two dependents share one table" + table_writes = [call for call in batch.calls if call[0] != "execute_raw"] + assert len(table_writes) == len(writes), "a dependent bound to a batch of its own would not land here" + assert len({call[0] for call in table_writes}) == len(writes), "two dependents share one table" + assert [call[0] for call in batch.calls].count("execute_raw") > 0, "the row locks share the cascade's batch" assert batch.commit_count == 1