mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): enforce settled customer budgets across workers (#51568)
* fix(proxy): enforce settled customer budgets across workers Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(proxy): sync customer budget reset API types Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(proxy): reseed customer spend through shared counters Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
This commit is contained in:
ishaan-j-berri
GitHub
parent
6a1f81b9f8
commit
6eda17aadb
14 changed files with 500 additions and 6 deletions
|
|
@ -22,6 +22,7 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
|||
from litellm.proxy.auth.auth_checks import can_object_call_model
|
||||
from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET
|
||||
from litellm.proxy.auth.customer_member_auto_create import ensure_customer_member
|
||||
from litellm.proxy.auth.customer_spend import current_customer_spend
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl
|
||||
|
|
@ -61,6 +62,7 @@ class _OrganizationRow(BaseModel):
|
|||
|
||||
customer_org_id: str
|
||||
spend: float = 0.0
|
||||
budget_reset_at: datetime | None = None
|
||||
budget: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
|
|
@ -69,6 +71,7 @@ class _TeamRow(BaseModel):
|
|||
|
||||
customer_team_id: str
|
||||
spend: float = 0.0
|
||||
budget_reset_at: datetime | None = None
|
||||
budget: Mapping[str, object] | None = None
|
||||
parent_org: _OrganizationRow
|
||||
|
||||
|
|
@ -78,6 +81,7 @@ class _MemberRow(BaseModel):
|
|||
|
||||
customer_member_id: str
|
||||
spend: float = 0.0
|
||||
budget_reset_at: datetime | None = None
|
||||
budget: Mapping[str, object] | None = None
|
||||
budget_override: Mapping[str, object] | None = None
|
||||
budget_override_expires_at: datetime | None = None
|
||||
|
|
@ -129,13 +133,16 @@ def customer_identity_from_headers(headers: Mapping[str, object] | None) -> Cust
|
|||
return identity
|
||||
|
||||
|
||||
def _scope(key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None) -> CustomerScopeAuth:
|
||||
def _scope(
|
||||
key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None, budget_reset_at: datetime | None = None
|
||||
) -> CustomerScopeAuth:
|
||||
if budget is None:
|
||||
return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend)
|
||||
return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend, budget_reset_at=budget_reset_at)
|
||||
return CustomerScopeAuth(
|
||||
key=key,
|
||||
entity_id=entity_id,
|
||||
spend=spend,
|
||||
budget_reset_at=budget_reset_at,
|
||||
max_budget=budget.max_budget,
|
||||
tpm_limit=budget.tpm_limit,
|
||||
rpm_limit=budget.rpm_limit,
|
||||
|
|
@ -163,13 +170,21 @@ def customer_hierarchy_from_member_row(row: Mapping[str, object], now: datetime)
|
|||
organization.customer_org_id,
|
||||
organization.spend,
|
||||
related_budget(organization.budget, now),
|
||||
organization.budget_reset_at,
|
||||
),
|
||||
_scope(
|
||||
"customer_team",
|
||||
team.customer_team_id,
|
||||
team.spend,
|
||||
related_budget(team.budget, now),
|
||||
team.budget_reset_at,
|
||||
),
|
||||
_scope("customer_team", team.customer_team_id, team.spend, related_budget(team.budget, now)),
|
||||
_scope(
|
||||
"customer_member",
|
||||
member.customer_member_id,
|
||||
member.spend,
|
||||
effective_member_budget(template, override, member.budget_override_expires_at, now),
|
||||
member.budget_reset_at,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
|
@ -186,6 +201,7 @@ def customer_hierarchy_from_organization_row(row: Mapping[str, object], now: dat
|
|||
organization.customer_org_id,
|
||||
organization.spend,
|
||||
related_budget(organization.budget, now),
|
||||
organization.budget_reset_at,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
|
@ -300,7 +316,9 @@ def _not_found(message: str, param: str) -> ProxyException:
|
|||
|
||||
|
||||
async def _organization_named_by_alias(
|
||||
customer_org_alias: str, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache
|
||||
customer_org_alias: str,
|
||||
prisma_client: "PrismaClient | None",
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> CustomerHierarchyAuth:
|
||||
organization: Final = (
|
||||
await get_customer_organization_hierarchy(customer_org_alias, prisma_client, user_api_key_cache)
|
||||
|
|
@ -411,7 +429,10 @@ async def resolve_customer_hierarchy(
|
|||
) -> CustomerHierarchyAuth:
|
||||
"""Load the hierarchy the headers name and admit the request against it. ``request_model`` is the model the
|
||||
route will actually serve, which may differ from the body."""
|
||||
hierarchy: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register)
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
cached: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register)
|
||||
hierarchy: Final = await current_customer_spend(cached, get_current_spend)
|
||||
check_customer_hierarchy_key_scope(hierarchy, key_metadata or {})
|
||||
if not RouteChecks.is_info_route(route):
|
||||
check_customer_hierarchy_budget(hierarchy)
|
||||
|
|
|
|||
53
litellm/proxy/auth/customer_spend.py
Normal file
53
litellm/proxy/auth/customer_spend.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
import asyncio
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.redis.stores.spend_counters import SpendCounterStore
|
||||
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
|
||||
|
||||
|
||||
def customer_spend_counter(scope: CustomerScopeAuth) -> str:
|
||||
period: Final = scope.budget_reset_at.isoformat() if scope.budget_reset_at is not None else "lifetime"
|
||||
return f"spend:{scope.key}:{scope.entity_id}:{period}"
|
||||
|
||||
|
||||
async def _current_scope(
|
||||
scope: CustomerScopeAuth, current: float | None, store: SpendCounterStore
|
||||
) -> CustomerScopeAuth:
|
||||
spend: Final = (
|
||||
await store.seed(customer_spend_counter(scope), scope.spend, increment_if_lost=False)
|
||||
if current is None
|
||||
else current
|
||||
)
|
||||
return scope.model_copy(update={"spend": max(scope.spend, spend)})
|
||||
|
||||
|
||||
async def current_customer_spend(
|
||||
hierarchy: CustomerHierarchyAuth,
|
||||
read_spend: Callable[[str, float, float | None], Awaitable[float]],
|
||||
) -> CustomerHierarchyAuth:
|
||||
readings: Final = await asyncio.gather(
|
||||
*(read_spend(customer_spend_counter(scope), scope.spend, scope.max_budget) for scope in hierarchy.scopes)
|
||||
)
|
||||
scopes: Final = tuple(
|
||||
scope.model_copy(update={"spend": max(scope.spend, spend)})
|
||||
for scope, spend in zip(hierarchy.scopes, readings, strict=True)
|
||||
)
|
||||
return hierarchy.model_copy(update={"scopes": tuple(scopes)})
|
||||
|
||||
|
||||
async def settle_customer_spend(metadata: Mapping[str, object], cost: float, store: SpendCounterStore) -> None:
|
||||
snapshot: Final = metadata.get("user_api_key_customer_hierarchy")
|
||||
if snapshot is None or cost == 0:
|
||||
return
|
||||
try:
|
||||
hierarchy: Final = CustomerHierarchyAuth.model_validate(snapshot)
|
||||
except ValidationError:
|
||||
return
|
||||
readings: Final = await asyncio.gather(*(store.read(customer_spend_counter(scope)) for scope in hierarchy.scopes))
|
||||
_ = await asyncio.gather(
|
||||
*(_current_scope(scope, reading[0], store) for scope, reading in zip(hierarchy.scopes, readings, strict=True))
|
||||
)
|
||||
_ = await store.increment_many(tuple((customer_spend_counter(scope), cost) for scope in hierarchy.scopes))
|
||||
|
|
@ -13,6 +13,7 @@ memory in long-lived deployments.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -34,6 +35,9 @@ from litellm.repositories.organization_repository import OrganizationRepository
|
|||
from litellm.repositories.project_repository import ProjectRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
BudgetWindowSpendRepository,
|
||||
CustomerMemberRepository,
|
||||
CustomerOrganizationRepository,
|
||||
CustomerTeamRepository,
|
||||
EndUserRepository,
|
||||
SpendLogsRepository,
|
||||
TeamMembershipRepository,
|
||||
|
|
@ -69,6 +73,23 @@ def _as_utc(value: datetime) -> datetime:
|
|||
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
_CUSTOMER_COUNTER: Final = re.compile(
|
||||
r"spend:(customer_org|customer_team|customer_member):(.+):"
|
||||
r"(lifetime|\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:\d{2})?)"
|
||||
)
|
||||
|
||||
|
||||
def _customer_counter_period_matches(row: object, period: str) -> bool:
|
||||
try:
|
||||
expected: Final = None if period == "lifetime" else datetime.fromisoformat(period)
|
||||
except ValueError:
|
||||
return False
|
||||
actual: Final = getattr(row, "budget_reset_at", None)
|
||||
if expected is None or actual is None:
|
||||
return expected is actual
|
||||
return isinstance(actual, datetime) and _as_utc(actual) == _as_utc(expected)
|
||||
|
||||
|
||||
class SpendCounterReseed:
|
||||
"""
|
||||
Reseeds spend counters from the authoritative DB and warms the cache,
|
||||
|
|
@ -147,11 +168,28 @@ class SpendCounterReseed:
|
|||
return None
|
||||
if row is None:
|
||||
return None
|
||||
customer: Final = _CUSTOMER_COUNTER.fullmatch(counter_key)
|
||||
if customer is not None and not _customer_counter_period_matches(row, customer.group(3)):
|
||||
return None
|
||||
return float(getattr(row, "spend", 0.0) or 0.0)
|
||||
|
||||
@staticmethod
|
||||
async def _counter_row(prisma_client: "PrismaClient", counter_key: str) -> object | None:
|
||||
async with db_lookup_gate.current():
|
||||
customer: Final = _CUSTOMER_COUNTER.fullmatch(counter_key)
|
||||
if customer is not None:
|
||||
family, entity_id, _ = customer.groups()
|
||||
if family == "customer_org":
|
||||
return await CustomerOrganizationRepository(prisma_client).table.find_unique(
|
||||
where={"customer_org_id": entity_id}
|
||||
)
|
||||
if family == "customer_team":
|
||||
return await CustomerTeamRepository(prisma_client).table.find_unique(
|
||||
where={"customer_team_id": entity_id}
|
||||
)
|
||||
return await CustomerMemberRepository(prisma_client).table.find_unique(
|
||||
where={"customer_member_id": entity_id}
|
||||
)
|
||||
if counter_key.startswith("spend:key:"):
|
||||
token: Final = counter_key[len("spend:key:") :]
|
||||
return await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token})
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
get_key_object,
|
||||
get_team_object,
|
||||
)
|
||||
from litellm.proxy.auth.customer_spend import settle_customer_spend
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded
|
||||
from litellm.proxy.db.db_spend_update_writer import (
|
||||
|
|
@ -1440,7 +1441,14 @@ async def _update_database_and_spend_counters_in_batch(
|
|||
await _release_budget_reservation(budget_reservation=budget_reservation)
|
||||
return False
|
||||
|
||||
from litellm.proxy.proxy_server import spend_counter_store
|
||||
|
||||
try:
|
||||
await settle_customer_spend(
|
||||
_OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs)),
|
||||
response_cost,
|
||||
spend_counter_store,
|
||||
)
|
||||
await increment_spend_counters(
|
||||
token=user_api_key,
|
||||
team_id=team_id,
|
||||
|
|
|
|||
|
|
@ -2140,6 +2140,7 @@ class LiteLLMProxyRequestSetup:
|
|||
user_api_key_customer_member_id=customer.customer_member_id if customer else None,
|
||||
user_api_key_customer_team_id=customer.customer_team_id if customer else None,
|
||||
user_api_key_customer_org_id=customer.customer_org_id if customer else None,
|
||||
user_api_key_customer_hierarchy=customer.model_dump(mode="json") if customer else None,
|
||||
user_api_key_user_email=user_api_key_dict.user_email,
|
||||
user_api_key_request_route=user_api_key_dict.request_route,
|
||||
user_api_key_budget_reset_at=(
|
||||
|
|
|
|||
|
|
@ -23,6 +23,9 @@ pub enum SpendCounter {
|
|||
Project(String),
|
||||
ModelAccessGroup(String),
|
||||
ShadowEval(String),
|
||||
CustomerOrg(String),
|
||||
CustomerTeam(String),
|
||||
CustomerMember(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
|
|
@ -51,6 +54,9 @@ impl SpendCounter {
|
|||
Self::Project(id) => format!("spend:project:{id}"),
|
||||
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
|
||||
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
|
||||
Self::CustomerOrg(id) => format!("spend:customer_org:{id}"),
|
||||
Self::CustomerTeam(id) => format!("spend:customer_team:{id}"),
|
||||
Self::CustomerMember(id) => format!("spend:customer_member:{id}"),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -95,7 +101,24 @@ impl SpendCounter {
|
|||
"project" => Self::Project(id.to_owned()),
|
||||
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
|
||||
"shadow_eval" => Self::ShadowEval(id.to_owned()),
|
||||
"customer_org" => Self::CustomerOrg(id.to_owned()),
|
||||
"customer_team" => Self::CustomerTeam(id.to_owned()),
|
||||
"customer_member" => Self::CustomerMember(id.to_owned()),
|
||||
_ => return Err(invalid()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SpendCounter;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::organization("spend:customer_org:tenant:lifetime")]
|
||||
#[case::team("spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00")]
|
||||
#[case::member("spend:customer_member:user@example.com:lifetime")]
|
||||
fn customer_counter_keys_round_trip(#[case] key: &str) {
|
||||
assert_eq!(SpendCounter::parse(key).unwrap().key(), key);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -65,7 +65,10 @@ impl SpendCounters {
|
|||
|
||||
fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration {
|
||||
match counter {
|
||||
SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis(
|
||||
SpendCounter::Committed(_)
|
||||
| SpendCounter::CustomerOrg(_)
|
||||
| SpendCounter::CustomerTeam(_)
|
||||
| SpendCounter::CustomerMember(_) => default_ttl.max(Duration::from_millis(
|
||||
self.committed_ttl_millis.load(Ordering::Relaxed),
|
||||
)),
|
||||
_ => default_ttl,
|
||||
|
|
@ -418,3 +421,23 @@ impl SpendCounters {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{SpendCounter, SpendCounters};
|
||||
use rstest::rstest;
|
||||
use std::time::Duration;
|
||||
|
||||
#[rstest]
|
||||
#[case::organization(SpendCounter::CustomerOrg("tenant:lifetime".into()))]
|
||||
#[case::team(SpendCounter::CustomerTeam("team:lifetime".into()))]
|
||||
#[case::member(SpendCounter::CustomerMember("member:lifetime".into()))]
|
||||
fn customer_counters_outlive_the_pending_database_flush(#[case] counter: SpendCounter) {
|
||||
let store = SpendCounters::new(None);
|
||||
store.set_committed_ttl(Duration::from_secs(120));
|
||||
assert_eq!(
|
||||
store.ttl(&counter, Duration::from_secs(1)),
|
||||
Duration::from_secs(120)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
|
|
@ -26,6 +27,7 @@ class CustomerScopeAuth(BaseModel):
|
|||
key: CustomerScopeKey
|
||||
entity_id: str
|
||||
spend: float = 0.0
|
||||
budget_reset_at: datetime | None = None
|
||||
max_budget: float | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
|
|
|
|||
|
|
@ -3148,6 +3148,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict):
|
|||
user_api_key_customer_member_id: ReadOnly[str | None]
|
||||
user_api_key_customer_team_id: ReadOnly[str | None]
|
||||
user_api_key_customer_org_id: ReadOnly[str | None]
|
||||
user_api_key_customer_hierarchy: NotRequired[ReadOnly[Mapping[str, object] | None]]
|
||||
user_api_key_request_route: str | None
|
||||
user_api_key_auth_metadata: dict[str, str] | None
|
||||
|
||||
|
|
|
|||
2
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -33236,6 +33236,8 @@ export interface components {
|
|||
CustomerScopeAuth: {
|
||||
/** Allowed Models */
|
||||
allowed_models?: string[] | null;
|
||||
/** Budget Reset At */
|
||||
budget_reset_at?: string | null;
|
||||
/** Entity Id */
|
||||
entity_id: string;
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -23,16 +23,126 @@ from litellm.proxy.auth.customer_hierarchy_auth import (
|
|||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.management_helpers.customer_hierarchy import ErrorCode, Failure
|
||||
from litellm.redis.stores.connection import connect
|
||||
from litellm.redis.stores.spend_counters import SpendCounterStore
|
||||
from litellm.types.proxy.auth.customer_hierarchy import (
|
||||
CustomerHierarchyAuth,
|
||||
CustomerRequestIdentity,
|
||||
CustomerScopeAuth,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def customer_spend_store(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_store", SpendCounterStore(connect(None).spend_counters()))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("alias_only", (False, True))
|
||||
async def test_a_missing_shared_counter_reseeds_spend_without_reloading_a_stale_cached_snapshot(
|
||||
alias_only: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
cache: Final = UserApiKeyCache()
|
||||
old: Final = _Prisma(
|
||||
_MemberTable({"member-1": _row(org_budget={"max_budget": 1})}),
|
||||
_OrganizationTable((_organization("tenant", max_budget=1),)),
|
||||
)
|
||||
persisted: Final = _Prisma(
|
||||
_MemberTable({"member-1": _row(org_budget={"max_budget": 1}, spend=(2, 0, 0))}),
|
||||
_OrganizationTable((_organization("tenant", max_budget=1, spend=2),)),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", persisted)
|
||||
if alias_only:
|
||||
_ = await get_customer_organization_hierarchy(
|
||||
"tenant",
|
||||
old, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
)
|
||||
else:
|
||||
_ = await get_customer_hierarchy(
|
||||
"member-1",
|
||||
old, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
)
|
||||
identity: Final = CustomerRequestIdentity(
|
||||
customer_org_alias="tenant" if alias_only else None,
|
||||
customer_member_id=None if alias_only else "member-1",
|
||||
)
|
||||
with pytest.raises(litellm.BudgetExceededError) as rejected:
|
||||
await resolve_customer_hierarchy(
|
||||
identity,
|
||||
persisted, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
None,
|
||||
"model",
|
||||
"/responses",
|
||||
)
|
||||
assert rejected.value.current_cost == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("level", ("org", "team", "member"))
|
||||
@pytest.mark.parametrize("alias_only", (False, True))
|
||||
async def test_settled_spend_reaches_another_workers_cached_hierarchy(
|
||||
monkeypatch: pytest.MonkeyPatch, level: str, alias_only: bool
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
shared: Final[dict[str, object]] = {}
|
||||
settling: Final = FakeSpendCounters(redis=shared).store()
|
||||
reading: Final = FakeSpendCounters(redis=shared).store()
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_store", reading)
|
||||
cache: Final = UserApiKeyCache()
|
||||
row: Final = _row(template={"max_budget": 1.0}, team_budget={"max_budget": 1.0}, org_budget={"max_budget": 1.0})
|
||||
prisma: Final = _Prisma(
|
||||
_MemberTable({"member-1": row}), _OrganizationTable((_organization("tenant", max_budget=1),))
|
||||
)
|
||||
identity: Final = CustomerRequestIdentity(
|
||||
customer_org_alias="tenant" if alias_only else None,
|
||||
customer_member_id=None if alias_only else "member-1",
|
||||
)
|
||||
await resolve_customer_hierarchy(
|
||||
identity,
|
||||
prisma, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
None,
|
||||
"model",
|
||||
"/responses",
|
||||
)
|
||||
counter: Final = f"spend:customer_{level}:{level}-1:lifetime"
|
||||
await settling.increment(counter, 1.25)
|
||||
if alias_only and level != "org":
|
||||
allowed: Final = await resolve_customer_hierarchy(
|
||||
identity,
|
||||
prisma, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
None,
|
||||
"model",
|
||||
"/responses",
|
||||
)
|
||||
assert allowed.customer_member_id is None
|
||||
return
|
||||
with pytest.raises(litellm.BudgetExceededError) as rejected:
|
||||
await resolve_customer_hierarchy(
|
||||
identity,
|
||||
prisma, # pyright: ignore[reportArgumentType] # database boundary fixture
|
||||
cache,
|
||||
None,
|
||||
"model",
|
||||
"/responses",
|
||||
)
|
||||
assert (rejected.value.entity_type, rejected.value.current_cost) == (f"customer_{level}", 1.25)
|
||||
|
||||
|
||||
def _row(
|
||||
*,
|
||||
template: Mapping[str, object] | None = None,
|
||||
|
|
@ -187,6 +297,10 @@ class _Record:
|
|||
def model_dump(self) -> Mapping[str, object]:
|
||||
return self.values
|
||||
|
||||
@property
|
||||
def spend(self):
|
||||
return self.values.get("spend", 0.0)
|
||||
|
||||
|
||||
class _OrganizationTable:
|
||||
def __init__(self, rows: Sequence[Mapping[str, object]] = ()) -> None:
|
||||
|
|
@ -200,6 +314,9 @@ class _OrganizationTable:
|
|||
matching: Final = [_Record(row) for row in self.rows if row.get("alias") == where["alias"]]
|
||||
return matching[:take] if take is not None else matching
|
||||
|
||||
async def find_unique(self, where: Mapping[str, object]):
|
||||
return next((_Record(row) for row in self.rows if row.get("customer_org_id") == where["customer_org_id"]), None)
|
||||
|
||||
|
||||
class _Db:
|
||||
def __init__(self, table: _MemberTable, organizations: _OrganizationTable) -> None:
|
||||
|
|
|
|||
157
tests/unit/proxy/auth/test_customer_spend.py
Normal file
157
tests/unit/proxy/auth/test_customer_spend.py
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.auth.customer_hierarchy_auth import check_customer_hierarchy_budget
|
||||
from litellm.proxy.auth.customer_spend import current_customer_spend as read_customer_spend
|
||||
from litellm.proxy.auth.customer_spend import customer_spend_counter, settle_customer_spend
|
||||
from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
|
||||
async def current_customer_spend(hierarchy, store, prisma=None):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
with pytest.MonkeyPatch.context() as context:
|
||||
context.setattr(proxy_server, "spend_counter_store", store)
|
||||
context.setattr(proxy_server, "prisma_client", prisma)
|
||||
return await read_customer_spend(hierarchy, proxy_server.get_current_spend)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scope", ("customer_org", "customer_team", "customer_member"))
|
||||
async def test_a_lost_counter_reseeds_committed_spend_without_refreshing_the_customer_snapshot(scope: str) -> None:
|
||||
from tests.unit.proxy.db.test_spend_counter_reseed import _FakeFindUniqueTable
|
||||
|
||||
row: Final = SimpleNamespace(spend=2.5, budget_reset_at=None)
|
||||
table: Final = _FakeFindUniqueTable(row)
|
||||
table_name: Final = {
|
||||
"customer_org": "litellm_customerorganizationtable",
|
||||
"customer_team": "litellm_customerteamtable",
|
||||
"customer_member": "litellm_customermembertable",
|
||||
}[scope]
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(**{table_name: table}))
|
||||
snapshot: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant", scopes=(CustomerScopeAuth(key=scope, entity_id="id:with:colons", max_budget=2),)
|
||||
)
|
||||
shared: Final[dict[str, object]] = {}
|
||||
first: Final = FakeSpendCounters(redis=shared).store()
|
||||
second: Final = FakeSpendCounters(redis=shared).store()
|
||||
for store in (first, second):
|
||||
loaded: Final = await current_customer_spend(snapshot, store, prisma)
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
check_customer_hierarchy_budget(loaded)
|
||||
assert snapshot.scopes[0].spend == 0
|
||||
assert len(table.where_clauses) == 1
|
||||
await first.delete((customer_spend_counter(snapshot.scopes[0]),))
|
||||
reseeded: Final = await current_customer_spend(snapshot, second, prisma)
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
check_customer_hierarchy_budget(reseeded)
|
||||
assert reseeded.scopes[0].spend == row.spend
|
||||
assert len(table.where_clauses) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_database_snapshot_ahead_of_settlement_does_not_count_the_pending_charge_twice() -> None:
|
||||
store: Final = FakeSpendCounters().store()
|
||||
initial: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant", scopes=(CustomerScopeAuth(key="customer_org", entity_id="tenant"),)
|
||||
)
|
||||
_ = await current_customer_spend(initial, store)
|
||||
persisted: Final = initial.model_copy(update={"scopes": (initial.scopes[0].model_copy(update={"spend": 0.5}),)})
|
||||
_ = await current_customer_spend(persisted, store)
|
||||
await settle_customer_spend({"user_api_key_customer_hierarchy": initial.model_dump(mode="json")}, 0.5, store)
|
||||
assert (await current_customer_spend(persisted, store)).scopes[0].spend == 0.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_member_b_settlement_blocks_cached_member_a_and_an_alias_only_request_on_another_worker() -> None:
|
||||
shared: Final[dict[str, object]] = {}
|
||||
settling: Final = FakeSpendCounters(redis=shared).store()
|
||||
reading: Final = FakeSpendCounters(redis=shared).store()
|
||||
organization: Final = CustomerScopeAuth(key="customer_org", entity_id="tenant", max_budget=1)
|
||||
member_a: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant",
|
||||
customer_member_id="a",
|
||||
scopes=(organization, CustomerScopeAuth(key="customer_member", entity_id="a")),
|
||||
)
|
||||
member_b: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant",
|
||||
customer_member_id="b",
|
||||
scopes=(organization, CustomerScopeAuth(key="customer_member", entity_id="b")),
|
||||
)
|
||||
alias_only: Final = CustomerHierarchyAuth(customer_org_id="tenant", scopes=(organization,))
|
||||
_ = await current_customer_spend(member_a, reading)
|
||||
_ = await current_customer_spend(alias_only, reading)
|
||||
await settle_customer_spend({"user_api_key_customer_hierarchy": member_b.model_dump(mode="json")}, 1.25, settling)
|
||||
for cached in (member_a, alias_only):
|
||||
current: Final = await current_customer_spend(cached, reading)
|
||||
with pytest.raises(litellm.BudgetExceededError) as rejected:
|
||||
check_customer_hierarchy_budget(current)
|
||||
assert rejected.value.entity_type == "customer_org"
|
||||
assert rejected.value.entity_id == "tenant"
|
||||
assert (await current_customer_spend(member_a, reading)).scopes[1].spend == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("redis", (True, False))
|
||||
async def test_settlement_counts_once_on_each_scope_without_modifying_the_cached_snapshot(redis: bool) -> None:
|
||||
handle: Final = FakeSpendCounters(redis={} if redis else None)
|
||||
store: Final = handle.store()
|
||||
hierarchy: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant",
|
||||
scopes=(
|
||||
CustomerScopeAuth(key="customer_org", entity_id="tenant", spend=1, max_budget=5),
|
||||
CustomerScopeAuth(key="customer_team", entity_id="team", spend=2),
|
||||
CustomerScopeAuth(key="customer_member", entity_id="member", spend=3),
|
||||
),
|
||||
)
|
||||
await settle_customer_spend({"user_api_key_customer_hierarchy": hierarchy.model_dump(mode="json")}, 0.5, store)
|
||||
current: Final = await current_customer_spend(hierarchy, store)
|
||||
assert tuple(scope.spend for scope in current.scopes) == (1.5, 2.5, 3.5)
|
||||
assert tuple(scope.spend for scope in hierarchy.scopes) == (1.0, 2.0, 3.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_budget_reset_uses_a_new_counter_even_when_an_old_request_settles_late() -> None:
|
||||
store: Final = FakeSpendCounters().store()
|
||||
old: Final = CustomerHierarchyAuth(
|
||||
customer_org_id="tenant",
|
||||
scopes=(CustomerScopeAuth(key="customer_org", entity_id="tenant", spend=2, max_budget=1),),
|
||||
)
|
||||
reset: Final = old.model_copy(
|
||||
update={
|
||||
"scopes": (
|
||||
old.scopes[0].model_copy(
|
||||
update={"spend": 0.0, "budget_reset_at": datetime(2026, 1, 1, tzinfo=timezone.utc)}
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
await settle_customer_spend({"user_api_key_customer_hierarchy": old.model_dump(mode="json")}, 0.5, store)
|
||||
current: Final = await current_customer_spend(reset, store)
|
||||
check_customer_hierarchy_budget(current)
|
||||
assert current.scopes[0].spend == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"metadata", ({}, {"user_api_key_customer_hierarchy": {}}, {"user_api_key_customer_hierarchy": "bad"})
|
||||
)
|
||||
async def test_missing_or_invalid_customer_metadata_does_not_charge_a_counter(metadata: dict[str, object]) -> None:
|
||||
handle: Final = FakeSpendCounters()
|
||||
await settle_customer_spend(metadata, 1, handle.store())
|
||||
assert handle.redis == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cost_does_not_initialize_customer_counters() -> None:
|
||||
handle: Final = FakeSpendCounters()
|
||||
hierarchy: Final = CustomerHierarchyAuth(customer_org_id="tenant", scopes=())
|
||||
await settle_customer_spend(
|
||||
{"user_api_key_customer_hierarchy": hierarchy.model_dump(mode="json")}, 0, handle.store()
|
||||
)
|
||||
assert handle.redis == {}
|
||||
|
|
@ -27,6 +27,48 @@ from litellm.redis.stores.spend_counters import SpendCounterStore
|
|||
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"family,table_name,id_field",
|
||||
(
|
||||
("customer_org", "litellm_customerorganizationtable", "customer_org_id"),
|
||||
("customer_team", "litellm_customerteamtable", "customer_team_id"),
|
||||
("customer_member", "litellm_customermembertable", "customer_member_id"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"period,reset_at,expected",
|
||||
(
|
||||
("lifetime", None, 4.5),
|
||||
(WINDOW_START.isoformat(), WINDOW_START, 4.5),
|
||||
(WINDOW_START.isoformat(), WINDOW_START.replace(tzinfo=None), 4.5),
|
||||
(WINDOW_START.isoformat(), WINDOW_START + timedelta(days=1), None),
|
||||
(WINDOW_START.isoformat(), WINDOW_START - timedelta(days=1), None),
|
||||
("lifetime", WINDOW_START, None),
|
||||
(WINDOW_START.isoformat(), None, None),
|
||||
("not-a-period", None, None),
|
||||
("2026-99-99T00:00:00+00:00", WINDOW_START, None),
|
||||
),
|
||||
)
|
||||
async def test_customer_counter_reseed_reads_only_the_matching_budget_period(
|
||||
family, table_name, id_field, period, reset_at, expected
|
||||
):
|
||||
table: Final = _FakeFindUniqueTable(SimpleNamespace(spend=4.5, budget_reset_at=reset_at))
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(**{table_name: table}))
|
||||
actual: Final = await SpendCounterReseed.from_db(prisma, f"spend:{family}:id:with:colons:{period}")
|
||||
assert actual == expected
|
||||
if period != "not-a-period":
|
||||
assert table.where_clauses == [{id_field: "id:with:colons"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", (None, ConnectionError("database is unavailable")))
|
||||
async def test_missing_or_unavailable_customer_spend_does_not_seed_a_counter(error):
|
||||
table: Final = _FakeFindUniqueTable(None, error)
|
||||
prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_customermembertable=table))
|
||||
assert await SpendCounterReseed.from_db(prisma, "spend:customer_member:member:lifetime") is None
|
||||
|
||||
|
||||
class _FakeFindUniqueTable:
|
||||
def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None:
|
||||
self._row = row
|
||||
|
|
|
|||
|
|
@ -47,6 +47,12 @@ _KEYS: Final = (
|
|||
(project_counter("proj-1"), "spend:project:proj-1"),
|
||||
(model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"),
|
||||
(shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"),
|
||||
("spend:customer_org:tenant:lifetime", "spend:customer_org:tenant:lifetime"),
|
||||
(
|
||||
"spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00",
|
||||
"spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00",
|
||||
),
|
||||
("spend:customer_member:user@example.com:lifetime", "spend:customer_member:user@example.com:lifetime"),
|
||||
(committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"),
|
||||
(committed_counter(team_counter("team-a")), "spend:committed:team:team-a"),
|
||||
(committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue