mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
fix(proxy): do not requeue a daily spend batch whose commit already left for postgres (#42786)
* fix(proxy): do not requeue a daily spend batch whose commit already left for postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): settle an interrupted daily spend commit from the shutdown flush instead of blocking the cancelled tick Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): burst two workers and SIGTERM during daily spend COMMIT, expect exactly once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
1519032d90
commit
093ceb576d
4 changed files with 486 additions and 12 deletions
|
|
@ -12,7 +12,8 @@ import os
|
|||
import random
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Coroutine, Mapping, Sequence
|
||||
from contextvars import ContextVar
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload
|
||||
|
|
@ -269,6 +270,80 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager:
|
|||
return tx
|
||||
|
||||
|
||||
_daily_spend_commit_started: Final[ContextVar[asyncio.Event | None]] = ContextVar(
|
||||
"_daily_spend_commit_started", default=None
|
||||
)
|
||||
|
||||
|
||||
def _mark_daily_spend_commit_started() -> None:
|
||||
started: Final = _daily_spend_commit_started.get()
|
||||
if started is not None:
|
||||
started.set()
|
||||
|
||||
|
||||
def _mark_daily_spend_commit_finished() -> None:
|
||||
started: Final = _daily_spend_commit_started.get()
|
||||
if started is not None:
|
||||
started.clear()
|
||||
|
||||
|
||||
def _start_daily_spend_commit(
|
||||
commit_started: asyncio.Event, commit: Callable[[], Coroutine[object, object, None]]
|
||||
) -> "asyncio.Task[None]":
|
||||
token: Final = _daily_spend_commit_started.set(commit_started)
|
||||
try:
|
||||
return asyncio.ensure_future(commit())
|
||||
finally:
|
||||
_daily_spend_commit_started.reset(token)
|
||||
|
||||
|
||||
def _track_interrupted_commit(commits: set[asyncio.Task[None]], settle: Coroutine[object, object, None]) -> None:
|
||||
task: Final = asyncio.ensure_future(settle)
|
||||
commits.add(task)
|
||||
task.add_done_callback(commits.discard)
|
||||
|
||||
|
||||
async def _settle_interrupted_commits(commits: set[asyncio.Task[None]]) -> None:
|
||||
while commits:
|
||||
await asyncio.wait(tuple(commits))
|
||||
|
||||
|
||||
async def _restore_tag_spend_the_commit_left_behind(
|
||||
commit_task: "asyncio.Task[None]",
|
||||
redis_update_buffer: RedisUpdateBuffer,
|
||||
transactions: dict[str, DailyTagSpendTransaction],
|
||||
) -> None:
|
||||
await asyncio.wait({commit_task})
|
||||
if commit_task.cancelled() or commit_task.exception() is None:
|
||||
return
|
||||
await redis_update_buffer.restore_transactions_to_redis(
|
||||
daily_tag_spend_update_transactions=transactions,
|
||||
)
|
||||
|
||||
|
||||
async def _requeue_daily_spend_the_commit_left_behind(
|
||||
commit_task: "asyncio.Task[None]",
|
||||
queue: DailySpendUpdateQueue,
|
||||
entity_type: str,
|
||||
transactions: dict[str, BaseDailySpendTransaction],
|
||||
) -> None:
|
||||
await asyncio.wait({commit_task})
|
||||
if commit_task.cancelled() or not transactions:
|
||||
return
|
||||
failure: Final = commit_task.exception()
|
||||
if failure is None:
|
||||
return
|
||||
spend_log_error(
|
||||
"Spend tracking - daily %s spend commit interrupted by shutdown failed. Re-queued %d rows for the "
|
||||
"shutdown flush. Error: %s",
|
||||
entity_type,
|
||||
len(transactions),
|
||||
str(failure),
|
||||
exc=failure,
|
||||
)
|
||||
await queue.add_update(transactions)
|
||||
|
||||
|
||||
# The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL),
|
||||
# so the roster check below cannot interleave with their writes. A row lock would deadlock with the
|
||||
# access-group endpoints, which lock a team row after an access-group lock.
|
||||
|
|
@ -391,6 +466,9 @@ class DBSpendUpdateWriter:
|
|||
self.daily_org_spend_update_queue = DailySpendUpdateQueue()
|
||||
self.daily_tag_spend_update_queue = DailySpendUpdateQueue()
|
||||
self.window_spend_update_queue = WindowSpendUpdateQueue()
|
||||
self.interrupted_tag_commits: set[asyncio.Task[None]] = (
|
||||
set()
|
||||
) # mutable-ok: same registry as DailySpendUpdateQueue.interrupted_commits
|
||||
|
||||
async def update_database(
|
||||
# LiteLLM management object fields
|
||||
|
|
@ -1606,17 +1684,24 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
commit_task: Final = asyncio.ensure_future(
|
||||
commit(
|
||||
commit_started: Final = asyncio.Event()
|
||||
commit_task: Final = _start_daily_spend_commit(
|
||||
commit_started,
|
||||
lambda: commit(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions),
|
||||
)
|
||||
),
|
||||
)
|
||||
try:
|
||||
await asyncio.shield(commit_task)
|
||||
except asyncio.CancelledError:
|
||||
if commit_started.is_set():
|
||||
queue.track_interrupted_commit(
|
||||
_requeue_daily_spend_the_commit_left_behind(commit_task, queue, entity_type, transactions)
|
||||
)
|
||||
raise
|
||||
commit_task.cancel()
|
||||
if transactions:
|
||||
await queue.add_update(transactions)
|
||||
|
|
@ -1841,23 +1926,36 @@ class DBSpendUpdateWriter:
|
|||
The drain is destructive, so a failed commit must push the transactions back for the next tick
|
||||
or their spend is lost permanently.
|
||||
"""
|
||||
await _settle_interrupted_commits(self.interrupted_tag_commits)
|
||||
daily_tag_spend_update_transactions: Final = (
|
||||
await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if not daily_tag_spend_update_transactions:
|
||||
return
|
||||
|
||||
commit_task: Final = asyncio.ensure_future(
|
||||
DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
commit_started: Final = asyncio.Event()
|
||||
commit_task: Final = _start_daily_spend_commit(
|
||||
commit_started,
|
||||
lambda: DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
),
|
||||
)
|
||||
try:
|
||||
await asyncio.shield(commit_task)
|
||||
except BaseException: # noqa: BLE001 # a cancel must restore the drained rows before its rollback returns
|
||||
if commit_started.is_set():
|
||||
_track_interrupted_commit(
|
||||
self.interrupted_tag_commits,
|
||||
_restore_tag_spend_the_commit_left_behind(
|
||||
commit_task,
|
||||
self.redis_update_buffer,
|
||||
daily_tag_spend_update_transactions,
|
||||
),
|
||||
)
|
||||
raise
|
||||
commit_task.cancel()
|
||||
await self.redis_update_buffer.restore_transactions_to_redis(
|
||||
daily_tag_spend_update_transactions=daily_tag_spend_update_transactions,
|
||||
|
|
@ -2382,6 +2480,8 @@ class DBSpendUpdateWriter:
|
|||
sql, params = build_bulk_upsert(table=table, batch=merged_batch)
|
||||
async with _spend_update_tx(prisma_client) as transaction:
|
||||
await transaction.execute_raw(sql, *params)
|
||||
_mark_daily_spend_commit_started()
|
||||
_mark_daily_spend_commit_finished()
|
||||
except Exception as batch_error:
|
||||
if _spend_commit_failure_is_requeue_safe(batch_error):
|
||||
spend_log_error(
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
from collections.abc import Coroutine
|
||||
from copy import deepcopy
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -57,6 +58,18 @@ class DailySpendUpdateQueue(BaseUpdateQueue):
|
|||
self.update_queue: asyncio.Queue[dict[str, BaseDailySpendTransaction]] = asyncio.Queue(
|
||||
maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE
|
||||
)
|
||||
self.interrupted_commits: set[asyncio.Task[None]] = (
|
||||
set()
|
||||
) # mutable-ok: registry of in-flight commit outcomes, entries leave via their done callback
|
||||
|
||||
def track_interrupted_commit(self, settle: Coroutine[object, object, None]) -> None:
|
||||
task: Final = asyncio.ensure_future(settle)
|
||||
self.interrupted_commits.add(task)
|
||||
task.add_done_callback(self.interrupted_commits.discard)
|
||||
|
||||
async def settle_interrupted_commits(self) -> None:
|
||||
while self.interrupted_commits:
|
||||
await asyncio.wait(tuple(self.interrupted_commits))
|
||||
|
||||
async def add_update(self, update: dict[str, BaseDailySpendTransaction]):
|
||||
"""Enqueue an update."""
|
||||
|
|
@ -81,6 +94,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue):
|
|||
self,
|
||||
) -> dict[str, BaseDailySpendTransaction]:
|
||||
"""Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates."""
|
||||
await self.settle_interrupted_commits()
|
||||
updates: Final = await self.flush_all_updates_from_in_memory_queue()
|
||||
if len(updates) > 0:
|
||||
verbose_proxy_logger.info(
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import signal
|
|||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
|
@ -13,16 +14,18 @@ import httpx
|
|||
import psycopg
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from integration._support.client import Gateway, delete_key_if_present, eventually, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from psycopg import sql
|
||||
|
||||
REQUESTS_WHILE_BLOCKED: Final = 6
|
||||
CANCEL_LOG_LINE: Final = "in-flight scheduled job(s) for shutdown"
|
||||
BATCH_DRAINED_LOG_LINE: Final = f"flushed {REQUESTS_WHILE_BLOCKED} daily spend update items from in-memory queue"
|
||||
MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new"
|
||||
COMMIT_DELAY_SECONDS: Final = 15
|
||||
BURST_REQUESTS: Final = 30
|
||||
|
||||
|
||||
def _api_requests(table: str, column: str, identity: str) -> int:
|
||||
|
|
@ -44,6 +47,50 @@ def _waiting_on(table: str) -> int:
|
|||
return waiting
|
||||
|
||||
|
||||
def _committing_daily_user_spend() -> int:
|
||||
rows: Final = read_rows(
|
||||
"SELECT count(*)::int AS committing FROM pg_stat_activity "
|
||||
"WHERE query='COMMIT' AND state='active' AND wait_event='PgSleep' AND pid IN "
|
||||
"(SELECT pid FROM pg_locks WHERE relation = %s::regclass AND mode='RowExclusiveLock')",
|
||||
('"LiteLLM_DailyUserSpend"',),
|
||||
)
|
||||
committing: Final = rows[0]["committing"]
|
||||
assert isinstance(committing, int)
|
||||
return committing
|
||||
|
||||
|
||||
def _install_slow_commit(user_id: str, fails_once: bool) -> str:
|
||||
suffix: Final = f"slow_commit_{uuid.uuid4().hex}"
|
||||
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
|
||||
connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(suffix)))
|
||||
connection.execute(
|
||||
sql.SQL(
|
||||
"CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $slow$ "
|
||||
"BEGIN PERFORM pg_sleep({}); "
|
||||
"IF {} AND nextval({}) = 1 THEN RAISE EXCEPTION 'integration: first COMMIT fails'; END IF; "
|
||||
"RETURN NULL; END $slow$"
|
||||
).format(
|
||||
sql.Identifier(suffix), sql.Literal(COMMIT_DELAY_SECONDS), sql.Literal(fails_once), sql.Literal(suffix)
|
||||
)
|
||||
)
|
||||
connection.execute(
|
||||
sql.SQL(
|
||||
'CREATE CONSTRAINT TRIGGER {} AFTER INSERT OR UPDATE ON "LiteLLM_DailyUserSpend" '
|
||||
"DEFERRABLE INITIALLY DEFERRED FOR EACH ROW WHEN (NEW.user_id = {}) EXECUTE FUNCTION {}()"
|
||||
).format(sql.Identifier(suffix), sql.Literal(user_id), sql.Identifier(suffix))
|
||||
)
|
||||
return suffix
|
||||
|
||||
|
||||
def _drop_slow_commit(suffix: str) -> None:
|
||||
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
|
||||
connection.execute(
|
||||
sql.SQL('DROP TRIGGER IF EXISTS {} ON "LiteLLM_DailyUserSpend"').format(sql.Identifier(suffix))
|
||||
)
|
||||
connection.execute(sql.SQL("DROP FUNCTION IF EXISTS {}()").format(sql.Identifier(suffix)))
|
||||
connection.execute(sql.SQL("DROP SEQUENCE IF EXISTS {}").format(sql.Identifier(suffix)))
|
||||
|
||||
|
||||
def _provider(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply(status=404, body=b'{"error":"not scripted"}')
|
||||
|
|
@ -77,6 +124,17 @@ class _Shutdown:
|
|||
def daily_user_requests(self) -> int:
|
||||
return _api_requests("LiteLLM_DailyUserSpend", "user_id", self.owner)
|
||||
|
||||
def spend_logs(self) -> int:
|
||||
rows: Final = read_rows('SELECT count(*)::int AS total FROM "LiteLLM_SpendLogs" WHERE "user"=%s', (self.owner,))
|
||||
total: Final = rows[0]["total"]
|
||||
assert isinstance(total, int)
|
||||
return total
|
||||
|
||||
def burst(self, requests: int) -> None:
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
for outcome in pool.map(lambda _: self.chat(), range(requests)):
|
||||
assert outcome is None
|
||||
|
||||
def logged(self, line: str, times: int = 1) -> bool:
|
||||
return self.owned.log.read_text(errors="replace").count(line) >= times
|
||||
|
||||
|
|
@ -121,7 +179,15 @@ def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
|
|||
|
||||
|
||||
@contextmanager
|
||||
def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int) -> Iterator[_Shutdown]:
|
||||
def _proxy_with_one_seeded_row(
|
||||
gateway: Gateway,
|
||||
tmp_path: Path,
|
||||
pool_limit: int,
|
||||
cancel_timeout_seconds: int = 5,
|
||||
settle_seconds: int = 0,
|
||||
requests: int = REQUESTS_WHILE_BLOCKED,
|
||||
workers: int = 1,
|
||||
) -> Iterator[_Shutdown]:
|
||||
owner: Final = f"integration-owner-{uuid.uuid4().hex}"
|
||||
with gateway.scenario() as scenario, wire_server(_provider) as wire:
|
||||
model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0)
|
||||
|
|
@ -133,9 +199,10 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int
|
|||
"LITELLM_LOG": "DEBUG",
|
||||
"GRACEFUL_SHUTDOWN_TIMEOUT": "1",
|
||||
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1",
|
||||
"SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": "5",
|
||||
"SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds),
|
||||
},
|
||||
config=_config_with_pool_limit(tmp_path, pool_limit),
|
||||
workers=workers,
|
||||
) as owned:
|
||||
key: Final = string_value(
|
||||
owned.gateway.post("/key/generate", {"user_id": owner, "team_id": team, "models": [model]})["key"]
|
||||
|
|
@ -145,8 +212,18 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int
|
|||
shutdown.chat()
|
||||
eventually(shutdown.daily_user_requests, lambda total: total == 1, seconds=60)
|
||||
yield shutdown
|
||||
assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == 1 + REQUESTS_WHILE_BLOCKED
|
||||
assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == 1 + REQUESTS_WHILE_BLOCKED
|
||||
written: Final = 1 + requests
|
||||
if settle_seconds:
|
||||
eventually(
|
||||
lambda: (
|
||||
_api_requests("LiteLLM_DailyUserSpend", "user_id", owner),
|
||||
_api_requests("LiteLLM_DailyTeamSpend", "team_id", team),
|
||||
),
|
||||
lambda totals: totals == (written, written),
|
||||
seconds=settle_seconds,
|
||||
)
|
||||
assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == written
|
||||
assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == written
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch")
|
||||
|
|
@ -187,3 +264,48 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exa
|
|||
lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _waiting_on("LiteLLM_DailyUserSpend") == 1,
|
||||
holder.rollback,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cancel_timeout_seconds", "commit_fails_once"),
|
||||
[
|
||||
pytest.param(60, False, id="cancel_budget_outlives_commit"),
|
||||
pytest.param(5, False, id="commit_outlives_cancel_budget"),
|
||||
pytest.param(60, True, id="commit_fails_within_cancel_budget"),
|
||||
pytest.param(5, True, id="commit_fails_after_cancel_budget"),
|
||||
],
|
||||
)
|
||||
def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_exactly_once(
|
||||
gateway: Gateway, tmp_path: Path, cancel_timeout_seconds: int, commit_fails_once: bool
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(
|
||||
gateway, tmp_path, pool_limit=10, cancel_timeout_seconds=cancel_timeout_seconds, settle_seconds=90
|
||||
) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
):
|
||||
suffix: Final = _install_slow_commit(shutdown.owner, fails_once=commit_fails_once)
|
||||
try:
|
||||
shutdown.chat_while_spend_update_is_blocked(memberships, "LiteLLM_TeamMembership")
|
||||
memberships.rollback()
|
||||
shutdown.terminate_once(
|
||||
lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _committing_daily_user_spend() == 1,
|
||||
lambda: None,
|
||||
)
|
||||
finally:
|
||||
_drop_slow_commit(suffix)
|
||||
|
||||
|
||||
def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_exactly_once(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with _proxy_with_one_seeded_row(
|
||||
gateway, tmp_path, pool_limit=10, settle_seconds=120, requests=BURST_REQUESTS, workers=2
|
||||
) as shutdown:
|
||||
suffix: Final = _install_slow_commit(shutdown.owner, fails_once=False)
|
||||
try:
|
||||
shutdown.burst(BURST_REQUESTS)
|
||||
eventually(shutdown.spend_logs, lambda total: total == 1 + BURST_REQUESTS, seconds=60)
|
||||
shutdown.terminate_once(lambda: _committing_daily_user_spend() >= 1, lambda: None)
|
||||
finally:
|
||||
_drop_slow_commit(suffix)
|
||||
|
|
|
|||
|
|
@ -4501,3 +4501,241 @@ async def test_tag_batch_drained_from_redis_and_cancelled_mid_flight_is_restored
|
|||
await asyncio.wait_for(db.rolled_back.wait(), timeout=5)
|
||||
assert db.transaction_outcomes == ["rollback"]
|
||||
assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == []
|
||||
|
||||
|
||||
class _CommittingDailySpendFakeDB(_DailySpendFakeDB):
|
||||
"""Runs the daily upsert at once but holds the COMMIT until released. A COMMIT that has left
|
||||
the client lands on the server whether or not the client keeps waiting for the reply."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(failing_table=None)
|
||||
self.committing = asyncio.Event()
|
||||
self.commit_release = asyncio.Event()
|
||||
self.transaction_outcomes: list[str] = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self) -> AsyncIterator["_CommittingDailySpendFakeDB"]:
|
||||
try:
|
||||
yield self
|
||||
except BaseException:
|
||||
self.transaction_outcomes.append("rollback")
|
||||
raise
|
||||
self.committing.set()
|
||||
try:
|
||||
await self.commit_release.wait()
|
||||
finally:
|
||||
self.transaction_outcomes.append("commit")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("queue_name", "entity_type", "entity_id_field", "table"), _DAILY_SPEND_ENTITIES)
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_that_lands_while_the_daily_batch_is_committing_waits_for_the_commit_and_does_not_requeue_it(
|
||||
queue_name: str, entity_type: str, entity_id_field: str, table: str
|
||||
):
|
||||
"""Shutdown cancels the tick after the COMMIT has left for Postgres. The server finishes that
|
||||
commit whatever the client does, so putting the batch back on the queue makes the final flush
|
||||
write the same spend a second time. The tick has to wait for the commit's outcome instead."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
queue = _DAILY_SPEND_QUEUES[queue_name](db_writer)
|
||||
await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)})
|
||||
await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)})
|
||||
db = _CommittingDailySpendFakeDB()
|
||||
|
||||
def flush(prisma_db: _DailySpendFakeDB):
|
||||
return db_writer._flush_daily_spend_queue(
|
||||
queue=queue,
|
||||
entity_type=entity_type,
|
||||
commit=_DAILY_SPEND_COMMITS[entity_type],
|
||||
n_retry_times=0,
|
||||
prisma_client=_WindowSpendFakePrisma(prisma_db),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
tick = asyncio.ensure_future(flush(db))
|
||||
await asyncio.wait_for(db.committing.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
finished, _ = await asyncio.wait({tick}, timeout=0.2)
|
||||
assert finished == {tick}, "the cancelled tick must hand the in-flight commit's outcome to the next flush"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
tick.result()
|
||||
assert len(queue.interrupted_commits) == 1
|
||||
|
||||
db.commit_release.set()
|
||||
await queue.settle_interrupted_commits()
|
||||
|
||||
assert db.transaction_outcomes == ["commit"]
|
||||
(upsert,) = _daily_upserts(db, table)
|
||||
assert _row_values(upsert, "api_requests") == [2]
|
||||
assert queue.update_queue.empty(), "a batch whose COMMIT already left for the server must not be requeued"
|
||||
|
||||
final_db = _DailySpendFakeDB(failing_table=None)
|
||||
await flush(final_db)
|
||||
assert _daily_upserts(final_db, table) == [], "the final flush must not write the committed batch again"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_batch_drained_from_redis_and_cancelled_while_committing_is_not_restored():
|
||||
"""Same in-flight COMMIT as the in-memory path, but the drained rows live in Redis. Restoring
|
||||
them after the server committed writes the tag spend twice on the next tick."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))}
|
||||
redis_buffer = _DrainedTagRedisBuffer(drained)
|
||||
db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer)
|
||||
db = _CommittingDailySpendFakeDB()
|
||||
|
||||
tick = asyncio.ensure_future(
|
||||
db_writer._drain_and_commit_daily_tag_spend_from_redis(
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
)
|
||||
await asyncio.wait_for(db.committing.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
finished, _ = await asyncio.wait({tick}, timeout=0.2)
|
||||
assert finished == {tick}, "the cancelled drain must hand the in-flight commit's outcome to the next drain"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
tick.result()
|
||||
assert len(db_writer.interrupted_tag_commits) == 1
|
||||
|
||||
db.commit_release.set()
|
||||
(settle,) = tuple(db_writer.interrupted_tag_commits)
|
||||
await settle
|
||||
|
||||
assert db.transaction_outcomes == ["commit"]
|
||||
assert redis_buffer.restored == [], (
|
||||
"a tag batch whose COMMIT already left for the server must not be restored to Redis"
|
||||
)
|
||||
|
||||
|
||||
class _CommitFailingDailySpendFakeDB(_CommittingDailySpendFakeDB):
|
||||
"""COMMIT leaves for the server but the reply comes back as a failure."""
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self) -> AsyncIterator["_CommitFailingDailySpendFakeDB"]:
|
||||
yield self
|
||||
self.committing.set()
|
||||
await self.commit_release.wait()
|
||||
self.transaction_outcomes.append("commit_failed")
|
||||
raise Exception("connection reset")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_while_committing_requeues_the_batch_when_the_commit_itself_fails():
|
||||
"""Waiting for the in-flight commit's outcome must not swallow a real commit failure:
|
||||
the batch still goes back on the queue and the next flush writes it once."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
queue = db_writer.daily_spend_update_queue
|
||||
await queue.add_update({"key-a": _daily_txn()})
|
||||
await queue.add_update({"key-a": _daily_txn()})
|
||||
db = _CommitFailingDailySpendFakeDB()
|
||||
|
||||
def flush(prisma_db: _DailySpendFakeDB):
|
||||
return db_writer._flush_daily_spend_queue(
|
||||
queue=queue,
|
||||
entity_type="user",
|
||||
commit=DBSpendUpdateWriter.update_daily_user_spend,
|
||||
n_retry_times=0,
|
||||
prisma_client=_WindowSpendFakePrisma(prisma_db),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
tick = asyncio.ensure_future(flush(db))
|
||||
await asyncio.wait_for(db.committing.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
finished, _ = await asyncio.wait({tick}, timeout=0.2)
|
||||
assert finished == {tick}, "the cancelled tick must not eat the shutdown budget waiting on the commit"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
tick.result()
|
||||
|
||||
db.commit_release.set()
|
||||
await queue.settle_interrupted_commits()
|
||||
|
||||
assert db.transaction_outcomes == ["commit_failed"]
|
||||
assert not queue.update_queue.empty(), "a batch whose COMMIT came back failed must be requeued"
|
||||
|
||||
final_db = _DailySpendFakeDB(failing_table=None)
|
||||
await flush(final_db)
|
||||
(upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend")
|
||||
assert _row_values(upsert, "api_requests") == [2]
|
||||
assert queue.update_queue.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_flush_that_lands_before_the_interrupted_commit_resolves_still_writes_a_failed_batch_once():
|
||||
"""The cancelled tick returns right away, so a COMMIT can still be in flight when the
|
||||
shutdown flush runs. If that commit later fails, the flush must first settle it, pick the
|
||||
requeued rows back up, and write them exactly once instead of losing them."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
queue = db_writer.daily_spend_update_queue
|
||||
await queue.add_update({"key-a": _daily_txn()})
|
||||
await queue.add_update({"key-a": _daily_txn()})
|
||||
db = _CommitFailingDailySpendFakeDB()
|
||||
|
||||
def flush(prisma_db: _DailySpendFakeDB):
|
||||
return db_writer._flush_daily_spend_queue(
|
||||
queue=queue,
|
||||
entity_type="user",
|
||||
commit=DBSpendUpdateWriter.update_daily_user_spend,
|
||||
n_retry_times=0,
|
||||
prisma_client=_WindowSpendFakePrisma(prisma_db),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
tick = asyncio.ensure_future(flush(db))
|
||||
await asyncio.wait_for(db.committing.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(tick, timeout=5)
|
||||
assert db.transaction_outcomes == [], "the COMMIT is still on the wire when the shutdown flush starts"
|
||||
|
||||
final_db = _DailySpendFakeDB(failing_table=None)
|
||||
shutdown_flush = asyncio.ensure_future(flush(final_db))
|
||||
finished, _ = await asyncio.wait({shutdown_flush}, timeout=0.2)
|
||||
assert finished == set(), "the shutdown flush must wait for the interrupted commit's outcome"
|
||||
assert _daily_upserts(final_db, "LiteLLM_DailyUserSpend") == []
|
||||
|
||||
db.commit_release.set()
|
||||
await asyncio.wait_for(shutdown_flush, timeout=5)
|
||||
|
||||
(upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend")
|
||||
assert _row_values(upsert, "api_requests") == [2]
|
||||
assert queue.update_queue.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_drain_that_lands_before_the_interrupted_tag_commit_resolves_restores_a_failed_batch():
|
||||
"""Same ordering for the Redis tag path: the shutdown drain must settle the interrupted
|
||||
commit before the destructive drain, or a commit that fails late is never restored."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))}
|
||||
redis_buffer = _DrainedTagRedisBuffer(drained)
|
||||
db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer)
|
||||
db = _CommitFailingDailySpendFakeDB()
|
||||
|
||||
def drain(prisma_db: _DailySpendFakeDB):
|
||||
return db_writer._drain_and_commit_daily_tag_spend_from_redis(
|
||||
prisma_client=_WindowSpendFakePrisma(prisma_db),
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
tick = asyncio.ensure_future(drain(db))
|
||||
await asyncio.wait_for(db.committing.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(tick, timeout=5)
|
||||
assert db.transaction_outcomes == []
|
||||
|
||||
final_db = _DailySpendFakeDB(failing_table=None)
|
||||
shutdown_drain = asyncio.ensure_future(drain(final_db))
|
||||
finished, _ = await asyncio.wait({shutdown_drain}, timeout=0.2)
|
||||
assert finished == set(), "the shutdown drain must wait for the interrupted commit's outcome"
|
||||
assert _daily_upserts(final_db, "LiteLLM_DailyTagSpend") == []
|
||||
|
||||
db.commit_release.set()
|
||||
await asyncio.wait_for(shutdown_drain, timeout=5)
|
||||
|
||||
assert redis_buffer.restored == [drained], "a tag batch whose COMMIT came back failed must be restored to Redis"
|
||||
(upsert,) = _daily_upserts(final_db, "LiteLLM_DailyTagSpend")
|
||||
assert _row_values(upsert, "api_requests") == [1]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue