Files
sientia-dataops-laborious-l…/e2e/test_import_model_gates.py
vitor-aignosi dac4f1718f QTZPOC-20
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.
2026-08-25 17:07:42 -03:00

676 lines
27 KiB
Python

"""One rejection per gate, against a real object, a real archive and a real filesystem.
The gates are unit-tested already. What only a real run can show is that a rejection is *recorded*
correctly and that it provisions nothing: the row ends `ORCHESTRATOR_VALIDATION_ERROR` naming that
gate's step, its catalog code and its catalog sentence, and no experiment, run, artifact, registry
version, model document or extracted file survives.
Every case here is hostile **input** — design E7's first preference. Nothing is patched, no
dependency is broken; the bundles are built by the shared factory and the gates do the rest. The
one exception to "input only" is the object-size ceiling, where the ceiling itself is narrowed
rather than a gigabyte being uploaded; that is a configuration change, not a patched code path, and
it is scoped to the one class that needs it.
`_assert_rejected` is where the shared claim of task 5.8 lives — terminal status, recorded step,
recorded code, catalog sentence, and `assert_nothing_provisioned` — so no case can quietly assert
four fifths of it.
"""
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,
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 {})
async def _run_rejected_import(
*,
payload: bytes,
password: str,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute: Any,
import_envelope: Any,
envelope_password: str | None = None,
) -> tuple[Any, int, str]:
"""Upload a hostile bundle, insert the row, run the import and expect it to fail."""
object_key, digest = import_upload_bundle(payload)
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(
envelope_password if envelope_password is not None else password
),
)
)
return handle, import_run_id, object_key
async def _assert_rejected(
*,
handle: Any,
import_run_id: int,
object_key: str,
step: ImportStep,
import_read_status_row: Any,
import_assert_nothing_provisioned: Any,
import_names: dict[str, str],
experiment_name: str | None = None,
model_name: str | None = None,
) -> dict[str, Any]:
"""The whole of task 5.8, in one call. Returns the row so a case can add to it."""
with pytest.raises(WorkflowFailureError):
await handle.result()
row = import_read_status_row(import_run_id)
assert row is not None, 'the rejection lost the row it was recorded on'
assert row['status'] == VALIDATION_ERROR_STATUS, (
f'terminal status is {row["status"]!r}, expected {VALIDATION_ERROR_STATUS!r}'
)
expected_code, expected_reason = 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}'
)
assert detail['code'] == expected_code.value, (
f'recorded code is {detail["code"]!r}, expected {expected_code.value!r}'
)
assert detail['workflow_id'] == handle.id
assert row['error_message'] == expected_reason, (
'error_message is not the catalog sentence for this step'
)
import_assert_nothing_provisioned(
experiment_name=experiment_name or import_names['experiment_name'],
model_name=model_name or import_names['model_name'],
object_key=object_key,
)
return row
class TestGate1ObjectLimits:
"""Gate 1 rejects before the object's body is fetched at all."""
@pytest.fixture
def import_limits_override(self) -> dict[str, str]:
"""A 1 KiB ceiling, so a factory bundle is over it.
The production ceiling is 1 GiB. Narrowing it is the only way to cross gate 1 without
uploading a gigabyte, and it changes a configured limit rather than a code path — the same
comparison runs, against a different number.
"""
return {'IMPORT_MAX_OBJECT_BYTES': '1024'}
async def test_an_oversized_object_is_rejected_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:
payload = import_bundle_factory()
assert len(payload) > 1024, 'the fixture bundle must exceed the narrowed ceiling'
handle, import_run_id, object_key = await _run_rejected_import(
payload=payload,
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.DOWNLOAD,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
class TestGate2Header:
"""Four header rejections, all landing on `download` — the step that says "not these bytes"."""
@pytest.mark.parametrize(
('case', 'header_kwargs'),
[
('wrong magic', {'magic': b'NOTSIENT'}),
('unsupported version', {'format_version': 9}),
('kdf memory above the ceiling', {'declared_memlimit': 512 * 1024 * 1024}),
('kdf passes above the ceiling', {'declared_opslimit': 9}),
('digest not of the chunk stream', {'digest_override': bytes(32)}),
],
)
async def test_a_malformed_header_is_rejected_at_download(
self,
case: str,
header_kwargs: dict[str, Any],
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:
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(**header_kwargs),
password=FACTORY_PASSWORD,
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,
)
row = await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.DOWNLOAD,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
# None of these may borrow the decryption sentence: the password is not implicated.
assert 'password' not in (row['error_message'] or '').lower(), case
class TestGate3Decryption:
"""The password, the tag and the framing — the three ways a stream refuses to open."""
async def test_a_wrong_password_is_rejected_at_decryption(
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 envelope opens — it is the *bundle* password inside it that is wrong.
Worth separating: an envelope the worker cannot open is also a decryption-step failure, but
for a different reason, and only this one proves the KDF actually ran and produced a key
the stream rejected.
"""
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(password=FACTORY_PASSWORD),
password=FACTORY_PASSWORD,
envelope_password='Wrong1Password!Entirely',
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.DECRYPTION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
async def test_a_tampered_chunk_is_rejected_at_decryption(
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:
"""One flipped byte inside the stream, with the header digest repaired to match.
Repairing the digest is what makes this gate 3 rather than gate 2: the file is
self-consistent, so only the AEAD can tell it was altered.
"""
from tests.helpers.bundle_factory import tamper_ciphertext
handle, import_run_id, object_key = await _run_rejected_import(
payload=tamper_ciphertext(import_bundle_factory()),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.DECRYPTION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
async def test_a_stream_without_its_final_tag_is_rejected_at_decryption(
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:
"""Every chunk framed as MESSAGE, so the stream claims more is coming and never ends.
The digest cannot see this — the bytes on disk are exactly the ones it covers. Only the
`secretstream` framing knows the file is incomplete.
"""
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(drop_final_tag=True),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.DECRYPTION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
class TestGate4ArchiveInspection:
"""The central directory alone decides — nothing has been written to disk yet."""
@pytest.fixture
def import_limits_override(self) -> dict[str, str]:
"""A 1 MiB declared-expansion ceiling, for the oversized-expansion case in this class."""
return {'IMPORT_MAX_UNCOMPRESSED_BYTES': str(1024 * 1024)}
@pytest.mark.parametrize(
('case', 'kwargs'),
[
('entry escaping the tree', {'traversal_entry': 'artifacts/../../escaped.csv'}),
('absolute entry path', {'absolute_entry': '/opt/absolute.csv'}),
('symbolic link', {'symlinks': {'artifacts/link.csv': 'artifacts/train_data.csv'}}),
('declared expansion over the ceiling', {'bomb_bytes': 4 * 1024 * 1024}),
],
)
async def test_a_hostile_archive_is_rejected_before_anything_is_extracted(
self,
case: str,
kwargs: dict[str, Any],
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:
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(**kwargs),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.ARCHIVE_INSPECTION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
# "Nothing extracted" is stronger than "nothing provisioned": not one member reached disk.
work_dir = tmp_path / 'import-work'
assert not list(work_dir.rglob('*')) if work_dir.exists() else True, case
class TestGate5Extraction:
"""A member that gate 4 passed and extraction still cannot place.
Gate 5 is about *placement*, not size: containment on the resolved path, and never overwriting
a file already written. A zip carrying the same name twice passes gate 4 — which checks names
and modes, not duplicates — and is refused here on the second member.
"""
async def test_a_duplicate_member_is_rejected_at_extraction(
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:
from tests.helpers.bundle_factory import default_tree
# `bomb_bytes` writes `artifacts/bomb.csv`, and `files` already carries that name, so the
# archive declares it twice. Small enough that no size ceiling is involved.
files = dict(default_tree())
files['artifacts/bomb.csv'] = b'first\n'
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(files=files, bomb_bytes=64),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.EXTRACTION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
class TestGate6Structure:
"""The layout and the metadata schema, checked on the extracted tree."""
@pytest.mark.parametrize(
('case', 'kwargs'),
[
('no metadata.json', {'include_metadata': False}),
('metadata missing run_id', {'drop_metadata_field': 'run_id'}),
('parameters missing target_variable', {'drop_parameter': 'target_variable'}),
],
)
async def test_an_incomplete_bundle_is_rejected_at_structure_validation(
self,
case: str,
kwargs: dict[str, Any],
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:
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(**kwargs),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.STRUCTURE_VALIDATION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
async def test_a_bundle_without_prediction_model_is_rejected(
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:
"""Without `prediction_model` there is nothing to register a version against."""
from tests.helpers.bundle_factory import default_tree
files = {
name: content
for name, content in default_tree().items()
if not name.startswith('artifacts/prediction_model/')
}
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(files=files),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.STRUCTURE_VALIDATION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
async def test_an_over_long_experiment_name_never_reaches_the_varchar_column(
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,
import_bff_engine: Any,
) -> None:
"""A 60-character experiment name is refused at gate 6, not by Postgres three steps later.
`experiment_name` is `VARCHAR(50)` and the import writes the bundle's own name onto the row
*before* provisioning. If the gate did not enforce the column's limit, the names write would
be a `DataError` — non-retryable, recorded as a status failure, on a row whose import was
otherwise fine. Asserting the row's `experiment_name` is untouched is what proves the write
never happened.
"""
from sqlalchemy import text
long_name = 'x' * 60
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(metadata={'experiment_name': long_name}),
password=FACTORY_PASSWORD,
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,
)
row = await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.STRUCTURE_VALIDATION,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
experiment_name=long_name,
)
assert row['experiment_name'] == 'import-e2e-placeholder', (
'the over-long name reached the row, so the gate did not hold the column'
)
with import_bff_engine.begin() as connection:
width = connection.execute(
text(
'SELECT character_maximum_length FROM information_schema.columns '
"WHERE table_name = 'experiment_run' AND column_name = 'experiment_name'"
)
).scalar_one()
assert width == 50, 'the column this gate exists to protect is no longer VARCHAR(50)'
class TestGate7ContentPolicy:
"""A file kind an MLflow artifact tree has no reason to carry."""
async def test_a_forbidden_file_type_is_rejected_at_content_policy(
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:
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(extra_files={'artifacts/payload.exe': b'MZ\x00\x00'}),
password=FACTORY_PASSWORD,
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,
)
await _assert_rejected(
handle=handle,
import_run_id=import_run_id,
object_key=object_key,
step=ImportStep.CONTENT_POLICY,
import_read_status_row=import_read_status_row,
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
import_names=import_names,
)
class TestTheRejectionSaysNothingTechnical:
"""`error_message` is a sentence for a person: no path, no object key, no traceback.
One case is enough to make the point structurally — the catalog is keyed by step and
`classify_import_failure` discards the exception's text entirely — and the per-gate cases above
each assert the sentence is *exactly* the catalog's, which is the same guarantee stated
positively.
"""
async def test_no_technical_detail_reaches_the_recorded_reason(
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,
) -> None:
handle, import_run_id, object_key = await _run_rejected_import(
payload=import_bundle_factory(include_metadata=False),
password=FACTORY_PASSWORD,
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) as failure:
await handle.result()
row = import_read_status_row(import_run_id)
assert row is not None
message = row['error_message'] or ''
for forbidden in (
object_key,
IMPORT_BUCKET,
IMPORT_PREFIX,
'experiment_run',
'Traceback',
'BundleGateError',
tempfile.gettempdir(),
'gate',
):
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
# The technical truth is not lost — it is on the Temporal error, where support can read it.
assert 'metadata' in failure_text(failure.value).lower()