vitor-aignosi
dac4f1718f
QTZPOC-20: Implement model import workflow for .sientia bundles Introduces the import_model workflow and its associated activities to enable importing .sientia bundles into MLflow and MongoDB. This includes: * New core workflow and activity logic for handling bundle download, decryption, validation, MLflow provisioning (experiments, runs, registered versions), and MongoDB model document creation. * Comprehensive E2E and unit test suites, including a dedicated test harness and a manual run script for development. * Updated .env.example and README.md with new configurations and usage instructions. * Ensures correct MLflow model source resolution for legacy loader compatibility and robust MongoDB document management. * Refactors existing utility scripts into a scripts/ directory.
951 lines
37 KiB
Python
951 lines
37 KiB
Python
"""Fixtures for the `.sientia` model-import e2e suite. Registered from `e2e/conftest.py`.
|
|
|
|
**Why this is a module of its own.** `e2e/conftest.py` is shared by every existing e2e suite and
|
|
long enough already; four containers' worth of setup, an MLflow server subprocess, the import
|
|
environment variables and an `ImportModel` worker do not belong in it. Every fixture here is
|
|
prefixed `import_`, so nothing collides with the predictions harness, and no existing fixture is
|
|
edited — the one exception being `log_statement=all` on the shared `postgres_container`, which the
|
|
absence assertions need and which changes no test's behaviour.
|
|
|
|
**The fidelity contract, which is the whole point of this file.** On the import path *nothing is
|
|
mocked*:
|
|
|
|
| Dependency | Here | Why real |
|
|
|---|---|---|
|
|
| MinIO | the shared `minio_container`, plus an import bucket | gate 1 reads a real object |
|
|
| BFF Postgres | a container of its own, `sientia-core-mlops-bff` | the real `CHECK` constraints are half of what is under test, and two tests stop the server |
|
|
| MongoDB | a container, `models` seeded with gapped ids and a unique index on `id` | `max(int(id)) + 1` against a real unique index |
|
|
| MLflow | `mlflow server` from *this* venv, sqlite-backed, artifacts proxied | experiment/run/artifact/registry semantics are what the provisioning activities are tested for |
|
|
| the crypto | the production reader, real Argon2id, real `secretstream` | a bundle the producer wrote must open |
|
|
|
|
What stays stubbed, deliberately: `CoreNotificationHandler`'s Mongo client keeps the existing
|
|
`mock_mongo_client`. Notifications are a side channel this suite counts, not a dependency it tests,
|
|
and a real notification collection would add a second write path to every absence assertion for
|
|
nothing.
|
|
|
|
The import `Activities` is **never** composed with `patch_mlflow`, `patch_minio_repository` or
|
|
`patch_create_engine`. `mlflow.set_tracking_uri` is process-global, so a leak would silently turn a
|
|
real assertion into an assertion about a `MagicMock`; `import_activities` asserts the live tracking
|
|
URI names the test server before it hands the object over.
|
|
|
|
**Bundles come from the unit suite.** `tests/helpers/bundle_factory.py` writes them and
|
|
`tests/fixtures/bundle/golden_model_v3.sientia` is the producer-emitted one. No third mirrored copy
|
|
of the producer's crypto is introduced here.
|
|
|
|
**Isolation** (the containers are session-scoped): per test the status table is rebuilt, the `models`
|
|
collection is dropped and reseeded, and the import prefix is emptied. MLflow state is *not* reset —
|
|
experiment and model names are made unique per test instead, because deleted registered models
|
|
leave tombstones `get_registered_model` still finds, which would make "second import of the same
|
|
model" unreadable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable, Iterator
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import mlflow
|
|
import pytest
|
|
import pytest_asyncio
|
|
import requests
|
|
from pymongo import MongoClient
|
|
from sqlalchemy import Engine, create_engine, text
|
|
from temporalio.client import WorkflowHandle
|
|
from temporalio.worker import Worker
|
|
from testcontainers.mongodb import MongoDbContainer
|
|
from testcontainers.postgres import PostgresContainer
|
|
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.import_model import ImportModel
|
|
from tests.helpers.bundle_factory import (
|
|
DEFAULT_ENVELOPE_KEY,
|
|
DEFAULT_PASSWORD,
|
|
ZipSpec,
|
|
build_envelope,
|
|
build_metadata_document,
|
|
default_tree,
|
|
envelope_key_b64,
|
|
make_bundle,
|
|
)
|
|
|
|
# --- constants -----------------------------------------------------------------------------------
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
FIXTURES_DIR = Path(__file__).resolve().parent / 'fixtures'
|
|
STATUS_TABLE_DDL = FIXTURES_DIR / 'experiment_run.sql'
|
|
STATUS_TABLE_DDL_V12 = FIXTURES_DIR / 'experiment_run_v12.sql'
|
|
GOLDEN_BUNDLE = REPO_ROOT / 'tests' / 'fixtures' / 'bundle' / 'golden_model_v3.sientia'
|
|
GOLDEN_BUNDLE_PASSWORD = 'Import1Model!Fixture'
|
|
|
|
BFF_DATABASE = 'sientia-core-mlops-bff'
|
|
BFF_USER = 'bff'
|
|
BFF_PASSWORD = 'bff'
|
|
|
|
# Lowercase and hyphens only: `chk_experiment_run_bucket_name_pattern` on the real table.
|
|
IMPORT_BUCKET = 'sientia-import'
|
|
IMPORT_PREFIX = 'imported_models/'
|
|
MODELS_COLLECTION = 'models'
|
|
MONGO_DATABASE = 'sientia'
|
|
IMPORT_TASK_QUEUE = 'import-model-e2e-queue'
|
|
|
|
# Deliberately gapped, and deliberately strings: the baseline migration stores `id` as a string and
|
|
# `max(int(id)) + 1` must not be confused with a count or with a lexicographic maximum ("9" > "30").
|
|
SEEDED_MODEL_IDS = ('1', '2', '7')
|
|
|
|
MLFLOW_HEALTH_TIMEOUT_SECONDS = 90.0
|
|
MLFLOW_HEALTH_POLL_SECONDS = 0.5
|
|
|
|
IMPORT_ACTIVITY_NAMES = (
|
|
'claim_import_status',
|
|
'record_import_status',
|
|
'record_import_names',
|
|
'record_import_terminal_status',
|
|
'report_import_status_write_failure',
|
|
'download_import_bundle',
|
|
'open_import_bundle',
|
|
'create_import_experiment',
|
|
'upload_import_artifacts',
|
|
'register_import_model_version',
|
|
'write_import_model_document',
|
|
'cleanup_import_files',
|
|
'ensure_import_bundle_retention',
|
|
)
|
|
|
|
|
|
# --- the BFF Postgres container ------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_bff_container() -> Iterator[PostgresContainer]:
|
|
"""The import log database, on a server of its own.
|
|
|
|
Separate from the shared `postgres_container` for two independent reasons, either sufficient:
|
|
the retry-recovery tests **stop and start this server**, which must not disturb the predictions
|
|
database every other e2e suite shares session scope with; and "the import touches no Postgres
|
|
database other than `sientia-core-mlops-bff`" is a stronger statement when the other database
|
|
lives on a different server.
|
|
|
|
`log_statement=all` is what makes the absence assertions possible: the server log is the only
|
|
vantage point that sees every client, including the psycopg pool `Storage` holds.
|
|
"""
|
|
container = PostgresContainer(
|
|
'postgres:15',
|
|
username=BFF_USER,
|
|
password=BFF_PASSWORD,
|
|
dbname=BFF_DATABASE,
|
|
).with_command('postgres -c log_statement=all')
|
|
container.start()
|
|
try:
|
|
yield container
|
|
finally:
|
|
container.stop()
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_bff_engine(import_bff_container: PostgresContainer) -> Iterator[Engine]:
|
|
"""A SQLAlchemy engine on the BFF container.
|
|
|
|
`pool_pre_ping` matters here rather than being habit: two tests stop the container underneath
|
|
this engine, and a stale pooled connection would fail the *fixture* instead of the assertion.
|
|
"""
|
|
engine = create_engine(import_bff_container.get_connection_url(), pool_pre_ping=True)
|
|
try:
|
|
yield engine
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
def _apply_ddl(engine: Engine, files: tuple[Path, ...]) -> None:
|
|
"""Drop `public.experiment_run` and rebuild it from `files`, in order."""
|
|
with engine.begin() as connection:
|
|
connection.execute(text('DROP TABLE IF EXISTS public.experiment_run CASCADE'))
|
|
for path in files:
|
|
# One transaction per file, because each file is one migration's worth of statements.
|
|
with engine.begin() as connection:
|
|
connection.execute(text(path.read_text(encoding='utf-8')))
|
|
|
|
|
|
@pytest.fixture
|
|
def import_status_table(import_bff_engine: Engine) -> Engine:
|
|
"""`public.experiment_run` as `V1`…`V12` leave it, rebuilt for this test.
|
|
|
|
Rebuilt rather than truncated: the pre-`V12` variant below has to be able to alternate with
|
|
this one inside a single session, and `TRUNCATE` cannot take a constraint away. The effect a
|
|
test sees is the one the spec asks for — the table holds only the row that test inserted.
|
|
"""
|
|
_apply_ddl(import_bff_engine, (STATUS_TABLE_DDL, STATUS_TABLE_DDL_V12))
|
|
return import_bff_engine
|
|
|
|
|
|
@pytest.fixture
|
|
def import_status_table_pre_v12(import_bff_engine: Engine) -> Engine:
|
|
"""The same table without `V12`: `run_type` present, the import statuses and pairing absent.
|
|
|
|
There is no second copy of the base schema — this applies the same file, and simply does not
|
|
apply the delta.
|
|
"""
|
|
_apply_ddl(import_bff_engine, (STATUS_TABLE_DDL,))
|
|
return import_bff_engine
|
|
|
|
|
|
InsertStatusRow = Callable[..., int]
|
|
|
|
|
|
@pytest.fixture
|
|
def import_insert_status_row(import_bff_engine: Engine) -> InsertStatusRow:
|
|
"""Insert one import log row the way the BFF will, and return its `id`.
|
|
|
|
Defaults mirror what the frontend sends: `PENDING`, `run_type='IMPORT'`, the object key in
|
|
`file_name`, and the uploader's digest inside `request_data` — which is where
|
|
`claim_import_status` looks for it.
|
|
"""
|
|
|
|
def insert(
|
|
*,
|
|
object_key: str,
|
|
expected_digest: str | None = None,
|
|
username: str = 'import.e2e@example.com',
|
|
experiment_name: str = 'import-e2e-placeholder',
|
|
status: str = 'PENDING',
|
|
run_type: str = 'IMPORT',
|
|
bucket_name: str = IMPORT_BUCKET,
|
|
request_data: dict[str, Any] | None = None,
|
|
orchestrator_response_data: dict[str, Any] | None = None,
|
|
) -> int:
|
|
payload = request_data if request_data is not None else {}
|
|
if expected_digest is not None and request_data is None:
|
|
payload = {'expected_digest': expected_digest}
|
|
now = datetime.now(UTC).replace(tzinfo=None)
|
|
statement = text(
|
|
'INSERT INTO public.experiment_run '
|
|
'(experiment_name, username, status, created_at, updated_at, bucket_name, file_name, '
|
|
' request_data, orchestrator_response_data, run_type) '
|
|
'VALUES (:experiment_name, :username, :status, :created_at, :updated_at, '
|
|
' :bucket_name, :file_name, CAST(:request_data AS JSONB), '
|
|
' CAST(:detail AS JSONB), :run_type) '
|
|
'RETURNING id'
|
|
)
|
|
with import_bff_engine.begin() as connection:
|
|
return int(
|
|
connection.execute(
|
|
statement,
|
|
{
|
|
'experiment_name': experiment_name,
|
|
'username': username,
|
|
'status': status,
|
|
'created_at': now,
|
|
'updated_at': now,
|
|
'bucket_name': bucket_name,
|
|
'file_name': object_key,
|
|
'request_data': json.dumps(payload),
|
|
'detail': (
|
|
json.dumps(orchestrator_response_data)
|
|
if orchestrator_response_data is not None
|
|
else None
|
|
),
|
|
'run_type': run_type,
|
|
},
|
|
).scalar_one()
|
|
)
|
|
|
|
return insert
|
|
|
|
|
|
ReadStatusRow = Callable[[int], dict[str, Any] | None]
|
|
|
|
|
|
@pytest.fixture
|
|
def import_read_status_row(import_bff_engine: Engine) -> ReadStatusRow:
|
|
"""Read one import log row as a plain mapping, or `None` when it is gone."""
|
|
|
|
def read(import_run_id: int) -> dict[str, Any] | None:
|
|
with import_bff_engine.begin() as connection:
|
|
row = (
|
|
connection.execute(
|
|
text('SELECT * FROM public.experiment_run WHERE id = :id'),
|
|
{'id': import_run_id},
|
|
)
|
|
.mappings()
|
|
.one_or_none()
|
|
)
|
|
return dict(row) if row is not None else None
|
|
|
|
return read
|
|
|
|
|
|
# --- MongoDB ------------------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_mongo_container() -> Iterator[MongoDbContainer]:
|
|
"""A real MongoDB for the `models` collection the import writes its document into."""
|
|
container = MongoDbContainer('mongo:7')
|
|
container.start()
|
|
try:
|
|
yield container
|
|
finally:
|
|
container.stop()
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_mongo_client(import_mongo_container: MongoDbContainer) -> Iterator[MongoClient]:
|
|
"""A client on the Mongo container, for seeding and for reading assertions back."""
|
|
client: MongoClient = MongoClient(import_mongo_container.get_connection_url())
|
|
try:
|
|
yield client
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def import_models_collection(import_mongo_client: MongoClient) -> Any:
|
|
"""The `models` collection, dropped and reseeded with gapped ids and a unique index on `id`.
|
|
|
|
The index is the assertion's teeth: without it "the id continues the collection" and "a name
|
|
already present is left alone" would both pass against a collection that quietly took a
|
|
duplicate.
|
|
"""
|
|
collection = import_mongo_client[MONGO_DATABASE][MODELS_COLLECTION]
|
|
collection.drop()
|
|
collection.insert_many(
|
|
[
|
|
{
|
|
'id': model_id,
|
|
'name': f'seeded_model_{model_id}',
|
|
'active': True,
|
|
'model_config': {
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'joblib',
|
|
'retention_minutes': 60,
|
|
'target': 'SE',
|
|
},
|
|
}
|
|
for model_id in SEEDED_MODEL_IDS
|
|
]
|
|
)
|
|
collection.create_index('id', unique=True)
|
|
return collection
|
|
|
|
|
|
# --- MinIO --------------------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_bucket(minio_container) -> Iterator[Any]:
|
|
"""The import bucket in the shared MinIO container, with its prefix emptied for this test.
|
|
|
|
The bucket is kept rather than recreated — a bucket carries the retention lifecycle rule the
|
|
import sets, and dropping it between tests would make that rule untestable.
|
|
"""
|
|
client = minio_container.get_client()
|
|
if not client.bucket_exists(IMPORT_BUCKET):
|
|
client.make_bucket(IMPORT_BUCKET)
|
|
_empty_prefix(client, IMPORT_BUCKET, IMPORT_PREFIX)
|
|
yield client
|
|
_empty_prefix(client, IMPORT_BUCKET, IMPORT_PREFIX)
|
|
|
|
|
|
def _empty_prefix(client: Any, bucket: str, prefix: str) -> None:
|
|
"""Remove every object under `prefix`, so a leftover cannot satisfy the next test."""
|
|
for obj in client.list_objects(bucket, prefix=prefix, recursive=True):
|
|
client.remove_object(bucket, obj.object_name)
|
|
|
|
|
|
UploadBundle = Callable[..., tuple[str, str]]
|
|
|
|
|
|
@pytest.fixture
|
|
def import_upload_bundle(import_bucket: Any) -> UploadBundle:
|
|
"""Upload bundle bytes under the import prefix; return `(object_key, hex_digest)`.
|
|
|
|
The digest is computed over the bytes handed in, which is what the frontend promises and what
|
|
gate 1 holds it to. A test that wants the digest to *disagree* passes its own.
|
|
"""
|
|
import io
|
|
|
|
def upload(payload: bytes, name: str | None = None) -> tuple[str, str]:
|
|
object_name = name or f'e2e_{uuid.uuid4().hex}.sientia'
|
|
object_key = f'{IMPORT_PREFIX}{object_name}'
|
|
import_bucket.put_object(
|
|
IMPORT_BUCKET,
|
|
object_key,
|
|
io.BytesIO(payload),
|
|
length=len(payload),
|
|
content_type='application/octet-stream',
|
|
)
|
|
return object_key, hashlib.sha256(payload).hexdigest()
|
|
|
|
return upload
|
|
|
|
|
|
# --- building bundles ----------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_envelope() -> Callable[[str], str]:
|
|
"""Wrap a bundle password the way the frontend will, under the key the worker is configured with.
|
|
|
|
The key is `bundle_factory.DEFAULT_ENVELOPE_KEY`, which `import_activities` puts in
|
|
`IMPORT_PASSWORD_KEY` — so an envelope built here is one the worker can actually open, and an
|
|
envelope built under any other key is a decryption-step failure, which is a test of its own.
|
|
"""
|
|
|
|
def build(password: str = DEFAULT_PASSWORD) -> str:
|
|
return build_envelope(password, DEFAULT_ENVELOPE_KEY)
|
|
|
|
return build
|
|
|
|
|
|
@pytest.fixture
|
|
def import_bundle_factory(import_names: dict[str, str]) -> Callable[..., bytes]:
|
|
"""Build a `.sientia` bundle carrying *this test's* unique model and experiment names.
|
|
|
|
Every knob the gate tests need is passed straight through to the shared factory — there is no
|
|
second implementation of the producer's header, KDF or framing here, which is the whole point of
|
|
reusing `tests/helpers/bundle_factory.py`.
|
|
|
|
The names default to `import_names` so two tests provisioning models in the same session never
|
|
observe each other's registry state (MLflow is deliberately not reset — design E13). A test that
|
|
*wants* the collision, like the second-import case, passes the same names twice.
|
|
"""
|
|
|
|
def build(
|
|
*,
|
|
password: str = DEFAULT_PASSWORD,
|
|
metadata: dict[str, Any] | None = None,
|
|
parameters: dict[str, Any] | None = None,
|
|
drop_metadata_field: str | None = None,
|
|
drop_top_level_key: str | None = None,
|
|
drop_parameter: str | None = None,
|
|
files: dict[str, bytes] | None = None,
|
|
include_metadata: bool = True,
|
|
symlinks: dict[str, str] | None = None,
|
|
absolute_entry: str | None = None,
|
|
traversal_entry: str | None = None,
|
|
bomb_bytes: int = 0,
|
|
extra_files: dict[str, bytes] | None = None,
|
|
**header_kwargs: Any,
|
|
) -> bytes:
|
|
document = build_metadata_document(
|
|
metadata={
|
|
'model_name': import_names['model_name'],
|
|
'experiment_name': import_names['experiment_name'],
|
|
**(metadata or {}),
|
|
},
|
|
parameters={'model_name': import_names['model_name'], **(parameters or {})},
|
|
drop_metadata_field=drop_metadata_field,
|
|
drop_top_level_key=drop_top_level_key,
|
|
drop_parameter=drop_parameter,
|
|
)
|
|
spec = ZipSpec(
|
|
files=default_tree() if files is None else files,
|
|
document=document,
|
|
include_metadata=include_metadata,
|
|
symlinks=symlinks or {},
|
|
absolute_entry=absolute_entry,
|
|
traversal_entry=traversal_entry,
|
|
bomb_bytes=bomb_bytes,
|
|
extra_files=extra_files or {},
|
|
)
|
|
return make_bundle(spec, password, **header_kwargs)
|
|
|
|
return build
|
|
|
|
|
|
# --- the MLflow tracking server -----------------------------------------------------------------
|
|
|
|
|
|
def _free_port() -> int:
|
|
"""A loopback port nothing is listening on, taken from an ephemeral bind."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
sock.bind(('127.0.0.1', 0))
|
|
return int(sock.getsockname()[1])
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_mlflow_server(tmp_path_factory: pytest.TempPathFactory) -> Iterator[dict[str, Any]]:
|
|
"""`mlflow server` from *this* venv, on a free loopback port.
|
|
|
|
A subprocess rather than a container because the venv pin **is** the production pin (2.10.1,
|
|
the same version the producer pins), so this tests the exact client/server pair that will run
|
|
in production — no image tag to drift out of step with `requirements.txt`.
|
|
|
|
sqlite for the backend store because the model registry is not available on a file-backed
|
|
tracking store, and `create_model_version` is half of what the provisioning activities do.
|
|
`--serve-artifacts` sends `log_artifacts` over HTTP through the server, as it goes in
|
|
production, while leaving the uploaded tree readable straight off the local filesystem.
|
|
|
|
Yields the host, the port and the artifacts directory. A server that never answers `/health`
|
|
fails *here*, with its own stderr attached, rather than making the first test fail obscurely.
|
|
"""
|
|
root = tmp_path_factory.mktemp('import_mlflow')
|
|
artifacts = root / 'mlartifacts'
|
|
artifacts.mkdir()
|
|
log_path = root / 'mlflow-server.log'
|
|
port = _free_port()
|
|
host = '127.0.0.1'
|
|
tracking_uri = f'http://{host}:{port}'
|
|
|
|
server: dict[str, Any] = {
|
|
'host': f'http://{host}',
|
|
'port': str(port),
|
|
'tracking_uri': tracking_uri,
|
|
'artifacts_dir': artifacts,
|
|
'root': root,
|
|
'log_path': log_path,
|
|
'process': _launch_mlflow(root, host, port, artifacts, log_path),
|
|
}
|
|
try:
|
|
_wait_for_mlflow(server['process'], tracking_uri, log_path)
|
|
yield server
|
|
finally:
|
|
_terminate_group(server['process'])
|
|
|
|
|
|
def _launch_mlflow(
|
|
root: Path, host: str, port: int, artifacts: Path, log_path: Path
|
|
) -> subprocess.Popen:
|
|
"""Start the server process. Separate from the fixture so one test can restart it.
|
|
|
|
The sqlite file and the artifacts directory are reused across a restart on purpose: a test that
|
|
stops the server to fail one activity must not also erase what earlier tests provisioned.
|
|
"""
|
|
with log_path.open('ab') as log_handle:
|
|
return subprocess.Popen( # noqa: S603 - fixed argv, no shell, all values built here
|
|
[
|
|
sys.executable,
|
|
'-m',
|
|
'mlflow',
|
|
'server',
|
|
'--host',
|
|
host,
|
|
'--port',
|
|
str(port),
|
|
'--backend-store-uri',
|
|
f'sqlite:///{root / "mlflow.db"}',
|
|
'--artifacts-destination',
|
|
str(artifacts),
|
|
'--serve-artifacts',
|
|
],
|
|
stdout=log_handle,
|
|
stderr=subprocess.STDOUT,
|
|
cwd=str(root),
|
|
# Its own process group: `mlflow server` is gunicorn, so it has children, and killing
|
|
# only the parent would leave workers holding the port for the rest of the session.
|
|
start_new_session=True,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def import_mlflow_stopped(import_mlflow_server: dict[str, Any]) -> Any:
|
|
"""A context manager that stops the tracking server and **always** brings it back.
|
|
|
|
This is design E7's second preference in its purest form: no patch, no mock — the dependency is
|
|
genuinely gone, which is the thing that actually happens in production when a pod is replaced.
|
|
|
|
The restart is in a `finally` because the server is session-scoped: a test that left it down
|
|
would fail every test after it, and the failure would point at the wrong place.
|
|
"""
|
|
|
|
@contextlib.contextmanager
|
|
def stopped() -> Iterator[None]:
|
|
_terminate_group(import_mlflow_server['process'])
|
|
try:
|
|
yield
|
|
finally:
|
|
import_mlflow_server['process'] = _launch_mlflow(
|
|
import_mlflow_server['root'],
|
|
'127.0.0.1',
|
|
int(import_mlflow_server['port']),
|
|
import_mlflow_server['artifacts_dir'],
|
|
import_mlflow_server['log_path'],
|
|
)
|
|
_wait_for_mlflow(
|
|
import_mlflow_server['process'],
|
|
import_mlflow_server['tracking_uri'],
|
|
import_mlflow_server['log_path'],
|
|
)
|
|
|
|
return stopped
|
|
|
|
|
|
def _wait_for_mlflow(process: subprocess.Popen, tracking_uri: str, log_path: Path) -> None:
|
|
"""Poll `/health` until it answers, or fail with what the server actually said."""
|
|
deadline = time.monotonic() + MLFLOW_HEALTH_TIMEOUT_SECONDS
|
|
while time.monotonic() < deadline:
|
|
if process.poll() is not None:
|
|
_terminate_group(process)
|
|
raise RuntimeError(
|
|
f'mlflow server exited with {process.returncode} before answering /health:\n'
|
|
f'{log_path.read_text(encoding="utf-8", errors="replace")}'
|
|
)
|
|
try:
|
|
if requests.get(f'{tracking_uri}/health', timeout=2).status_code == 200:
|
|
return
|
|
except requests.RequestException:
|
|
pass
|
|
time.sleep(MLFLOW_HEALTH_POLL_SECONDS)
|
|
_terminate_group(process)
|
|
raise RuntimeError(
|
|
f'mlflow server did not answer /health within {MLFLOW_HEALTH_TIMEOUT_SECONDS}s:\n'
|
|
f'{log_path.read_text(encoding="utf-8", errors="replace")}'
|
|
)
|
|
|
|
|
|
def _terminate_group(process: subprocess.Popen) -> None:
|
|
"""Stop the server and every gunicorn worker it forked."""
|
|
if process.poll() is not None:
|
|
return
|
|
with contextlib.suppress(ProcessLookupError, PermissionError):
|
|
os.killpg(os.getpgid(process.pid), signal.SIGTERM)
|
|
try:
|
|
process.wait(timeout=15)
|
|
except subprocess.TimeoutExpired:
|
|
with contextlib.suppress(ProcessLookupError, PermissionError):
|
|
os.killpg(os.getpgid(process.pid), signal.SIGKILL)
|
|
process.wait(timeout=15)
|
|
|
|
|
|
# --- per-test naming ----------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_names(request: pytest.FixtureRequest) -> dict[str, str]:
|
|
"""Experiment and model names unique to this test.
|
|
|
|
MLflow state is deliberately not reset between tests (design E13): deleting a registered model
|
|
leaves a tombstone `get_registered_model` still finds, which would make "second import of the
|
|
same model" read as a leak rather than as the thing it is testing. Unique names per test buy
|
|
the same isolation without the tombstones.
|
|
"""
|
|
token = uuid.uuid4().hex[:8]
|
|
stem = request.node.name.replace('[', '_').replace(']', '')[:24]
|
|
return {
|
|
'token': token,
|
|
'model_name': f'm_{stem}_{token}'[:40],
|
|
'experiment_name': f'e_{stem}_{token}'[:40],
|
|
}
|
|
|
|
|
|
# --- the import Activities ----------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_limits_override() -> dict[str, str]:
|
|
"""Extra import environment variables for this test, applied last so they win.
|
|
|
|
Empty by default. A test that needs a ceiling it can actually cross — gate 1's object size is
|
|
1 GiB, and uploading a gigabyte to prove a limit is not a test anybody runs twice — overrides
|
|
this fixture inside its own class, so the narrowed ceiling is scoped to the cases that need it
|
|
and every other test keeps the production defaults.
|
|
"""
|
|
return {}
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def import_activities(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
import_limits_override: dict[str, str],
|
|
tmp_path: Path,
|
|
postgres_container,
|
|
import_bff_container: PostgresContainer,
|
|
import_mongo_container: MongoDbContainer,
|
|
import_mlflow_server: dict[str, Any],
|
|
import_bucket: Any,
|
|
minio_container,
|
|
mock_logger,
|
|
notification_handler,
|
|
) -> Any:
|
|
"""An `Activities` whose import path talks to nothing fake.
|
|
|
|
Note what is *not* in the fixture's parameter list: `patch_mlflow`, `patch_minio_repository`
|
|
and `patch_create_engine`. Composing any of them would replace exactly the thing under test,
|
|
and because `mlflow.set_tracking_uri` is process-global the substitution would be silent — so
|
|
the live tracking URI is asserted below rather than assumed.
|
|
|
|
The import configuration is built by the production `build_import_config` /
|
|
`build_import_status_config` from environment variables set here, so the variable *names* are
|
|
covered too and not just the dictionaries they produce.
|
|
"""
|
|
from laborious.utils.connectors_config import build_import_config, build_import_status_config
|
|
|
|
minio_port = minio_container.get_exposed_port(9000)
|
|
|
|
monkeypatch.setenv('IMPORT_PASSWORD_KEY', envelope_key_b64(DEFAULT_ENVELOPE_KEY))
|
|
monkeypatch.setenv('IMPORT_BUNDLE_BUCKET', IMPORT_BUCKET)
|
|
monkeypatch.setenv('IMPORT_BUNDLE_PREFIX', IMPORT_PREFIX)
|
|
monkeypatch.setenv('IMPORT_MODELS_COLLECTION', MODELS_COLLECTION)
|
|
monkeypatch.setenv('IMPORT_WORK_DIR', str(tmp_path / 'import-work'))
|
|
monkeypatch.setenv('IMPORT_STATUS_DB_HOST', import_bff_container.get_container_host_ip())
|
|
monkeypatch.setenv('IMPORT_STATUS_DB_PORT', str(import_bff_container.get_exposed_port(5432)))
|
|
monkeypatch.setenv('IMPORT_STATUS_DB_NAME', BFF_DATABASE)
|
|
monkeypatch.setenv('IMPORT_STATUS_DB_USER', BFF_USER)
|
|
monkeypatch.setenv('IMPORT_STATUS_DB_PASSWORD', BFF_PASSWORD)
|
|
for name, value in import_limits_override.items():
|
|
monkeypatch.setenv(name, value)
|
|
|
|
activities = Activities(
|
|
# The predictions database, which the import must never touch. It is pointed at a real
|
|
# server on purpose: "no statement reached it" is only worth asserting when there is a log
|
|
# that would have shown one.
|
|
postgres_config={
|
|
'host': 'localhost',
|
|
'port': postgres_container.get_exposed_port(5432),
|
|
'user': 'test',
|
|
'password': 'test',
|
|
'dbname': 'test',
|
|
'min_connections': 1,
|
|
'max_connections': 5,
|
|
},
|
|
mlflow_config={
|
|
'host': import_mlflow_server['host'],
|
|
'port': import_mlflow_server['port'],
|
|
'username': 'e2e',
|
|
'password': 'e2e',
|
|
},
|
|
minio_config={
|
|
'endpoint_url': f'localhost:{minio_port}',
|
|
'access_key': 'minioadmin',
|
|
'secret_key': 'minioadmin',
|
|
'default_bucket': IMPORT_BUCKET,
|
|
'retention_hours': 24,
|
|
'secure': False,
|
|
},
|
|
opc_config={},
|
|
pi_web_api_config={
|
|
'base_url': 'http://localhost:8080',
|
|
'auth_type': 'bearer',
|
|
'auth_token': 'e2e',
|
|
},
|
|
logger=mock_logger,
|
|
notification_handler=notification_handler,
|
|
mongo_config={
|
|
'connection_string': import_mongo_container.get_connection_url(),
|
|
'database_name': MONGO_DATABASE,
|
|
},
|
|
import_config=build_import_config(),
|
|
import_status_config=build_import_status_config(),
|
|
)
|
|
|
|
# If a mocked-MLflow fixture ever leaks into an import test, this is where it stops.
|
|
assert mlflow.get_tracking_uri() == import_mlflow_server['tracking_uri'], (
|
|
f'MLflow tracking URI is {mlflow.get_tracking_uri()!r}, not the test server '
|
|
f'{import_mlflow_server["tracking_uri"]!r} — a mocked MLflow fixture has leaked in'
|
|
)
|
|
|
|
try:
|
|
yield activities
|
|
finally:
|
|
await activities.shutdown()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def import_worker(temporal_test_env, import_activities: Any) -> Any:
|
|
"""An `ImportModel` worker with the thirteen import activities, on the import queue."""
|
|
async with Worker(
|
|
temporal_test_env.client,
|
|
task_queue=IMPORT_TASK_QUEUE,
|
|
workflows=[ImportModel],
|
|
activities=[getattr(import_activities, name) for name in IMPORT_ACTIVITY_NAMES],
|
|
) as worker:
|
|
yield worker
|
|
|
|
|
|
# --- running an import --------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_build_input() -> Callable[..., dict[str, Any]]:
|
|
"""The workflow input, as the frontend assembles it."""
|
|
|
|
def build(
|
|
*,
|
|
import_run_id: int,
|
|
object_key: str,
|
|
expected_digest: str | None,
|
|
password_envelope: str,
|
|
bucket: str | None = IMPORT_BUCKET,
|
|
**extra: Any,
|
|
) -> dict[str, Any]:
|
|
payload: dict[str, Any] = {
|
|
'import_run_id': import_run_id,
|
|
'object_key': object_key,
|
|
'password_envelope': password_envelope,
|
|
}
|
|
if expected_digest is not None:
|
|
payload['expected_digest'] = expected_digest
|
|
if bucket is not None:
|
|
payload['bucket'] = bucket
|
|
payload.update(extra)
|
|
return payload
|
|
|
|
return build
|
|
|
|
|
|
@pytest.fixture
|
|
def import_execute(temporal_test_env, import_worker: Any) -> Callable[..., Any]:
|
|
"""Run one import to completion and hand back the handle, whatever the outcome.
|
|
|
|
The handle rather than the result, because half this suite asserts on *how* a run failed —
|
|
the error type Temporal recorded, the event history — and `execute_workflow` would have thrown
|
|
that away.
|
|
"""
|
|
|
|
async def execute(payload: dict[str, Any], workflow_id: str | None = None) -> WorkflowHandle:
|
|
handle = await temporal_test_env.client.start_workflow(
|
|
ImportModel.run,
|
|
payload,
|
|
id=workflow_id or f'import-e2e-{uuid.uuid4().hex}',
|
|
task_queue=IMPORT_TASK_QUEUE,
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
# The exception is the assertion's subject in the negative cases; the caller reads it
|
|
# back off the handle so a failure here is never mistaken for a harness error.
|
|
await handle.result()
|
|
return handle
|
|
|
|
return execute
|
|
|
|
|
|
@pytest.fixture
|
|
def import_history_events() -> Callable[[WorkflowHandle], Any]:
|
|
"""Every event of a run's history, flattened to the text the secrecy assertions scan.
|
|
|
|
Returns `(events, blob)`: the event list for structural assertions, and one string holding
|
|
every payload in the history for "this value appears nowhere" assertions.
|
|
"""
|
|
|
|
async def fetch(handle: WorkflowHandle) -> tuple[list[Any], str]:
|
|
events = [event async for event in handle.fetch_history_events()]
|
|
return events, '\n'.join(str(event) for event in events)
|
|
|
|
return fetch
|
|
|
|
|
|
# --- reading a failure ---------------------------------------------------------------------------
|
|
#
|
|
# Temporal wraps what actually went wrong: `WorkflowFailureError` -> `ActivityError` ->
|
|
# `ApplicationError`, and only the innermost one carries the sentence and the declared error type.
|
|
# `str(exception)` on the outermost is "Workflow execution failed", which is true and useless, so
|
|
# every negative assertion in this suite goes through these two functions instead.
|
|
|
|
|
|
def failure_chain(error: BaseException) -> list[BaseException]:
|
|
"""The exception and everything it was caused by, outermost first."""
|
|
chain: list[BaseException] = []
|
|
current: BaseException | None = error
|
|
while current is not None and not any(current is seen for seen in chain):
|
|
chain.append(current)
|
|
current = current.__cause__
|
|
return chain
|
|
|
|
|
|
def failure_text(error: BaseException) -> str:
|
|
"""Every link of the chain as `Type: message`, joined — what a substring assertion reads."""
|
|
return '\n'.join(f'{type(part).__name__}: {part}' for part in failure_chain(error))
|
|
|
|
|
|
def failure_types(error: BaseException) -> set[str]:
|
|
"""Every Python class name in the chain, plus every `ApplicationError.type` it declares.
|
|
|
|
The declared `type` is the one that survives serialisation into the event history, so it is
|
|
what "the run failed with the not-claimable error type" actually means.
|
|
"""
|
|
names: set[str] = set()
|
|
for part in failure_chain(error):
|
|
names.add(type(part).__name__)
|
|
declared = getattr(part, 'type', None)
|
|
if isinstance(declared, str):
|
|
names.add(declared)
|
|
return names
|
|
|
|
|
|
# --- the absence assertion ----------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def import_assert_nothing_provisioned(
|
|
import_mlflow_server: dict[str, Any],
|
|
import_models_collection: Any,
|
|
import_bucket: Any,
|
|
tmp_path: Path,
|
|
) -> Callable[..., None]:
|
|
"""Assert a rejected import left nothing behind, in one call over all four places.
|
|
|
|
One helper rather than four assertions per test, because "nothing was provisioned" is a single
|
|
claim and a test that checks three of its four halves reads as if it checked all four.
|
|
"""
|
|
|
|
def assert_nothing(
|
|
*,
|
|
experiment_name: str,
|
|
model_name: str,
|
|
object_key: str | None = None,
|
|
seeded_ids: tuple[str, ...] = SEEDED_MODEL_IDS,
|
|
) -> None:
|
|
client = mlflow.tracking.MlflowClient(tracking_uri=import_mlflow_server['tracking_uri'])
|
|
|
|
experiment = client.get_experiment_by_name(experiment_name)
|
|
assert experiment is None, f'experiment {experiment_name!r} was created'
|
|
|
|
versions = client.search_model_versions(f"name='{model_name}'")
|
|
assert not list(versions), f'registry holds versions of {model_name!r}'
|
|
|
|
documents = list(import_models_collection.find({}, {'_id': False}))
|
|
assert {str(document['id']) for document in documents} == set(seeded_ids), (
|
|
f'the models collection changed: {documents}'
|
|
)
|
|
assert not [d for d in documents if d.get('name') == model_name], (
|
|
f'a model document for {model_name!r} was written'
|
|
)
|
|
|
|
# The uploaded object itself stays — the frontend put it there and retention removes it.
|
|
# What must not be there is anything the *import* wrote under the prefix.
|
|
stored = {
|
|
obj.object_name
|
|
for obj in import_bucket.list_objects(
|
|
IMPORT_BUCKET, prefix=IMPORT_PREFIX, recursive=True
|
|
)
|
|
}
|
|
assert stored <= ({object_key} if object_key else set()), (
|
|
f'the import wrote objects under {IMPORT_PREFIX}: {stored}'
|
|
)
|
|
|
|
work_dir = tmp_path / 'import-work'
|
|
leftovers = list(work_dir.rglob('*')) if work_dir.exists() else []
|
|
assert not leftovers, f'the worker filesystem still holds {leftovers}'
|
|
|
|
return assert_nothing
|
|
|
|
|
|
# --- reading the golden fixture -----------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def import_golden_bundle() -> dict[str, Any]:
|
|
"""The producer-emitted bundle: its bytes, its digest and the password it was written with.
|
|
|
|
This is the only place in the suite that pays the producer's real Argon2id cost (256 MiB,
|
|
3 passes). Everything else runs the factory's cheap cost, which is the right split — the
|
|
reader honours whatever the header declares, up to its ceiling.
|
|
"""
|
|
payload = GOLDEN_BUNDLE.read_bytes()
|
|
return {
|
|
'bytes': payload,
|
|
'digest': hashlib.sha256(payload).hexdigest(),
|
|
'password': GOLDEN_BUNDLE_PASSWORD,
|
|
'path': GOLDEN_BUNDLE,
|
|
}
|