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/snapshot.py b/litellm/accounting/snapshot.py index d5652f79ed0..ab85df08320 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from litellm.router import Router _FIELDS: Final = TypeAdapter(dict[str, object]) +_JSON_SAFE_INTEGER: Final = 2**53 - 1 def _mapping(value: object) -> Mapping[str, object] | None: @@ -86,6 +87,7 @@ def build_accounting_payload( "metadata": accounting_metadata, }, non_finite_as_null=True, + integer_limit=_JSON_SAFE_INTEGER, ) ) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index df71f2da31c..be18d8ac6b1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, + DeliveryBatchResult, DeliveryProjection, LedgerDatabase, LedgerTransaction, @@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol): _JOB: Final = "accounting_delivery" +def _filled(result: DeliveryBatchResult, batch_size: int) -> bool: + return result.applied + len(result.parked) >= batch_size + + +def _progressed(result: DeliveryBatchResult) -> bool: + return bool(result.applied or result.failed) + + async def flush_accounting_deliveries( writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" ) -> None: @@ -54,9 +63,9 @@ 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: + async def run_batch(*, retry_only: bool) -> DeliveryBatchResult: result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) if result.applied: await writer.invalidate_accounting_entities(applied, proxy_logging_obj) @@ -68,17 +77,25 @@ async def flush_accounting_deliveries( proxy_logging_obj, message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", ) - return bool(result.applied or result.failed) + return result - await run_batch(retry_only=False) + fresh_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline: + break retry_deadline: Final = time.monotonic() + 5 for _ in range(batch_size): - if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)): + break + retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")) + cleanup_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if ( + await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2) + < batch_size * 2 + or time.monotonic() >= cleanup_deadline + ): break - await repository.cleanup_applied( - retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")), - batch_size=batch_size * 2, - ) metrics: Final = await repository.metrics() record_deliveries(metrics, time.monotonic() - started) if metrics.batch_count >= int( diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 909e1222ae6..ea26d549326 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -19,6 +19,12 @@ def _finite_or_none(value: float | None) -> float | None: return None if isinstance(value, float) and not math.isfinite(value) else value +def _bounded(value: float | None, integer_limit: int | None) -> float | None: + if integer_limit is None or isinstance(value, bool) or not isinstance(value, int): + return value + return max(-integer_limit, min(integer_limit, value)) + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, @@ -26,12 +32,14 @@ def safe_json_structure( key: str | None = None, *, non_finite_as_null: bool = False, + integer_limit: int | None = None, ) -> object: """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them. + With integer_limit, integers are clamped to that magnitude. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -51,7 +59,8 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return _finite_or_none(obj) if non_finite_as_null else obj + bounded: Final = _bounded(obj, integer_limit) + return _finite_or_none(bounded) if non_finite_as_null else bounded # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b5c0c5f8d3..d3ac264518e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2507,8 +2507,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ use_spend_logs_ledger: bool = Field( - False, - description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.", + True, + description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", ) request_tag_mappings: Sequence[RequestTagMapping] | None = Field( 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..ff57596232f 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()) @@ -712,7 +730,7 @@ class DBSpendUpdateWriter: proxy_settings: Final = cast( # cast-ok: proxy settings use string keys Mapping[str, object], general_settings ) - use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True + use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) if kwargs is not None: kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker @@ -863,8 +881,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) if "accounting_payload" in payload: + from litellm.proxy.utils import accounting_only_spend_log + await self._insert_spend_log_to_db( - payload={**payload, "accounting_log_disabled": disable_spend_logs}, + payload=accounting_only_spend_log(payload) + if disable_spend_logs + else {**payload, "accounting_log_disabled": False}, prisma_client=prisma_client, llm_router=llm_router, ) @@ -2294,6 +2316,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 +2339,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 +2483,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 +2533,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 +2567,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..0e65e98c557 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -6,18 +6,27 @@ 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 +from pydantic import ConfigDict + 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.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel +class _ListingTeam(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + team_id: str + + class _TouchedRow(LiteLLMBaseModel): kind: str object_id: str @@ -80,11 +89,56 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str: return f"WITH {ctes} {rows}" +_TEAMS_LISTING_SQL: Final = ( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C"' +) + +_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 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' +) +_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"))') +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, *, @@ -98,8 +152,8 @@ 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 + 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( 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 a2928bba287..00187790901 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5545,6 +5545,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() @@ -6845,24 +6859,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 7dd92aecb2d..d5bbe891089 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -5,7 +5,7 @@ from itertools import chain from types import TracebackType from typing import TYPE_CHECKING, Final, Literal, Protocol -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.repositories.prisma_protocols import TableActions @@ -41,7 +41,10 @@ class _Delivery(BaseModel): delivery_id: str attempts: int - accounting_payload: JsonValue + + +class _ClaimedDelivery(_Delivery): + accounting_payload: str @dataclass(frozen=True, slots=True) @@ -68,9 +71,11 @@ class DeliveryMetrics(BaseModel): _DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...]) +_CLAIMED: Final = TypeAdapter(tuple[_ClaimedDelivery, ...]) _PAYLOAD: Final = TypeAdapter(dict[str, object]) _TIMEOUT: Final = timedelta(seconds=60) -_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery" +_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload::text AS accounting_payload +FROM "LiteLLM_SpendLogDelivery" WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP AND (attempts > 0) = $2::boolean ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED""" @@ -79,10 +84,20 @@ 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, accounting_payload""" +RETURNING delivery_id, attempts""" +_UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload' +WHERE delivery_id = ANY($1::text[]) AND state = 'pending'""" + + +def _decoded(delivery: _ClaimedDelivery) -> DeliveryProjection | None: + try: + return DeliveryProjection(_PAYLOAD.validate_json(delivery.accounting_payload)) + except ValueError: + return None def _is_permanent_failure(error: Exception) -> bool: @@ -138,26 +153,25 @@ class AccountingDeliveryRepository: try: async with self._database.tx(timeout=_TIMEOUT) as transaction: retries: Final = ( - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True)) + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, 1, True)) if retry_only is not False else () ) - claimed = retries or ( # rebind-ok: retain claimed identities after rollback - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) + rows: Final = retries or ( + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) if retry_only is not True else () ) - if not claimed: - return DeliveryBatchResult() - await apply( - transaction, - tuple( - DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload)) - for delivery in claimed - ), - ) - await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed)) - return DeliveryBatchResult(applied=len(claimed)) + decoded: Final = tuple((row, _decoded(row)) for row in rows) + undecodable: Final = tuple(row.delivery_id for row, projection in decoded if projection is None) + applicable: Final = tuple((row, projection) for row, projection in decoded if projection is not None) + if undecodable: + await transaction.execute_raw(_UNDECODABLE, undecodable) + claimed = tuple(row for row, _ in applicable) # rebind-ok: retain claimed identities after rollback + if claimed: + await apply(transaction, tuple(projection for _, projection in applicable)) + await transaction.execute_raw(_APPLIED, tuple(row.delivery_id for row in claimed)) + return DeliveryBatchResult(applied=len(claimed), parked=undecodable) except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry if not claimed: raise 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_spend_lock_order.py b/tests/integration/database/test_spend_lock_order.py new file mode 100644 index 00000000000..a7e00b7fa87 --- /dev/null +++ b/tests/integration/database/test_spend_lock_order.py @@ -0,0 +1,197 @@ +import asyncio +import subprocess +import sys +from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma + +from litellm.proxy._types import DBSpendUpdateTransactions +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.repositories import unit_of_work +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] +_KEY: Final = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" +_CHARGES: Final = DBSpendUpdateTransactions( + user_list_transactions={"user": 1.0}, + end_user_list_transactions={}, + key_list_transactions={_KEY: 1.0}, + team_list_transactions={"team": 1.0}, + team_member_list_transactions={"team_id::team::user_id::user": 1.0}, + org_list_transactions={"org": 1.0}, + org_member_list_transactions={"org_id::org::user_id::user": 1.0}, + project_list_transactions={"project": 1.0}, + tag_list_transactions={"tag": 1.0}, + agent_list_transactions={}, + model_access_group_list_transactions={"group": 1.0}, +) +_SEED: Final = ( + "INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 't', 't')", + 'INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) ' + "VALUES ('org', 'org', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')", + 'INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles) ' + """VALUES ('team', '[{"user_id": "user", "role": "user"}]')""", + "INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id, budget_id) VALUES ('team', 'user', 'budget')", + "INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')", + f"INSERT INTO \"LiteLLM_VerificationToken\" (token, budget_id) VALUES ('{_KEY}', 'budget')", + "INSERT INTO \"LiteLLM_ProjectTable\" (project_id, budget_id, created_by, updated_by) " + "VALUES ('project', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_TagTable\" (tag_name, budget_id) VALUES ('tag', 'budget')", + "INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name, budget_id) VALUES ('group', 'budget')", +) + + +@pytest.fixture(scope="module") +def schema_sql() -> str: + return subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), "--script"], + check=True, + capture_output=True, + text=True, + cwd=_ROOT, + ).stdout + + +@pytest.fixture +def database(schema_sql: str) -> Iterator[str]: + with scratch_database("spend_lock_order") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(schema_sql) + for statement in _SEED: + connection.execute(statement) + yield url + + +@pytest_asyncio.fixture(loop_scope="function") +async def ledger(database: str) -> AsyncIterator[Prisma]: + client: Final = Prisma(datasource={"url": database}) + await client.connect() + try: + yield client + finally: + await client.disconnect() + + +async def _apply_charges(client: Prisma) -> None: + async with client.tx() as transaction: + await DBSpendUpdateWriter().apply_accounting_entities(client, None, transaction, _CHARGES) + + +async def _ledger_waits_for_a_lock(url: str) -> None: + with psycopg.connect(url, autocommit=True) as observer: + for _ in range(200): + if observer.execute("SELECT 1 FROM pg_locks WHERE NOT granted").fetchone(): + return + await asyncio.sleep(0.05) + raise AssertionError("the ledger never waited on the other writer") + + +def _spend(url: str) -> Sequence[object]: + with psycopg.connect(url) as connection: + return connection.execute( + 'SELECT (SELECT spend FROM "LiteLLM_UserTable"), (SELECT spend FROM "LiteLLM_TeamTable"), ' + '(SELECT spend FROM "LiteLLM_TeamMembership"), (SELECT spend FROM "LiteLLM_VerificationToken"), ' + '(SELECT spend FROM "LiteLLM_OrganizationTable"), (SELECT spend FROM "LiteLLM_ProjectTable")' + ).fetchone() or () + + +_Writer = Callable[[psycopg.Connection[tuple[object, ...]]], None] + + +def _roster_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT pg_advisory_xact_lock(hashtext('team'))") + + +def _roster_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET members_with_roles = members_with_roles WHERE team_id = 'team'") + + +def _access_group_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET access_group_ids = ARRAY['ag'] WHERE team_id = 'team'") + + +def _access_group_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(f"UPDATE \"LiteLLM_VerificationToken\" SET access_group_ids = ARRAY['ag'] WHERE token = '{_KEY}'") + + +def _cascade_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(unit_of_work._LOCK_BUDGET_ORGS_SQL.replace("$1", "%s"), (["budget"],)) + + +def _cascade_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + for statement in ( + unit_of_work._LOCK_BUDGET_MEMBERSHIPS_SQL, + unit_of_work._LOCK_BUDGET_KEYS_SQL, + unit_of_work._LOCK_BUDGET_PROJECTS_SQL, + unit_of_work._LOCK_BUDGET_TAGS_SQL, + unit_of_work._LOCK_BUDGET_ACCESS_GROUPS_SQL, + ): + connection.execute(statement.replace("$1", "%s"), (["budget"],)) + connection.execute("UPDATE \"LiteLLM_TagTable\" SET spend = 0 WHERE budget_id = 'budget'") + + +@pytest.mark.parametrize( + ("first", "rest"), + [(_roster_first, _roster_rest), (_access_group_first, _access_group_rest), (_cascade_first, _cascade_rest)], + ids=["team_roster", "access_group", "budget_cascade"], +) +@pytest.mark.asyncio(loop_scope="function") +async def test_charges_wait_for_a_concurrent_writer_instead_of_deadlocking( + database: str, ledger: Prisma, first: _Writer, rest: _Writer +) -> None: + """The other writer takes its first lock, the ledger starts and must block holding nothing the other writer + needs next, so the other writer finishes and the charge lands once after it.""" + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + first(writer) + charges: Final = asyncio.create_task(_apply_charges(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(rest, writer) + writer.commit() + await asyncio.wait_for(charges, timeout=20) + + assert _spend(database) == (1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + + +async def _rename(client: Prisma) -> None: + await sync_model_allowlists_for_renamed_model( + SimpleNamespace(db=client), + model_id="deployment", + old_name="old-model", + new_name="new-model", + llm_router=None, + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio(loop_scope="function") +async def test_a_model_rename_waits_for_a_roster_change_instead_of_deadlocking(database: str, ledger: Prisma) -> None: + """The roster change holds the team lock and the user row, so the rename must wait before it locks the team + row the roster change updates next.""" + with psycopg.connect(database, autocommit=True) as seed: + seed.execute("UPDATE \"LiteLLM_TeamTable\" SET models = ARRAY['old-model']") + seed.execute("UPDATE \"LiteLLM_UserTable\" SET models = ARRAY['old-model']") + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + _roster_first(writer) + writer.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + rename: Final = asyncio.create_task(_rename(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(_roster_rest, writer) + writer.commit() + await asyncio.wait_for(rename, timeout=20) + + with psycopg.connect(database) as connection: + assert connection.execute( + 'SELECT (SELECT models FROM "LiteLLM_TeamTable"), (SELECT models FROM "LiteLLM_UserTable")' + ).fetchone() == (["new-model"], ["new-model"]) diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index 19f01a6b902..d1227e85565 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -131,9 +131,9 @@ async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(deliver "accounting_log_disabled": mode == "disabled", } assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1 - assert delivery_database.execute( - 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"' - ).fetchone() == (first["accounting_payload"],) + assert delivery_database.execute('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + first["accounting_payload"], + ) assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( [(first["spend"],)] if mode == "local" else [] ) @@ -415,11 +415,28 @@ def test_delivery_migration_preserves_unused_old_receipts(state): assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) +@pytest.mark.parametrize( + ("payload", "huge_state"), + [ + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 401) || '}')::jsonb)", + "applied", + ), + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 5000) || '}')::jsonb)", + "parked", + ), + ("'[]'::jsonb", "parked"), + ], + ids=["beyond_float", "beyond_python_int_parse", "not_an_object"], +) @pytest.mark.asyncio -async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client): - _insert(delivery_database, (_row("bad"), _row("good"))) +async def test_a_payload_the_engine_cannot_decode_never_blocks_other_deliveries( + delivery_database, delivery_client, payload, huge_state +): + _insert(delivery_database, (_row("huge"), _row("good"))) delivery_database.execute( - 'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$' + f'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload = {payload} WHERE request_id = $$huge$$' ) async def apply(transaction, deliveries): @@ -428,12 +445,16 @@ async def test_malformed_payload_is_parked_without_losing_other_deliveries(deliv ) repository = AccountingDeliveryRepository(delivery_client) - assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) - delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') - assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",) - assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + results = [await repository.apply_pending(apply, max_attempts=1) for _ in range(3)] + _insert(delivery_database, (_row("later"),)) + results.append(await repository.apply_pending(apply, max_attempts=1)) + + applied: Final = 3 if huge_state == "applied" else 2 + assert sum(result.applied for result in results) == applied + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5 * applied,) + assert delivery_database.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("good", "applied"), ("huge", huge_state), ("later", "applied")] @pytest.mark.asyncio @@ -463,19 +484,29 @@ async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode( from litellm.proxy.utils import jsonify_object completed = { - **_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + **_row("batch_completed"), + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": '{"status":"success","private":"request detail"}', "accounting_payload": { - "request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "request_id": "batch_completed", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": {"accounting_savings": {"compression": 0.2}}, }, } monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}}) client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object) writer = DBSpendUpdateWriter() - first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled) + first = await writer._claim_settle_once_cost_row( + payload=completed, prisma_client=client, disable_spend_logs=disabled + ) again = await writer._claim_settle_once_cost_row( - payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled + payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, + prisma_client=client, + disable_spend_logs=disabled, ) assert (first, again) == (True, False) assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == ( diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index e1894aa66bf..60f3c074cca 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -63,10 +63,10 @@ _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,transactions,distinct_models", [(None, 251, 1, False), (None, 1001, 3, False), ("200", 251, 2, False), ("3001", 3001, 1, 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 +async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, transactions: int, distinct_models: bool ): from litellm.accounting.worker import flush_accounting_deliveries @@ -92,13 +92,37 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction( )) == row_count - 1 started = time.monotonic() await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + print(f"Ledger backlog queued={row_count} seconds={time.monotonic() - started:.3f}") with psycopg.connect(parity_database) as connection: - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) - assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) - assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (row_count,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (0,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (transactions,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (row_count * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (row_count if distinct_models else 1,) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + "FROM generate_series(1, 501) AS n" + ) + await client.db.connect() + try: + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) finally: await client.db.disconnect() @@ -404,6 +428,7 @@ async def _assert_legacy_database_parity( assert len(rows) == 1 if path == "ledger" and retention == "lean": assert rows[0]["accounting_log_disabled"] is True + assert not {"messages", "response", "proxy_server_request", "metadata"}.intersection(rows[0]) frozen = rows[0]["accounting_payload"] if isinstance(frozen, str): frozen = json.loads(frozen) diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 17186a89922..c78d6bec757 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -150,6 +150,7 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) monkeypatch.setattr(proxy_server, "prisma_client", client) monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) monkeypatch.setattr(client, "spend_log_transactions", []) monkeypatch.setattr(client, "model_usage_transactions", []) writer: Final = DBSpendUpdateWriter() @@ -521,7 +522,9 @@ def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> N assert dict(daily_user.increments)["compression_saved_tokens"] == 7 -def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None: +def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_the_charge( + monkeypatch: pytest.MonkeyPatch, +) -> None: import litellm monkeypatch.setitem( @@ -552,6 +555,8 @@ def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(m ) restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0 - assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [ - 0.5 - ] + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert [ + dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken" + ] == [0.5] diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1b8bdd21326..740ce4f433a 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -230,9 +230,7 @@ def test_pydantic_base_model(): inner: InnerModel tags: list - outer = OuterModel( - name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"] - ) + outer = OuterModel(name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"]) # Test a pydantic model at the top level result = json.loads(safe_dumps(outer)) @@ -261,3 +259,9 @@ def test_safe_json_structure_keeps_tuples_and_drops_non_string_keys(): assert structure == {"models": ("A", "B"), "tags": ["X", "Y"], "nested": {"deep": ("C",)}} assert type(structure["models"]) is tuple assert json.loads(safe_dumps(data)) == {"models": ["a", "b"], "tags": ["x", "y"], "nested": {"deep": ["c"]}} + + +def test_integer_limit_clamps_both_signs_and_leaves_bools_and_floats_alone() -> None: + assert safe_json_structure( + {"big": 10**400, "small": -(10**400), "inside": 7, "flag": True, "ratio": 1e300}, integer_limit=100 + ) == {"big": 100, "small": -100, "inside": 7, "flag": True, "ratio": 1e300} diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 92114b00398..eee5d96511b 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3324,9 +3324,13 @@ class TestMultiPodBatchCostClaim: @pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external")) @pytest.mark.parametrize("collector_enabled", (False, True)) async def test_builtin_accounting_retries_only_before_a_possible_charge( - self, failure_phase: str, collector_enabled: bool + self, failure_phase: str, collector_enabled: bool, monkeypatch: pytest.MonkeyPatch ) -> None: import base64 + + from litellm.proxy import proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) from datetime import datetime import httpx 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/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index c6366a40504..0c7292674a6 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2851,9 +2851,12 @@ async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_requ live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch ): """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row - it already logged from being counted twice.""" + it already logged from being counted twice. Only the legacy writers queue this row, so the ledger is off.""" + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import proxy_logging_obj + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue await window_queue.flush_and_get_aggregated_window_spend_transactions() key_cache: Final = DualCache() 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..9665b15bd60 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", 'SELECT team_id FROM "LiteLLM_TeamTable"')): + 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) @@ -11420,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"] 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/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index e119acd5ffc..c3e916d581d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3819,7 +3819,10 @@ class TestSpendLogsPayload: await _wait_for_mock_call(mock_client) kwargs = mock_client.call_args.kwargs - payload: SpendLogsPayload = kwargs["payload"] + payload: SpendLogsPayload = { + key: value for key, value in kwargs["payload"].items() if not key.startswith("accounting_") + } + assert kwargs["payload"]["accounting_payload"]["spend"] == kwargs["payload"]["spend"] expected_payload = SpendLogsPayload( **{ "request_id": "chatcmpl-34df56d5-4807-45c1-bb99-61e52586b802", diff --git a/tests/unit/proxy/test_unit_test_proxy_hooks.py b/tests/unit/proxy/test_unit_test_proxy_hooks.py index 6ce7154f2d1..7990ec12037 100644 --- a/tests/unit/proxy/test_unit_test_proxy_hooks.py +++ b/tests/unit/proxy/test_unit_test_proxy_hooks.py @@ -11,7 +11,7 @@ import litellm @pytest.mark.asyncio async def test_disable_spend_logs(): """ - Test that the spend logs are not written to the database when disable_spend_logs is True + With disable_spend_logs the request is queued only for its charge, flagged so no SpendLogs row is written """ # Mock the necessary components @@ -41,5 +41,4 @@ async def test_disable_spend_logs(): org_id="org_id", kwargs={}, ) - # Verify no spend logs were added - assert len(mock_prisma_client.spend_log_transactions) == 0 + assert [row["accounting_log_disabled"] for row in mock_prisma_client.spend_log_transactions] == [True] diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index afc282a14d4..888f664b638 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -142,7 +142,9 @@ def mock_prisma_client() -> MagicMock: client.db.stop_token_refresh_task = AsyncMock() client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) client.db.execute_raw = AsyncMock() - client.db.tx = MagicMock() + client.db.tx = MagicMock( + return_value=MagicMock(__aenter__=AsyncMock(return_value=client.db), __aexit__=AsyncMock(return_value=False)) + ) client.db.batch_ = MagicMock() for table_name in _PRISMA_TABLES: setattr(client.db, table_name, _make_table_mock()) 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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f3c465bcf85..d14e490f606 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31345,8 +31345,8 @@ export interface components { use_google_kms?: boolean | null; /** * Use Spend Logs Ledger - * @description Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass. - * @default false + * @description Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues. + * @default true */ use_spend_logs_ledger: boolean; /**