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.
315 lines
13 KiB
Python
315 lines
13 KiB
Python
"""An import that works, twice over: once from the producer's own bytes and once from the factory.
|
|
|
|
Two origins, because they prove different things. The **golden** bundle is bytes the producer's code
|
|
emitted — it decrypts at the producer's real Argon2id cost (256 MiB, 3 passes) and carries the tree
|
|
the producer really writes, so it is the only case that would notice the producer's layout moving.
|
|
The **factory** bundle is synthesised here, cheaply, with names unique to the test; it is what the
|
|
rest of the suite runs on, so a happy path over it is what makes the negative cases' "and nothing
|
|
was provisioned" meaningful.
|
|
|
|
What "it worked" means, in full: the experiment exists, a run exists inside it carrying the derived
|
|
run name, the artifact tree is under the run with `prediction_model` and `data_model` at its root,
|
|
every parameter the bundle declared is logged plus the origin's identity under `source.*`, a
|
|
registry version exists, the MongoDB document exists, and the status row ends `IMPORT_SUCCESS`.
|
|
|
|
Two things a successful import must **not** have done are asserted here too, because they are the
|
|
difference between "imported" and "live": no version is in `Production`, and the document is
|
|
`active: false`. Promotion and activation are manual operator steps (QTZPOC-21).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import mlflow
|
|
import pytest
|
|
|
|
from e2e.import_harness import (
|
|
IMPORT_BUCKET,
|
|
IMPORT_PREFIX,
|
|
SEEDED_MODEL_IDS,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
|
|
|
GOLDEN_MODEL_NAME = 'fixture_model'
|
|
GOLDEN_EXPERIMENT_NAME = 'fixture_experiment'
|
|
GOLDEN_MODEL_VERSION = '3'
|
|
GOLDEN_TARGET = 'SE'
|
|
|
|
REQUIRED_ARTIFACT_ROOTS = {'prediction_model', 'data_model'}
|
|
|
|
|
|
def _mlflow_client(server: dict[str, Any]) -> Any:
|
|
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
|
|
|
|
|
def _as_dict(value: Any) -> dict[str, Any]:
|
|
"""`request_data` / `orchestrator_response_data` come back as a dict or as JSON text."""
|
|
return json.loads(value) if isinstance(value, str) else (value or {})
|
|
|
|
|
|
async def _run_import(
|
|
*,
|
|
payload: bytes,
|
|
import_upload_bundle: Any,
|
|
import_insert_status_row: Any,
|
|
import_build_input: Any,
|
|
import_execute: Any,
|
|
password: str,
|
|
import_envelope: Any,
|
|
) -> tuple[Any, int, str]:
|
|
"""Upload, insert the row the BFF would have inserted, run the workflow. Return the outcome."""
|
|
object_key, digest = import_upload_bundle(payload)
|
|
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
|
handle = await import_execute(
|
|
import_build_input(
|
|
import_run_id=import_run_id,
|
|
object_key=object_key,
|
|
expected_digest=digest,
|
|
password_envelope=import_envelope(password),
|
|
)
|
|
)
|
|
return handle, import_run_id, object_key
|
|
|
|
|
|
def _assert_provisioned(
|
|
*,
|
|
client: Any,
|
|
experiment_name: str,
|
|
model_name: str,
|
|
run_name: str,
|
|
expected_parameters: dict[str, str],
|
|
expected_source_fields: dict[str, str],
|
|
) -> tuple[str, Any]:
|
|
"""The whole positive claim about MLflow, in one place. Returns `(run_id, version)`."""
|
|
experiment = client.get_experiment_by_name(experiment_name)
|
|
assert experiment is not None, f'no experiment named {experiment_name!r}'
|
|
|
|
runs = [
|
|
run
|
|
for run in client.search_runs([experiment.experiment_id])
|
|
if run.info.run_name == run_name
|
|
]
|
|
assert len(runs) == 1, f'expected exactly one run named {run_name!r}, found {len(runs)}'
|
|
run = runs[0]
|
|
assert run.info.status == 'FINISHED'
|
|
|
|
artifact_roots = {artifact.path for artifact in client.list_artifacts(run.info.run_id)}
|
|
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, (
|
|
f'the run holds {sorted(artifact_roots)}, missing {sorted(REQUIRED_ARTIFACT_ROOTS - artifact_roots)}'
|
|
)
|
|
# `prediction_model` is a directory the legacy loader walks, not an opaque blob.
|
|
prediction_files = {
|
|
artifact.path.rsplit('/', 1)[-1]
|
|
for artifact in client.list_artifacts(run.info.run_id, 'prediction_model')
|
|
}
|
|
assert {'MLmodel', 'model.pkl'} <= prediction_files, prediction_files
|
|
|
|
for name, value in expected_parameters.items():
|
|
assert run.data.params.get(name) == value, (
|
|
f'parameter {name!r} is {run.data.params.get(name)!r}, expected {value!r}'
|
|
)
|
|
for name, value in expected_source_fields.items():
|
|
assert run.data.params.get(f'source.{name}') == value, (
|
|
f'source.{name} is {run.data.params.get(f"source.{name}")!r}, expected {value!r}'
|
|
)
|
|
|
|
versions = list(client.search_model_versions(f"name='{model_name}'"))
|
|
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
|
|
version = versions[0]
|
|
assert version.run_id == run.info.run_id
|
|
return run.info.run_id, version
|
|
|
|
|
|
class TestGoldenBundle:
|
|
"""The producer's own bytes, at the producer's own KDF cost."""
|
|
|
|
async def test_the_producer_artefact_imports_end_to_end(
|
|
self,
|
|
import_status_table: Any,
|
|
import_models_collection: Any,
|
|
import_golden_bundle: dict[str, Any],
|
|
import_mlflow_server: dict[str, Any],
|
|
import_upload_bundle: Any,
|
|
import_insert_status_row: Any,
|
|
import_build_input: Any,
|
|
import_execute: Any,
|
|
import_envelope: Any,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
handle, import_run_id, object_key = await _run_import(
|
|
payload=import_golden_bundle['bytes'],
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
password=import_golden_bundle['password'],
|
|
import_envelope=import_envelope,
|
|
)
|
|
result = await handle.result()
|
|
|
|
assert result['model_name'] == GOLDEN_MODEL_NAME
|
|
assert result['import_run_id'] == import_run_id
|
|
assert result['version'] == '1', 'the destination numbering starts at 1, not at the bundle'
|
|
|
|
run_id, version = _assert_provisioned(
|
|
client=_mlflow_client(import_mlflow_server),
|
|
experiment_name=GOLDEN_EXPERIMENT_NAME,
|
|
model_name=GOLDEN_MODEL_NAME,
|
|
run_name=f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}',
|
|
expected_parameters={'target_variable': GOLDEN_TARGET, 'nan_treatment': 'drop'},
|
|
expected_source_fields={
|
|
'model_name': GOLDEN_MODEL_NAME,
|
|
'experiment_name': GOLDEN_EXPERIMENT_NAME,
|
|
'model_version': GOLDEN_MODEL_VERSION,
|
|
'model_project': 'fixture_project',
|
|
},
|
|
)
|
|
assert result['run_id'] == run_id
|
|
|
|
document = import_models_collection.find_one({'name': GOLDEN_MODEL_NAME}, {'_id': False})
|
|
assert document is not None, 'no model document was written'
|
|
assert document['id'] == str(max(int(i) for i in SEEDED_MODEL_IDS) + 1)
|
|
assert document['model_config']['target'] == GOLDEN_TARGET
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
assert row['status'] == 'IMPORT_SUCCESS'
|
|
assert row['error_message'] is None
|
|
assert row['experiment_name'] == GOLDEN_EXPERIMENT_NAME
|
|
assert row['run_name'] == f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}'
|
|
assert row['file_name'] == object_key
|
|
assert row['bucket_name'] == IMPORT_BUCKET
|
|
assert object_key.startswith(IMPORT_PREFIX)
|
|
|
|
# Imported, not live.
|
|
assert version.current_stage in (None, 'None'), version.current_stage
|
|
assert document['active'] is False
|
|
|
|
|
|
class TestFactoryBundle:
|
|
"""A synthesised tree with names unique to this test — the shape the rest of the suite runs on."""
|
|
|
|
async def test_a_factory_bundle_imports_end_to_end(
|
|
self,
|
|
import_status_table: Any,
|
|
import_models_collection: Any,
|
|
import_names: dict[str, str],
|
|
import_bundle_factory: Any,
|
|
import_mlflow_server: dict[str, Any],
|
|
import_upload_bundle: Any,
|
|
import_insert_status_row: Any,
|
|
import_build_input: Any,
|
|
import_execute: Any,
|
|
import_envelope: Any,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
payload = import_bundle_factory(parameters={'target_variable': 'SE'})
|
|
handle, import_run_id, _ = await _run_import(
|
|
payload=payload,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
password='Factory1Model!Test',
|
|
import_envelope=import_envelope,
|
|
)
|
|
result = await handle.result()
|
|
|
|
assert result['model_name'] == import_names['model_name']
|
|
run_id, version = _assert_provisioned(
|
|
client=_mlflow_client(import_mlflow_server),
|
|
experiment_name=import_names['experiment_name'],
|
|
model_name=import_names['model_name'],
|
|
# DEFAULT_METADATA declares model_version '2'.
|
|
run_name=f'import-{import_names["model_name"]}-v2',
|
|
expected_parameters={'target_variable': 'SE', 'retrain': 'False'},
|
|
expected_source_fields={
|
|
'model_name': import_names['model_name'],
|
|
'experiment_name': import_names['experiment_name'],
|
|
'model_version': '2',
|
|
},
|
|
)
|
|
assert result['run_id'] == run_id
|
|
assert result['version'] == '1'
|
|
|
|
document = import_models_collection.find_one(
|
|
{'name': import_names['model_name']}, {'_id': False}
|
|
)
|
|
assert document is not None
|
|
assert document['model_config'] == {
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'joblib',
|
|
'retention_minutes': 60,
|
|
'target': 'SE',
|
|
}
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
assert row['status'] == 'IMPORT_SUCCESS'
|
|
detail = _as_dict(row['orchestrator_response_data'])
|
|
assert detail['import']['workflow_id'] == handle.id
|
|
|
|
assert version.current_stage in (None, 'None')
|
|
assert document['active'] is False
|
|
|
|
|
|
class TestLegacyLoaderResolution:
|
|
"""The registry source an import writes has to be one the *existing* loader can parse.
|
|
|
|
This is the assertion that could only ever be made against a real registry: the importer writes
|
|
the run's resolved artifact URI with `prediction_model` appended, because `get_model_run_id`
|
|
extracts the run id as `source.split('/')[2]` — an index that is the run id only in the resolved
|
|
form. A literal `runs:/<run_id>/prediction_model` would put `prediction_model` at index 2, and a
|
|
`models:/m-*` source would have none of the parts at all.
|
|
|
|
It is called with the stage the version is actually in rather than `Production`, because an
|
|
import deliberately promotes nothing — passing `Production` here would test the promotion step
|
|
this feature does not have.
|
|
"""
|
|
|
|
async def test_the_imported_source_resolves_to_the_import_run(
|
|
self,
|
|
import_status_table: Any,
|
|
import_models_collection: Any,
|
|
import_names: dict[str, str],
|
|
import_bundle_factory: Any,
|
|
import_activities: Any,
|
|
import_mlflow_server: dict[str, Any],
|
|
import_upload_bundle: Any,
|
|
import_insert_status_row: Any,
|
|
import_build_input: Any,
|
|
import_execute: Any,
|
|
import_envelope: Any,
|
|
) -> None:
|
|
payload = import_bundle_factory()
|
|
handle, _, _ = await _run_import(
|
|
payload=payload,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
password='Factory1Model!Test',
|
|
import_envelope=import_envelope,
|
|
)
|
|
result = await handle.result()
|
|
model_name = import_names['model_name']
|
|
|
|
client = _mlflow_client(import_mlflow_server)
|
|
version = next(iter(client.search_model_versions(f"name='{model_name}'")))
|
|
assert version.source.endswith('/prediction_model'), version.source
|
|
|
|
# The production method, on the production repository object, against the real registry.
|
|
resolved = import_activities.import_mlflow_repository.get_model_run_id(
|
|
model_name, stage=version.current_stage
|
|
)
|
|
assert resolved == result['run_id'], (
|
|
f'get_model_run_id resolved {resolved!r} from source {version.source!r}, '
|
|
f'but the import created run {result["run_id"]!r}'
|
|
)
|
|
|
|
artifact_roots = {artifact.path for artifact in client.list_artifacts(resolved)}
|
|
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, sorted(artifact_roots)
|