vitor-aignosi
192834bde3
QTZPOC-20: Refactor model import workflow, introduce atomic model ID generation, and clarify MinIO expiry. Fuses bundle download, decryption, MLflow experiment creation, and artifact upload into a single activity to prevent cross-worker issues. Implements an atomic counter for MongoDB model IDs, replacing the max(id) + 1 scan to avoid collisions during concurrent imports. Adds comprehensive documentation for MinIO object expiry mechanisms and refines lifecycle rule application. Improves error handling with clearer responsibility classification and centralizes MLflow repository injection.
316 lines
13 KiB
Python
316 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',
|
|
'skip_predict_timestamp': True,
|
|
}
|
|
|
|
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)
|