diff --git a/.circleci/config.yml b/.circleci/config.yml index cc9aa7fe1c4..83acf0ac1c8 100644 --- a/.circleci/config.yml +++ b/.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: diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py index 56029c406fb..5a73c54e3f5 100644 --- a/.circleci/scripts/run_migration_tests.py +++ b/.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, diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py index 735adeedbdb..b7604a4fdda 100644 --- a/tests/e2e/migrations/conftest.py +++ b/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 diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py index 0f5793b81dd..dd126b994d3 100644 --- a/tests/e2e/migrations/containers.py +++ b/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, diff --git a/tests/e2e/migrations/test_rolling_upgrade.py b/tests/e2e/migrations/test_rolling_upgrade.py new file mode 100644 index 00000000000..4d60d20df4a --- /dev/null +++ b/tests/e2e/migrations/test_rolling_upgrade.py @@ -0,0 +1,55 @@ +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") + 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" diff --git a/tests/e2e/migrations/test_shaped_database.py b/tests/e2e/migrations/test_shaped_database.py new file mode 100644 index 00000000000..20c4368ae33 --- /dev/null +++ b/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" diff --git a/tests/e2e/migrations/test_upgrade.py b/tests/e2e/migrations/test_upgrade.py new file mode 100644 index 00000000000..23f0bbe9124 --- /dev/null +++ b/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" diff --git a/tests/e2e/migrations/upgrade.py b/tests/e2e/migrations/upgrade.py new file mode 100644 index 00000000000..2758d0fb974 --- /dev/null +++ b/tests/e2e/migrations/upgrade.py @@ -0,0 +1,117 @@ +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" + + +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" + + +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