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.
385 lines
17 KiB
Python
385 lines
17 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 real unique index refuses the activity's own insert | broken dependency |
|
|
|
|
Two departures from what the tasks anticipated, both recorded in the test that found them:
|
|
|
|
- **Registration needs a patch too**, not just artifact upload. One process serves both tracking and
|
|
the registry, so stopping it fails experiment creation three steps earlier; and MLflow validates
|
|
nothing about a registered model name except that it is non-empty, which gate 6 already requires.
|
|
Both patches are on the single client call the step *is*, never on the activity or the workflow, so
|
|
what is under test — the workflow attributing this step's failure to this step — is untouched.
|
|
- **The model document's unique-index violation is on `model_config.target`**, not on `id`. The
|
|
duplicate-`id` state the task describes cannot be constructed at all; `TestModelDocument`'s
|
|
docstring shows why, and the mechanism the task wanted — a real index refusing a real insert — is
|
|
kept intact.
|
|
"""
|
|
|
|
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.utils.bundle.steps import ImportStep
|
|
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
|
|
|
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, _responsibility = 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,
|
|
)
|
|
|
|
# Download, decryption, experiment creation and artifact upload are one fused activity
|
|
# attempt (`model-import-fix-bugs`, bug 8): the workflow only learns the bundle's names once
|
|
# that whole attempt returns, and an experiment-creation failure means it never does. The
|
|
# placeholder names inserted with the row are therefore still there — narrowed from the old
|
|
# "names before MLflow is touched" invariant, per `model-import-status-record`'s updated spec.
|
|
assert row is not None
|
|
assert row['experiment_name'] == 'import-e2e-placeholder'
|
|
|
|
# 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 refusing the real insert.
|
|
|
|
**Why the index is on `model_config.target` and not on `id`.** Task 7.4 asks for "a document
|
|
already holding the id the activity is about to compute". That state cannot be constructed, and
|
|
why is worth more than the test would have been: `_next_model_document_id` computes
|
|
`max(int(id)) + 1` over `distinct('id')` of the *same* collection, so every value that
|
|
stringifies to the colliding id is counted by that maximum. Put a document at `'8'` and the next
|
|
id becomes `'9'`. Make the id unparseable and it is ignored by the numbering *and* does not
|
|
collide. There is no third case — the numbering cannot collide with itself, and the only real
|
|
duplicate-key path on `id` is two imports racing, which this change lists as a non-goal. That is
|
|
a good property of the implementation, not a gap.
|
|
|
|
So the mechanism the task asked for is kept — a real unique index refusing the activity's own
|
|
insert — on a key that can actually collide. `model_config.target` is `'SE'` on the seeded
|
|
documents and `'SE'` on the one this import writes, so a unique index there makes the insert a
|
|
genuine `DuplicateKeyError` from the server. No patch, no mock: the same class of failure, from
|
|
the same enforcement mechanism, at the same step.
|
|
"""
|
|
|
|
async def test_a_unique_index_violation_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:
|
|
# All three seeded documents carry target 'SE', so the index has to be buildable first:
|
|
# give two of them their own target and leave exactly one holding 'SE'. That one is what the
|
|
# import's own document — also target 'SE' — collides with.
|
|
for ordinal, seeded_id in enumerate(sorted(SEEDED_MODEL_IDS)[1:], start=1):
|
|
import_models_collection.update_one(
|
|
{'id': seeded_id}, {'$set': {'model_config.target': f'OTHER_{ordinal}'}}
|
|
)
|
|
import_models_collection.create_index('model_config.target', unique=True, name='e2e_target')
|
|
try:
|
|
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,
|
|
parameters={'target_variable': 'SE'},
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
row = import_read_status_row(import_run_id)
|
|
finally:
|
|
# The collection is dropped and reseeded per test, but dropping the index here keeps the
|
|
# failure local even if that ever changes.
|
|
import_models_collection.drop_index('e2e_target')
|
|
|
|
_assert_platform_failure(
|
|
row=row,
|
|
step=ImportStep.MODEL_DOCUMENT,
|
|
workflow_id=handle.id,
|
|
object_key=object_key,
|
|
)
|
|
|
|
# No document for this model, and the seeded ones are exactly as they were.
|
|
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
|
assert {
|
|
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
|
} == set(SEEDED_MODEL_IDS)
|
|
|
|
# Everything before it stands — registered, and unusable, which is the point of the split
|
|
# between this step's sentence and the registration step's.
|
|
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'])
|