QTZPOC-20

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.
This commit is contained in:
2026-09-08 17:06:25 -03:00
parent e4fc12e529
commit 192834bde3
34 changed files with 1485 additions and 531 deletions

View File

@@ -102,6 +102,9 @@ IMPORT_BUCKET = 'sientia-import'
IMPORT_PREFIX = 'imported_models/'
MODELS_COLLECTION = 'models'
MONGO_DATABASE = 'sientia'
# Must track `IMPORT_MODEL_ID_COUNTERS_COLLECTION`'s default in connectors_config.py — the harness
# never sets that env var, so production's default is what the worker under test actually reads from.
MODEL_ID_COUNTERS_COLLECTION = 'model_id_counters'
IMPORT_TASK_QUEUE = 'import-model-e2e-queue'
# Deliberately gapped, and deliberately strings: the baseline migration stores `id` as a string and
@@ -119,10 +122,7 @@ IMPORT_ACTIVITY_NAMES = (
'record_import_terminal_status',
'report_import_status_write_failure',
'report_import_pipeline_failure',
'download_import_bundle',
'open_import_bundle',
'create_import_experiment',
'upload_import_artifacts',
'prepare_and_upload_import_bundle',
'register_import_model_version',
'write_import_model_document',
'cleanup_import_files',
@@ -331,6 +331,12 @@ def import_models_collection(import_mongo_client: MongoClient) -> Any:
The index is the assertion's teeth: without it "the id continues the collection" and "a name
already present is left alone" would both pass against a collection that quietly took a
duplicate.
New ids no longer come from scanning this collection — `_next_model_document_id` reads an
atomic counter in a separate `model_id_counters` collection instead (a different database, no
less), so it survives this drop untouched. Reset it here too, to the same seeded high-water
mark, or the next id handed out silently carries over from whatever an earlier test in the same
run already claimed.
"""
collection = import_mongo_client[MONGO_DATABASE][MODELS_COLLECTION]
collection.drop()
@@ -351,6 +357,13 @@ def import_models_collection(import_mongo_client: MongoClient) -> Any:
]
)
collection.create_index('id', unique=True)
counters = import_mongo_client[MONGO_DATABASE][MODEL_ID_COUNTERS_COLLECTION]
counters.replace_one(
{'_id': MODELS_COLLECTION},
{'_id': MODELS_COLLECTION, 'seq': max(int(seeded) for seeded in SEEDED_MODEL_IDS)},
upsert=True,
)
return collection

View File

@@ -39,7 +39,7 @@ def _as_dict(value: Any) -> dict[str, Any]:
def _assert_step(row: dict[str, Any], step: ImportStep, workflow_id: str) -> None:
code, reason = IMPORT_FAILURE_CATALOG[step]
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
detail = _as_dict(row['orchestrator_response_data'])['import']
assert row['status'] == VALIDATION_ERROR_STATUS, row['status']
assert detail['step'] == step.value, detail

View File

@@ -92,7 +92,7 @@ async def _assert_rejected(
f'terminal status is {row["status"]!r}, expected {VALIDATION_ERROR_STATUS!r}'
)
expected_code, expected_reason = IMPORT_FAILURE_CATALOG[step]
expected_code, expected_reason, _expected_responsibility = IMPORT_FAILURE_CATALOG[step]
detail = _as_dict(row['orchestrator_response_data'])['import']
assert detail['step'] == step.value, (
f'recorded step is {detail["step"]!r}, expected {step.value!r}'

View File

@@ -244,6 +244,7 @@ class TestFactoryBundle:
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'SE',
'skip_predict_timestamp': True,
}
row = import_read_status_row(import_run_id)

View File

@@ -69,7 +69,7 @@ def _assert_platform_failure(
f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
)
code, reason = IMPORT_FAILURE_CATALOG[step]
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
@@ -152,11 +152,13 @@ class TestExperimentCreation:
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.
# 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_names['experiment_name']
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
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}))

View File

@@ -34,6 +34,7 @@ CONTRACTED_MODEL_CONFIG_KEYS = {
'predict_flavor',
'retention_minutes',
'target',
'skip_predict_timestamp',
}
@@ -203,7 +204,7 @@ class TestModelDocumentContent:
assert document['active'] is False
assert set(document['model_config']) == CONTRACTED_MODEL_CONFIG_KEYS, (
f'model_config holds {sorted(document["model_config"])}, not exactly the four '
f'model_config holds {sorted(document["model_config"])}, not exactly the '
f'contracted keys'
)
assert document['model_config'] == {
@@ -211,6 +212,7 @@ class TestModelDocumentContent:
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'PRODUCT_QUALITY',
'skip_predict_timestamp': True,
}
# The document has exactly these fields and no others — nothing was smuggled in alongside.

View File

@@ -155,24 +155,30 @@ class TestOrdering:
import_envelope: Any,
import_read_status_row: Any,
) -> None:
"""Observed at the first call experiment creation makes, so nothing exists in MLflow yet.
"""Observed at the first call registration makes, so the experiment and run already exist.
Download, decryption, experiment creation and artifact upload are fused into one activity
attempt (`model-import-status-record`'s design, decision on the fused step): the names are
not known to the workflow, and so cannot be recorded, until that whole activity returns —
the guarantee is anchored to the two steps that follow it, registration and the model
document, not to the experiment or run, which by construction already exist by then.
This is the stronger form of the same claim task 7.1 makes from a failure: here the import
*succeeds*, and the names were still already written — so their presence is not an artefact
of the run having stopped.
*succeeds*, and the names were still already written before registration ran — so their
presence is not an artefact of the run having stopped.
"""
repository = import_activities.import_mlflow_repository
real_get_experiment = repository.get_experiment
client = import_activities.import_mlflow_repository.client
real_get_registered_model = client.get_registered_model
snapshots: list[dict[str, Any] | None] = []
observed_run_id: list[int] = [0]
def observing_get_experiment(*args: Any, **kwargs: Any) -> Any:
def observing_get_registered_model(*args: Any, **kwargs: Any) -> Any:
snapshots.append(import_read_status_row(observed_run_id[0]))
return real_get_experiment(*args, **kwargs)
return real_get_registered_model(*args, **kwargs)
object_key, digest = import_upload_bundle(import_bundle_factory())
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
repository.get_experiment = observing_get_experiment
client.get_registered_model = observing_get_registered_model
try:
handle = await import_execute(
import_build_input(
@@ -184,9 +190,9 @@ class TestOrdering:
)
await handle.result()
finally:
repository.get_experiment = real_get_experiment
client.get_registered_model = real_get_registered_model
assert snapshots, 'experiment creation never ran'
assert snapshots, 'registration never ran'
mid_flight = snapshots[0]
assert mid_flight is not None
assert mid_flight['status'] == RUNNING_STATUS
@@ -436,7 +442,7 @@ class TestInputAgainstTheRow:
row = import_read_status_row(import_run_id)
assert row is not None
code, reason = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
assert row['status'] == VALIDATION_ERROR_STATUS
assert row['error_message'] == reason
detail = _as_dict(row['orchestrator_response_data'])['import']