mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(proxy): keep the in-flight daily spend batch when shutdown cancels the flush (#42593)
* fix(proxy): keep the in-flight daily spend batch when shutdown cancels the flush A daily spend batch drained from the in-memory queue was dropped for good when the scheduler tick was cancelled by shutdown, because asyncio.CancelledError bypasses the except Exception requeue. The flush now requeues the drained rows on cancellation and re-raises, and each daily batch upsert runs in an interactive transaction so a statement that already reached Postgres is rolled back with the cancel instead of committing behind the requeue Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): requeue the cancelled daily spend batch before its rollback returns Behind a lock the rollback of the cancelled interactive transaction only returns once the blocked statement does, which is after the shutdown flush has already run. The commit now runs as a shielded task so the cancelled tick requeues the batch at once and lets the rollback finish in the background. The final flush then finds the rows and writes them exactly once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): give the recording db a transaction seam for the bulk upsert tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): route the mocked daily tag spend upsert through the transaction seam Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): restore the drained Redis tag batch when shutdown cancels its commit 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
21a2d828df
commit
a80379baf8
7 changed files with 520 additions and 17 deletions
|
|
@ -1606,13 +1606,21 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
try:
|
||||
await commit(
|
||||
commit_task: Final = asyncio.ensure_future(
|
||||
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:
|
||||
commit_task.cancel()
|
||||
if transactions:
|
||||
await queue.add_update(transactions)
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush
|
||||
if not transactions:
|
||||
return
|
||||
|
|
@ -1839,14 +1847,18 @@ class DBSpendUpdateWriter:
|
|||
if not daily_tag_spend_update_transactions:
|
||||
return
|
||||
|
||||
try:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
commit_task: Final = asyncio.ensure_future(
|
||||
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,
|
||||
)
|
||||
except Exception:
|
||||
)
|
||||
try:
|
||||
await asyncio.shield(commit_task)
|
||||
except BaseException: # noqa: BLE001 # a cancel must restore the drained rows before its rollback returns
|
||||
commit_task.cancel()
|
||||
await self.redis_update_buffer.restore_transactions_to_redis(
|
||||
daily_tag_spend_update_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
|
|
@ -2368,7 +2380,8 @@ class DBSpendUpdateWriter:
|
|||
table=table, transactions=tuple(transactions_to_process.values())
|
||||
)
|
||||
sql, params = build_bulk_upsert(table=table, batch=merged_batch)
|
||||
await prisma_client.db.execute_raw(sql, *params)
|
||||
async with _spend_update_tx(prisma_client) as transaction:
|
||||
await transaction.execute_raw(sql, *params)
|
||||
except Exception as batch_error:
|
||||
if _spend_commit_failure_is_requeue_safe(batch_error):
|
||||
spend_log_error(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import time
|
|||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -45,8 +46,37 @@ def stop_root_process(process: subprocess.Popen[bytes]) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class OwnedProxy:
|
||||
gateway: Gateway
|
||||
process: subprocess.Popen[bytes]
|
||||
log: Path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str], *, config: Path | None = None, remove_environment: tuple[str, ...] = ()) -> Iterator[Gateway]:
|
||||
def owned_proxy(
|
||||
gateway: Gateway,
|
||||
directory: Path,
|
||||
overrides: Mapping[str, str],
|
||||
*,
|
||||
config: Path | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
) -> Iterator[Gateway]:
|
||||
with owned_proxy_process(
|
||||
gateway, directory, overrides, config=config, remove_environment=remove_environment
|
||||
) as owned:
|
||||
yield owned.gateway
|
||||
|
||||
|
||||
@contextmanager
|
||||
def owned_proxy_process(
|
||||
gateway: Gateway,
|
||||
directory: Path,
|
||||
overrides: Mapping[str, str],
|
||||
*,
|
||||
config: Path | None = None,
|
||||
remove_environment: tuple[str, ...] = (),
|
||||
) -> Iterator[OwnedProxy]:
|
||||
with socket.socket() as reserve:
|
||||
reserve.bind(("127.0.0.1", 0))
|
||||
port: Final = reserve.getsockname()[1]
|
||||
|
|
@ -60,7 +90,8 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str],
|
|||
}
|
||||
output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
with (output / f"owned-proxy-{uuid.uuid4().hex}.log").open("w") as log:
|
||||
log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log"
|
||||
with log_path.open("w") as log:
|
||||
process: Final = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
|
|
@ -95,7 +126,7 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str],
|
|||
pass
|
||||
assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
|
||||
time.sleep(0.1)
|
||||
yield Gateway(client, gateway.key, gateway.upstream_url)
|
||||
yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, log_path)
|
||||
finally:
|
||||
root_stopped: Final = stop_root_process(process)
|
||||
residual: Final = group_members(process.pid)
|
||||
|
|
|
|||
|
|
@ -281,6 +281,12 @@
|
|||
"tests/integration/spend/test_filtered_ledger.py::test_rotated_keys_users_and_model_groups_preserve_success_failure_cache_ledger": [
|
||||
"quota_management.spend_tracking.filtered_ledger_preserves_owner_identity_and_totals"
|
||||
],
|
||||
"tests/integration/spend/test_shutdown_flush.py::test_daily_spend_batch_cancelled_while_waiting_for_a_pool_connection_is_written_by_the_final_flush": [
|
||||
"quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch"
|
||||
],
|
||||
"tests/integration/spend/test_shutdown_flush.py::test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exactly_once": [
|
||||
"quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch"
|
||||
],
|
||||
"tests/integration/spend/test_spend_calculate.py::test_spend_calculate_rejects_unpriced_model_with_400": [
|
||||
"quota_management.spend_tracking.spend_calculate.rejects_unpriced_model"
|
||||
],
|
||||
|
|
|
|||
189
tests/integration/spend/test_shutdown_flush.py
Normal file
189
tests/integration/spend/test_shutdown_flush.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
import json
|
||||
import os
|
||||
import signal
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
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
|
||||
|
||||
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"
|
||||
|
||||
|
||||
def _api_requests(table: str, column: str, identity: str) -> int:
|
||||
rows: Final = read_rows(
|
||||
f'SELECT coalesce(sum(api_requests), 0)::int AS total FROM "{table}" WHERE {column}=%s', (identity,)
|
||||
)
|
||||
total: Final = rows[0]["total"]
|
||||
assert isinstance(total, int)
|
||||
return total
|
||||
|
||||
|
||||
def _waiting_on(table: str) -> int:
|
||||
rows: Final = read_rows(
|
||||
"SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE wait_event_type='Lock' AND query LIKE %s",
|
||||
(f'%"{table}"%',),
|
||||
)
|
||||
waiting: Final = rows[0]["waiting"]
|
||||
assert isinstance(waiting, int)
|
||||
return waiting
|
||||
|
||||
|
||||
def _provider(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply(status=404, body=b'{"error":"not scripted"}')
|
||||
assert request.target == "/v1/chat/completions"
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-" + uuid.uuid4().hex,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Shutdown:
|
||||
owner: str
|
||||
team: str
|
||||
owned: OwnedProxy
|
||||
key: str
|
||||
model: str
|
||||
|
||||
def chat(self) -> None:
|
||||
body: Final = {"model": self.model, "messages": [{"role": "user", "content": f"spend {uuid.uuid4().hex}"}]}
|
||||
assert self.owned.gateway.request("POST", "/v1/chat/completions", body, key=self.key).status_code == 200
|
||||
|
||||
def daily_user_requests(self) -> int:
|
||||
return _api_requests("LiteLLM_DailyUserSpend", "user_id", self.owner)
|
||||
|
||||
def logged(self, line: str, times: int = 1) -> bool:
|
||||
return self.owned.log.read_text(errors="replace").count(line) >= times
|
||||
|
||||
def chat_while_spend_update_is_blocked(self, blocker: psycopg.Connection, table: str) -> None:
|
||||
blocker.execute(f'LOCK TABLE "{table}" IN EXCLUSIVE MODE')
|
||||
for _ in range(REQUESTS_WHILE_BLOCKED):
|
||||
self.chat()
|
||||
eventually(lambda: _waiting_on(table), lambda waiting: waiting == 1, seconds=30)
|
||||
|
||||
def start_blocked_model_insert(self) -> threading.Thread:
|
||||
body: Final = {
|
||||
"model_name": f"integration-blocked-{uuid.uuid4().hex}",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "integration-provider-key"},
|
||||
"model_info": {},
|
||||
}
|
||||
|
||||
def insert() -> None:
|
||||
try:
|
||||
self.owned.gateway.request("POST", "/model/new", body)
|
||||
except httpx.TransportError:
|
||||
pass
|
||||
|
||||
thread: Final = threading.Thread(target=insert, daemon=True)
|
||||
thread.start()
|
||||
return thread
|
||||
|
||||
def terminate_once(self, blocked: Callable[[], bool], release: Callable[[], None]) -> None:
|
||||
eventually(blocked, lambda state: state, seconds=60)
|
||||
self.owned.process.send_signal(signal.SIGTERM)
|
||||
eventually(lambda: self.logged(CANCEL_LOG_LINE), lambda seen: seen, seconds=60)
|
||||
release()
|
||||
self.owned.process.wait(timeout=120)
|
||||
|
||||
|
||||
def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["general_settings"]["database_connection_pool_limit"] = pool_limit
|
||||
config["general_settings"]["database_connection_pool_timeout"] = 60
|
||||
path: Final = tmp_path / f"pool-{pool_limit}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int) -> 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)
|
||||
team: Final = scenario.team(models=[model])
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"LITELLM_LOG": "DEBUG",
|
||||
"GRACEFUL_SHUTDOWN_TIMEOUT": "1",
|
||||
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1",
|
||||
"SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": "5",
|
||||
},
|
||||
config=_config_with_pool_limit(tmp_path, pool_limit),
|
||||
) as owned:
|
||||
key: Final = string_value(
|
||||
owned.gateway.post("/key/generate", {"user_id": owner, "team_id": team, "models": [model]})["key"]
|
||||
)
|
||||
scenario.cleanups.callback(delete_key_if_present, gateway, key)
|
||||
shutdown: Final = _Shutdown(owner, team, owned, key, model)
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch")
|
||||
def test_daily_spend_batch_cancelled_while_waiting_for_a_pool_connection_is_written_by_the_final_flush(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=2) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as models,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
):
|
||||
models.execute('LOCK TABLE "LiteLLM_ProxyModelTable" IN EXCLUSIVE MODE')
|
||||
first: Final = shutdown.start_blocked_model_insert()
|
||||
eventually(lambda: _waiting_on("LiteLLM_ProxyModelTable"), lambda waiting: waiting == 1, seconds=30)
|
||||
shutdown.chat_while_spend_update_is_blocked(memberships, "LiteLLM_TeamMembership")
|
||||
second: Final = shutdown.start_blocked_model_insert()
|
||||
eventually(lambda: shutdown.logged(MODEL_INSERT_ARRIVED_LOG_LINE, times=2), lambda seen: seen, seconds=30)
|
||||
memberships.rollback()
|
||||
eventually(lambda: _waiting_on("LiteLLM_ProxyModelTable"), lambda waiting: waiting == 2, seconds=30)
|
||||
shutdown.terminate_once(lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE), models.rollback)
|
||||
first.join(timeout=30)
|
||||
second.join(timeout=30)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch")
|
||||
def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exactly_once(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
with (
|
||||
_proxy_with_one_seeded_row(gateway, tmp_path, pool_limit=10) as shutdown,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as holder,
|
||||
psycopg.connect(os.environ["DATABASE_URL"]) as memberships,
|
||||
):
|
||||
holder.execute('SELECT 1 FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s FOR UPDATE', (shutdown.owner,))
|
||||
shutdown.chat_while_spend_update_is_blocked(memberships, "LiteLLM_TeamMembership")
|
||||
memberships.rollback()
|
||||
shutdown.terminate_once(
|
||||
lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _waiting_on("LiteLLM_DailyUserSpend") == 1,
|
||||
holder.rollback,
|
||||
)
|
||||
|
|
@ -100,6 +100,7 @@ async def test_daily_tag_spend_retries_then_succeeds():
|
|||
1,
|
||||
]
|
||||
)
|
||||
prisma_client.db.tx.return_value.__aenter__.return_value.execute_raw = prisma_client.db.execute_raw
|
||||
|
||||
daily_spend_transactions: Dict[str, DailyTagSpendTransaction] = {
|
||||
"k": {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
"""Tests for the single-statement daily spend upsert (LIT-5291)."""
|
||||
|
||||
import re
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -149,6 +151,13 @@ class _RecordingDb:
|
|||
self.statements.append((query, args))
|
||||
return len(args)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self) -> AsyncIterator["_RecordingDb"]:
|
||||
yield self
|
||||
|
||||
def tx(self, timeout: object = None) -> AbstractAsyncContextManager["_RecordingDb"]:
|
||||
return self._tx()
|
||||
|
||||
|
||||
class _RecordingPrismaClient:
|
||||
def __init__(self) -> None:
|
||||
|
|
|
|||
|
|
@ -5,11 +5,11 @@ import logging
|
|||
import re
|
||||
|
||||
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -20,7 +20,7 @@ from redis.exceptions import DataError
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
|
||||
from litellm.proxy._types import DailyTagSpendTransaction, Litellm_EntityType, SpendUpdateQueueItem
|
||||
from litellm.proxy.db.db_spend_update_writer import (
|
||||
_TEAM_ADVISORY_LOCK_SQL,
|
||||
_TEAM_MEMBER_SPEND_SQL,
|
||||
|
|
@ -28,6 +28,8 @@ from litellm.proxy.db.db_spend_update_writer import (
|
|||
_SpendTableName,
|
||||
_spend_tables_left_to_send,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
build_window_spend_transaction,
|
||||
|
|
@ -303,6 +305,13 @@ class _RecordingDb:
|
|||
return self._execute_raw()
|
||||
return len(args)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self) -> AsyncIterator["_RecordingDb"]:
|
||||
yield self
|
||||
|
||||
def tx(self, timeout: timedelta | None = None) -> AbstractAsyncContextManager["_RecordingDb"]:
|
||||
return self._tx()
|
||||
|
||||
|
||||
class _RecordingPrisma:
|
||||
def __init__(self, execute_raw: Callable[[], int] | None = None) -> None:
|
||||
|
|
@ -3770,8 +3779,15 @@ async def test_commit_spend_updates_does_not_retry_non_deadlock_data_error(monke
|
|||
@pytest.mark.asyncio
|
||||
async def test_update_daily_spend_retries_deadlock(monkeypatch):
|
||||
"""The daily-spend upsert path retries a deadlock on the bulk upsert and then drains successfully."""
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[_deadlock_error(), None])
|
||||
outcomes = iter([_deadlock_error(), None])
|
||||
|
||||
def first_attempt_deadlocks():
|
||||
outcome = next(outcomes)
|
||||
if outcome is not None:
|
||||
raise outcome
|
||||
return 1
|
||||
|
||||
mock_prisma_client = _RecordingPrisma(execute_raw=first_attempt_deadlocks)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
|
|
@ -3786,7 +3802,7 @@ async def test_update_daily_spend_retries_deadlock(monkeypatch):
|
|||
entity_id_field="user_id",
|
||||
)
|
||||
|
||||
assert mock_prisma_client.db.execute_raw.call_count == 2
|
||||
assert len(mock_prisma_client.db.statements) == 2
|
||||
assert daily_spend_transactions == {}
|
||||
proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
|
@ -4247,3 +4263,241 @@ async def test_daily_transaction_attributes_caching_savings_only_with_an_injecti
|
|||
assert transaction["cache_creation_input_tokens"] == 1111
|
||||
assert transaction["prompt_caching_savings_spend"] != 0.0
|
||||
assert transaction["gateway_injected_caching_savings_spend"] == 0.0
|
||||
|
||||
|
||||
class _StallingDailySpendFakeDB(_DailySpendFakeDB):
|
||||
"""Holds the daily upsert aimed at one table until it is cancelled, like a starved pool does.
|
||||
|
||||
The rollback of that transaction waits for ``rollback_release``: the query engine only
|
||||
rolls back once the statement it is running has returned, which behind a lock takes
|
||||
as long as the lock is held."""
|
||||
|
||||
def __init__(self, stalled_table: str) -> None:
|
||||
super().__init__(failing_table=None)
|
||||
self.stalled_table = stalled_table
|
||||
self.stalled = asyncio.Event()
|
||||
self.rollback_release = asyncio.Event()
|
||||
self.rolled_back = asyncio.Event()
|
||||
self.transaction_outcomes: list[str] = []
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
if self.stalled_table in query:
|
||||
self.stalled.set()
|
||||
await asyncio.Event().wait()
|
||||
return await super().execute_raw(query, *args)
|
||||
|
||||
@asynccontextmanager
|
||||
async def _tx(self) -> AsyncIterator["_StallingDailySpendFakeDB"]:
|
||||
try:
|
||||
yield self
|
||||
except BaseException:
|
||||
await self.rollback_release.wait()
|
||||
self.transaction_outcomes.append("rollback")
|
||||
self.rolled_back.set()
|
||||
raise
|
||||
self.transaction_outcomes.append("commit")
|
||||
|
||||
|
||||
def _daily_entity_txn(entity_id_field: str) -> dict:
|
||||
return {key: value for key, value in _daily_txn().items() if key != "user_id"} | {entity_id_field: "entity-1"}
|
||||
|
||||
|
||||
_DAILY_SPEND_ENTITIES: Final = [
|
||||
pytest.param("daily_spend_update_queue", "user", "user_id", "LiteLLM_DailyUserSpend", id="user"),
|
||||
pytest.param("daily_team_spend_update_queue", "team", "team_id", "LiteLLM_DailyTeamSpend", id="team"),
|
||||
pytest.param("daily_org_spend_update_queue", "org", "organization_id", "LiteLLM_DailyOrganizationSpend", id="org"),
|
||||
pytest.param("daily_tag_spend_update_queue", "tag", "tag", "LiteLLM_DailyTagSpend", id="tag"),
|
||||
pytest.param(
|
||||
"daily_end_user_spend_update_queue", "end_user", "end_user_id", "LiteLLM_DailyEndUserSpend", id="end_user"
|
||||
),
|
||||
pytest.param("daily_agent_spend_update_queue", "agent", "agent_id", "LiteLLM_DailyAgentSpend", id="agent"),
|
||||
]
|
||||
|
||||
_DAILY_SPEND_QUEUES: Final[dict[str, Callable[[DBSpendUpdateWriter], DailySpendUpdateQueue]]] = {
|
||||
"daily_spend_update_queue": lambda writer: writer.daily_spend_update_queue,
|
||||
"daily_team_spend_update_queue": lambda writer: writer.daily_team_spend_update_queue,
|
||||
"daily_org_spend_update_queue": lambda writer: writer.daily_org_spend_update_queue,
|
||||
"daily_tag_spend_update_queue": lambda writer: writer.daily_tag_spend_update_queue,
|
||||
"daily_end_user_spend_update_queue": lambda writer: writer.daily_end_user_spend_update_queue,
|
||||
"daily_agent_spend_update_queue": lambda writer: writer.daily_agent_spend_update_queue,
|
||||
}
|
||||
|
||||
_DAILY_SPEND_COMMITS: Final = {
|
||||
"user": DBSpendUpdateWriter.update_daily_user_spend,
|
||||
"team": DBSpendUpdateWriter.update_daily_team_spend,
|
||||
"org": DBSpendUpdateWriter.update_daily_org_spend,
|
||||
"tag": DBSpendUpdateWriter.update_daily_tag_spend,
|
||||
"end_user": DBSpendUpdateWriter.update_daily_end_user_spend,
|
||||
"agent": DBSpendUpdateWriter.update_daily_agent_spend,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("queue_name", "entity_type", "entity_id_field", "table"), _DAILY_SPEND_ENTITIES)
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_spend_batch_cancelled_mid_flight_is_rolled_back_requeued_and_written_once_by_the_next_flush(
|
||||
queue_name: str, entity_type: str, entity_id_field: str, table: str
|
||||
):
|
||||
"""Shutdown cancels the scheduler tick while a drained batch waits on the database. The
|
||||
batch has left the queue, so unless the cancellation puts it back, the final flush finds
|
||||
nothing and the spend is gone (F2). The upsert runs in an interactive transaction so a
|
||||
statement that did reach Postgres is rolled back with the cancel and the requeued rows
|
||||
land exactly once."""
|
||||
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 = _StallingDailySpendFakeDB(stalled_table=table)
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.failure_handler = AsyncMock()
|
||||
|
||||
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=proxy_logging_obj,
|
||||
)
|
||||
|
||||
tick = asyncio.ensure_future(flush(db))
|
||||
await asyncio.wait_for(db.stalled.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
finished, _ = await asyncio.wait({tick}, timeout=1)
|
||||
assert finished == {tick}, "the cancelled tick must return before the rolled-back statement unwinds"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
tick.result()
|
||||
|
||||
assert not queue.update_queue.empty(), "the cancelled batch must go back on the queue before the rollback lands"
|
||||
assert db.transaction_outcomes == []
|
||||
db.rollback_release.set()
|
||||
await asyncio.wait_for(db.rolled_back.wait(), timeout=5)
|
||||
assert db.transaction_outcomes == ["rollback"]
|
||||
assert _daily_upserts(db, table) == []
|
||||
|
||||
final_db = _DailySpendFakeDB(failing_table=None)
|
||||
await flush(final_db)
|
||||
|
||||
(upsert,) = _daily_upserts(final_db, table)
|
||||
assert _row_values(upsert, entity_id_field) == ["entity-1"]
|
||||
assert _row_values(upsert, "spend") == [pytest.approx(0.2)]
|
||||
assert _row_values(upsert, "api_requests") == [2]
|
||||
assert queue.update_queue.empty()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_flush_of_an_empty_daily_queue_requeues_nothing():
|
||||
"""A cancel that lands with nothing drained must not push an empty batch onto the queue."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
db = _StallingDailySpendFakeDB(stalled_table="LiteLLM_DailyUserSpend")
|
||||
|
||||
class _CancellingQueue(type(db_writer.daily_spend_update_queue)):
|
||||
async def flush_and_get_aggregated_daily_spend_update_transactions(self):
|
||||
drained = await super().flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
asyncio.current_task().cancel()
|
||||
await asyncio.sleep(0)
|
||||
return drained
|
||||
|
||||
queue = _CancellingQueue()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await db_writer._flush_daily_spend_queue(
|
||||
queue=queue,
|
||||
entity_type="user",
|
||||
commit=DBSpendUpdateWriter.update_daily_user_spend,
|
||||
n_retry_times=0,
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert queue.update_queue.empty()
|
||||
|
||||
|
||||
class _AnnouncingDailySpendFakeDB(_DailySpendFakeDB):
|
||||
"""Signals ``written`` the moment the daily upsert has been committed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(failing_table=None)
|
||||
self.written = asyncio.Event()
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
rows = await super().execute_raw(query, *args)
|
||||
self.written.set()
|
||||
return rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_that_lands_after_the_daily_batch_committed_does_not_requeue_it():
|
||||
"""The commit has returned but the tick has not resumed yet when the cancel arrives.
|
||||
Putting the batch back now would write the same spend twice on the final flush."""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
queue = db_writer.daily_spend_update_queue
|
||||
await queue.add_update({"key-a": _daily_txn()})
|
||||
db = _AnnouncingDailySpendFakeDB()
|
||||
|
||||
tick = asyncio.ensure_future(
|
||||
db_writer._flush_daily_spend_queue(
|
||||
queue=queue,
|
||||
entity_type="user",
|
||||
commit=DBSpendUpdateWriter.update_daily_user_spend,
|
||||
n_retry_times=0,
|
||||
prisma_client=_WindowSpendFakePrisma(db),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
)
|
||||
await db.written.wait()
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await tick
|
||||
|
||||
assert len(_daily_upserts(db, "LiteLLM_DailyUserSpend")) == 1
|
||||
assert queue.update_queue.empty(), "a batch that already committed must not be requeued"
|
||||
|
||||
|
||||
class _DrainedTagRedisBuffer:
|
||||
"""Hands out one drained tag batch and records whatever is restored."""
|
||||
|
||||
def __init__(self, drained: dict[str, DailyTagSpendTransaction]) -> None:
|
||||
self.drained = drained
|
||||
self.restored: list[dict[str, DailyTagSpendTransaction]] = []
|
||||
|
||||
async def get_all_daily_tag_spend_update_transactions_from_redis_buffer(
|
||||
self,
|
||||
) -> dict[str, DailyTagSpendTransaction]:
|
||||
return self.drained
|
||||
|
||||
async def restore_transactions_to_redis(
|
||||
self, daily_tag_spend_update_transactions: dict[str, DailyTagSpendTransaction]
|
||||
) -> None:
|
||||
self.restored.append(daily_tag_spend_update_transactions)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_batch_drained_from_redis_and_cancelled_mid_flight_is_restored_before_its_rollback_returns():
|
||||
"""The Redis tag drain is destructive. A shutdown cancel used to leave the batch nowhere:
|
||||
Redis no longer had it and the interactive transaction rolled the statement back."""
|
||||
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 = _StallingDailySpendFakeDB(stalled_table="LiteLLM_DailyTagSpend")
|
||||
|
||||
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.stalled.wait(), timeout=5)
|
||||
tick.cancel()
|
||||
finished, _ = await asyncio.wait({tick}, timeout=1)
|
||||
assert finished == {tick}, "the cancelled drain must return before the rolled-back statement unwinds"
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
tick.result()
|
||||
|
||||
assert redis_buffer.restored == [drained], "the drained tag batch must be back in Redis before the rollback lands"
|
||||
assert db.transaction_outcomes == []
|
||||
db.rollback_release.set()
|
||||
await asyncio.wait_for(db.rolled_back.wait(), timeout=5)
|
||||
assert db.transaction_outcomes == ["rollback"]
|
||||
assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue