Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_benchmark_runs_postgres.py

@@ -0,0 +1,446 @@
+"""Managed PostgreSQL qualification for benchmark-run persistence and restart behavior.
+
+This test is enabled only by the disposable PostgreSQL boundary in managed-validation.yml.
+It creates one uniquely named database on that boundary, then drops only that database if it
+still exists and is still owned by the expected validation role.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from datetime import UTC, datetime
+from pathlib import Path
+
+import pytest
+import sqlalchemy as sa
+from alembic.config import Config
+from sqlalchemy import create_engine, func, inspect, select, text
+from sqlalchemy.engine import URL, make_url
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import Session
+
+from alembic import command
+from modelforge_api.domain.benchmarking import BenchmarkRunStartRequest
+from modelforge_api.persistence.models import (
+    ArtifactSet,
+    ArtifactSetMember,
+    BenchmarkRun,
+    Deployment,
+    GatewayRequest,
+    LabExecutionRun,
+)
+from modelforge_api.services.benchmark_runs import BenchmarkRunService
+from modelforge_api.services.serving import ServingService
+from modelforge_api.services.transient_payloads import MemoryPayloadStore
+from tests.test_dispatch_context import operator_fixture
+
+_PRIOR = "20260924_0038"
+_CURRENT = "20260924_0039"
+_LIFECYCLE_OPT_IN = "managed-validation-v1"
+
+
+def _managed_test_urls() -> tuple[URL, URL, URL]:
+    """Fail closed unless the dedicated disposable PostgreSQL boundary opted in."""
+    if os.getenv("MODELFORGE_TEST_POSTGRES_DATABASE_LIFECYCLE") is None:
+        pytest.skip("requires the managed-validation disposable PostgreSQL lifecycle opt-in")
+    if os.getenv("MODELFORGE_TEST_POSTGRES_DATABASE_LIFECYCLE") != _LIFECYCLE_OPT_IN:
+        pytest.fail("managed PostgreSQL database lifecycle opt-in has an unexpected value")
+
+    raw_urls = (
+        os.getenv("MODELFORGE_TEST_POSTGRES_URL"),
+        os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL"),
+        os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL"),
+    )
+    if not all(raw_urls):
+        pytest.fail("managed PostgreSQL lifecycle opt-in requires all three test role URLs")
+
+    urls = tuple(make_url(value) for value in raw_urls if value is not None)
+    expected_users = ("postgres", "modelforge", "modelforge_runtime")
+    expected_passwords = (
+        os.getenv("MODELFORGE_POSTGRES_ADMIN_PASSWORD"),
+        os.getenv("MODELFORGE_MIGRATION_DB_PASSWORD"),
+        os.getenv("MODELFORGE_RUNTIME_DB_PASSWORD"),
+    )
+    labels = ("admin", "migration owner", "runtime")
+    if not all(expected_passwords):
+        pytest.fail("managed PostgreSQL role passwords are unavailable in this test process")
+
+    for url, username, password, label in zip(
+        urls, expected_users, expected_passwords, labels, strict=True
+    ):
+        if (
+            url.drivername != "postgresql+psycopg"
+            or url.host != "postgres"
+            or url.port != 5432
+            or url.database != "modelforge"
+            or url.username != username
+            or url.query
+        ):
+            pytest.fail(
+                f"managed PostgreSQL {label} URL does not match the isolated CI test boundary"
+            )
+        if url.password != password:
+            pytest.fail(
+                f"managed PostgreSQL {label} URL does not match its dedicated test role password"
+            )
+    return urls  # type: ignore[return-value]
+
+
+def _current_revision(engine) -> str:
+    with engine.connect() as connection:
+        return connection.scalar(text("SELECT version_num FROM alembic_version"))
+
+
+def _alembic_config(database_url: URL) -> Config:
+    config = Config()
+    config.set_main_option("script_location", str(Path(__file__).resolve().parents[1] / "alembic"))
+    rendered_url = database_url.render_as_string(hide_password=False).replace("%", "%%")
+    config.set_main_option("sqlalchemy.url", rendered_url)
+    return config
+
+
+def _insert_legacy_run(
+    connection: sa.Connection,
+    *,
+    run_id: uuid.UUID,
+    deployment_id: uuid.UUID,
+    suite_id: uuid.UUID,
+) -> None:
+    """Seed the exact pre-0039 row shape to prove the additive migration preserves history."""
+    legacy_runs = sa.table(
+        "benchmark_runs",
+        sa.column("id", sa.Uuid()),
+        sa.column("deployment_id", sa.Uuid()),
+        sa.column("suite_id", sa.Uuid()),
+        sa.column("model_revision_sha", sa.String(64)),
+        sa.column("artifact_sha256", sa.String(64)),
+        sa.column("runtime_type", sa.String(32)),
+        sa.column("runtime_version", sa.String(128)),
+        sa.column("runtime_image_digest", sa.String(71)),
+        sa.column("launch_arguments", sa.JSON()),
+        sa.column("suite_revision", sa.String(128)),
+        sa.column("dataset_revision", sa.String(255)),
+        sa.column("environment_fingerprint", sa.JSON()),
+        sa.column("environment_digest", sa.String(64)),
+        sa.column("seed", sa.Integer()),
+        sa.column("results", sa.JSON()),
+        sa.column("status", sa.String(32)),
+        sa.column("comparable", sa.Boolean()),
+        sa.column("stale_reason", sa.Text()),
+        sa.column("started_at", sa.DateTime(timezone=True)),
+        sa.column("completed_at", sa.DateTime(timezone=True)),
+    )
+    now = datetime.now(UTC)
+    connection.execute(
+        legacy_runs.insert().values(
+            id=run_id,
+            deployment_id=deployment_id,
+            suite_id=suite_id,
+            model_revision_sha="a" * 40,
+            artifact_sha256="b" * 64,
+            runtime_type="onnx",
+            runtime_version="1.0",
+            runtime_image_digest=None,
+            launch_arguments={},
+            suite_revision="1",
+            dataset_revision="builtin:legacy-benchmark-fixture-v1",
+            environment_fingerprint={},
+            environment_digest="c" * 64,
+            seed=None,
+            results={"legacy_fixture": True},
+            status="completed",
+            comparable=True,
+            stale_reason=None,
+            started_at=now,
+            completed_at=now,
+        )
+    )
+
+
+def test_managed_postgresql_benchmark_migration_race_and_restart(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Exercise migration 0039 and concurrent/restarted starts against real PostgreSQL."""
+    admin_url, owner_base_url, runtime_base_url = _managed_test_urls()
+    database_name = f"test_benchmark_run_{uuid.uuid4().hex}"
+    admin_engine = create_engine(admin_url, isolation_level="AUTOCOMMIT", pool_pre_ping=True)
+    created_database = False
+    owner_engine = None
+    runtime_engine = None
+
+    try:
+        with admin_engine.connect() as connection:
+            already_exists = connection.scalar(
+                text("SELECT 1 FROM pg_database WHERE datname = :database_name"),
+                {"database_name": database_name},
+            )
+            assert already_exists is None
+            connection.execute(text(f'CREATE DATABASE "{database_name}" OWNER modelforge'))
+            created_database = True
+
+        owner_url = owner_base_url.set(database=database_name)
+        runtime_url = runtime_base_url.set(database=database_name)
+        owner_engine = create_engine(owner_url, pool_pre_ping=True)
+        runtime_engine = create_engine(
+            runtime_url,
+            pool_pre_ping=True,
+            connect_args={
+                "connect_timeout": 15,
+                "options": "-c statement_timeout=30000 -c lock_timeout=15000",
+            },
+        )
+        migration = _alembic_config(owner_url)
+        command.upgrade(migration, _PRIOR)
+
+        legacy_run_id = uuid.uuid4()
+        with Session(owner_engine) as setup_session:
+            serving, target_request, selected = operator_fixture(setup_session)
+            benchmark_service = BenchmarkRunService(serving)
+            advertised_suite = benchmark_service.suites()[0]
+            suite = benchmark_service._suite(advertised_suite.id)
+            request = BenchmarkRunStartRequest(
+                capability="rag.embedding@1",
+                suite_id=suite.id,
+                target=target_request.target,
+                deployment_id=selected.capability_deployment_id,
+                experiment_route_id=selected.id,
+            )
+            artifact_set = setup_session.get(ArtifactSet, selected.deployment.artifact_set_id)
+            artifact_id = setup_session.scalar(
+                select(ArtifactSetMember.artifact_id)
+                .where(ArtifactSetMember.artifact_set_id == selected.deployment.artifact_set_id)
+                .limit(1)
+            )
+            assert artifact_set is not None and artifact_id is not None
+            legacy_deployment = Deployment(
+                name=f"benchmark-migration-fixture-{uuid.uuid4().hex}",
+                revision_id=artifact_set.revision_id,
+                artifact_id=artifact_id,
+                derived_artifact_id=None,
+                runtime_profile_id=selected.deployment.runtime_profile_id,
+                compute_node_id=selected.deployment.compute_node_id,
+                channel="candidate",
+                status="draft",
+                capability_contract_ids=[],
+            )
+            setup_session.add(legacy_deployment)
+            setup_session.flush()
+            legacy_deployment_id = legacy_deployment.id
+            setup_session.commit()
+            deployment_settings = serving._deployment_settings
+            manifests = serving.manifests
+            with owner_engine.begin() as connection:
+                _insert_legacy_run(
+                    connection,
+                    run_id=legacy_run_id,
+                    deployment_id=legacy_deployment_id,
+                    suite_id=suite.id,
+                )
+
+        command.upgrade(migration, "head")
+        assert _current_revision(owner_engine) == _CURRENT
+
+        column_metadata = {
+            column["name"]: column for column in inspect(owner_engine).get_columns("benchmark_runs")
+        }
+        upgraded_columns = set(column_metadata)
+        assert column_metadata["deployment_id"]["nullable"] is True
+        assert {
+            "capability_deployment_id",
+            "target_snapshot",
+            "actor_subject",
+            "idempotency_key",
+            "request_fingerprint",
+            "failure_code",
+            "cancellation_requested_at",
+            "created_at",
+            "generation",
+        } <= upgraded_columns
+        constraints = inspect(owner_engine)
+        assert "ck_benchmark_runs_operator_identity" in {
+            constraint["name"] for constraint in constraints.get_check_constraints("benchmark_runs")
+        }
+        assert "uq_benchmark_runs_actor_idempotency" in {
+            constraint["name"]
+            for constraint in constraints.get_unique_constraints("benchmark_runs")
+        }
+        capability_fk = next(
+            foreign_key
+            for foreign_key in constraints.get_foreign_keys("benchmark_runs")
+            if foreign_key["name"] == "fk_benchmark_runs_capability_deployment"
+        )
+        assert capability_fk["referred_table"] == "capability_deployments"
+        assert capability_fk["options"]["ondelete"] == "RESTRICT"
+
+        with Session(owner_engine) as session:
+            historical = session.get(BenchmarkRun, legacy_run_id)
+            assert historical is not None
+            assert historical.deployment_id == legacy_deployment_id
+            assert historical.actor_subject is None
+            assert historical.capability_deployment_id is None
+            assert historical.created_at is not None
+            assert historical.generation == 1
+
+        with owner_engine.connect() as connection:
+            historical_values = dict(
+                connection.execute(
+                    select(BenchmarkRun.__table__).where(BenchmarkRun.id == legacy_run_id)
+                )
+                .mappings()
+                .one()
+            )
+        historical_values["id"] = uuid.uuid4()
+        historical_values["actor_subject"] = "partial-identity-must-fail"
+        with pytest.raises(IntegrityError), owner_engine.begin() as connection:
+            connection.execute(BenchmarkRun.__table__.insert().values(historical_values))
+
+        with runtime_engine.connect() as connection:
+            assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _CURRENT
+            for privilege in ("SELECT", "INSERT", "UPDATE"):
+                assert connection.scalar(
+                    text("SELECT has_table_privilege(current_user, 'benchmark_runs', :privilege)"),
+                    {"privilege": privilege},
+                )
+            for privilege in ("DELETE", "TRUNCATE"):
+                assert not connection.scalar(
+                    text("SELECT has_table_privilege(current_user, 'benchmark_runs', :privilege)"),
+                    {"privilege": privilege},
+                )
+
+        original_provenance = BenchmarkRunService._provenance
+        provenance_barrier = threading.Barrier(2)
+
+        def synchronized_provenance(service, benchmark_request):
+            provenance = original_provenance(service, benchmark_request)
+            provenance_barrier.wait(timeout=30)
+            return provenance
+
+        monkeypatch.setattr(BenchmarkRunService, "_provenance", synchronized_provenance)
+
+        def start_once() -> uuid.UUID:
+            with Session(runtime_engine) as session:
+                serving = ServingService(
+                    session,
+                    deployment_settings,
+                    manifests,
+                    MemoryPayloadStore(),
+                )
+                return (
+                    BenchmarkRunService(serving)
+                    .start(
+                        request,
+                        actor_subject="managed-postgres-operator",
+                        idempotency_key="same-benchmark-start",
+                    )
+                    .id
+                )
+
+        with ThreadPoolExecutor(max_workers=2) as pool:
+            run_ids = list(pool.map(lambda _index: start_once(), range(2)))
+        assert run_ids[0] == run_ids[1]
+
+        benchmark_run_id = run_ids[0]
+        with Session(runtime_engine) as session:
+            benchmark = session.get(BenchmarkRun, benchmark_run_id)
+            assert benchmark is not None
+            assert benchmark.actor_subject == "managed-postgres-operator"
+            assert benchmark.idempotency_key == "same-benchmark-start"
+            assert benchmark.capability_deployment_id == request.deployment_id
+            assert benchmark.target_snapshot["deployment_id"] == str(request.deployment_id)
+            assert (
+                session.scalar(
+                    select(func.count())
+                    .select_from(BenchmarkRun)
+                    .where(
+                        BenchmarkRun.actor_subject == "managed-postgres-operator",
+                        BenchmarkRun.idempotency_key == "same-benchmark-start",
+                    )
+                )
+                == 1
+            )
+            child_id = uuid.UUID(benchmark.results["execution"]["active_lab_run_id"])
+            child = session.get(LabExecutionRun, child_id)
+            assert child is not None
+            assert child.actor_subject == "managed-postgres-operator"
+            assert child.idempotency_key == f"benchmark-{benchmark_run_id.hex}-00"
+            assert session.get(GatewayRequest, child.gateway_request_id) is not None
+            assert (
+                session.scalar(
+                    select(func.count())
+                    .select_from(LabExecutionRun)
+                    .where(
+                        LabExecutionRun.actor_subject == "managed-postgres-operator",
+                        LabExecutionRun.idempotency_key == f"benchmark-{benchmark_run_id.hex}-00",
+                    )
+                )
+                == 1
+            )
+
+        # Recreate the engine and service after the race: the committed parent/child link must
+        # remain readable and replaying the same actor/key must not create another child.
+        runtime_engine.dispose()
+        runtime_engine = create_engine(
+            runtime_url,
+            pool_pre_ping=True,
+            connect_args={
+                "connect_timeout": 15,
+                "options": "-c statement_timeout=30000 -c lock_timeout=15000",
+            },
+        )
+        with Session(runtime_engine) as session:
+            serving = ServingService(
+                session,
+                deployment_settings,
+                manifests,
+                MemoryPayloadStore(),
+            )
+            service = BenchmarkRunService(serving)
+            recovered = service.get(benchmark_run_id, actor_subject="managed-postgres-operator")
+            replay = service.start(
+                request,
+                actor_subject="managed-postgres-operator",
+                idempotency_key="same-benchmark-start",
+            )
+            assert recovered.id == replay.id == benchmark_run_id
+            assert replay.contract.deployment_id == str(request.deployment_id)
+            assert len(service.list(actor_subject="managed-postgres-operator")) == 1
+            assert (
+                session.scalar(
+                    select(func.count())
+                    .select_from(LabExecutionRun)
+                    .where(LabExecutionRun.actor_subject == "managed-postgres-operator")
+                )
+                == 1
+            )
+
+        # The downgrade must refuse to erase operator history and leave the migrated schema live.
+        with pytest.raises(RuntimeError, match="Operator benchmark history exists"):
+            command.downgrade(migration, _PRIOR)
+        assert _current_revision(owner_engine) == _CURRENT
+        with Session(owner_engine) as session:
+            assert session.get(BenchmarkRun, benchmark_run_id) is not None
+            assert session.get(BenchmarkRun, legacy_run_id) is not None
+    finally:
+        if runtime_engine is not None:
+            runtime_engine.dispose()
+        if owner_engine is not None:
+            owner_engine.dispose()
+        if created_database:
+            with admin_engine.connect() as connection:
+                database_owner = connection.scalar(
+                    text(
+                        "SELECT pg_get_userbyid(datdba) FROM pg_database "
+                        "WHERE datname = :database_name"
+                    ),
+                    {"database_name": database_name},
+                )
+                if database_owner is not None:
+                    if database_owner != "modelforge":
+                        raise RuntimeError(
+                            "refusing to remove the qualification database after its owner changed"
+                        )
+                    connection.execute(text(f'DROP DATABASE "{database_name}" WITH (FORCE)'))
+        admin_engine.dispose()