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:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}'
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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}))
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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']
|
||||
|
||||
Reference in New Issue
Block a user