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.
365 lines
15 KiB
Python
365 lines
15 KiB
Python
"""Each provisioning step fails as itself: its own code, its own sentence, nothing half-built.
|
|
|
|
The four steps after the gates — experiment creation, artifact upload, registration, the model
|
|
document — each get one failure, and every one lands `IMPORT_ERROR` rather than
|
|
`ORCHESTRATOR_VALIDATION_ERROR`. That split is not cosmetic: it is the table's existing distinction
|
|
between "the person who submitted this can fix it" and "only support can", and a provisioning
|
|
failure is always the second kind.
|
|
|
|
**How each failure is caused**, in design E7's order of preference:
|
|
|
|
| Step | Cause | Preference |
|
|
|---|---|---|
|
|
| `experiment_creation` | the MLflow server is genuinely stopped | broken dependency |
|
|
| `artifact_upload` | targeted patch on `log_artifacts` | patch, justified in place |
|
|
| `registration` | targeted patch on `create_model_version` | patch, justified in place |
|
|
| `model_document` | a document already holding the id the activity will compute, against the real unique index | broken dependency |
|
|
|
|
Two patches rather than the one design E7 anticipated. The reason is the same for both and is
|
|
recorded in each test: with one server serving both tracking and registry, and MLflow validating
|
|
nothing about a registered model name except that it is non-empty (which gate 6 already requires),
|
|
there is no way to make the world hostile to *just* one of these two steps. Stopping the server
|
|
fails experiment creation first; a malformed name is not malformed to MLflow. The patch is on the
|
|
client call the step makes, not on the activity or the workflow, so what is being tested is still
|
|
"the workflow attributes this step's failure to this step".
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from temporalio.client import WorkflowFailureError
|
|
|
|
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX, SEEDED_MODEL_IDS
|
|
from laborious.activities.model_import_errors import IMPORT_FAILURE_CATALOG
|
|
from laborious.utils.bundle.steps import ImportStep
|
|
|
|
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
|
|
|
PLATFORM_ERROR_STATUS = 'IMPORT_ERROR'
|
|
FACTORY_PASSWORD = 'Factory1Model!Test'
|
|
|
|
|
|
def _as_dict(value: Any) -> dict[str, Any]:
|
|
return json.loads(value) if isinstance(value, str) else (value or {})
|
|
|
|
|
|
def _assert_platform_failure(
|
|
*,
|
|
row: dict[str, Any] | None,
|
|
step: ImportStep,
|
|
workflow_id: str,
|
|
object_key: str,
|
|
) -> None:
|
|
"""The shared claim of task 7.5, asserted by every case in this module.
|
|
|
|
`error_message` is the catalog sentence *exactly* — not "contains" — which is the strongest
|
|
possible form of "carries no path, table name, object key or traceback": there is nothing in it
|
|
that did not come from the catalog.
|
|
"""
|
|
assert row is not None, 'the failure lost the row it was recorded on'
|
|
assert row['status'] == PLATFORM_ERROR_STATUS, (
|
|
f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
|
|
)
|
|
|
|
code, reason = IMPORT_FAILURE_CATALOG[step]
|
|
detail = _as_dict(row['orchestrator_response_data'])['import']
|
|
assert detail['step'] == step.value, detail
|
|
assert detail['code'] == code.value, detail
|
|
assert detail['workflow_id'] == workflow_id, detail
|
|
assert row['error_message'] == reason
|
|
|
|
message = row['error_message'] or ''
|
|
for forbidden in (
|
|
object_key,
|
|
IMPORT_BUCKET,
|
|
IMPORT_PREFIX,
|
|
'experiment_run',
|
|
'Traceback',
|
|
'Error',
|
|
'mlflow',
|
|
'Mongo',
|
|
tempfile.gettempdir(),
|
|
):
|
|
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
|
|
|
|
|
|
async def _upload_and_start(
|
|
*,
|
|
import_bundle_factory: Any,
|
|
import_upload_bundle: Any,
|
|
import_insert_status_row: Any,
|
|
import_build_input: Any,
|
|
import_execute: Any,
|
|
import_envelope: Any,
|
|
**factory_kwargs: Any,
|
|
) -> tuple[Any, int, str]:
|
|
object_key, digest = import_upload_bundle(import_bundle_factory(**factory_kwargs))
|
|
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(FACTORY_PASSWORD),
|
|
)
|
|
)
|
|
return handle, import_run_id, object_key
|
|
|
|
|
|
class TestExperimentCreation:
|
|
"""The tracking server is not there. No patch, no mock — the process is stopped."""
|
|
|
|
async def test_a_stopped_tracking_server_fails_at_experiment_creation(
|
|
self,
|
|
import_status_table: Any,
|
|
import_models_collection: Any,
|
|
import_names: dict[str, str],
|
|
import_bundle_factory: 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,
|
|
import_mlflow_stopped: Any,
|
|
import_assert_nothing_provisioned: Any,
|
|
) -> None:
|
|
with import_mlflow_stopped():
|
|
handle, import_run_id, object_key = await _upload_and_start(
|
|
import_bundle_factory=import_bundle_factory,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
import_envelope=import_envelope,
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
row = import_read_status_row(import_run_id)
|
|
|
|
_assert_platform_failure(
|
|
row=row,
|
|
step=ImportStep.EXPERIMENT_CREATION,
|
|
workflow_id=handle.id,
|
|
object_key=object_key,
|
|
)
|
|
|
|
# The bundle *was* opened before MLflow was reached, so the names are on the row — which is
|
|
# what task 9.3 asks for, observed from a failure that stopped exactly there.
|
|
assert row is not None
|
|
assert row['experiment_name'] == import_names['experiment_name']
|
|
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
|
|
|
|
# Nothing in Mongo: the document is the last step and never ran.
|
|
documents = list(import_models_collection.find({}, {'_id': False}))
|
|
assert {str(d['id']) for d in documents} == set(SEEDED_MODEL_IDS)
|
|
import_assert_nothing_provisioned(
|
|
experiment_name=import_names['experiment_name'],
|
|
model_name=import_names['model_name'],
|
|
object_key=object_key,
|
|
)
|
|
|
|
|
|
class TestArtifactUpload:
|
|
"""The step design E7 names as the one a patch is needed for."""
|
|
|
|
async def test_a_failing_artifact_upload_is_recorded_as_its_own_step(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
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,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
"""Why a patch: with the server up and the bundle valid, nothing else fails only this step.
|
|
|
|
Hostile input cannot reach here — a bundle that gets past gate 7 has a readable tree, and
|
|
one that does not never reaches artifact upload. A broken dependency cannot isolate it
|
|
either: the only dependency this step has is the tracking server, and stopping that fails
|
|
experiment creation one step earlier. So the patch is on the single client call the step
|
|
makes, `log_artifacts`; everything around it — the run creation, the parameters, the
|
|
workflow's attribution of the failure — is the real code.
|
|
"""
|
|
client = import_activities.import_mlflow_repository.client
|
|
|
|
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
|
raise OSError('artifact store is unreachable')
|
|
|
|
monkeypatch.setattr(client, 'log_artifacts', refuse)
|
|
|
|
handle, import_run_id, object_key = await _upload_and_start(
|
|
import_bundle_factory=import_bundle_factory,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
import_envelope=import_envelope,
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
|
|
_assert_platform_failure(
|
|
row=import_read_status_row(import_run_id),
|
|
step=ImportStep.ARTIFACT_UPLOAD,
|
|
workflow_id=handle.id,
|
|
object_key=object_key,
|
|
)
|
|
|
|
# The experiment exists — it was the previous step and it succeeded. Nothing after does.
|
|
mlflow_client = _fresh_client(import_mlflow_server)
|
|
assert mlflow_client.get_experiment_by_name(import_names['experiment_name']) is not None
|
|
assert not list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
|
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
|
|
|
|
|
class TestRegistration:
|
|
"""The artifacts are stored and the model still cannot be registered."""
|
|
|
|
async def test_a_failing_registration_is_recorded_as_its_own_step(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
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,
|
|
import_read_status_row: Any,
|
|
) -> None:
|
|
"""Why a patch here too, contrary to design E7's expectation that this step needs none.
|
|
|
|
Hostile input was the first thing tried: a registered model name MLflow refuses. It does not
|
|
exist — `_validate_model_name` rejects only the empty string, which gate 6 already refuses,
|
|
and the sqlite backend store does not enforce the `VARCHAR(256)` the schema declares, so
|
|
even a 300-character name is accepted. A broken dependency was the second: one process
|
|
serves both tracking and the registry, so stopping it fails experiment creation three steps
|
|
earlier. The patch is on `create_model_version`, the one call this step is.
|
|
"""
|
|
client = import_activities.import_mlflow_repository.client
|
|
|
|
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
|
raise RuntimeError('registry write rejected')
|
|
|
|
monkeypatch.setattr(client, 'create_model_version', refuse)
|
|
|
|
handle, import_run_id, object_key = await _upload_and_start(
|
|
import_bundle_factory=import_bundle_factory,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
import_envelope=import_envelope,
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
|
|
_assert_platform_failure(
|
|
row=import_read_status_row(import_run_id),
|
|
step=ImportStep.REGISTRATION,
|
|
workflow_id=handle.id,
|
|
object_key=object_key,
|
|
)
|
|
|
|
# The run and its artifacts survive — "the model's files were stored, but the model itself
|
|
# could not be registered" is exactly what the catalog sentence for this step says.
|
|
mlflow_client = _fresh_client(import_mlflow_server)
|
|
experiment = mlflow_client.get_experiment_by_name(import_names['experiment_name'])
|
|
assert experiment is not None
|
|
runs = mlflow_client.search_runs([experiment.experiment_id])
|
|
assert len(runs) == 1
|
|
assert {a.path for a in mlflow_client.list_artifacts(runs[0].info.run_id)} >= {
|
|
'prediction_model',
|
|
'data_model',
|
|
}
|
|
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
|
|
|
|
|
class TestModelDocument:
|
|
"""The last step, failed by a real unique index on a real collection."""
|
|
|
|
async def test_a_duplicate_id_fails_at_model_document(
|
|
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:
|
|
"""A document already holding the id `max(int(id)) + 1` will compute, under another name.
|
|
|
|
Under another name on purpose: a document with the *same* name would make the activity skip
|
|
the insert entirely, which is a different behaviour (task 8.6) and would not fail. This one
|
|
collides only on `id`, so the insert is attempted and the index refuses it.
|
|
|
|
No patch and no mock: the failure is `pymongo`'s duplicate-key error from a real index,
|
|
which is the thing that would happen in production if two imports raced.
|
|
"""
|
|
collision_id = str(max(int(seeded) for seeded in SEEDED_MODEL_IDS) + 1)
|
|
import_models_collection.insert_one(
|
|
{'id': collision_id, 'name': 'a_different_model', 'active': True, 'model_config': {}}
|
|
)
|
|
|
|
handle, import_run_id, object_key = await _upload_and_start(
|
|
import_bundle_factory=import_bundle_factory,
|
|
import_upload_bundle=import_upload_bundle,
|
|
import_insert_status_row=import_insert_status_row,
|
|
import_build_input=import_build_input,
|
|
import_execute=import_execute,
|
|
import_envelope=import_envelope,
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
|
|
_assert_platform_failure(
|
|
row=import_read_status_row(import_run_id),
|
|
step=ImportStep.MODEL_DOCUMENT,
|
|
workflow_id=handle.id,
|
|
object_key=object_key,
|
|
)
|
|
|
|
# No document for this model, and the colliding one is untouched.
|
|
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
|
blocker = import_models_collection.find_one({'id': collision_id}, {'_id': False})
|
|
assert blocker == {
|
|
'id': collision_id,
|
|
'name': 'a_different_model',
|
|
'active': True,
|
|
'model_config': {},
|
|
}
|
|
|
|
# Everything before it stands: "the model was registered, but the settings ... could not be
|
|
# saved" is the catalog sentence, and this is what makes it true rather than merely stored.
|
|
mlflow_client = _fresh_client(import_mlflow_server)
|
|
versions = list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
|
assert len(versions) == 1
|
|
|
|
|
|
def _fresh_client(server: dict[str, Any]) -> Any:
|
|
"""An MlflowClient bound to the test server, independent of the one the activities hold."""
|
|
import mlflow
|
|
|
|
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|