Merge 482508c675 into aeede53112

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()

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()

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()