diff --git a/tests/integration/_support/scheduled_write.py b/tests/integration/_support/scheduled_write.py new file mode 100644 index 00000000000..021e94c891e --- /dev/null +++ b/tests/integration/_support/scheduled_write.py @@ -0,0 +1,102 @@ +import asyncio +import time +from collections.abc import AsyncGenerator, Callable, Coroutine +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import Final + +from apscheduler.schedulers.asyncio import ( # pyright: ignore[reportMissingTypeStubs] # APScheduler has no stubs + AsyncIOScheduler, +) + +from litellm.proxy.shutdown.scheduled_jobs import AwaitableAsyncIOExecutor, stop_in_flight_scheduler_jobs +from tests.integration._support.client import eventually +from tests.integration._support.database import read_rows + + +async def wait_until(predicate: Callable[[], bool], description: str) -> None: + deadline: Final = time.monotonic() + 15 + while not predicate(): + assert time.monotonic() < deadline, description + await asyncio.sleep(0.01) + + +async def wait_for_blocked_update(database_url: str, blocker_pid: int, table: str) -> None: + await asyncio.to_thread( + eventually, + lambda: read_rows( + "SELECT pid FROM pg_stat_activity WHERE datname=current_database() " + "AND state='active' AND wait_event_type='Lock' AND %s::int=ANY(pg_blocking_pids(pid)) " + "AND query ILIKE '%%UPDATE%%' AND query LIKE %s", + (str(blocker_pid), f'%"{table}"%'), + database_url=database_url, + ), + lambda rows: len(rows) == 1, + seconds=15, + ) + + +async def wait_for_database_idle(database_url: str) -> None: + await asyncio.to_thread( + eventually, + lambda: read_rows( + "SELECT pid, state, query FROM pg_stat_activity WHERE datname=current_database() " + "AND pid<>pg_backend_pid() AND backend_type='client backend' " + "AND (state<>'idle' OR xact_start IS NOT NULL)", + (), + database_url=database_url, + ), + lambda rows: not rows, + seconds=15, + ) + + +@dataclass(frozen=True, slots=True) +class ScheduledWrite: + scheduler: AsyncIOScheduler + executor: AwaitableAsyncIOExecutor + outcome: asyncio.Future[None] + + async def cancel(self, release: Callable[[], None]) -> None: + jobs: Final = self.executor.in_flight_jobs() + assert len(jobs) == 1, "Expected one blocked application job before shutdown" + shutdown: Final = asyncio.create_task( + stop_in_flight_scheduler_jobs( + self.scheduler, self.executor, finish_timeout_seconds=0.01, cancel_timeout_seconds=15 + ) + ) + try: + await wait_until( + lambda: any(isinstance(job, asyncio.Task) and job.cancelling() > 0 for job in jobs), + "Scheduler never requested cancellation of the blocked write", + ) + finally: + release() + await asyncio.wait_for(shutdown, timeout=20) + + +@asynccontextmanager +async def scheduled_write(operation: Callable[[], Coroutine[object, object, None]]) -> AsyncGenerator[ScheduledWrite]: + executor: Final = AwaitableAsyncIOExecutor() + scheduler: Final = AsyncIOScheduler(executors={"default": executor}) + outcome: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future() + + async def run() -> None: + try: + await operation() + except BaseException as error: + outcome.set_exception(error) + raise + else: + outcome.set_result(None) + + scheduler.start() + scheduler.add_job(run, "date") # pyright: ignore[reportUnknownMemberType] # APScheduler's job arguments lack types + try: + yield ScheduledWrite(scheduler, executor, outcome) + finally: + await stop_in_flight_scheduler_jobs( + scheduler, executor, finish_timeout_seconds=0, cancel_timeout_seconds=15 + ) + if outcome.done() and not outcome.cancelled(): + outcome.exception() diff --git a/tests/integration/spend/conftest.py b/tests/integration/spend/conftest.py new file mode 100644 index 00000000000..7c902ddf92d --- /dev/null +++ b/tests/integration/spend/conftest.py @@ -0,0 +1,55 @@ +import os +import subprocess +import sys +from collections.abc import AsyncGenerator, Iterator +from pathlib import Path +from typing import Final + +import pytest +import pytest_asyncio +from redis import Redis + +from litellm.caching.redis_cache import RedisCache +from litellm.proxy import proxy_server +from litellm.proxy.utils import PrismaClient +from tests.integration._support.database import scratch_database +from tests.integration._support.prisma_clients import connected_proxy_prisma +from tests.integration._support.redis_process import owned_redis + + +@pytest.fixture(scope="module") +def key_spend_database() -> Iterator[str]: + with scratch_database("key_spend") as database_url: + schema: Final = Path(__file__).resolve().parents[3] / "litellm/proxy/schema.prisma" + result: Final = subprocess.run( + [sys.executable, "-I", "-m", "prisma", "db", "push", "--skip-generate", "--schema", str(schema)], + env={**os.environ, "DATABASE_URL": database_url, "DIRECT_URL": database_url}, + capture_output=True, + text=True, + timeout=90, + ) + assert result.returncode == 0, result.stdout + result.stderr + yield database_url + + +@pytest_asyncio.fixture(loop_scope="function") +async def buffered_spend_worker( + key_spend_database: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> AsyncGenerator[tuple[PrismaClient, "Redis[str]"]]: + monkeypatch.setattr(proxy_server, "general_settings", {"use_redis_transaction_buffer": True}) + monkeypatch.setenv("DATABASE_URL", key_spend_database) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + with owned_redis(tmp_path) as redis_server: + cache: Final = RedisCache(host=redis_server.host, port=redis_server.port) + try: + async with connected_proxy_prisma(key_spend_database) as database: + database.proxy_logging_obj.update_values( # pyright: ignore[reportUnknownMemberType] # Untyped alerts + redis_cache=cache + ) + observer_client: Final[Redis[str]] = Redis( + host=redis_server.host, port=redis_server.port, decode_responses=True + ) + with observer_client as observer: + yield database, observer + finally: + await cache.disconnect() diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index 174403d1873..f01a408d362 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,3 +1,4 @@ +import asyncio import hashlib import json import os @@ -9,18 +10,25 @@ from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace +from functools import partial from pathlib import Path -from typing import Final +from typing import Final, Literal import httpx import psycopg import pytest import yaml from integration._support.client import Gateway, delete_key_if_present, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, write_rows from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.scheduled_write import scheduled_write, wait_for_blocked_update, wait_for_database_idle from integration._support.wire import Reply, Request, wire_server from psycopg import sql +from redis import Redis + +from litellm.constants import REDIS_UPDATE_BUFFER_KEY +from litellm.proxy._types import Litellm_EntityType +from litellm.proxy.utils import PrismaClient, update_spend REQUESTS_WHILE_BLOCKED: Final = 6 CANCEL_LOG_LINE: Final = "in-flight scheduled job(s) for shutdown" @@ -538,3 +546,70 @@ def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_ex shutdown.terminate_once(lambda: _committing_daily_user_spend() >= 1, lambda: None) finally: _drop_slow_commit(suffix) + + +@pytest.mark.parametrize("schedule", ("normal", "locked", "cancelled")) +async def test_key_spend_worker_recovery_records_buffered_increment_exactly_once( + key_spend_database: str, + buffered_spend_worker: tuple[PrismaClient, "Redis[str]"], + schedule: Literal["normal", "locked", "cancelled"], +) -> None: + database, observer = buffered_spend_worker + alpha: Final = f"alpha-{uuid.uuid4().hex}" + beta: Final = f"beta-{uuid.uuid4().hex}" + write_rows( + 'INSERT INTO "LiteLLM_VerificationToken" (token, spend, total_spend, models) ' + "VALUES (%s, 100, 200, '{}'), (%s, 100, 200, '{}')", + (alpha, beta), + database_url=key_spend_database, + ) + queue: Final = database.proxy_logging_obj.db_spend_update_writer.spend_update_queue + worker: Final = partial( + update_spend, prisma_client=database, db_writer_client=None, proxy_logging_obj=database.proxy_logging_obj + ) + other_key: Final = partial( + read_rows, + 'SELECT row_to_json(k) AS row FROM "LiteLLM_VerificationToken" k WHERE token=%s', + (beta,), + database_url=key_spend_database, + ) + untouched: Final = other_key() + + def assert_recovered() -> None: + assert read_rows( + 'SELECT token, spend, total_spend FROM "LiteLLM_VerificationToken" ' + "WHERE token IN (%s, %s) ORDER BY token", + (alpha, beta), + database_url=key_spend_database, + ) == [ + {"token": alpha, "spend": 103.0, "total_spend": 203.0}, + {"token": beta, "spend": 100.0, "total_spend": 200.0}, + ], "Settled worker and replay must record the original +3 once and leave the other key alone" + assert observer.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [] + assert queue.update_queue.empty() + assert other_key() == untouched + + await queue.add_update({"entity_type": Litellm_EntityType.KEY, "entity_id": alpha, "response_cost": 3.0}) + with psycopg.connect(key_spend_database) as blocker: + if schedule != "normal": + blocker.execute('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s FOR UPDATE', (alpha,)) + async with scheduled_write(worker) as job: + try: + if schedule != "normal": + await wait_for_blocked_update( + key_spend_database, blocker.info.backend_pid, "LiteLLM_VerificationToken" + ) + assert observer.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [] + assert queue.update_queue.empty() + if schedule == "cancelled": + await job.cancel(blocker.rollback) + finally: + blocker.rollback() + outcome, = await asyncio.wait_for(asyncio.gather(job.outcome, return_exceptions=True), timeout=20) + assert outcome is None or (schedule == "cancelled" and isinstance(outcome, asyncio.CancelledError)), outcome + await wait_for_database_idle(key_spend_database) + if schedule != "cancelled": + assert_recovered() + await asyncio.wait_for(worker(), timeout=20) + await wait_for_database_idle(key_spend_database) + assert_recovered()