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.
197 lines
8.3 KiB
Python
197 lines
8.3 KiB
Python
"""Does the harness itself come up? Diagnosable separately from the suite that runs on it.
|
|
|
|
Four containers, an `mlflow server` subprocess and a Temporal worker is a lot of moving parts, and
|
|
when one of them fails to start the first *real* test fails somewhere deep inside an activity with
|
|
a message about something else. These tests exist so that a broken harness says so in its own
|
|
words: each one checks one piece is alive and does nothing with the import path.
|
|
|
|
The last test is the exception — it runs a real workflow, but with an input designed to be rejected
|
|
at the very first step, so what it proves is only that the wiring reaches the first activity.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import mlflow
|
|
import pytest
|
|
from sqlalchemy import text
|
|
from temporalio.client import WorkflowFailureError
|
|
|
|
from e2e.import_harness import (
|
|
BFF_DATABASE,
|
|
IMPORT_ACTIVITY_NAMES,
|
|
IMPORT_BUCKET,
|
|
IMPORT_PREFIX,
|
|
IMPORT_TASK_QUEUE,
|
|
SEEDED_MODEL_IDS,
|
|
failure_text,
|
|
failure_types,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
|
|
|
|
|
class TestContainers:
|
|
"""Each dependency is real, reachable, and the one the harness says it is."""
|
|
|
|
async def test_bff_container_holds_the_status_table(self, import_status_table: Any) -> None:
|
|
with import_status_table.begin() as connection:
|
|
database = connection.execute(text('SELECT current_database()')).scalar_one()
|
|
columns = {
|
|
row[0]
|
|
for row in connection.execute(
|
|
text(
|
|
'SELECT column_name FROM information_schema.columns '
|
|
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
|
)
|
|
)
|
|
}
|
|
assert database == BFF_DATABASE
|
|
assert {'id', 'status', 'run_type', 'request_data', 'orchestrator_response_data'} <= columns
|
|
|
|
async def test_bff_container_logs_every_statement(self, import_bff_container: Any) -> None:
|
|
"""`log_statement=all` — without it the absence assertions have no vantage point."""
|
|
result = import_bff_container.exec(
|
|
['psql', '-U', 'bff', '-d', BFF_DATABASE, '-tAc', 'SHOW log_statement']
|
|
)
|
|
assert b'all' in result.output
|
|
|
|
async def test_the_v12_delta_is_what_separates_the_two_tables(
|
|
self, import_status_table_pre_v12: Any
|
|
) -> None:
|
|
"""The pre-`V12` table has `run_type` and refuses the import terminal statuses."""
|
|
with import_status_table_pre_v12.begin() as connection:
|
|
constraints = {
|
|
row[0]
|
|
for row in connection.execute(
|
|
text(
|
|
'SELECT conname FROM pg_constraint '
|
|
"WHERE conrelid = 'public.experiment_run'::regclass"
|
|
)
|
|
)
|
|
}
|
|
columns = {
|
|
row[0]
|
|
for row in connection.execute(
|
|
text(
|
|
'SELECT column_name FROM information_schema.columns '
|
|
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
|
)
|
|
)
|
|
}
|
|
assert 'run_type' in columns
|
|
assert 'chk_experiment_run_terminal_status_run_type' not in constraints
|
|
|
|
async def test_the_seeded_collection_has_gapped_ids_and_a_unique_index(
|
|
self, import_models_collection: Any
|
|
) -> None:
|
|
ids = {
|
|
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
|
}
|
|
assert ids == set(SEEDED_MODEL_IDS)
|
|
unique_keys = {
|
|
tuple(index['key'].keys())
|
|
for index in import_models_collection.list_indexes()
|
|
if index.get('unique')
|
|
}
|
|
assert ('id',) in unique_keys
|
|
|
|
async def test_the_import_bucket_exists_and_its_prefix_is_empty(
|
|
self, import_bucket: Any
|
|
) -> None:
|
|
assert import_bucket.bucket_exists(IMPORT_BUCKET)
|
|
assert not list(
|
|
import_bucket.list_objects(IMPORT_BUCKET, prefix=IMPORT_PREFIX, recursive=True)
|
|
)
|
|
|
|
async def test_uploading_a_bundle_returns_a_key_and_its_digest(
|
|
self, import_upload_bundle: Any, import_bucket: Any
|
|
) -> None:
|
|
object_key, digest = import_upload_bundle(b'not a real bundle, just bytes')
|
|
assert object_key.startswith(IMPORT_PREFIX)
|
|
assert len(digest) == 64
|
|
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
|
|
|
|
async def test_the_mlflow_server_answers_and_serves_a_registry(
|
|
self, import_mlflow_server: dict[str, Any]
|
|
) -> None:
|
|
client = mlflow.tracking.MlflowClient(tracking_uri=import_mlflow_server['tracking_uri'])
|
|
experiment_id = client.create_experiment('harness_smoke_experiment')
|
|
run = client.create_run(experiment_id, run_name='harness-smoke-run')
|
|
client.create_registered_model('harness_smoke_model')
|
|
version = client.create_model_version(
|
|
'harness_smoke_model', f'runs:/{run.info.run_id}/prediction_model', run.info.run_id
|
|
)
|
|
assert version.version == '1'
|
|
assert version.current_stage in (None, 'None')
|
|
|
|
|
|
class TestActivitiesAndWorker:
|
|
"""The object under test is built from the real dependencies, and the worker registers it."""
|
|
|
|
async def test_activities_point_at_the_real_dependencies(
|
|
self, import_activities: Any, import_mlflow_server: dict[str, Any]
|
|
) -> None:
|
|
assert mlflow.get_tracking_uri() == import_mlflow_server['tracking_uri']
|
|
assert import_activities.import_config['bucket'] == IMPORT_BUCKET
|
|
assert import_activities.import_config['prefix'] == IMPORT_PREFIX
|
|
assert import_activities.status_config['dbname'] == BFF_DATABASE
|
|
assert import_activities.mongo_repository is not None
|
|
# Not a mock: a MagicMock would satisfy any attribute access, so the type is the assertion.
|
|
assert type(import_activities.import_minio_repository).__name__ == 'MinioRepository'
|
|
|
|
async def test_every_import_activity_is_registered(self, import_activities: Any) -> None:
|
|
for name in IMPORT_ACTIVITY_NAMES:
|
|
assert callable(getattr(import_activities, name)), f'{name} is missing'
|
|
|
|
async def test_the_worker_polls_the_import_queue(self, import_worker: Any) -> None:
|
|
assert import_worker.task_queue == IMPORT_TASK_QUEUE
|
|
|
|
async def test_a_workflow_input_reaches_the_first_activity(
|
|
self,
|
|
import_status_table: Any,
|
|
import_execute: Any,
|
|
import_build_input: Any,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
"""The claim runs, and the row it could not claim proves the wiring is live.
|
|
|
|
An id no row carries is the cheapest input that gets past workflow validation and into
|
|
`claim_import_status`, so a run that ends failed *there* means the client, the worker, the
|
|
workflow, the activity and the BFF connection are all in place. Nothing else is exercised.
|
|
"""
|
|
handle = await import_execute(
|
|
import_build_input(
|
|
import_run_id=987654321,
|
|
object_key=f'{IMPORT_PREFIX}nothing.sientia',
|
|
expected_digest='0' * 64,
|
|
password_envelope='not-a-real-envelope',
|
|
)
|
|
)
|
|
with pytest.raises(WorkflowFailureError) as failure:
|
|
await handle.result()
|
|
text_of_failure = failure_text(failure.value)
|
|
assert 'claim' in text_of_failure.lower()
|
|
assert '987654321' in text_of_failure
|
|
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
|
assert import_read_status_row(987654321) is None
|
|
|
|
async def test_the_insert_helper_writes_a_claimable_row(
|
|
self,
|
|
import_status_table: Any,
|
|
import_insert_status_row: Any,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
object_key = f'{IMPORT_PREFIX}claimable.sientia'
|
|
row_id = import_insert_status_row(object_key=object_key, expected_digest='a' * 64)
|
|
row = import_read_status_row(row_id)
|
|
assert row is not None
|
|
assert row['status'] == 'PENDING'
|
|
assert row['run_type'] == 'IMPORT'
|
|
assert row['file_name'] == object_key
|
|
request_data = row['request_data']
|
|
payload = json.loads(request_data) if isinstance(request_data, str) else request_data
|
|
assert payload['expected_digest'] == 'a' * 64
|