vitor-aignosi
dac4f1718f
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.
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.activities.model_import_errors import IMPORT_FAILURE_CATALOG
|
|
from laborious.utils.bundle.steps import ImportStep
|
|
|
|
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 = 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,
|
|
)
|