mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix: make scheduled key, user, and team spend resets idempotent (#50298)
* fix: make scheduled spend resets idempotent Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: cover reset roundoff with shared Redis workers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: isolate reset workers and type raw adapters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: keep updated_at moving on zero-spend budget resets * fix: never leave spend below zero after a budget reset * fix(proxy): keep the stored budget reset time when budget_duration is resent unchanged Key, user, and team updates recomputed budget_reset_at whenever the body carried budget_duration, and the Admin UI key form sends it on every save. Under the guarded reset, a resend during a sweep skipped that period's reset, and a resend after the boundary skipped it on main too. An unchanged duration now keeps the stored reset time * fix(proxy): keep a resent budget reset time only once it is due A save that resends the same budget_duration keeps the stored reset time only when that time has already passed, so the scheduled sweep still resets the period it is due for. A row that is not due yet realigns its reset time on save, the same as before. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
3daeb8daf2
commit
ac08eb5663
12 changed files with 704 additions and 106 deletions
|
|
@ -8,7 +8,14 @@ from datetime import datetime, timedelta, timezone
|
|||
from enum import Enum
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Generic, Literal, Protocol, TypeVar
|
||||
from typing import (
|
||||
Final,
|
||||
Generic,
|
||||
Literal,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # get_data returns validated rows selected by table_name
|
||||
)
|
||||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
|
|
@ -94,6 +101,7 @@ _RowT = TypeVar("_RowT")
|
|||
class _RowReset(Generic[_RowT]):
|
||||
row: _RowT
|
||||
spend_decrement: float
|
||||
expected_reset_at: datetime | None
|
||||
|
||||
|
||||
_LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}})
|
||||
|
|
@ -1050,6 +1058,8 @@ class ResetBudgetJob:
|
|||
token=k.row.token,
|
||||
budget_reset_at=k.row.budget_reset_at,
|
||||
spend_decrement=k.spend_decrement,
|
||||
expected_reset_at=k.expected_reset_at,
|
||||
spend_after_reset=float(k.row.spend or 0.0),
|
||||
)
|
||||
|
||||
async def _write_user_reset_updates(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None:
|
||||
|
|
@ -1075,6 +1085,8 @@ class ResetBudgetJob:
|
|||
user_id=u.row.user_id,
|
||||
budget_reset_at=u.row.budget_reset_at,
|
||||
spend_decrement=u.spend_decrement,
|
||||
expected_reset_at=u.expected_reset_at,
|
||||
spend_after_reset=float(u.row.spend or 0.0),
|
||||
)
|
||||
|
||||
async def _write_team_reset_updates(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None:
|
||||
|
|
@ -1100,6 +1112,8 @@ class ResetBudgetJob:
|
|||
team_id=t.row.team_id,
|
||||
budget_reset_at=t.row.budget_reset_at,
|
||||
spend_decrement=t.spend_decrement,
|
||||
expected_reset_at=t.expected_reset_at,
|
||||
spend_after_reset=float(t.row.spend or 0.0),
|
||||
)
|
||||
|
||||
def _emit_phase_failure(
|
||||
|
|
@ -1156,7 +1170,12 @@ class ResetBudgetJob:
|
|||
if keys_to_reset is not None and len(keys_to_reset) > 0:
|
||||
for key in keys_to_reset:
|
||||
try:
|
||||
pre_reset_spend = float(key.spend or 0.0)
|
||||
pre_reset_spend, expected_reset_at = (
|
||||
float(key.spend or 0.0),
|
||||
cast( # cast-ok: get_data's key branch returns validated key rows
|
||||
LiteLLM_VerificationToken, key
|
||||
).budget_reset_at,
|
||||
)
|
||||
updated_key = await ResetBudgetJob._reset_budget_for_key(
|
||||
key=key,
|
||||
current_time=now,
|
||||
|
|
@ -1167,6 +1186,7 @@ class ResetBudgetJob:
|
|||
_RowReset(
|
||||
row=updated_key,
|
||||
spend_decrement=pre_reset_spend - float(updated_key.spend or 0.0),
|
||||
expected_reset_at=expected_reset_at,
|
||||
)
|
||||
)
|
||||
else:
|
||||
|
|
@ -1263,7 +1283,12 @@ class ResetBudgetJob:
|
|||
if users_to_reset is not None and len(users_to_reset) > 0:
|
||||
for user in users_to_reset:
|
||||
try:
|
||||
pre_reset_spend = float(user.spend or 0.0)
|
||||
pre_reset_spend, expected_reset_at = (
|
||||
float(user.spend or 0.0),
|
||||
cast( # cast-ok: get_data's user branch returns validated user rows
|
||||
LiteLLM_UserTable, user
|
||||
).budget_reset_at,
|
||||
)
|
||||
updated_user = await ResetBudgetJob._reset_budget_for_user(
|
||||
user=user,
|
||||
current_time=now,
|
||||
|
|
@ -1274,6 +1299,7 @@ class ResetBudgetJob:
|
|||
_RowReset(
|
||||
row=updated_user,
|
||||
spend_decrement=pre_reset_spend - float(updated_user.spend or 0.0),
|
||||
expected_reset_at=expected_reset_at,
|
||||
)
|
||||
)
|
||||
else:
|
||||
|
|
@ -1378,7 +1404,12 @@ class ResetBudgetJob:
|
|||
if teams_to_reset is not None and len(teams_to_reset) > 0:
|
||||
for team in teams_to_reset:
|
||||
try:
|
||||
pre_reset_spend = float(team.spend or 0.0)
|
||||
pre_reset_spend, expected_reset_at = (
|
||||
float(team.spend or 0.0),
|
||||
cast( # cast-ok: get_data's team branch returns validated team rows
|
||||
LiteLLM_TeamTable, team
|
||||
).budget_reset_at,
|
||||
)
|
||||
updated_team = await ResetBudgetJob._reset_budget_for_team(
|
||||
team=team,
|
||||
current_time=now,
|
||||
|
|
@ -1389,6 +1420,7 @@ class ResetBudgetJob:
|
|||
_RowReset(
|
||||
row=updated_team,
|
||||
spend_decrement=pre_reset_spend - float(updated_team.spend or 0.0),
|
||||
expected_reset_at=expected_reset_at,
|
||||
)
|
||||
)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from datetime import datetime, time, timezone
|
||||
from typing import Final
|
||||
|
||||
from pydantic import ConfigDict
|
||||
|
||||
|
|
@ -81,6 +82,15 @@ def get_budget_reset_time(budget_duration: str) -> datetime:
|
|||
return compute_budget_reset_at(budget_duration, get_budget_reset_settings())
|
||||
|
||||
|
||||
def keeps_due_budget_reset(
|
||||
requested_duration: str, stored_duration: str | None, stored_reset_at: datetime | None
|
||||
) -> bool:
|
||||
if stored_reset_at is None or requested_duration != stored_duration:
|
||||
return False
|
||||
stored_utc: Final = stored_reset_at if stored_reset_at.tzinfo else stored_reset_at.replace(tzinfo=timezone.utc)
|
||||
return stored_utc <= datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _is_persistable_budget_duration(budget_duration: str) -> bool:
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.proxy.auth.password_policy import (
|
|||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
|
||||
from litellm.proxy.common_utils.openai_error_payload import openai_shaped_proxy_exception, proxy_exception_for
|
||||
from litellm.proxy.common_utils.timezone_utils import keeps_due_budget_reset
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
object_permission_cache_key,
|
||||
|
|
@ -1576,6 +1577,16 @@ async def _update_single_user_helper(
|
|||
|
||||
if existing_user_row is not None:
|
||||
existing_user_row = LiteLLM_UserTable.model_validate(existing_user_row.model_dump(exclude_none=True))
|
||||
if (
|
||||
"budget_reset_at" in non_default_values
|
||||
and isinstance(non_default_values["budget_duration"], str)
|
||||
and keeps_due_budget_reset(
|
||||
non_default_values["budget_duration"],
|
||||
existing_user_row.budget_duration,
|
||||
existing_user_row.budget_reset_at,
|
||||
)
|
||||
):
|
||||
del non_default_values["budget_reset_at"]
|
||||
|
||||
# Prevent budget self-escalation (GHSA-wvg4-6222-3q4r): non-admin callers
|
||||
# must not be able to raise their own budget/spend fields.
|
||||
|
|
|
|||
|
|
@ -84,7 +84,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
|
|||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for
|
||||
from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset
|
||||
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET, UserApiKeyCache, is_user_key_cache_key
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage
|
||||
|
|
@ -3116,11 +3116,11 @@ async def prepare_key_update_data(
|
|||
non_default_values["budget_duration"] = None
|
||||
non_default_values["budget_reset_at"] = None
|
||||
elif isinstance(budget_duration, str) and len(budget_duration) > 0:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
key_reset_at: Final = get_budget_reset_time(budget_duration=budget_duration)
|
||||
non_default_values["budget_reset_at"] = key_reset_at
|
||||
non_default_values["budget_duration"] = budget_duration
|
||||
if not keeps_due_budget_reset(
|
||||
budget_duration, existing_key_row.budget_duration, existing_key_row.budget_reset_at
|
||||
):
|
||||
non_default_values["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
|
||||
|
||||
if "budget_limits" in non_default_values:
|
||||
requested_windows: Final = _budget_limit_windows(data.budget_limits)
|
||||
|
|
|
|||
|
|
@ -135,6 +135,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
|||
from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars
|
||||
from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch
|
||||
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
team_membership_auth_cache_key,
|
||||
|
|
@ -2774,7 +2775,7 @@ async def update_team(
|
|||
)
|
||||
|
||||
# Check budget_duration and budget_reset_at
|
||||
_set_budget_reset_at(data, updated_kv)
|
||||
_set_budget_reset_at(data, updated_kv, existing_team)
|
||||
|
||||
_team_member_fields_in_request: Final = {
|
||||
field
|
||||
|
|
@ -3002,19 +3003,16 @@ async def patch_team(
|
|||
raise handle_exception_on_proxy(e)
|
||||
|
||||
|
||||
def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
|
||||
def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict, existing_team: LiteLLM_TeamTable) -> None:
|
||||
"""Set budget_reset_at in updated_kv if budget_duration is provided."""
|
||||
if data.budget_duration is not None:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
reset_at: Final = get_budget_reset_time(budget_duration=data.budget_duration)
|
||||
updated_kv["budget_reset_at"] = reset_at
|
||||
if data.budget_duration is not None and not keeps_due_budget_reset(
|
||||
data.budget_duration, existing_team.budget_duration, existing_team.budget_reset_at
|
||||
):
|
||||
updated_kv["budget_reset_at"] = get_budget_reset_time(budget_duration=data.budget_duration)
|
||||
elif "budget_duration" in updated_kv and updated_kv["budget_duration"] is None:
|
||||
updated_kv["budget_reset_at"] = None
|
||||
|
||||
if data.budget_limits is not None and len(data.budget_limits) > 0:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
initialized_windows: Final = []
|
||||
for window in data.budget_limits:
|
||||
w = window if isinstance(window, dict) else window.model_dump()
|
||||
|
|
|
|||
|
|
@ -18,7 +18,8 @@ carrying fields the update input type rejects (see #27730).
|
|||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from math import isfinite, ulp
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
|
||||
from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch
|
||||
|
|
@ -29,13 +30,79 @@ def _spend_reset_data(budget_reset_at: datetime | None, spend_decrement: float)
|
|||
return {"spend": spend, "budget_reset_at": budget_reset_at}
|
||||
|
||||
|
||||
_SPEND_RESET_SET: Final = (
|
||||
" SET spend = CASE WHEN ABS(spend - $3::text::double precision) <= $5::text::double precision"
|
||||
" THEN 0 ELSE GREATEST(spend - $3::text::double precision, 0) END, budget_reset_at = $2::timestamp,"
|
||||
" updated_at = (NOW() AT TIME ZONE 'UTC')"
|
||||
)
|
||||
_SPEND_RESET_WHERE: Final = " AND budget_reset_at IS NOT DISTINCT FROM $4::timestamp"
|
||||
_KEY_SPEND_RESET_SQL: Final = (
|
||||
'UPDATE "LiteLLM_VerificationToken"' + _SPEND_RESET_SET + " WHERE token = $1" + _SPEND_RESET_WHERE
|
||||
)
|
||||
_USER_SPEND_RESET_SQL: Final = (
|
||||
'UPDATE "LiteLLM_UserTable"' + _SPEND_RESET_SET + " WHERE user_id = $1" + _SPEND_RESET_WHERE
|
||||
)
|
||||
_TEAM_SPEND_RESET_SQL: Final = (
|
||||
'UPDATE "LiteLLM_TeamTable"' + _SPEND_RESET_SET + " WHERE team_id = $1" + _SPEND_RESET_WHERE
|
||||
)
|
||||
|
||||
|
||||
def _reset_timestamp(value: datetime | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
reset_at: Final = (value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value).astimezone(timezone.utc)
|
||||
return reset_at.replace(tzinfo=None).isoformat(timespec="microseconds")
|
||||
|
||||
|
||||
def _queue_guarded_spend_reset(
|
||||
batch: PrismaBatch,
|
||||
query: str,
|
||||
identity: str,
|
||||
budget_reset_at: datetime | None,
|
||||
spend_decrement: float,
|
||||
expected_reset_at: datetime | None,
|
||||
spend_after_reset: float,
|
||||
) -> bool:
|
||||
if budget_reset_at is None or spend_decrement == 0 or not isfinite(spend_decrement):
|
||||
return False
|
||||
# Prisma rounds Float reads, so exact subtraction can still leave a few ULPs.
|
||||
tolerance: Final = 4 * ulp(spend_decrement) if spend_after_reset == 0.0 else 0.0
|
||||
batch.execute_raw(
|
||||
query,
|
||||
identity,
|
||||
_reset_timestamp(budget_reset_at),
|
||||
repr(spend_decrement),
|
||||
_reset_timestamp(expected_reset_at),
|
||||
str(tolerance),
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KeySpendResetWrites:
|
||||
table: BatchTable
|
||||
batch: PrismaBatch
|
||||
|
||||
def queue_spend_reset(self, token: str, budget_reset_at: datetime | None, spend_decrement: float) -> None:
|
||||
self.table.update(
|
||||
where={"token": token},
|
||||
def queue_spend_reset(
|
||||
self,
|
||||
token: str,
|
||||
budget_reset_at: datetime | None,
|
||||
spend_decrement: float,
|
||||
expected_reset_at: datetime | None,
|
||||
spend_after_reset: float = 0.0,
|
||||
) -> None:
|
||||
if _queue_guarded_spend_reset(
|
||||
self.batch,
|
||||
_KEY_SPEND_RESET_SQL,
|
||||
token,
|
||||
budget_reset_at,
|
||||
spend_decrement,
|
||||
expected_reset_at,
|
||||
spend_after_reset,
|
||||
):
|
||||
return
|
||||
self.table.update_many(
|
||||
where={"token": token, "budget_reset_at": expected_reset_at},
|
||||
data=_spend_reset_data(budget_reset_at, spend_decrement),
|
||||
)
|
||||
|
||||
|
|
@ -43,10 +110,28 @@ class KeySpendResetWrites:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendResetWrites:
|
||||
table: BatchTable
|
||||
batch: PrismaBatch
|
||||
|
||||
def queue_spend_reset(self, user_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None:
|
||||
self.table.update(
|
||||
where={"user_id": user_id},
|
||||
def queue_spend_reset(
|
||||
self,
|
||||
user_id: str,
|
||||
budget_reset_at: datetime | None,
|
||||
spend_decrement: float,
|
||||
expected_reset_at: datetime | None,
|
||||
spend_after_reset: float = 0.0,
|
||||
) -> None:
|
||||
if _queue_guarded_spend_reset(
|
||||
self.batch,
|
||||
_USER_SPEND_RESET_SQL,
|
||||
user_id,
|
||||
budget_reset_at,
|
||||
spend_decrement,
|
||||
expected_reset_at,
|
||||
spend_after_reset,
|
||||
):
|
||||
return
|
||||
self.table.update_many(
|
||||
where={"user_id": user_id, "budget_reset_at": expected_reset_at},
|
||||
data=_spend_reset_data(budget_reset_at, spend_decrement),
|
||||
)
|
||||
|
||||
|
|
@ -54,10 +139,28 @@ class UserSpendResetWrites:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class TeamSpendResetWrites:
|
||||
table: BatchTable
|
||||
batch: PrismaBatch
|
||||
|
||||
def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None:
|
||||
self.table.update(
|
||||
where={"team_id": team_id},
|
||||
def queue_spend_reset(
|
||||
self,
|
||||
team_id: str,
|
||||
budget_reset_at: datetime | None,
|
||||
spend_decrement: float,
|
||||
expected_reset_at: datetime | None,
|
||||
spend_after_reset: float = 0.0,
|
||||
) -> None:
|
||||
if _queue_guarded_spend_reset(
|
||||
self.batch,
|
||||
_TEAM_SPEND_RESET_SQL,
|
||||
team_id,
|
||||
budget_reset_at,
|
||||
spend_decrement,
|
||||
expected_reset_at,
|
||||
spend_after_reset,
|
||||
):
|
||||
return
|
||||
self.table.update_many(
|
||||
where={"team_id": team_id, "budget_reset_at": expected_reset_at},
|
||||
data=_spend_reset_data(budget_reset_at, spend_decrement),
|
||||
)
|
||||
|
||||
|
|
@ -191,9 +294,9 @@ class BudgetCascadeUnitOfWork:
|
|||
async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> AsyncGenerator[SpendResetUnitOfWork, None]:
|
||||
batch = new_batch()
|
||||
yield SpendResetUnitOfWork(
|
||||
keys=KeySpendResetWrites(table=batch.litellm_verificationtoken),
|
||||
users=UserSpendResetWrites(table=batch.litellm_usertable),
|
||||
teams=TeamSpendResetWrites(table=batch.litellm_teamtable),
|
||||
keys=KeySpendResetWrites(table=batch.litellm_verificationtoken, batch=batch),
|
||||
users=UserSpendResetWrites(table=batch.litellm_usertable, batch=batch),
|
||||
teams=TeamSpendResetWrites(table=batch.litellm_teamtable, batch=batch),
|
||||
)
|
||||
await batch.commit()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,16 +1,25 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.process import OwnedState, owned_proxy
|
||||
from prisma import Prisma
|
||||
from psycopg import sql
|
||||
from redis import Redis
|
||||
|
||||
from litellm.repositories.unit_of_work import spend_reset_unit_of_work
|
||||
|
||||
pytestmark: Final = pytest.mark.shares_worker_state
|
||||
|
||||
RESET_LEASE_KEY: Final = "cronjob_lock:reset_budget_job"
|
||||
|
|
@ -20,6 +29,21 @@ FAST_RESET_TICK: Final = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _peer_reset_lease() -> Generator["Redis[bytes]", None, None]:
|
||||
cache: Final["Redis[bytes]"] = Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
with cache:
|
||||
eventually(
|
||||
lambda: cache.set(RESET_LEASE_KEY, PEER_POD_LEASE, ex=300, nx=True),
|
||||
lambda claimed: claimed is True,
|
||||
seconds=60,
|
||||
)
|
||||
try:
|
||||
yield cache
|
||||
finally:
|
||||
cache.delete(RESET_LEASE_KEY)
|
||||
|
||||
|
||||
def _team_spend(team: str) -> float:
|
||||
rows: Final = read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team,))
|
||||
assert len(rows) == 1, rows
|
||||
|
|
@ -38,25 +62,26 @@ def _make_team_budget_due(team: str, spend: float) -> None:
|
|||
|
||||
|
||||
@pytest.mark.covers("spend.budget_reset.one_pod_sweeps_per_tick")
|
||||
@pytest.mark.parametrize("starting_spend", (0.5, 40.080000000000005, 1009.9442791267694))
|
||||
def test_reset_sweep_skips_ticks_while_another_pod_holds_the_lease_and_resumes_after_release(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
gateway: Gateway, tmp_path: Path, starting_spend: float
|
||||
) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache,
|
||||
):
|
||||
model: Final = scenario.model(input_cost_per_token=0.0, output_cost_per_token=0.0)
|
||||
team: Final = scenario.team(max_budget=1.0, budget_duration="30d")
|
||||
team: Final = scenario.team(max_budget=2000000.0, budget_duration="30d")
|
||||
key: Final = scenario.key(team_id=team)
|
||||
_make_team_budget_due(team, spend=0.5)
|
||||
assert _team_spend(team) == 0.5
|
||||
_make_team_budget_due(team, spend=starting_spend)
|
||||
assert _team_spend(team) == starting_spend
|
||||
eventually(
|
||||
lambda: cache.set(RESET_LEASE_KEY, PEER_POD_LEASE, ex=120, nx=True),
|
||||
lambda claimed: claimed is True,
|
||||
seconds=60,
|
||||
)
|
||||
try:
|
||||
with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as replica:
|
||||
with owned_proxy(gateway, tmp_path, FAST_RESET_TICK, workers=2) as replica:
|
||||
response: Final = replica.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
|
|
@ -65,13 +90,272 @@ def test_reset_sweep_skips_ticks_while_another_pod_holds_the_lease_and_resumes_a
|
|||
)
|
||||
assert response.status_code == 200, response.text
|
||||
held: Final = eventually(
|
||||
lambda: _team_spend(team), lambda spend: spend != 0.5, seconds=8, return_last_on_timeout=True
|
||||
lambda: _team_spend(team),
|
||||
lambda spend: spend != starting_spend,
|
||||
seconds=8,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert held == starting_spend, (
|
||||
f"team {team} was swept while another pod held the reset lease: spend={held}"
|
||||
)
|
||||
assert held == 0.5, f"team {team} was swept while another pod held the reset lease: spend={held}"
|
||||
assert cache.get(RESET_LEASE_KEY) == PEER_POD_LEASE.encode()
|
||||
cache.delete(RESET_LEASE_KEY)
|
||||
swept: Final = eventually(lambda: _team_spend(team), lambda spend: spend == 0.0, seconds=15)
|
||||
swept: Final = eventually(
|
||||
lambda: _team_spend(team), lambda spend: spend == 0.0, seconds=15, return_last_on_timeout=True
|
||||
)
|
||||
info: Final = replica.request("GET", f"/team/info?team_id={team}")
|
||||
assert info.status_code == 200, info.text
|
||||
assert info.json()["team_info"]["spend"] == 0.0
|
||||
assert swept == 0.0
|
||||
eventually(lambda: cache.get(RESET_LEASE_KEY), lambda value: value is None, seconds=15)
|
||||
finally:
|
||||
cache.delete(RESET_LEASE_KEY)
|
||||
|
||||
|
||||
def _blocked_reset_updates(table: str) -> int:
|
||||
rows: Final = read_rows(
|
||||
"SELECT count(*) AS blocked FROM pg_stat_activity WHERE datname = current_database() "
|
||||
"AND wait_event_type = 'Lock' AND query LIKE %s AND query LIKE '%%UPDATE%%'",
|
||||
(f'%"{table}"%',),
|
||||
)
|
||||
return int(str(rows[0]["blocked"]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ("key", "user", "team"))
|
||||
@pytest.mark.parametrize(
|
||||
("workers", "spend", "increment", "rollover", "reset_timezone"),
|
||||
(
|
||||
(3, 40.0, 0.25, False, "UTC"),
|
||||
(1, 40.080000000000005, 0.0, False, "UTC"),
|
||||
(1, 1234567.8901234567, 0.0, False, "UTC"),
|
||||
(1, 1009.9442791267694, 0.0, False, "UTC"),
|
||||
(3, 40.000000000000014, 0.0, True, "UTC"),
|
||||
(1, 0.30000000000000004, 0.0, False, "America/Los_Angeles"),
|
||||
),
|
||||
ids=(
|
||||
"overlap-preserves-new-spend",
|
||||
"fractional-roundoff",
|
||||
"large-roundoff",
|
||||
"four-ulp-roundoff",
|
||||
"rollover-keeps-small-carry",
|
||||
"non-utc-roundoff",
|
||||
),
|
||||
)
|
||||
def test_scheduled_reset_preserves_new_spend_and_never_leaves_roundoff(
|
||||
gateway: Gateway,
|
||||
tmp_path: Path,
|
||||
entity: Literal["key", "user", "team"],
|
||||
workers: int,
|
||||
spend: float,
|
||||
increment: float,
|
||||
rollover: bool,
|
||||
reset_timezone: str,
|
||||
) -> None:
|
||||
max_budget: Final = 40.0 if rollover else 2000000.0
|
||||
expected_spend: Final = spend - max_budget + increment if rollover else increment
|
||||
config: Final = tmp_path / "without-redis.yaml"
|
||||
config.write_text(
|
||||
"model_list: []\ngeneral_settings:\n"
|
||||
" master_key: os.environ/LITELLM_MASTER_KEY\n"
|
||||
" database_url: os.environ/DATABASE_URL\n"
|
||||
" store_model_in_db: true\n"
|
||||
"litellm_settings:\n cache: false\n"
|
||||
f" budget_rollover: {str(rollover).lower()}\n"
|
||||
f" timezone: {reset_timezone}\n"
|
||||
)
|
||||
with gateway.scenario() as scenario, _peer_reset_lease() as peer:
|
||||
identity: Final = (
|
||||
sha256(scenario.key(max_budget=max_budget, budget_duration="1d").encode()).hexdigest()
|
||||
if entity == "key"
|
||||
else scenario.user(max_budget=max_budget, budget_duration="1d")
|
||||
if entity == "user"
|
||||
else scenario.team(max_budget=max_budget, budget_duration="1d")
|
||||
)
|
||||
table, column = {
|
||||
"key": ("LiteLLM_VerificationToken", "token"),
|
||||
"user": ("LiteLLM_UserTable", "user_id"),
|
||||
"team": ("LiteLLM_TeamTable", "team_id"),
|
||||
}[entity]
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as gate:
|
||||
gate.execute(
|
||||
sql.SQL(
|
||||
"UPDATE {} SET spend = %s, budget_reset_at = now() - interval '1 day',"
|
||||
" updated_at = now() - interval '1 day' WHERE {} = %s"
|
||||
).format(sql.Identifier(table), sql.Identifier(column)),
|
||||
(spend, identity),
|
||||
)
|
||||
gate.commit()
|
||||
gate.execute(
|
||||
sql.SQL("SELECT 1 FROM {} WHERE {} = %s FOR UPDATE").format(
|
||||
sql.Identifier(table), sql.Identifier(column)
|
||||
),
|
||||
(identity,),
|
||||
)
|
||||
with owned_proxy(
|
||||
gateway,
|
||||
tmp_path,
|
||||
FAST_RESET_TICK,
|
||||
config=config,
|
||||
remove_environment=("REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_URL"),
|
||||
workers=workers,
|
||||
state=OwnedState.SHARED_WITH_WORKER,
|
||||
) as replica:
|
||||
try:
|
||||
blocked: Final = eventually(
|
||||
lambda: _blocked_reset_updates(table), lambda count: count == workers, seconds=30
|
||||
)
|
||||
assert blocked == workers
|
||||
gate.execute(
|
||||
sql.SQL("UPDATE {} SET spend = spend + %s WHERE {} = %s").format(
|
||||
sql.Identifier(table), sql.Identifier(column)
|
||||
),
|
||||
(increment, identity),
|
||||
)
|
||||
gate.commit()
|
||||
eventually(lambda: _blocked_reset_updates(table), lambda count: count == 0, seconds=30)
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as observer:
|
||||
result: Final = observer.execute(
|
||||
sql.SQL(
|
||||
"SELECT spend, budget_reset_at > now(), updated_at > now() - interval '1 hour'"
|
||||
" FROM {} WHERE {} = %s"
|
||||
).format(sql.Identifier(table), sql.Identifier(column)),
|
||||
(identity,),
|
||||
).fetchall()
|
||||
parameter, info = {
|
||||
"key": ("key", "info"),
|
||||
"user": ("user_id", "user_info"),
|
||||
"team": ("team_id", "team_info"),
|
||||
}[entity]
|
||||
response: Final = replica.request("GET", f"/{entity}/info?{parameter}={identity}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()[info]["spend"] == pytest.approx(expected_spend, rel=1e-14, abs=0.0)
|
||||
assert result == [(expected_spend, True, True)], result
|
||||
assert peer.get(RESET_LEASE_KEY) == PEER_POD_LEASE.encode()
|
||||
finally:
|
||||
gate.rollback()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ("key", "user", "team"))
|
||||
def test_stale_zero_reset_cannot_erase_another_workers_rollover(gateway: Gateway, entity: str) -> None:
|
||||
carry: Final = 1.4210854715202004e-14
|
||||
previous_reset: Final = datetime(2020, 1, 1, tzinfo=timezone.utc)
|
||||
next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1)
|
||||
with gateway.scenario() as scenario:
|
||||
identity: Final = (
|
||||
sha256(scenario.key(max_budget=40.0, budget_duration="1d").encode()).hexdigest()
|
||||
if entity == "key"
|
||||
else scenario.user(max_budget=40.0, budget_duration="1d")
|
||||
if entity == "user"
|
||||
else scenario.team(max_budget=40.0, budget_duration="1d")
|
||||
)
|
||||
table, column = {
|
||||
"key": ("LiteLLM_VerificationToken", "token"),
|
||||
"user": ("LiteLLM_UserTable", "user_id"),
|
||||
"team": ("LiteLLM_TeamTable", "team_id"),
|
||||
}[entity]
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
||||
connection.execute(
|
||||
sql.SQL("UPDATE {} SET spend = %s, budget_reset_at = %s WHERE {} = %s").format(
|
||||
sql.Identifier(table), sql.Identifier(column)
|
||||
),
|
||||
(40.0 + carry, previous_reset, identity),
|
||||
)
|
||||
|
||||
async def commit_winner_then_stale_reset() -> None:
|
||||
async with Prisma() as client:
|
||||
for spend_after_reset in (carry, 0.0):
|
||||
async with spend_reset_unit_of_work(client.batch_) as writes:
|
||||
writer: Final = {"key": writes.keys, "user": writes.users, "team": writes.teams}[entity]
|
||||
writer.queue_spend_reset(identity, next_reset, 40.0, previous_reset, spend_after_reset)
|
||||
assert read_rows(
|
||||
sql.SQL("SELECT spend FROM {} WHERE {} = %s")
|
||||
.format(sql.Identifier(table), sql.Identifier(column))
|
||||
.as_string(),
|
||||
(identity,),
|
||||
) == [{"spend": carry}]
|
||||
|
||||
asyncio.run(commit_winner_then_stale_reset())
|
||||
parameter, info = {
|
||||
"key": ("key", "info"),
|
||||
"user": ("user_id", "user_info"),
|
||||
"team": ("team_id", "team_info"),
|
||||
}[entity]
|
||||
response: Final = gateway.request("GET", f"/{entity}/info?{parameter}={identity}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()[info]["spend"] == pytest.approx(carry, rel=1e-14, abs=0.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ("user", "team"))
|
||||
def test_reset_starts_the_window_of_a_row_that_never_had_a_reset_time(gateway: Gateway, entity: str) -> None:
|
||||
next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1)
|
||||
with gateway.scenario() as scenario:
|
||||
identity: Final = (
|
||||
scenario.user(max_budget=100.0, budget_duration="1d")
|
||||
if entity == "user"
|
||||
else scenario.team(max_budget=100.0, budget_duration="1d")
|
||||
)
|
||||
table, column, info = {
|
||||
"user": ("LiteLLM_UserTable", "user_id", "user_info"),
|
||||
"team": ("LiteLLM_TeamTable", "team_id", "team_info"),
|
||||
}[entity]
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
||||
connection.execute(
|
||||
sql.SQL("UPDATE {} SET spend = 40, budget_reset_at = NULL WHERE {} = %s").format(
|
||||
sql.Identifier(table), sql.Identifier(column)
|
||||
),
|
||||
(identity,),
|
||||
)
|
||||
|
||||
async def reset() -> None:
|
||||
async with Prisma() as client:
|
||||
async with spend_reset_unit_of_work(client.batch_) as writes:
|
||||
(writes.users if entity == "user" else writes.teams).queue_spend_reset(
|
||||
identity, next_reset, 40.0, None
|
||||
)
|
||||
|
||||
asyncio.run(reset())
|
||||
response: Final = gateway.request("GET", f"/{entity}/info?{column}={identity}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()[info]["spend"] == 0.0
|
||||
assert response.json()[info]["budget_reset_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entity", ("key", "user", "team"))
|
||||
@pytest.mark.parametrize(("spend_decrement", "spend_after_reset"), ((40.0, 0.0), (30.0, 10.0)))
|
||||
def test_reset_after_a_manual_spend_reset_never_goes_negative(
|
||||
gateway: Gateway, entity: str, spend_decrement: float, spend_after_reset: float
|
||||
) -> None:
|
||||
previous_reset: Final = datetime(2020, 1, 1, tzinfo=timezone.utc)
|
||||
next_reset: Final = datetime.now(timezone.utc).replace(microsecond=0) + timedelta(days=1)
|
||||
with gateway.scenario() as scenario:
|
||||
identity: Final = (
|
||||
sha256(scenario.key(max_budget=100.0, budget_duration="1d").encode()).hexdigest()
|
||||
if entity == "key"
|
||||
else scenario.user(max_budget=100.0, budget_duration="1d")
|
||||
if entity == "user"
|
||||
else scenario.team(max_budget=100.0, budget_duration="1d")
|
||||
)
|
||||
table, column, parameter, info = {
|
||||
"key": ("LiteLLM_VerificationToken", "token", "key", "info"),
|
||||
"user": ("LiteLLM_UserTable", "user_id", "user_id", "user_info"),
|
||||
"team": ("LiteLLM_TeamTable", "team_id", "team_id", "team_info"),
|
||||
}[entity]
|
||||
with psycopg.connect(os.environ["DATABASE_URL"]) as connection:
|
||||
connection.execute(
|
||||
sql.SQL("UPDATE {} SET spend = 0, budget_reset_at = %s WHERE {} = %s").format(
|
||||
sql.Identifier(table), sql.Identifier(column)
|
||||
),
|
||||
(previous_reset, identity),
|
||||
)
|
||||
|
||||
async def reset_with_the_spend_read_before_the_manual_reset() -> None:
|
||||
async with Prisma() as client:
|
||||
async with spend_reset_unit_of_work(client.batch_) as writes:
|
||||
{"key": writes.keys, "user": writes.users, "team": writes.teams}[entity].queue_spend_reset(
|
||||
identity, next_reset, spend_decrement, previous_reset, spend_after_reset
|
||||
)
|
||||
|
||||
asyncio.run(reset_with_the_spend_read_before_the_manual_reset())
|
||||
response: Final = gateway.request("GET", f"/{entity}/info?{parameter}={identity}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()[info]["spend"] == 0.0
|
||||
|
|
|
|||
|
|
@ -77,6 +77,33 @@ class MockTable:
|
|||
return {"count": 1}
|
||||
|
||||
|
||||
def _raw_reset_write(query: str, args: Sequence[object]) -> dict[str, object] | None:
|
||||
if "CASE WHEN ABS(spend -" not in query:
|
||||
return None
|
||||
table, column = {
|
||||
"LiteLLM_VerificationToken": ("key", "token"),
|
||||
"LiteLLM_UserTable": ("user", "user_id"),
|
||||
"LiteLLM_TeamTable": ("team", "team_id"),
|
||||
}[query.split('"')[1]]
|
||||
identity, reset_at, decrement, expected_reset_at, _ = args
|
||||
assert isinstance(identity, str) and isinstance(reset_at, str) and isinstance(decrement, str)
|
||||
assert expected_reset_at is None or isinstance(expected_reset_at, str)
|
||||
return {
|
||||
"table": table,
|
||||
"op": "update_many",
|
||||
"where": {
|
||||
column: identity,
|
||||
"budget_reset_at": datetime.fromisoformat(expected_reset_at).replace(tzinfo=timezone.utc)
|
||||
if expected_reset_at is not None
|
||||
else None,
|
||||
},
|
||||
"data": {
|
||||
"spend": {"decrement": float(decrement)},
|
||||
"budget_reset_at": datetime.fromisoformat(reset_at).replace(tzinfo=timezone.utc),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class MockBatcher:
|
||||
"""Captures the writes queued on one `db.batch_()` and whether it committed.
|
||||
|
||||
|
|
@ -115,7 +142,9 @@ class MockBatcher:
|
|||
self.litellm_endusertable = _Table("enduser", self)
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None:
|
||||
self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args})
|
||||
self.calls.append(
|
||||
_raw_reset_write(query, args) or {"table": "raw", "op": "execute_raw", "query": query, "args": args}
|
||||
)
|
||||
|
||||
async def commit(self):
|
||||
self.committed = True
|
||||
|
|
@ -283,15 +312,15 @@ def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(rese
|
|||
|
||||
asyncio.run(
|
||||
reset_budget_job._write_key_reset_updates(
|
||||
updated_keys=[_RowReset(row=k, spend_decrement=(k.spend or 0.0)) for k in keys]
|
||||
updated_keys=[_RowReset(row=k, spend_decrement=(k.spend or 0.0), expected_reset_at=reset_at) for k in keys]
|
||||
)
|
||||
)
|
||||
|
||||
assert _batch_writes(mock_prisma_client, "key") == [
|
||||
{
|
||||
"table": "key",
|
||||
"op": "update",
|
||||
"where": {"token": "tok-ok"},
|
||||
"op": "update_many",
|
||||
"where": {"token": "tok-ok", "budget_reset_at": reset_at},
|
||||
"data": {"spend": {"decrement": 0.0}, "budget_reset_at": reset_at},
|
||||
}
|
||||
]
|
||||
|
|
@ -323,7 +352,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
|
|||
key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"]
|
||||
assert len(key_writes) == 1
|
||||
write = key_writes[0]
|
||||
assert write["where"] == {"token": "tok-key-1"}
|
||||
assert write["where"] == {"token": "tok-key-1", "budget_reset_at": now}
|
||||
assert write["data"]["spend"] == {"decrement": 100.0}
|
||||
assert write["data"]["budget_reset_at"] > now
|
||||
assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
|
@ -403,7 +432,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client):
|
|||
user_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "user"]
|
||||
assert len(user_writes) == 1
|
||||
write = user_writes[0]
|
||||
assert write["where"] == {"user_id": "uid-1"}
|
||||
assert write["where"] == {"user_id": "uid-1", "budget_reset_at": now}
|
||||
assert write["data"]["spend"] == {"decrement": 200.0}
|
||||
assert write["data"]["budget_reset_at"] > now
|
||||
assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
|
@ -432,7 +461,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
|
|||
team_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "team"]
|
||||
assert len(team_writes) == 1
|
||||
write = team_writes[0]
|
||||
assert write["where"] == {"team_id": "tid-1"}
|
||||
assert write["where"] == {"team_id": "tid-1", "budget_reset_at": now}
|
||||
assert write["data"]["spend"] == {"decrement": 500.0}
|
||||
assert write["data"]["budget_reset_at"] > now
|
||||
assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
|
@ -663,9 +692,13 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
|
|||
("user", {"user_id": "uid-all-1"}, 200.0),
|
||||
("team", {"team_id": "tid-all-1"}, 500.0),
|
||||
]:
|
||||
writes = _batch_writes(mock_prisma_client, table_name, op="update")
|
||||
writes = [
|
||||
write
|
||||
for write in _batch_writes(mock_prisma_client, table_name, op="update_many")
|
||||
if "budget_reset_at" in write["where"]
|
||||
]
|
||||
assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
|
||||
assert writes[0]["where"] == where
|
||||
assert writes[0]["where"] == {**where, "budget_reset_at": now}
|
||||
assert writes[0]["data"]["spend"] == {"decrement": decrement}
|
||||
assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
||||
|
|
@ -1839,7 +1872,6 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma
|
|||
assert evicted == {f"end_user_id:cust-{i:06d}" for i in range(population)}
|
||||
|
||||
|
||||
|
||||
def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_finish(
|
||||
mock_prisma_client, monkeypatch
|
||||
):
|
||||
|
|
@ -2379,8 +2411,8 @@ def test_key_reset_walks_the_due_rows_one_chunk_at_a_time(monkeypatch):
|
|||
assert _fetch_limits(client, "key") == [2, 2]
|
||||
assert len(client.db.batchers) == 2
|
||||
assert all(batcher.committed for batcher in client.db.batchers)
|
||||
assert [len(batcher.calls) for batcher in client.db.batchers] == [2, 1]
|
||||
assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2", "k3"]
|
||||
assert [sum(call["op"] == "update_many" for call in batcher.calls) for batcher in client.db.batchers] == [2, 1]
|
||||
assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update_many")] == ["k1", "k2", "k3"]
|
||||
|
||||
|
||||
def test_key_reset_stops_after_a_chunk_shorter_than_the_batch_size(monkeypatch):
|
||||
|
|
@ -2453,7 +2485,7 @@ def test_user_and_team_resets_are_chunked_too(monkeypatch, phase, table_name, ro
|
|||
assert client.fetches_by_table[table_name] == 2
|
||||
assert _fetch_limits(client, table_name) == [2, 2]
|
||||
assert len(client.db.batchers) == 2
|
||||
assert len(_batch_writes(client, table_name, op="update")) == 3
|
||||
assert len(_batch_writes(client, table_name, op="update_many")) == 3
|
||||
|
||||
|
||||
def test_budget_table_reset_walks_chunks_until_it_runs_dry(monkeypatch):
|
||||
|
|
@ -2510,7 +2542,7 @@ def test_key_reset_stops_when_the_new_reset_time_is_not_in_the_future(monkeypatc
|
|||
asyncio.run(job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert client.fetches_by_table["key"] == 1
|
||||
assert len(_batch_writes(client, "key", op="update")) == 2
|
||||
assert len(_batch_writes(client, "key", op="update_many")) == 2
|
||||
|
||||
|
||||
def test_budget_table_reset_stops_when_the_new_window_is_not_in_the_future(monkeypatch):
|
||||
|
|
@ -2577,7 +2609,7 @@ def test_key_reset_keeps_paging_when_some_rows_in_a_chunk_fail(monkeypatch):
|
|||
_run_and_drain_hooks(job.reset_budget_for_litellm_keys)
|
||||
|
||||
assert client.fetches_by_table["key"] == 2
|
||||
assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2"]
|
||||
assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update_many")] == ["k1", "k2"]
|
||||
assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == ["reset_budget_keys"]
|
||||
assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == {"num_keys_found"}
|
||||
assert [call["call_type"] for call in logging_obj.service_logging_obj.success_calls] == ["reset_budget_keys"]
|
||||
|
|
@ -2602,7 +2634,7 @@ def test_user_and_team_chunks_report_progress_despite_a_failed_row(
|
|||
_run_and_drain_hooks(getattr(job, phase))
|
||||
|
||||
assert client.fetches_by_table[table_name] == 2
|
||||
assert len(_batch_writes(client, table_name, op="update")) == 2
|
||||
assert len(_batch_writes(client, table_name, op="update_many")) == 2
|
||||
assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == [call_type]
|
||||
|
||||
|
||||
|
|
@ -3025,7 +3057,7 @@ def test_transient_transport_error_on_read_reconnects_and_still_resets(phase, ta
|
|||
asyncio.run(getattr(job, phase)())
|
||||
|
||||
assert client.reconnect_reasons == [reason]
|
||||
assert len(_batch_writes(client, table_name, op="update")) == 1
|
||||
assert len(_batch_writes(client, table_name, op="update_many")) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -3049,7 +3081,7 @@ def test_connect_error_on_write_reconnects_and_commits(phase, table_name, reason
|
|||
|
||||
assert client.reconnect_reasons == [reason]
|
||||
assert client.commit_attempts == 2
|
||||
assert len(_batch_writes(client, table_name, op="update")) == 1
|
||||
assert len(_batch_writes(client, table_name, op="update_many")) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ambiguous_error_name", ["ReadError", "ReadTimeout"])
|
||||
|
|
@ -3081,7 +3113,7 @@ def test_ambiguous_transport_error_on_read_still_retries(ambiguous_error_name):
|
|||
asyncio.run(job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert client.reconnect_reasons == ["reset_budget_read_keys_failure"]
|
||||
assert len(_batch_writes(client, "key", op="update")) == 1
|
||||
assert len(_batch_writes(client, "key", op="update_many")) == 1
|
||||
|
||||
|
||||
def test_transport_error_on_budget_cascade_read_reconnects_and_commits():
|
||||
|
|
@ -3318,7 +3350,15 @@ def rollover_enabled(monkeypatch):
|
|||
],
|
||||
)
|
||||
def test_direct_reset_carries_overage_when_rollover_enabled(
|
||||
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch, run_phase, table, id_field, id_value, row_factory
|
||||
rollover_enabled,
|
||||
reset_budget_job,
|
||||
mock_prisma_client,
|
||||
monkeypatch,
|
||||
run_phase,
|
||||
table,
|
||||
id_field,
|
||||
id_value,
|
||||
row_factory,
|
||||
):
|
||||
"""spend=150 against max_budget=100 must decrement by the cap (leaving 50)
|
||||
rather than zero the row, and the spend counter must be seeded with 50."""
|
||||
|
|
@ -3330,7 +3370,7 @@ def test_direct_reset_carries_overage_when_rollover_enabled(
|
|||
|
||||
writes = _batch_writes(mock_prisma_client, table)
|
||||
assert len(writes) == 1
|
||||
assert writes[0]["where"] == {id_field: id_value}
|
||||
assert writes[0]["where"] == {id_field: id_value, "budget_reset_at": now}
|
||||
assert writes[0]["data"]["spend"] == {"decrement": 100.0}
|
||||
assert writes[0]["data"]["budget_reset_at"] > now
|
||||
counter_prefix = {"key": "spend:key", "user": "spend:user", "team": "spend:team"}[table]
|
||||
|
|
@ -3652,7 +3692,7 @@ def test_reset_decrement_preserves_spend_landed_after_read(
|
|||
|
||||
writes = _batch_writes(mock_prisma_client, table)
|
||||
assert len(writes) == 1
|
||||
assert writes[0]["where"] == {id_field: id_value}
|
||||
assert writes[0]["where"] == {id_field: id_value, "budget_reset_at": now}
|
||||
assert writes[0]["data"]["spend"] == {"decrement": 5.0}
|
||||
assert writes[0]["data"]["budget_reset_at"] > now
|
||||
assert _apply_spend_payload(db_spend=5.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(0.4)
|
||||
|
|
@ -3756,6 +3796,7 @@ async def test_a_budget_window_write_renders_a_postgres_update_span_for_its_tabl
|
|||
assert update.await_count == 1
|
||||
assert await postgres_span_names() == (span_name,)
|
||||
|
||||
|
||||
def _attrify(d: dict):
|
||||
"""
|
||||
Wrap a dict so that attribute access (`.token`, `.user_id`, `.team_id`,
|
||||
|
|
@ -3775,7 +3816,8 @@ def _attrify(d: dict):
|
|||
def __setattr__(self, k, v):
|
||||
self[k] = v
|
||||
|
||||
return _AttrDict(d)
|
||||
return _AttrDict({"budget_reset_at": None, **d})
|
||||
|
||||
|
||||
def _wire_batcher_for_test(prisma_client, fail_commit=False):
|
||||
"""
|
||||
|
|
@ -3834,12 +3876,19 @@ def _wire_batcher_for_test(prisma_client, fail_commit=False):
|
|||
batcher.litellm_organizationtable = _Table("org")
|
||||
batcher.litellm_tagtable = _Table("tag")
|
||||
batcher.litellm_endusertable = _Table("enduser")
|
||||
def execute_raw(query: str, *args: object) -> None:
|
||||
write: Final = _raw_reset_write(query, args)
|
||||
if write is not None:
|
||||
queued.append(write)
|
||||
|
||||
batcher.execute_raw = execute_raw
|
||||
batcher.commit = commit
|
||||
return batcher
|
||||
|
||||
prisma_client.db.batch_ = MagicMock(side_effect=make_batcher)
|
||||
return batch_calls
|
||||
|
||||
|
||||
def _wire_cascade_reads_for_test(prisma_client, endusers=()):
|
||||
"""
|
||||
The budget tier's cascade reads the rows it is about to zero, so their
|
||||
|
|
@ -3860,6 +3909,7 @@ def _wire_cascade_reads_for_test(prisma_client, endusers=()):
|
|||
getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[])
|
||||
prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=list(endusers))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_keys_partial_failure():
|
||||
"""
|
||||
|
|
@ -3919,9 +3969,7 @@ async def test_reset_budget_keys_partial_failure():
|
|||
# Simulate successful reset modification
|
||||
key["spend"] = 0.0
|
||||
# Compute a new reset time based on the budget duration
|
||||
key["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=key["budget_duration"])
|
||||
).isoformat()
|
||||
key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"])
|
||||
return key
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -3957,6 +4005,7 @@ async def test_reset_budget_keys_partial_failure():
|
|||
for call in failure_hook_calls
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_users_partial_failure():
|
||||
"""
|
||||
|
|
@ -4007,9 +4056,7 @@ async def test_reset_budget_users_partial_failure():
|
|||
raise Exception("Simulated failure for user1")
|
||||
else:
|
||||
user["spend"] = 0.0
|
||||
user["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=user["budget_duration"])
|
||||
).isoformat()
|
||||
user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"])
|
||||
return user
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4038,6 +4085,7 @@ async def test_reset_budget_users_partial_failure():
|
|||
for call in failure_hook_calls
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_endusers_cascade_failure_is_all_or_nothing():
|
||||
"""
|
||||
|
|
@ -4251,12 +4299,12 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c
|
|||
assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}, "spend": {"gt": 0}}
|
||||
assert enduser_writes[0]["data"] == {"spend": 0}
|
||||
|
||||
key_writes: Final = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"]
|
||||
key_writes: Final = [c for c in batch_calls if c["table"] == "key" and "token" in c["where"]]
|
||||
user_writes: Final = [c for c in batch_calls if c["table"] == "user"]
|
||||
team_writes: Final = [c for c in batch_calls if c["table"] == "team"]
|
||||
assert len(key_writes) == 2
|
||||
assert len(user_writes) == 1
|
||||
assert user_writes[0]["where"] == {"user_id": "user2"}
|
||||
assert user_writes[0]["where"] == {"user_id": "user2", "budget_reset_at": None}
|
||||
assert len(team_writes) == 2
|
||||
for c in key_writes + user_writes + team_writes:
|
||||
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
|
@ -4300,9 +4348,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
raise Exception("Simulated failure for team1")
|
||||
else:
|
||||
team["spend"] = 0.0
|
||||
team["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=team["budget_duration"])
|
||||
).isoformat()
|
||||
team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"])
|
||||
return team
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4315,7 +4361,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
prisma_client.update_data.assert_not_awaited()
|
||||
team_writes = [c for c in batch_calls if c["table"] == "team"]
|
||||
assert len(team_writes) == 1
|
||||
assert team_writes[0]["where"] == {"team_id": "team2"}
|
||||
assert team_writes[0]["where"] == {"team_id": "team2", "budget_reset_at": None}
|
||||
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
|
||||
|
||||
|
|
@ -4327,6 +4373,7 @@ async def test_reset_budget_teams_partial_failure():
|
|||
for call in failure_hook_calls
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_keys_success():
|
||||
"""
|
||||
|
|
@ -4355,9 +4402,7 @@ async def test_service_logger_keys_success():
|
|||
|
||||
async def fake_reset_key(key, current_time, reset_settings=None):
|
||||
key["spend"] = 0.0
|
||||
key["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=key["budget_duration"])
|
||||
).isoformat()
|
||||
key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"])
|
||||
return key
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4386,6 +4431,7 @@ async def test_service_logger_keys_success():
|
|||
# Failure hook should not be executed.
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_keys_failure():
|
||||
"""
|
||||
|
|
@ -4412,9 +4458,7 @@ async def test_service_logger_keys_failure():
|
|||
if key["id"] == "key1":
|
||||
raise Exception("Simulated failure for key1")
|
||||
key["spend"] = 0.0
|
||||
key["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=key["budget_duration"])
|
||||
).isoformat()
|
||||
key["budget_reset_at"] = current_time + timedelta(seconds=key["budget_duration"])
|
||||
return key
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4448,6 +4492,7 @@ async def test_service_logger_keys_failure():
|
|||
# Success hook should not be called.
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_users_success():
|
||||
"""
|
||||
|
|
@ -4476,9 +4521,7 @@ async def test_service_logger_users_success():
|
|||
|
||||
async def fake_reset_user(user, current_time, reset_settings=None):
|
||||
user["spend"] = 0.0
|
||||
user["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=user["budget_duration"])
|
||||
).isoformat()
|
||||
user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"])
|
||||
return user
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4504,6 +4547,7 @@ async def test_service_logger_users_success():
|
|||
assert event_metadata.get("num_users_failed") == 0
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_users_failure():
|
||||
"""
|
||||
|
|
@ -4529,9 +4573,7 @@ async def test_service_logger_users_failure():
|
|||
if user["id"] == "user1":
|
||||
raise Exception("Simulated failure for user1")
|
||||
user["spend"] = 0.0
|
||||
user["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=user["budget_duration"])
|
||||
).isoformat()
|
||||
user["budget_reset_at"] = current_time + timedelta(seconds=user["budget_duration"])
|
||||
return user
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4562,6 +4604,7 @@ async def test_service_logger_users_failure():
|
|||
assert "users_found" not in event_metadata
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_teams_success():
|
||||
"""
|
||||
|
|
@ -4590,9 +4633,7 @@ async def test_service_logger_teams_success():
|
|||
|
||||
async def fake_reset_team(team, current_time, reset_settings=None):
|
||||
team["spend"] = 0.0
|
||||
team["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=team["budget_duration"])
|
||||
).isoformat()
|
||||
team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"])
|
||||
return team
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4618,6 +4659,7 @@ async def test_service_logger_teams_success():
|
|||
assert event_metadata.get("num_teams_failed") == 0
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_teams_failure():
|
||||
"""
|
||||
|
|
@ -4643,9 +4685,7 @@ async def test_service_logger_teams_failure():
|
|||
if team["id"] == "team1":
|
||||
raise Exception("Simulated failure for team1")
|
||||
team["spend"] = 0.0
|
||||
team["budget_reset_at"] = (
|
||||
current_time + timedelta(seconds=team["budget_duration"])
|
||||
).isoformat()
|
||||
team["budget_reset_at"] = current_time + timedelta(seconds=team["budget_duration"])
|
||||
return team
|
||||
|
||||
with patch.object(
|
||||
|
|
@ -4676,6 +4716,7 @@ async def test_service_logger_teams_failure():
|
|||
assert "teams_found" not in event_metadata
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_endusers_success():
|
||||
"""
|
||||
|
|
@ -4740,6 +4781,7 @@ async def test_service_logger_endusers_success():
|
|||
assert event_metadata.get("num_endusers_failed") == 0
|
||||
proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_service_logger_endusers_failure():
|
||||
"""
|
||||
|
|
@ -4806,6 +4848,7 @@ async def test_service_logger_endusers_failure():
|
|||
assert "budgets_found" not in event_metadata
|
||||
proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_for_litellm_team_members_called():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2459,6 +2459,54 @@ async def test_user_status_update_refreshes_cached_user(
|
|||
broadcast.assert_awaited_once_with(cache_key=saved_user.user_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_fields", "role_default_duration", "stored_reset_at", "moves_reset"),
|
||||
[
|
||||
({"budget_duration": "30d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), False),
|
||||
({"budget_duration": "30d"}, None, datetime(2099, 1, 1, tzinfo=timezone.utc), True),
|
||||
({"budget_duration": "7d"}, None, datetime(2020, 1, 1, tzinfo=timezone.utc), True),
|
||||
({"user_role": LitellmUserRoles.INTERNAL_USER}, "30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False),
|
||||
],
|
||||
)
|
||||
async def test_user_update_keeps_a_due_budget_reset_when_the_duration_is_resent(
|
||||
mocker: MockerFixture,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_fields: dict[str, object],
|
||||
role_default_duration: str | None,
|
||||
stored_reset_at: datetime,
|
||||
moves_reset: bool,
|
||||
) -> None:
|
||||
import litellm
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import _update_single_user_helper
|
||||
|
||||
monkeypatch.setattr(litellm, "internal_user_budget_duration", role_default_duration)
|
||||
saved_user: Final = LiteLLM_UserTable(
|
||||
user_id="user-spruce",
|
||||
budget_duration="30d",
|
||||
budget_reset_at=stored_reset_at,
|
||||
)
|
||||
prisma_client: Final = mocker.MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_first = 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
|
||||
mocker.patch( # test-quality-ok: keep the Redis publication boundary offline
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
new_callable=mocker.AsyncMock,
|
||||
)
|
||||
|
||||
await _update_single_user_helper(
|
||||
user_request=UpdateUserRequest(user_id=saved_user.user_id, **request_fields),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="admin-spruce", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
written: Final = prisma_client.update_data.call_args.kwargs["data"]
|
||||
reset_at: Final = written.get("budget_reset_at")
|
||||
assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset
|
||||
assert ("budget_reset_at" in written) is moves_reset
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bulk_user_model_budget_clear_serializes_and_refreshes_cache(mocker: MockerFixture) -> None:
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
|
|
|
|||
|
|
@ -2109,6 +2109,34 @@ async def test_update_key_end_user_budget_id_folds_into_metadata_and_survives_om
|
|||
assert untouched["metadata"]["end_user_budget_id"] == "svc-a-budget"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("stored_duration", "stored_reset_at", "moves_reset"),
|
||||
[
|
||||
("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False),
|
||||
("30d", datetime(2020, 1, 1), False),
|
||||
("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True),
|
||||
("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True),
|
||||
("30d", None, True),
|
||||
],
|
||||
)
|
||||
async def test_update_key_keeps_a_due_budget_reset_when_the_duration_is_resent(
|
||||
stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool
|
||||
) -> None:
|
||||
existing_key: Final = LiteLLM_VerificationToken(
|
||||
token="hashed", budget_duration=stored_duration, budget_reset_at=stored_reset_at
|
||||
)
|
||||
|
||||
updated: Final = await prepare_key_update_data(
|
||||
data=UpdateKeyRequest(key="sk-1", budget_duration="30d"), existing_key_row=existing_key
|
||||
)
|
||||
|
||||
assert updated["budget_duration"] == "30d"
|
||||
reset_at: Final = updated.get("budget_reset_at")
|
||||
assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset
|
||||
assert ("budget_reset_at" in updated) is moves_reset
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_key_clears_end_user_budget_id_with_empty_string():
|
||||
"""Sending an empty end_user_budget_id detaches the key default without touching any budget row,
|
||||
|
|
|
|||
|
|
@ -13155,7 +13155,7 @@ def test_set_budget_reset_at_clears_when_budget_duration_null():
|
|||
data = UpdateTeamRequest(team_id="test-team", budget_duration=None)
|
||||
updated_kv = {"team_id": "test-team", "budget_duration": None}
|
||||
|
||||
_set_budget_reset_at(data, updated_kv)
|
||||
_set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team"))
|
||||
|
||||
assert "budget_reset_at" in updated_kv
|
||||
assert updated_kv["budget_reset_at"] is None
|
||||
|
|
@ -13172,7 +13172,7 @@ def test_set_budget_reset_at_noop_when_budget_duration_not_sent():
|
|||
data = UpdateTeamRequest(team_id="test-team")
|
||||
updated_kv = {"team_id": "test-team"}
|
||||
|
||||
_set_budget_reset_at(data, updated_kv)
|
||||
_set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team"))
|
||||
|
||||
assert "budget_reset_at" not in updated_kv
|
||||
|
||||
|
|
@ -13188,12 +13188,41 @@ def test_set_budget_reset_at_sets_value_when_budget_duration_provided():
|
|||
data = UpdateTeamRequest(team_id="test-team", budget_duration="30d")
|
||||
updated_kv = {"team_id": "test-team", "budget_duration": "30d"}
|
||||
|
||||
_set_budget_reset_at(data, updated_kv)
|
||||
_set_budget_reset_at(data, updated_kv, LiteLLM_TeamTable(team_id="test-team"))
|
||||
|
||||
assert "budget_reset_at" in updated_kv
|
||||
assert updated_kv["budget_reset_at"] is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stored_duration", "stored_reset_at", "moves_reset"),
|
||||
[
|
||||
("30d", datetime(2020, 1, 1, tzinfo=timezone.utc), False),
|
||||
("30d", datetime(2020, 1, 1), False),
|
||||
("30d", datetime(2099, 1, 1, tzinfo=timezone.utc), True),
|
||||
("7d", datetime(2020, 1, 1, tzinfo=timezone.utc), True),
|
||||
("30d", None, True),
|
||||
],
|
||||
)
|
||||
def test_set_budget_reset_at_keeps_a_due_reset_when_the_duration_is_resent(
|
||||
stored_duration: str, stored_reset_at: datetime | None, moves_reset: bool
|
||||
) -> None:
|
||||
from litellm.proxy._types import UpdateTeamRequest
|
||||
from litellm.proxy.management_endpoints.team_endpoints import _set_budget_reset_at
|
||||
|
||||
updated_kv: Final[dict[str, object]] = {"team_id": "test-team", "budget_duration": "30d"}
|
||||
|
||||
_set_budget_reset_at(
|
||||
UpdateTeamRequest(team_id="test-team", budget_duration="30d"),
|
||||
updated_kv,
|
||||
LiteLLM_TeamTable(team_id="test-team", budget_duration=stored_duration, budget_reset_at=stored_reset_at),
|
||||
)
|
||||
|
||||
reset_at: Final = updated_kv.get("budget_reset_at")
|
||||
assert (isinstance(reset_at, datetime) and reset_at > datetime.now(timezone.utc)) is moves_reset
|
||||
assert ("budget_reset_at" in updated_kv) is moves_reset
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_team_member_budget_duration_calls_update_budget():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -52,16 +52,24 @@ async def test_updates_across_tables_share_one_batch_and_commit_once():
|
|||
reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
async with spend_reset_unit_of_work(lambda: batch) as uow:
|
||||
uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at, spend_decrement=1.5)
|
||||
uow.users.queue_spend_reset(user_id="user-1", budget_reset_at=reset_at, spend_decrement=2.5)
|
||||
uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None, spend_decrement=0.0)
|
||||
uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at, spend_decrement=1.5, expected_reset_at=None)
|
||||
uow.users.queue_spend_reset(
|
||||
user_id="user-1", budget_reset_at=reset_at, spend_decrement=2.5, expected_reset_at=None
|
||||
)
|
||||
uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None, spend_decrement=0.0, expected_reset_at=None)
|
||||
assert batch.commit_count == 0
|
||||
|
||||
assert batch.commit_count == 1
|
||||
assert batch.calls == [
|
||||
("litellm_verificationtoken", {"token": "tok-1"}, {"spend": {"decrement": 1.5}, "budget_reset_at": reset_at}),
|
||||
("litellm_usertable", {"user_id": "user-1"}, {"spend": {"decrement": 2.5}, "budget_reset_at": reset_at}),
|
||||
("litellm_teamtable", {"team_id": "team-1"}, {"spend": {"decrement": 0.0}, "budget_reset_at": None}),
|
||||
assert [call[2]["args"][:4] for call in batch.calls if call[0] == "execute_raw"] == [
|
||||
("tok-1", "2026-08-03T12:00:00.000000", "1.5", None),
|
||||
("user-1", "2026-08-03T12:00:00.000000", "2.5", None),
|
||||
]
|
||||
assert [call for call in batch.calls if call[0] != "execute_raw"] == [
|
||||
(
|
||||
"litellm_teamtable.update_many",
|
||||
{"team_id": "team-1", "budget_reset_at": None},
|
||||
{"spend": {"decrement": 0.0}, "budget_reset_at": None},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -70,7 +78,7 @@ async def test_raising_inside_block_skips_commit():
|
|||
|
||||
async def _blow_up_mid_transaction():
|
||||
async with spend_reset_unit_of_work(lambda: batch) as uow:
|
||||
uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None, spend_decrement=0.0)
|
||||
uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None, spend_decrement=0.0, expected_reset_at=None)
|
||||
raise RuntimeError("boom")
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
|
|
@ -133,7 +141,11 @@ async def test_enduser_page_commits_on_its_own_batch():
|
|||
{"user_id": {"in": ["enduser-1", "enduser-2"]}, "spend": {"gt": 0}},
|
||||
{"spend": 0},
|
||||
),
|
||||
("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}}, {"spend": {"decrement": 5.0}}),
|
||||
(
|
||||
"litellm_endusertable.update_many",
|
||||
{"user_id": {"in": ["enduser-3"]}, "spend": {"gt": 5.0}},
|
||||
{"spend": {"decrement": 5.0}},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue