447 lines
18 KiB
Python
447 lines
18 KiB
Python
"""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()
|