QTZPOC-20
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.
This commit is contained in:
364
e2e/test_import_model_provisioning_failures.py
Normal file
364
e2e/test_import_model_provisioning_failures.py
Normal file
@@ -0,0 +1,364 @@
|
||||
"""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'])
|
||||
Reference in New Issue
Block a user