vitor-aignosi
e95e65b72b
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
185 lines
7.7 KiB
Python
185 lines
7.7 KiB
Python
"""One table, one database — proven from the servers' own logs, not from the application's side.
|
|
|
|
Decision D22 says an import writes to `public.experiment_run` and to nothing else relational: no
|
|
project table, no model table, no row in the predictions database. That claim is only worth making
|
|
if it is checked from somewhere that sees *every* client, and there is exactly one such place.
|
|
|
|
Why not the obvious alternatives:
|
|
|
|
- a SQLAlchemy `before_cursor_execute` listener sees the status engine and is blind to the psycopg
|
|
pool `Storage` holds — which is precisely the client that would betray the claim;
|
|
- `pg_stat_statements` is aggregated, so it cannot be scoped to the window of one import.
|
|
|
|
So both containers run with `log_statement=all` and the window is bracketed by a marker statement on
|
|
each server. The assertions are about *which relation names appear*, not about exact statement text,
|
|
because the text is the driver's business and changes with it.
|
|
|
|
The predictions database is deliberately real and reachable here: `import_activities` is built with
|
|
working credentials for it, so "no statement reached it" is a claim about the import's behaviour and
|
|
not about a missing connection.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from temporalio.client import WorkflowFailureError
|
|
|
|
from e2e.import_harness import names_no_relation, strip_literals
|
|
|
|
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
|
|
|
FACTORY_PASSWORD = 'Factory1Model!Test'
|
|
|
|
# Relations that exist in the predictions database and in nobody's import path.
|
|
FOREIGN_RELATIONS = (
|
|
'predictions_schema',
|
|
'laborious_data',
|
|
'transformed_data',
|
|
'predictions',
|
|
)
|
|
|
|
# Tables a reader might expect an import to want, and which decision D22 says do not exist.
|
|
RELATIONS_D22_DENIES = ('project', 'model_registry', 'models', 'import_run', 'model_project')
|
|
|
|
|
|
class TestOneTableOneDatabase:
|
|
"""A full import, from claim to terminal write, watched from both servers."""
|
|
|
|
async def test_every_statement_the_import_issues_names_experiment_run(
|
|
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_read_status_row: Any,
|
|
import_log_window: 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)
|
|
|
|
with import_log_window() as captured:
|
|
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),
|
|
)
|
|
)
|
|
await handle.result()
|
|
|
|
row = import_read_status_row(import_run_id)
|
|
assert row is not None
|
|
assert row['status'] == 'IMPORT_SUCCESS', (
|
|
'the window has to bracket a *working* import, or it proves nothing'
|
|
)
|
|
|
|
bff_statements = captured['bff']
|
|
assert bff_statements, 'the status database logged nothing at all — the window is wrong'
|
|
|
|
touching_a_relation = [
|
|
statement for statement in bff_statements if not names_no_relation(statement)
|
|
]
|
|
assert touching_a_relation, 'no relational statement was logged, so nothing was checked'
|
|
for statement in touching_a_relation:
|
|
assert 'experiment_run' in statement, (
|
|
f'a statement naming another relation reached the status database: {statement!r}'
|
|
)
|
|
|
|
# And no relation a reader might expect the feature to want appears anywhere in the window.
|
|
# Literals are stripped first: the import writes the bundle's own `model_project` value into
|
|
# the JSON detail, and `factory_project` is data, not a reference to a `project` table.
|
|
window_text = '\n'.join(strip_literals(s) for s in bff_statements).lower()
|
|
for relation in RELATIONS_D22_DENIES:
|
|
assert relation not in window_text, (
|
|
f'{relation!r} appears in the import window: D22 says no such table is written'
|
|
)
|
|
for relation in FOREIGN_RELATIONS:
|
|
assert relation not in window_text, f'{relation!r} appears in the import window'
|
|
|
|
async def test_the_predictions_database_sees_no_statement_from_an_import(
|
|
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_read_status_row: Any,
|
|
import_log_window: Any,
|
|
) -> None:
|
|
"""Not "no statement against our tables" — no statement at all.
|
|
|
|
The import holds a working connection to this server (`Activities` is built with one), so
|
|
this is about what the import chooses to do, not about what it is able to do.
|
|
"""
|
|
object_key, digest = import_upload_bundle(import_bundle_factory())
|
|
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
|
|
|
with import_log_window() as captured:
|
|
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),
|
|
)
|
|
)
|
|
await handle.result()
|
|
|
|
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
|
|
|
|
predictions_statements = [
|
|
statement for statement in captured['predictions'] if not names_no_relation(statement)
|
|
]
|
|
assert not predictions_statements, (
|
|
f'the import issued statements against the predictions database: '
|
|
f'{predictions_statements}'
|
|
)
|
|
|
|
async def test_a_rejected_import_also_touches_nothing_else(
|
|
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_log_window: Any,
|
|
) -> None:
|
|
"""The failure path writes more status rows than the success path, so it is worth its own case.
|
|
|
|
A rejected import runs the claim, the progress hints, cleanup and the terminal write — more
|
|
statements than a success, and every one of them still has to name the same table.
|
|
"""
|
|
object_key, digest = import_upload_bundle(import_bundle_factory(include_metadata=False))
|
|
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
|
|
|
with import_log_window() as captured:
|
|
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()
|
|
|
|
for statement in captured['bff']:
|
|
if names_no_relation(statement):
|
|
continue
|
|
assert 'experiment_run' in statement, statement
|
|
assert not [s for s in captured['predictions'] if not names_no_relation(s)]
|