Merge pull request #42294 from BerriAI/litellm_/litellm-db-migration-scope-98b60d

test(migrations): cover the release-to-release upgrade path

.circleci/config.yml

@@ -6,6 +6,9 @@ parameters:
   migration_candidate_image:
     type: string
     default: ""
+  migration_baseline_image:
+    type: string
+    default: "ghcr.io/berriai/litellm-database:v1.102.0"
   migration_source_sha:
     type: string
     default: ""
@@ -2946,7 +2949,10 @@ jobs:
     parameters:
       suite:
         type: enum
-        enum: [startup, recovery, legacy]
+        enum: [startup, recovery, legacy, upgrade, shaped]
+      baseline:
+        type: boolean
+        default: false
     machine:
       image: ubuntu-2204:2024.04.1
     resource_class: large
@@ -2954,6 +2960,7 @@ jobs:
     environment:
       LITELLM_MIGRATION_TESTS: "1"
       LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci
+      LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
       MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres
       MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres
       MIGRATION_TEST_OUTPUT: /tmp/migration-results
@@ -2981,6 +2988,16 @@ jobs:
       - wait_for_service:
           url: tcp://localhost:5432
           timeout: "60"
+      - when:
+          condition: << parameters.baseline >>
+          steps:
+            - run:
+                name: Pull the baseline release the upgrade starts from
+                environment:
+                  BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >>
+                command: |
+                  [[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1
+                  docker pull "$BASELINE_IMAGE"
       - run:
           name: Run migration startup regressions
           environment:
@@ -3188,6 +3205,16 @@ workflows:
           name: migration-legacy-and-pooling
           suite: legacy
           requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-upgrade
+          suite: upgrade
+          baseline: true
+          requires: [build_docker_database_image]
+      - migration_startup_tests:
+          name: migration-upgrade-shaped
+          suite: shaped
+          baseline: true
+          requires: [build_docker_database_image]
   migration_startup_scheduled:
     triggers:
       - schedule:

.circleci/scripts/run_migration_tests.py

@@ -13,6 +13,8 @@ SUITES: Final = {
     "startup": (("test_startup.py",), 12),
     "recovery": (("test_recovery.py",), 15),
     "legacy": (("test_legacy.py", "test_pooling.py"), 11),
+    "upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),
+    "shaped": (("test_shaped_database.py",), 1),
 }
 
 
@@ -93,6 +95,7 @@ def main() -> int:
             {
                 **metadata,
                 "suite": suite,
+                "baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""),
                 "expected_cases": expected,
                 "passed": passed,
                 "pytest_exit_code": result.returncode,

tests/e2e/migrations/conftest.py

@@ -60,3 +60,32 @@ def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Con
     output: Final = Path(configured) / request.node.name if configured else tmp_path
     output.mkdir(parents=True, exist_ok=True)
     return Containers(migration_image, output)
+
+
+@pytest.fixture(scope="session")
+def baseline_image(tmp_path_factory: pytest.TempPathFactory) -> str:
+    configured: Final = os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE")
+    assert configured, "LITELLM_MIGRATION_BASELINE_IMAGE must name the released image the upgrade starts from"
+    image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}")
+    assert image.startswith("sha256:"), "Unable to identify the baseline image"
+    output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp())))
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "baseline-image.json").write_text(json.dumps({"requested": configured, "image_id": image}))
+    return image
+
+
+@pytest.fixture(scope="session")
+def baseline_template(
+    databases: Databases, baseline_image: str, tmp_path_factory: pytest.TempPathFactory
+) -> Iterator[Database]:
+    output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "baseline-seed"
+    with databases.create() as database:
+        with Containers(baseline_image, output).start(database) as replica:
+            ready((replica,), database)
+        yield database
+
+
+@pytest.fixture
+def baseline_database(databases: Databases, baseline_template: Database) -> Iterator[Database]:
+    with databases.create(baseline_template) as database:
+        yield database

tests/e2e/migrations/containers.py

@@ -5,7 +5,7 @@ import subprocess
 import time
 from collections.abc import Callable, Generator, Mapping
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Final
 from uuid import uuid4
@@ -123,6 +123,9 @@ class Containers:
     image: str
     output: Path
 
