diff --git a/litellm/proxy/management_endpoints/management_v1/customers.py b/litellm/proxy/management_endpoints/management_v1/customers.py index be538603d6f..a678a535385 100644 --- a/litellm/proxy/management_endpoints/management_v1/customers.py +++ b/litellm/proxy/management_endpoints/management_v1/customers.py @@ -52,6 +52,11 @@ T: Final = TypeVar("T") _ERROR_DETAILS: Final[Mapping[ErrorCode, tuple[int, str, str]]] = MappingProxyType( { + ErrorCode.CAPACITY_BUSY: ( + 503, + "Customer portal busy", + "Customer portal capacity is being updated. Retry the request.", + ), ErrorCode.FREE_TIER_LIMIT: ( 403, "Free customer portal limit reached", diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index 63b51267b6d..f55bb81899e 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -16,6 +16,7 @@ from litellm.proxy.common_utils.timezone_utils import ( get_budget_reset_time, ) from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.list_api.common import escape_like from litellm.proxy.management_helpers.customer_budget_usage import customer_budget_usage from litellm.repositories.budget_repository import create_budget_row, update_budget_row @@ -46,6 +47,7 @@ if TYPE_CHECKING: class ErrorCode(str, Enum): FREE_TIER_LIMIT = "free_tier_limit" + CAPACITY_BUSY = "capacity_busy" NOT_FOUND = "not_found" PARENT_NOT_FOUND = "parent_not_found" ALREADY_EXISTS = "already_exists" @@ -76,6 +78,7 @@ FREE_ORGANIZATION_LIMIT: Final = 5 FREE_TEAMS_PER_ORGANIZATION_LIMIT: Final = 5 FREE_MEMBERS_PER_TEAM_LIMIT: Final = 5 FREE_CAPACITY_ADVISORY_LOCK: Final = 714023581 +FREE_CAPACITY_LOCK_TIMEOUT_MS: Final = 1000 class CustomerPortalLimits(BaseModel): @@ -501,14 +504,22 @@ class CustomerHierarchyStore: self.budgets: TableActions[PrismaRecord] = db.litellm_budgettable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings self.window_spend: TableActions[LiteLLM_BudgetWindowSpend] = db.litellm_budgetwindowspend # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings - async def lock_free_capacity(self) -> None: - await self.db.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock($1)", FREE_CAPACITY_ADVISORY_LOCK) + async def lock_free_capacity(self, scope: str | None = None) -> None: + await self.db.query_raw("SELECT set_config('lock_timeout', $1, true)", f"{FREE_CAPACITY_LOCK_TIMEOUT_MS}ms") + if scope is None: + await self.db.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock($1)", FREE_CAPACITY_ADVISORY_LOCK) + return + await self.db.query_raw( + "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended($1, $2))", + scope, + FREE_CAPACITY_ADVISORY_LOCK, + ) async def member_capacity_reached(self, team_ids: Sequence[str]) -> bool: if self.premium_user or not team_ids: return False - await self.lock_free_capacity() - for team_id in team_ids: + for team_id in sorted(frozenset(team_ids)): + await self.lock_free_capacity(f"members:{team_id}") if await self.members.count(where={"parent_team_ids": {"has": team_id}}) >= FREE_MEMBERS_PER_TEAM_LIMIT: return True return False @@ -876,7 +887,7 @@ class CustomerHierarchyStore: if parent is None: return Failure(ErrorCode.PARENT_NOT_FOUND) if not self.premium_user: - await self.lock_free_capacity() + await self.lock_free_capacity(f"teams:{payload.parent_org_id}") if await self.teams.count( where={"parent_org_id": payload.parent_org_id} ) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT and ( @@ -1192,7 +1203,7 @@ class CustomerHierarchy: *, operation_type: TransactionOperation = TransactionOperation.READ_WRITE, ) -> T | Failure: - from prisma.errors import ForeignKeyViolationError, RecordNotFoundError, UniqueViolationError + from prisma.errors import DataError, ForeignKeyViolationError, RecordNotFoundError, UniqueViolationError try: async with self.db.tx() as tx: @@ -1205,6 +1216,10 @@ class CustomerHierarchy: return Failure( ErrorCode.HAS_CHILDREN if operation_type is TransactionOperation.DELETE else ErrorCode.PARENT_NOT_FOUND ) + except DataError as error: + if PrismaDBExceptionHandler.is_lock_timeout_error(error): + return Failure(ErrorCode.CAPACITY_BUSY) + raise async def create_organization(self, payload: OrganizationCreate, actor: str) -> Organization | Failure: return await self.transaction(lambda store: store.create_organization(payload, actor)) diff --git a/litellm/proxy/management_helpers/customer_upsert.py b/litellm/proxy/management_helpers/customer_upsert.py index a89049e34f8..6285a7580e5 100644 --- a/litellm/proxy/management_helpers/customer_upsert.py +++ b/litellm/proxy/management_helpers/customer_upsert.py @@ -15,6 +15,7 @@ from litellm.proxy.management_helpers.customer_hierarchy import ( from litellm.types.proxy.management_endpoints.customer_hierarchy import ( Budget, BudgetTemplate, + CustomerMember, MemberCreate, MemberPatch, MemberUpsert, @@ -159,7 +160,7 @@ async def customer_team( else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) ) if not store.premium_user: - await store.lock_free_capacity() + await store.lock_free_capacity(f"teams:{organization_id}") if ( await store.teams.find_unique(where={"customer_team_id": default_team_id(organization_id)}) is None and await store.teams.count(where={"parent_org_id": organization_id}) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT @@ -184,6 +185,23 @@ async def customer_team( return team if team.parent_org_id == organization_id else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH) +async def member_team_ids( + store: CustomerHierarchyStore, + organization_id: str, + existing: CustomerMember | Failure, + team_id: str | None, + actor: str, +) -> tuple[str, ...] | Failure: + if not isinstance(existing, Failure) and team_id is None: + return existing.parent_team_ids + team: Final = await customer_team(store, organization_id, team_id, actor) + if isinstance(team, Failure): + return team + if isinstance(existing, Failure): + return (team.customer_team_id,) + return (team.customer_team_id, *(tid for tid in existing.parent_team_ids if tid != team.customer_team_id)) + + async def upsert_member( store: CustomerHierarchyStore, payload: MemberUpsert, actor: str ) -> MemberUpsertResponse | Failure: @@ -226,23 +244,18 @@ async def upsert_member( ) if selected is None and override is None: return Failure(ErrorCode.BUDGET_REQUIRED) - team: Final = await customer_team(store, organization.customer_org_id, payload.parent_team_id, actor) - if isinstance(team, Failure): - return team - team_ids: Final = ( - (team.customer_team_id,) - if isinstance(existing, Failure) - else existing.parent_team_ids - if payload.parent_team_id is None - else (team.customer_team_id, *(tid for tid in existing.parent_team_ids if tid != team.customer_team_id)) + team_ids: Final = await member_team_ids( + store, organization.customer_org_id, existing, payload.parent_team_id, actor ) + if isinstance(team_ids, Failure): + return team_ids if isinstance(existing, Failure): member: Final = await store.create_member( MemberCreate.model_validate( { **patch.model_dump(exclude_unset=True), "customer_member_id": payload.customer_member_id, - "parent_team_ids": (team.customer_team_id,), + "parent_team_ids": team_ids, } ), actor, diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index 2b9eb2a8164..3c94c657e33 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Final, cast import pytest from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from prisma.errors import DataError, RawQueryError from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key, customer_org_alias_cache_key @@ -19,6 +20,7 @@ from litellm.types.proxy.management_endpoints.customer_hierarchy import Customer from tests.unit.proxy.management_helpers.test_customer_hierarchy import ( NOW, BoundedChildrenTable, + CapacityErrorDatabase, Database, DeleteConflictTable, Record, @@ -127,6 +129,15 @@ def test_member_upsert_respects_capacity_for_default_team_creation_and_added_mem "/management/v1/customer_members", json={"customer_member_id": "moving", "parent_team_ids": ["extra-0"]} ) assert original.status_code == 201, original.text + for patch in ({}, {"alias": "renamed"}): + kept: Final = http.post( + "/management/v1/customer_members/upsert", + json={"customer_org_id": "org", "customer_member_id": "moving", **patch}, + ) + assert kept.status_code == 200, kept.text + assert kept.json()["created"] is False + assert kept.json()["parent_team_ids"] == ["extra-0"] + assert kept.json()["alias"] == patch.get("alias") moved: Final = http.post( "/management/v1/customer_members/upsert", json={"customer_org_id": "org", "customer_member_id": "moving", "parent_team_id": "team"}, @@ -137,11 +148,63 @@ def test_member_upsert_respects_capacity_for_default_team_creation_and_added_mem assert http.get("/management/v1/customer_teams/org:default").status_code == 404 +@pytest.mark.parametrize("error_code", ("55P03", "XX000")) +def test_capacity_contention_is_retryable_without_hiding_other_database_errors(error_code: str) -> None: + http: Final = client(database=CapacityErrorDatabase(error_code, "members:team"), premium_user=False) + family(http) + payload: Final = {"customer_member_id": "blocked", "parent_team_ids": ["team"]} + if error_code != "55P03": + with pytest.raises(RawQueryError, match="advisory lock failed"): + http.post("/management/v1/customer_members", json=payload) + return + response: Final = http.post("/management/v1/customer_members", json=payload) + assert response.status_code == 503, response.text + assert response.json() == { + "type": "urn:litellm:error:customer-hierarchy-capacity-busy", + "title": "Customer portal busy", + "status": 503, + "detail": "Customer portal capacity is being updated. Retry the request.", + } + assert http.get("/management/v1/customer_members/blocked").status_code == 404 + + def test_customer_portal_limits_require_admin() -> None: response: Final = client(role=LitellmUserRoles.INTERNAL_USER).get("/management/v1/customer_portal/limits") assert response.status_code == 403, response.text +@pytest.mark.parametrize( + ("error", "retryable"), + ( + (DataError({"user_facing_error": {"meta": {"code": "55P03", "message": "lock unavailable"}}}), True), + (DataError({"user_facing_error": {"message": "canceling statement due to lock timeout"}}), True), + (DataError({"user_facing_error": {"message": "other database error", "meta": {"code": "XX000"}}}), False), + ), + ids=("sqlstate", "message-only", "other-error"), +) +def test_team_model_write_contention_is_retryable_and_rolls_back(error: DataError, retryable: bool) -> None: + database: Final = Database(team_create_error=error) + http: Final = client(database=database, premium_user=False) + created: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "org"}) + assert created.status_code == 201, created.text + payload: Final = {"customer_team_id": "blocked", "parent_org_id": "org", "budget": {"max_budget": 10}} + if not retryable: + with pytest.raises(DataError, match="other database error") as raised: + http.post("/management/v1/customer_teams", json=payload) + assert raised.value is error + else: + response: Final = http.post("/management/v1/customer_teams", json=payload) + assert response.status_code == 503, response.text + assert response.json() == { + "type": "urn:litellm:error:customer-hierarchy-capacity-busy", + "title": "Customer portal busy", + "status": 503, + "detail": "Customer portal capacity is being updated. Retry the request.", + } + assert http.get("/management/v1/customer_teams/blocked").status_code == 404 + assert not database.litellm_budgettable.rows + + def family(http: TestClient) -> dict[str, object]: org: Final = http.post( "/management/v1/customer_organizations", diff --git a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py index da73d4617bf..cb36781e9cb 100644 --- a/tests/unit/proxy/management_helpers/test_customer_hierarchy.py +++ b/tests/unit/proxy/management_helpers/test_customer_hierarchy.py @@ -1,3 +1,4 @@ +import asyncio import json import sqlite3 from collections.abc import AsyncIterator, Mapping, Sequence @@ -5,10 +6,12 @@ from contextlib import asynccontextmanager from datetime import datetime, timezone from functools import reduce from itertools import chain +from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast import pytest from prisma import Json +from prisma.errors import DataError, RawQueryError from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from litellm.models.budget import BudgetLimit, LiteLLM_BudgetTable @@ -23,8 +26,10 @@ from litellm.proxy.management_helpers.customer_hierarchy import ( ) from litellm.types.proxy.management_endpoints.customer_hierarchy import ( Budget, + CustomerMember, MemberCreate, OrganizationCreate, + Team, TeamCreate, ) @@ -84,9 +89,10 @@ def matches_where(row: Mapping[str, object], where: Mapping[str, object]) -> boo class Table: - def __init__(self, id_field: str, database: "Database") -> None: + def __init__(self, id_field: str, database: "Database", *, create_error: DataError | None = None) -> None: self.id_field = id_field self.database = database + self.create_error = create_error self.rows: Mapping[str, Record] = {} async def include(self, row: Record, include: Mapping[str, object] | None) -> Record: @@ -109,6 +115,8 @@ class Table: async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record: from prisma.errors import UniqueViolationError + if self.create_error is not None: + raise self.create_error if str(data[self.id_field]) in self.rows: raise UniqueViolationError({}) row: Final = Record.model_validate( @@ -237,16 +245,63 @@ class DeleteConflictTable(Table): raise ForeignKeyViolationError({}) -class Database: +class AdvisoryLocks: def __init__(self) -> None: - self.litellm_customerorganizationtable = Table("customer_org_id", self) - self.litellm_customerteamtable = Table("customer_team_id", self) - self.litellm_customermembertable = Table("customer_member_id", self) - self.litellm_budgettable = Table("budget_id", self) - self.litellm_budgetwindowspend = Table("entity_id", self) - self.litellm_verificationtoken = Table("token", self) - self.litellm_projecttable = Table("project_id", self) - self.litellm_customerbudgetincreaserequest = Table("request_id", self) + self.locks: Mapping[tuple[str, ...], asyncio.Lock] = MappingProxyType({}) + + def lock_for(self, query: str, args: Sequence[object]) -> asyncio.Lock: + key: Final = (query, *(str(arg) for arg in args)) + existing: Final = self.locks.get(key) + if existing is not None: + return existing + lock: Final = asyncio.Lock() + self.locks = MappingProxyType({**self.locks, key: lock}) + return lock + + +class PausedCreateTable(Table): + def __init__(self, id_field: str, database: "Database", entered: asyncio.Event, release: asyncio.Event) -> None: + super().__init__(id_field, database) + self.entered = entered + self.release = release + + async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record: + if data.get(self.id_field) == "blocked": + self.entered.set() + await self.release.wait() + return await super().create(data, include) + + +class Database: + def __init__( + self, + advisory_locks: AdvisoryLocks | None = None, + source: "Database | None" = None, + *, + team_create_error: DataError | None = None, + ) -> None: + self.litellm_customerorganizationtable = ( + Table("customer_org_id", self) if source is None else source.litellm_customerorganizationtable + ) + self.litellm_customerteamtable = ( + Table("customer_team_id", self, create_error=team_create_error) + if source is None + else source.litellm_customerteamtable + ) + self.litellm_customermembertable = ( + Table("customer_member_id", self) if source is None else source.litellm_customermembertable + ) + self.litellm_budgettable = Table("budget_id", self) if source is None else source.litellm_budgettable + self.litellm_budgetwindowspend = ( + Table("entity_id", self) if source is None else source.litellm_budgetwindowspend + ) + self.litellm_verificationtoken = Table("token", self) if source is None else source.litellm_verificationtoken + self.litellm_projecttable = Table("project_id", self) if source is None else source.litellm_projecttable + self.litellm_customerbudgetincreaserequest = ( + Table("request_id", self) if source is None else source.litellm_customerbudgetincreaserequest + ) + self.advisory_locks = advisory_locks + self.held_locks: tuple[asyncio.Lock, ...] = () self.locked_organization_ids: tuple[str, ...] = () def has_budget_reference(self, budget_id: str, relation: str) -> bool: @@ -265,9 +320,21 @@ class Database: table, field = reference return any(row.model_dump().get(field) == budget_id for row in table.rows.values()) - async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: - if "pg_advisory_xact_lock" in query: + async def acquire_advisory_lock(self, query: str, args: Sequence[object]) -> tuple[Mapping[str, object], ...]: + if self.advisory_locks is None: return () + lock: Final = self.advisory_locks.lock_for(query, args) + if lock in self.held_locks: + return () + await lock.acquire() + self.held_locks = (*self.held_locks, lock) + return () + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if "set_config('lock_timeout'" in query: + return () + if "pg_advisory_xact_lock" in query: + return await self.acquire_advisory_lock(query, args) if "SELECT EXISTS" in query: return ({"referenced": False},) if "FOR UPDATE" in query: @@ -334,6 +401,81 @@ class Database: for table, rows in zip(tables, snapshots): table.rows = rows raise + finally: + for lock in self.held_locks: + lock.release() + self.held_locks = () + + +class CapacityErrorDatabase(Database): + def __init__(self, error_code: str, scope: str) -> None: + super().__init__() + self.error_code = error_code + self.scope = scope + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + if "pg_advisory_xact_lock" in query and args and args[0] == self.scope: + raise RawQueryError( + {"user_facing_error": {"meta": {"code": self.error_code, "message": "advisory lock failed"}}} + ) + return await super().query_raw(query, *args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ("team", "member")) +@pytest.mark.parametrize("same_parent", (False, True)) +async def test_capacity_contention_is_isolated_to_each_parent(kind: str, same_parent: bool) -> None: + locks: Final = AdvisoryLocks() + database: Final = Database(locks) + hierarchy: Final = CustomerHierarchy( + cast("Prisma", database), premium_user=False, clock=lambda: NOW + ) # cast-ok: in-memory Prisma edge + for org_id in ("a", "b"): + assert not isinstance( + await hierarchy.create_organization(OrganizationCreate(customer_org_id=org_id), "admin"), Failure + ) + if kind == "member": + assert not isinstance( + await hierarchy.create_team(TeamCreate(customer_team_id=org_id, parent_org_id=org_id), "admin"), Failure + ) + entered: Final = asyncio.Event() + release: Final = asyncio.Event() + if kind == "team": + database.litellm_customerteamtable = PausedCreateTable("customer_team_id", database, entered, release) + else: + database.litellm_customermembertable = PausedCreateTable("customer_member_id", database, entered, release) + other: Final = CustomerHierarchy( + cast("Prisma", Database(locks, source=database)), premium_user=False, clock=lambda: NOW + ) # cast-ok: independent connection to the shared in-memory tables + + async def create(store: CustomerHierarchy, parent: str, identity: str) -> Team | CustomerMember | Failure: + if kind == "team": + return await store.create_team(TeamCreate(customer_team_id=identity, parent_org_id=parent), "admin") + return await store.create_member( + MemberCreate(customer_member_id=identity, parent_team_ids=(parent,), budget_override=Budget(max_budget=10)), + "admin", + ) + + for index in range(4): + assert not isinstance(await create(hierarchy, "a", f"seed-{index}"), Failure) + blocked: Final = asyncio.create_task(create(hierarchy, "a", "blocked")) + await entered.wait() + independent: Final = asyncio.create_task(create(other, "a" if same_parent else "b", "independent")) + try: + if same_parent: + await asyncio.sleep(0) + release.set() + result: Final = await asyncio.wait_for(independent, timeout=10) + if same_parent: + assert result == Failure(ErrorCode.FREE_TIER_LIMIT) + else: + assert not isinstance(result, Failure) + assert result == ( + await other.get_team("independent") if kind == "team" else await other.get_member("independent") + ) + finally: + release.set() + assert not isinstance(await blocked, Failure) @pytest.mark.asyncio