Files
sientia-dataops-laborious-l…/e2e/test_import_model_cleanup.py
vitor-aignosi e95e65b72b QTZPOC-20
QTZPOC-20: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests.

This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides.

Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2
2026-08-26 15:45:33 -03:00

248 lines
10 KiB
Python

"""What the worker is left holding, what the bucket is left promising, and what the history reveals.
Three things an import must be true about after it finishes, whichever way it finished:
- **The worker's disk is empty of it.** No extracted tree, no downloaded object, no decrypted zip.
The bundle is decrypted onto local disk, so a leak here is a plaintext model sitting in a scratch
directory for the next import to trip over — and after a *failure* is exactly when that would go
unnoticed.
- **The bucket carries the retention rule.** The import never deletes the uploaded object: it is the
record of an attempt, and the lifecycle rule is what eventually removes it. That rule is applied
by the importer rather than assumed from a Helm value, so it is checkable from here.
- **The status record outlives the object**, because it is the audit trail and the object is not.
And the one claim that can only be made against a real Temporal history: no plaintext password
anywhere in it. The workflow carries an envelope it cannot open, and the activity that opens it
never returns what it found. That is the design; this reads the history and checks it.
"""
from __future__ import annotations
from typing import Any
import pytest
from temporalio.client import WorkflowFailureError
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX
from tests.helpers.bundle_factory import DEFAULT_ENVELOPE_KEY, envelope_key_b64
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
FACTORY_PASSWORD = 'Factory1Model!Test'
RETENTION_RULE_ID = 'sientia-import-bundle-expiry'
async def _run(
*,
payload: bytes,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute: Any,
import_envelope: Any,
password: str = FACTORY_PASSWORD,
) -> tuple[Any, int, str]:
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(password),
)
)
return handle, import_run_id, object_key
class TestWorkerFilesystem:
"""Nothing of the import is left on the worker — on the success path or on either failure path."""
@pytest.mark.parametrize(
('case', 'factory_kwargs', 'expect_success'),
[
('success', {}, True),
('rejected at decryption', {'drop_final_tag': True}, False),
('rejected at structure validation', {'include_metadata': False}, False),
],
)
async def test_the_work_directory_is_empty_afterwards(
self,
case: str,
factory_kwargs: dict[str, Any],
expect_success: bool,
import_status_table: Any,
import_models_collection: Any,
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute: Any,
import_envelope: Any,
tmp_path: Any,
) -> None:
"""Two different failure steps, because they clean up at different points.
A decryption failure leaves a partially written zip in the staging directory; a structure
failure has already extracted a whole tree. The reader unwinds both, and the cleanup activity
removes the scratch directory after either — but they are separate code paths, so asserting
one would not cover the other.
"""
handle, _, _ = await _run(
payload=import_bundle_factory(**factory_kwargs),
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,
)
if expect_success:
await handle.result()
else:
with pytest.raises(WorkflowFailureError):
await handle.result()
work_dir = tmp_path / 'import-work'
leftovers = sorted(str(path) for path in work_dir.rglob('*')) if work_dir.exists() else []
assert not leftovers, f'{case}: the worker still holds {leftovers}'
class TestRetention:
"""The rule the uploaded object's eventual removal depends on."""
async def test_the_import_bucket_carries_the_retention_rule(
self,
import_status_table: Any,
import_models_collection: Any,
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,
) -> None:
handle, _, object_key = await _run(
payload=import_bundle_factory(),
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 handle.result()
lifecycle = import_bucket.get_bucket_lifecycle(IMPORT_BUCKET)
assert lifecycle is not None, 'the bucket carries no lifecycle configuration at all'
rules = {rule.rule_id: rule for rule in lifecycle.rules}
assert RETENTION_RULE_ID in rules, (
f'no {RETENTION_RULE_ID!r} rule; the bucket has {sorted(rules)}'
)
rule = rules[RETENTION_RULE_ID]
assert rule.status == 'Enabled'
assert rule.rule_filter.prefix == IMPORT_PREFIX
assert rule.expiration.days == 7
# The import did not delete the uploaded object: it is the record of an attempt, and the
# rule above is what removes it later.
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
async def test_the_status_record_outlives_the_uploaded_object(
self,
import_status_table: Any,
import_models_collection: Any,
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,
) -> None:
"""Deleting the object is what the retention rule will eventually do; the verdict must survive.
The record is the audit trail. If reading a finished import's outcome depended on the bundle
still being in the bucket, every import would become unreadable after seven days.
"""
handle, import_run_id, object_key = await _run(
payload=import_bundle_factory(),
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 handle.result()
import_bucket.remove_object(IMPORT_BUCKET, object_key)
with pytest.raises(Exception): # noqa: B017 - MinIO's own not-found type is not the subject
import_bucket.stat_object(IMPORT_BUCKET, object_key)
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS'
assert row['file_name'] == object_key
assert row['run_name'] is not None
class TestEventHistorySecrecy:
"""The claim that only a real event history can settle."""
async def test_no_plaintext_password_or_key_appears_in_the_history(
self,
import_status_table: Any,
import_models_collection: Any,
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute: Any,
import_envelope: Any,
import_history_events: Any,
) -> None:
"""Dump every event of a completed import and search the whole thing.
Three values must be absent, and they are absent for three different reasons:
- the **bundle password** never crosses an activity boundary — `open_import_bundle` decrypts
the envelope inside itself and returns metadata, never the password;
- the **envelope key** is configuration, read from the environment inside the activity, so it
is never an argument to anything;
- the **decrypted password** is the first of these two by another name, and is searched for
separately because a future refactor that returned it "just for logging" would pass a test
that only looked for the literal input.
The envelope itself *is* in the history, and must be: it is how the password reaches the
activity. That it is there is asserted too, so this test cannot pass by finding nothing at
all — which is the way an assertion like this usually rots.
"""
password = 'Secret1Password!InHistory'
handle, _, _ = await _run(
payload=import_bundle_factory(password=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,
password=password,
)
await handle.result()
events, blob = await import_history_events(handle)
assert events, 'the history is empty'
assert password not in blob, 'the bundle password appears in the event history'
assert envelope_key_b64(DEFAULT_ENVELOPE_KEY) not in blob, (
'the envelope key appears in the event history'
)
assert DEFAULT_ENVELOPE_KEY.hex() not in blob.lower(), (
'the envelope key appears in the event history, hex-encoded'
)
# The negative assertions above are only meaningful if the history is actually readable and
# actually carries the envelope. Without this, an empty or opaque history would "pass".
assert 'password_envelope' in blob, (
'the history does not even carry the envelope, so the search above proved nothing'
)