+    def using(self, image: str) -> "Containers":
+        return replace(self, image=image)
+
     @contextmanager
     def start(
         self,

tests/e2e/migrations/test_rolling_upgrade.py

@@ -0,0 +1,57 @@
+from typing import Final
+
+import pytest
+
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import (
+    CACHED_PLAN,
+    assert_history_clean,
+    assert_upgraded,
+    auth_traffic,
+    confirm,
+    keep_serving,
+    migration_names,
+    provision,
+)
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+class TestRollingUpgrade:
+    def test_baseline_replica_keeps_serving_while_the_candidate_migrates(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, _ = provision(old)
+            before: Final = migration_names(baseline_database)
+            with auth_traffic(old, key) as traffic:
+                keep_serving(traffic, "the baseline replica authenticating before the upgrade")
+                with containers.start(baseline_database) as new:
+                    ready((new,), baseline_database)
+                    assert_upgraded(before, migration_names(baseline_database))
+                    keep_serving(traffic, "the baseline replica authenticating after the schema moved")
+                    with auth_traffic(old, provision(new)[0]) as uncached:
+                        keep_serving(uncached, "the baseline replica resolving a key minted after the schema moved")
+            assert_history_clean(baseline_database)
+            assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement"
+            assert old.state().Running, "The baseline replica died during the upgrade"
+
+    def test_both_releases_serve_and_share_keys_during_the_overlap(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            old_key, old_alias = provision(old)
+            before: Final = migration_names(baseline_database)
+            with containers.start(baseline_database) as new:
+                ready((new,), baseline_database)
+                assert_upgraded(before, migration_names(baseline_database))
+                new_key, new_alias = provision(new)
+                with auth_traffic(old, old_key) as old_traffic, auth_traffic(new, new_key) as new_traffic:
+                    keep_serving(old_traffic, "the baseline replica serving through the overlap")
+                    keep_serving(new_traffic, "the candidate replica serving through the overlap")
+                confirm(old, new_key, new_alias)
+                confirm(new, old_key, old_alias)
+            assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement"

tests/e2e/migrations/test_shaped_database.py

@@ -0,0 +1,43 @@
+from typing import Final
+
+import pytest
+
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
+
+SPEND_ROWS: Final = 20_000
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+def seed_spend_logs(database: Database, rows: int) -> None:
+    database.execute(
+        'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime") '
+        "SELECT 'upgrade-shape-' || g, 'acompletion', now() - (g || ' seconds')::interval, "
+        "now() - (g || ' seconds')::interval FROM generate_series(1, %s) AS g",
+        (rows,),
+    )
+    assert database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((rows,),)
+
+
+class TestPopulatedDatabaseUpgrade:
+    def test_upgrade_completes_and_preserves_a_populated_spend_log(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, alias = provision(old)
+        seed_spend_logs(baseline_database, SPEND_ROWS)
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as new:
+            ready((new,), baseline_database)
+            assert_upgraded(before, migration_names(baseline_database))
+            confirm(new, key, alias)
+        assert_history_clean(baseline_database)
+        assert baseline_database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((SPEND_ROWS,),), (
+            "The upgrade lost spend rows"
+        )
+        assert baseline_database.query(
+            'SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE "startTime" IS NULL OR "endTime" IS NULL'
+        ) == ((0,),), "The upgrade nulled timestamps on existing spend rows"

tests/e2e/migrations/test_upgrade.py

@@ -0,0 +1,47 @@
+from contextlib import ExitStack
+from typing import Final
+
+import pytest
+
+from .checks import start_replicas
+from .containers import Containers, ready
+from .database import Database
+from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
+
+pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
+
+
+class TestReleaseUpgrade:
+    def test_candidate_applies_the_pending_release_migrations(
+        self, containers: Containers, baseline_database: Database
+    ) -> None:
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as replica:
+            ready((replica,), baseline_database)
+        assert_upgraded(before, migration_names(baseline_database))
+        assert_history_clean(baseline_database)
+
+    def test_upgrade_preserves_keys_minted_by_the_baseline_release(
+        self, containers: Containers, baseline_image: str, baseline_database: Database
+    ) -> None:
+        with containers.using(baseline_image).start(baseline_database) as old:
+            ready((old,), baseline_database)
+            key, alias = provision(old)
+            confirm(old, key, alias)
+        before: Final = migration_names(baseline_database)
+        with containers.start(baseline_database) as new:
+            ready((new,), baseline_database)
+            assert_upgraded(before, migration_names(baseline_database))
+            confirm(new, key, alias)
+
+    def test_concurrent_replicas_upgrade_a_baseline_database_once(
+        self, containers: Containers, baseline_database: Database
+    ) -> None:
+        before: Final = migration_names(baseline_database)
+        with ExitStack() as stack:
+            ready(start_replicas(stack, containers, baseline_database), baseline_database)
+        assert_upgraded(before, migration_names(baseline_database))
+        assert_history_clean(baseline_database)
+        assert baseline_database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count > 1") == (
+            (0,),
+        ), "A migration was executed more than once across the upgrading replicas"

tests/e2e/migrations/upgrade.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import threading
+from collections.abc import Generator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Final
+from uuid import uuid4
+
+from e2e_http import Result, Success, unwrap
+from models import (
+    KeyGenerateBody,
+    KeyGenerateResponse,
+    KeyInfoParams,
+    KeyInfoResponse,
+    ModelsListParams,
+    ModelsListResponse,
+)
+from pydantic import BaseModel
+
+from .containers import Replica, until
+from .database import Database
+
+CACHED_PLAN: Final = "cached plan must not change result type"
+
+
+def provision(replica: Replica) -> tuple[str, str]:
+    alias: Final = f"upgrade-{uuid4().hex}"
+    key: Final = unwrap(
+        replica.transport.post(
+            "/key/generate",
+            headers=replica.transport.master,
+            json=KeyGenerateBody(key_alias=alias),
+            response_type=KeyGenerateResponse,
+        )
+    ).key
+    return key, alias
+
+
+def confirm(replica: Replica, key: str, alias: str) -> None:
+    info: Final = unwrap(
+        replica.transport.get(
+            "/key/info",
+            headers=replica.transport.master,
+            params=KeyInfoParams(key=key),
+            response_type=KeyInfoResponse,
+        )
+    )
+    assert info.info.key_alias == alias, "Key minted on one release did not resolve on the other"
+
+
+@dataclass(slots=True)
+class Outcomes:
+    served: int = 0
+    failures: list[str] = field(default_factory=list)
+
+    def record(self, result: Result[BaseModel]) -> None:
+        match result:
+            case Success():
+                self.served += 1
+            case _:
+                self.failures.append(result.model_dump_json())
+
+
+@contextmanager
+def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generator[Outcomes]:
+    outcomes: Final = Outcomes()
+    stop: Final = threading.Event()
+
+    def drive() -> None:
+        while not stop.is_set():
+            outcomes.record(
+                replica.transport.get(
+                    "/v1/models",
+                    headers=replica.transport.bearer(key),
+                    params=ModelsListParams(),
+                    response_type=ModelsListResponse,
+                    timeout=10,
+                )
+            )
+            stop.wait(interval)
+
+    thread: Final = threading.Thread(target=drive, name="upgrade-auth-traffic", daemon=True)
+    thread.start()
+    try:
+        yield outcomes
+    finally:
+        stop.set()
+        thread.join(30)
+        assert not thread.is_alive(), "Auth traffic thread did not stop"
+    assert not outcomes.failures, (
+        f"Virtual-key auth failed on {replica.name} after the traffic window closed: {outcomes.failures[:5]}"
+    )
+
+
+def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int:
+    target: Final = outcomes.served + calls
+    until(description, lambda: outcomes.served >= target or bool(outcomes.failures))
+    assert not outcomes.failures, f"Virtual-key auth failed during {description}: {outcomes.failures[:5]}"
+    return outcomes.served
+
+
+def migration_names(database: Database) -> frozenset[str]:
+    return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations"))
+
+
+def assert_history_clean(database: Database) -> None:
+    assert database.query(
+        "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL"
+    ) == ((0,),), "The upgrade left an unfinished or rolled-back migration behind"
+    assert database.query(
+        "SELECT count(*) FROM (SELECT migration_name FROM _prisma_migrations GROUP BY migration_name "
+        "HAVING count(*) > 1) duplicated"
+    ) == ((0,),), "A migration was recorded more than once, so it ran on more than one replica"
+
+
+def assert_upgraded(before: frozenset[str], after: frozenset[str]) -> frozenset[str]:
+    applied: Final = after - before
+    assert applied, (
+        "The candidate applied no migrations the baseline release had not: the pinned "
+        "LITELLM_MIGRATION_BASELINE_IMAGE is at or ahead of the candidate, so this suite proves nothing"
+    )
+    assert not before - after, "The upgrade removed migration history the baseline release had already applied"
+    return applied