mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
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 <noreply@anthropic.com>
This commit is contained in:
Joshua Valluru
parent
0fa9ac3ce0
commit
bbb85ed48e
19 changed files with 395 additions and 238 deletions
|
|
@ -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)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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::<value>::user_id::<value>"; 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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue