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.
215 lines
8.4 KiB
Python
215 lines
8.4 KiB
Python
"""The digest is a promise about the bytes in the bucket, and it is kept against the bucket.
|
|
|
|
Three ways the promise can be broken, and where each is caught:
|
|
|
|
- the uploader computed the digest over *other* bytes — caught at `download`, after the object is
|
|
fetched and before its header is parsed, so a wrong file never costs an Argon2id pass;
|
|
- the object was replaced between the upload and the run — the same failure, which is the point:
|
|
the check is against what is stored now, not against what was stored when the row was written;
|
|
- no digest at all — caught at `received`, before anything is downloaded, because a request with no
|
|
digest is not a request this workflow can honour.
|
|
|
|
The first two matter because MinIO is real here. A mocked object store would return whatever the
|
|
test told it to, and "the digest is enforced against the bytes that were uploaded" would be a
|
|
statement about the mock.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import io
|
|
import json
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from temporalio.client import WorkflowFailureError
|
|
|
|
from e2e.import_harness import IMPORT_BUCKET, failure_text
|
|
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]
|
|
|
|
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_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_step(row: dict[str, Any], step: ImportStep, workflow_id: str) -> None:
|
|
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
|
|
assert detail['code'] == code.value, detail
|
|
assert detail['workflow_id'] == workflow_id
|
|
assert row['error_message'] == reason
|
|
|
|
|
|
class TestWrongDigest:
|
|
"""A digest computed over bytes that are not the ones in the bucket."""
|
|
|
|
async def test_a_digest_from_another_file_fails_at_download(
|
|
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_assert_nothing_provisioned: Any,
|
|
) -> None:
|
|
"""The bundle is perfectly valid; only the digest is another file's.
|
|
|
|
The row carries the same wrong digest as the input, so the request and the record agree and
|
|
the `received` guard passes — otherwise this would be a `received` failure about a
|
|
contradictory request instead of the download-step digest check it is meant to exercise.
|
|
"""
|
|
payload = import_bundle_factory()
|
|
other_digest = hashlib.sha256(
|
|
import_bundle_factory(metadata={'run_id': 'other'})
|
|
).hexdigest()
|
|
assert other_digest != hashlib.sha256(payload).hexdigest()
|
|
|
|
object_key, _ = import_upload_bundle(payload)
|
|
import_run_id = import_insert_status_row(
|
|
object_key=object_key, expected_digest=other_digest
|
|
)
|
|
handle = await import_execute(
|
|
import_build_input(
|
|
import_run_id=import_run_id,
|
|
object_key=object_key,
|
|
expected_digest=other_digest,
|
|
password_envelope=import_envelope(FACTORY_PASSWORD),
|
|
)
|
|
)
|
|
with pytest.raises(WorkflowFailureError) as failure:
|
|
await handle.result()
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
|
|
|
# No header was parsed: the failure names the digest, not the format. If the digest check
|
|
# ran *after* the header the message here would be about magic bytes or a version.
|
|
text_of_failure = failure_text(failure.value).lower()
|
|
assert 'digest' in text_of_failure, text_of_failure
|
|
assert 'marker' not in text_of_failure and 'version' not in text_of_failure
|
|
|
|
# The names are still unwritten, which is the other half of "no header was parsed": the
|
|
# bundle was never opened, so nothing knew what to call it.
|
|
assert row['run_name'] is None
|
|
assert row['experiment_name'] == 'import-e2e-placeholder'
|
|
|
|
import_assert_nothing_provisioned(
|
|
experiment_name=import_names['experiment_name'],
|
|
model_name=import_names['model_name'],
|
|
object_key=object_key,
|
|
)
|
|
|
|
|
|
class TestOverwrittenObject:
|
|
"""The bytes changed after the row was written. Only a real object store can show this."""
|
|
|
|
async def test_replacing_the_object_after_upload_fails_at_download(
|
|
self,
|
|
import_status_table: Any,
|
|
import_models_collection: Any,
|
|
import_names: dict[str, str],
|
|
import_bundle_factory: Any,
|
|
import_bucket: 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_assert_nothing_provisioned: Any,
|
|
) -> None:
|
|
original = import_bundle_factory()
|
|
object_key, digest = import_upload_bundle(original)
|
|
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
|
|
|
# Same key, different bytes — also a valid bundle, so nothing but the digest can tell.
|
|
replacement = import_bundle_factory(metadata={'run_id': 'ff00ff00ff00ff00ff00ff00ff00ff00'})
|
|
assert hashlib.sha256(replacement).hexdigest() != digest
|
|
import_bucket.put_object(
|
|
IMPORT_BUCKET,
|
|
object_key,
|
|
io.BytesIO(replacement),
|
|
length=len(replacement),
|
|
content_type='application/octet-stream',
|
|
)
|
|
|
|
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),
|
|
)
|
|
)
|
|
with pytest.raises(WorkflowFailureError):
|
|
await handle.result()
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
|
import_assert_nothing_provisioned(
|
|
experiment_name=import_names['experiment_name'],
|
|
model_name=import_names['model_name'],
|
|
object_key=object_key,
|
|
)
|
|
|
|
|
|
class TestMissingDigest:
|
|
"""No digest is refused at the door, before a single byte is fetched."""
|
|
|
|
async def test_an_absent_digest_fails_at_received_with_nothing_downloaded(
|
|
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_assert_nothing_provisioned: Any,
|
|
tmp_path: Any,
|
|
) -> None:
|
|
object_key, digest = import_upload_bundle(import_bundle_factory())
|
|
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=None,
|
|
password_envelope=import_envelope(FACTORY_PASSWORD),
|
|
)
|
|
)
|
|
with pytest.raises(WorkflowFailureError) as failure:
|
|
await handle.result()
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
_assert_step(row, ImportStep.RECEIVED, handle.id)
|
|
assert 'expected_digest' in failure_text(failure.value)
|
|
|
|
# Nothing downloaded: the work directory was never even created.
|
|
work_dir = tmp_path / 'import-work'
|
|
assert not work_dir.exists() or not list(work_dir.rglob('*'))
|
|
import_assert_nothing_provisioned(
|
|
experiment_name=import_names['experiment_name'],
|
|
model_name=import_names['model_name'],
|
|
object_key=object_key,
|
|
)
|