gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
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
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-26 15:45:33 -03:00
parent dac4f1718f
commit e95e65b72b
32 changed files with 4247 additions and 1040 deletions
Download Patch File Download Diff File Expand all files Collapse all files

182
e2e/conftest.py
View File

@@ -2,9 +2,19 @@
Pytest configuration and fixtures for E2E tests.
"""
import importlib
import sys
from unittest.mock import AsyncMock, MagicMock, patch
# Load every BSON codec now, for its side effect only. `pymongo` imports `bson.objectid`,
# `bson.timestamp`, `bson.decimal128` and a dozen more lazily, and the first time is while an
# activity is running — after Temporal's workflow sandbox has snapshotted `sys.modules`. The sandbox
# then warns once per submodule ("Module bson.timestamp was imported after initial workflow load"),
# which was ~20 lines of noise per test that touches MongoDB, enough to bury a warning worth reading.
# This has to happen **before** the workflow modules are imported below, which is why it is here and
# not in `import_harness.py` — pytest loads `pytest_plugins` only after this module's body has run.
import bson.json_util # noqa: F401
# E2E workflows under test do not run ModelAnalysis; stub before Activities import.
_model_analysis_module = MagicMock()
_model_analysis_module.ModelAnalysis = MagicMock
@@ -18,14 +28,11 @@ import pytest
import pytest_asyncio
from sientia_do.notifications.handlers import CoreNotificationHandler
from sientia_do.observability.metrics_controller import MetricsController
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sqlalchemy import create_engine, text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from testcontainers.minio import MinioContainer
from testcontainers.postgres import PostgresContainer
from testcontainers.community.minio import MinioContainer
from testcontainers.community.postgres import PostgresContainer
from e2e.opc_test_server import OpcE2ETestServer
from laborious.activities.activities import Activities
@@ -42,9 +49,56 @@ pytest_plugins = ['e2e.import_harness']
TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017'
TEST_DATABASE_NAME = 'test_db'
# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised.
JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model'
JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact'
# Dedicated model served through the explicit `joblib` flavor, so no MLflow flavor loader runs.
JOBLIB_FLAVOR_MODEL_NAME = 'joblib_flavor_model'
# Module carrying the classes of the joblib artifacts' models. It ships inside the artifact's
# `code/` directory and is importable only while that directory is on `sys.path` — which is what
# the joblib flavor arranges before deserializing.
JOBLIB_FLAVOR_CODE_MODULE = 'joblib_flavor_probe'
JOBLIB_FLAVOR_WEIGHTS = {'feature_1': 2.0, 'feature_2': 3.0}
JOBLIB_FLAVOR_CODE_SOURCE = '''"""Model classes that exist only inside the artifact's `code/` directory."""
import pandas as pd
FEATURES = ('feature_1', 'feature_2')
class JoblibFlavorPredictor:
"""Prediction model: a weighted sum of the transform output columns."""
def __init__(self, weights):
self.weights = weights
def fit(self, data):
return self
def predict(self, data):
frame = pd.DataFrame(data)
return [
sum(row[name] * self.weights[name] for name in FEATURES)
for _, row in frame.iterrows()
]
class JoblibFlavorTransformer:
"""Transform model: emits the fixed feature frame the prediction model expects."""
target_variable = 'target'
def fit(self, data):
return self.predict(data)
def predict(self, data):
frame = pd.DataFrame(data)
result = pd.DataFrame(
{'feature_1': [0.234] * len(frame), 'feature_2': [0.783] * len(frame)}
)
result.index = frame.index
return result
'''
@pytest_asyncio.fixture(scope='session')
@@ -389,61 +443,78 @@ def mock_mlflow_models():
}
def _build_joblib_fallback_estimator():
"""
Build a real sklearn estimator for the joblib fallback artifact.
def _write_artifact_code_dir(artifact_dir):
"""Write the artifact's `code/` directory, carrying the module with the model classes."""
code_dir = artifact_dir / 'code'
code_dir.mkdir()
(code_dir / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').write_text(JOBLIB_FLAVOR_CODE_SOURCE)
return code_dir
`request_predict` appends a string `timestamp` column to the transformed
data, so the estimator selects the transform output columns explicitly and
drops the rest.
def _build_joblib_flavor_artifact(artifact_dir, class_name, *args):
"""
features = ['feature_1', 'feature_2']
estimator = Pipeline(
steps=[
(
'select',
ColumnTransformer([('features', 'passthrough', features)], remainder='drop'),
),
('regressor', LinearRegression()),
]
)
train = pd.DataFrame(
{
'feature_1': [0.0, 1.0, 2.0, 3.0],
'feature_2': [0.0, 1.0, 0.0, 1.0],
'timestamp': ['2024-01-01 12:00:00'] * 4,
}
)
estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3)
return estimator
Build one joblib artifact: `code/` with the class module, plus the compressed `model.pkl`.
Instantiating **and** dumping happen while `code/` is on `sys.path` — `pickle` refuses to
serialize a class it cannot import. Both are undone on the way out, so from here the pickled
class is unreachable until something puts the artifact's `code/` back on the path. That
something is the flavor under test.
"""
code_dir = _write_artifact_code_dir(artifact_dir)
sys.path.insert(0, str(code_dir))
try:
module = importlib.import_module(JOBLIB_FLAVOR_CODE_MODULE)
model = getattr(module, class_name)(*args)
joblib.dump(model, artifact_dir / 'model.pkl', compress=3)
finally:
sys.path.remove(str(code_dir))
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
importlib.invalidate_caches()
return {'path': str(artifact_dir), 'code_dir': str(code_dir), 'model': model}
@pytest.fixture
def joblib_fallback_artifact(tmp_path):
def joblib_flavor_artifacts(tmp_path):
"""
Real artifact directory holding a joblib-compressed sklearn estimator.
Real artifact directories for the `joblib` flavor, one per model type.
Mirrors what the joblib fallback finds after downloading `prediction_model`:
a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an
`MLmodel` placeholder.
Each holds a top-level `model.pkl` written by `joblib.dump(..., compress=3)` and a real
`code/` directory with the module where the model's class lives. No `MLmodel` is written:
the artifact this flavor reads — and the one a joblib retrain writes — has none.
Return:
dict with 'path' (the artifact directory) and 'estimator' (the object
that was serialized, so tests can compute the expected prediction).
dict keyed by artifact name ('prediction_model', 'data_model'), each value carrying
'path', 'code_dir' and 'model' (the object serialized, so tests can compute the
expected prediction).
"""
artifact_dir = tmp_path / 'prediction_model'
artifact_dir.mkdir()
original_sys_path = list(sys.path)
estimator = _build_joblib_fallback_estimator()
joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3)
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
builds = {
'prediction_model': ('JoblibFlavorPredictor', (JOBLIB_FLAVOR_WEIGHTS,)),
'data_model': ('JoblibFlavorTransformer', ()),
}
return {'path': str(artifact_dir), 'estimator': estimator}
artifacts = {}
for artifact_name, (class_name, args) in builds.items():
artifact_dir = tmp_path / artifact_name
artifact_dir.mkdir()
artifacts[artifact_name] = _build_joblib_flavor_artifact(artifact_dir, class_name, *args)
yield artifacts
# The flavor prepends `code/` to `sys.path` and deliberately never removes it.
sys.path[:] = original_sys_path
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
importlib.invalidate_caches()
@pytest.fixture
def joblib_fallback_artifact_without_pickle(tmp_path):
"""Artifact directory with no top-level `.pkl`, so the joblib fallback also fails."""
def joblib_flavor_artifact_without_pickle(tmp_path):
"""Artifact directory with no top-level `.pkl`, so the joblib flavor fails on it."""
artifact_dir = tmp_path / 'prediction_model_without_pickle'
artifact_dir.mkdir()
(artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n')
@@ -451,14 +522,12 @@ def joblib_fallback_artifact_without_pickle(tmp_path):
@pytest_asyncio.fixture
def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts):
"""Patch mlflow module in repository with load_model mocks."""
mock_mlflow = MagicMock()
# Mock sklearn.load_model
def mock_sklearn_load_model(model_uri):
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
if 'data_model' in model_uri or 'transform' in model_uri.lower():
return mock_mlflow_models['transform_model']
return mock_mlflow_models['predict_model']
@@ -468,8 +537,6 @@ def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
# Mock pyfunc.load_model
def mock_pyfunc_load_model(model_uri):
if JOBLIB_FALLBACK_MODEL_NAME in model_uri:
raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR)
if 'artifacts' in model_uri or 'tmp' in model_uri:
return mock_mlflow_models['pyfunc_model']
if 'data_model' in model_uri or 'transform' in model_uri.lower():
@@ -499,8 +566,13 @@ def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact):
]
),
)
# Only the joblib fallback calls download_artifacts directly on the client.
mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path'])
# Only the joblib flavor calls download_artifacts directly on the client, one artifact
# per model type — a single return value would serve `data_model` from `prediction_model`.
def mock_download_artifacts(_run_id, artifact_name):
return joblib_flavor_artifacts[artifact_name]['path']
mock_mlflow_client.download_artifacts = MagicMock(side_effect=mock_download_artifacts)
mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client)
with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow):

401
e2e/import_harness.py
View File

@@ -41,10 +41,12 @@ model" unreadable.
from __future__ import annotations
import asyncio
import contextlib
import hashlib
import json
import os
import re
import signal
import socket
import subprocess
@@ -52,7 +54,7 @@ import sys
import time
import uuid
from collections.abc import Callable, Iterator
from datetime import UTC, datetime
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any
@@ -63,9 +65,11 @@ import requests
from pymongo import MongoClient
from sqlalchemy import Engine, create_engine, text
from temporalio.client import WorkflowHandle
from temporalio.common import RetryPolicy
from temporalio.worker import Worker
from testcontainers.mongodb import MongoDbContainer
from testcontainers.postgres import PostgresContainer
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner, SandboxRestrictions
from testcontainers.community.mongodb import MongoDbContainer
from testcontainers.community.postgres import PostgresContainer
from laborious.activities.activities import Activities
from laborious.workflows.import_model import ImportModel
@@ -106,6 +110,7 @@ SEEDED_MODEL_IDS = ('1', '2', '7')
MLFLOW_HEALTH_TIMEOUT_SECONDS = 90.0
MLFLOW_HEALTH_POLL_SECONDS = 0.5
MONGO_SERVER_SELECTION_MS = 5000
IMPORT_ACTIVITY_NAMES = (
'claim_import_status',
@@ -139,13 +144,22 @@ def import_bff_container() -> Iterator[PostgresContainer]:
`log_statement=all` is what makes the absence assertions possible: the server log is the only
vantage point that sees every client, including the psycopg pool `Storage` holds.
The host port is **bound explicitly** rather than left ephemeral. Docker re-allocates an
ephemeral published port when a container is started again, so a stop/start would move the
database to a new address and the retry-recovery tests would be asserting that a *different*
server came up. A fixed binding is what makes "it came back" mean what it says.
"""
container = PostgresContainer(
'postgres:15',
username=BFF_USER,
password=BFF_PASSWORD,
dbname=BFF_DATABASE,
).with_command('postgres -c log_statement=all')
container = (
PostgresContainer(
'postgres:15',
username=BFF_USER,
password=BFF_PASSWORD,
dbname=BFF_DATABASE,
)
.with_command('postgres -c log_statement=all')
.with_bind_ports(5432, _free_port())
)
container.start()
try:
yield container
@@ -732,7 +746,13 @@ async def import_activities(
logger=mock_logger,
notification_handler=notification_handler,
mongo_config={
'connection_string': import_mongo_container.get_connection_url(),
# A short server-selection timeout so the "MongoDB is down" case fails in seconds
# instead of spending pymongo's 30-second default. Everything else about the client is
# the production one.
'connection_string': (
f'{import_mongo_container.get_connection_url()}'
f'?serverSelectionTimeoutMS={MONGO_SERVER_SELECTION_MS}'
),
'database_name': MONGO_DATABASE,
},
import_config=build_import_config(),
@@ -753,16 +773,208 @@ async def import_activities(
@pytest_asyncio.fixture
async def import_worker(temporal_test_env, import_activities: Any) -> Any:
"""An `ImportModel` worker with the thirteen import activities, on the import queue."""
"""An `ImportModel` worker with the thirteen import activities, on the import queue.
`bson` and `pymongo` are declared **passthrough** on the sandbox. Without that, the workflow
module's `imports_passed_through()` block pulls in `Activities` → `MongoDBRepository` → `pymongo`
during the first workflow activation, and `pymongo` loads its BSON codecs lazily
(`bson.objectid`, `bson.timestamp`, `bson.decimal128`, …). Each one is an import the sandbox has
not seen, so it warns "Module bson.X was imported after initial workflow load" — about twenty
lines per test that touches MongoDB.
Declaring them passthrough is the supported fix and the accurate one: these modules are pure
serialisation code with no workflow-relevant global state, which is exactly what passthrough is
for. Eagerly importing them on the host does **not** help — the sandbox judges imports against
its own module table during an activation, not against the host's `sys.modules`.
"""
async with Worker(
temporal_test_env.client,
task_queue=IMPORT_TASK_QUEUE,
workflows=[ImportModel],
activities=[getattr(import_activities, name) for name in IMPORT_ACTIVITY_NAMES],
workflow_runner=SandboxedWorkflowRunner(
restrictions=SandboxRestrictions.default.with_passthrough_modules('bson', 'pymongo')
),
) as worker:
yield worker
# --- the real-clock environment, for the retry-envelope tests ------------------------------------
#
# Design E9: everything else in this suite runs on the shared time-skipping environment. Two claims
# cannot: "the database came back *during* the retry envelope" needs the restart to occupy real
# time, and "the failure landed on the *first* attempt" is an assertion about elapsed attempts. Both
# run here instead, and this is the one place the harness deliberately weakens fidelity — so it says
# so, rather than hiding it in a fixture.
@pytest.fixture
def import_status_retry_shortening() -> dict[str, Any]:
"""How much to shorten the status retry policies for a real-clock test.
The **real** values (10 attempts, 1 s doubling to 30 s, a 300 s envelope) stay asserted where
they belong — unit test 12.1 of `import-model-workflow` pins them. What the tests here prove is
the recovery *behaviour*, which is identical at any interval, so waiting 151 seconds to observe
it would be waiting for arithmetic.
A constant half-second interval rather than a shortened doubling: it makes the envelope a number
a reader can multiply out (`attempts × interval`) instead of a geometric series.
"""
return {
'initial_interval': 0.5,
'backoff_coefficient': 1.0,
'maximum_attempts': 24, # 12 s of retrying
'hint_maximum_attempts': 2,
'envelope_seconds': 25,
'write_timeout_seconds': 5,
}
@pytest_asyncio.fixture
async def import_real_clock_env() -> Any:
"""A `WorkflowEnvironment` on the real clock."""
from temporalio.testing import WorkflowEnvironment
env = await WorkflowEnvironment.start_local()
async with env:
yield env
@pytest_asyncio.fixture
async def import_worker_real_clock(
monkeypatch: pytest.MonkeyPatch,
import_real_clock_env: Any,
import_status_retry_shortening: dict[str, Any],
import_activities: Any,
) -> Any:
"""An `ImportModel` worker on the real clock, with the status retry policies shortened.
Two harness choices worth knowing about:
- The policies are patched **on the workflow module** before the worker starts, because the
workflow reads them as module globals when it schedules an activity.
- The worker runs `UnsandboxedWorkflowRunner`. The sandbox imports the workflow module into a
namespace of its own, so a patch applied to the host module would simply not be seen and the
test would silently run the production 151-second envelope. Unsandboxed is what makes the
patch effective; nothing else about the workflow changes.
"""
from temporalio.worker import UnsandboxedWorkflowRunner
from laborious.workflows import import_model as workflow_module
shortening = import_status_retry_shortening
monkeypatch.setattr(
workflow_module,
'IMPORT_STATUS_RETRY_POLICY',
RetryPolicy(
maximum_attempts=int(shortening['maximum_attempts']),
initial_interval=timedelta(seconds=shortening['initial_interval']),
backoff_coefficient=shortening['backoff_coefficient'],
maximum_interval=timedelta(seconds=shortening['initial_interval']),
non_retryable_error_types=[
'ImportLogRowNotClaimableError',
'ImportStatusRejectedError',
],
),
)
monkeypatch.setattr(
workflow_module,
'IMPORT_STATUS_HINT_RETRY_POLICY',
RetryPolicy(
maximum_attempts=int(shortening['hint_maximum_attempts']),
initial_interval=timedelta(seconds=shortening['initial_interval']),
backoff_coefficient=1.0,
maximum_interval=timedelta(seconds=shortening['initial_interval']),
non_retryable_error_types=[
'ImportLogRowNotClaimableError',
'ImportStatusRejectedError',
],
),
)
monkeypatch.setattr(
workflow_module,
'STATUS_WRITE_ENVELOPE',
timedelta(seconds=shortening['envelope_seconds']),
)
monkeypatch.setattr(
workflow_module,
'STATUS_WRITE_TIMEOUT',
timedelta(seconds=shortening['write_timeout_seconds']),
)
async with Worker(
import_real_clock_env.client,
task_queue=IMPORT_TASK_QUEUE,
workflows=[ImportModel],
activities=[getattr(import_activities, name) for name in IMPORT_ACTIVITY_NAMES],
workflow_runner=UnsandboxedWorkflowRunner(),
) as worker:
yield worker
@pytest.fixture
def import_execute_real_clock(import_real_clock_env: Any, import_worker_real_clock: Any) -> Any:
"""Start an import on the real-clock environment and hand back the handle."""
async def execute(payload: dict[str, Any], workflow_id: str | None = None) -> WorkflowHandle:
return await import_real_clock_env.client.start_workflow(
ImportModel.run,
payload,
id=workflow_id or f'import-e2e-rc-{uuid.uuid4().hex}',
task_queue=IMPORT_TASK_QUEUE,
)
return execute
@pytest.fixture
def import_bff_stopped(import_bff_container: PostgresContainer, import_bff_engine: Engine) -> Any:
"""An **async** context manager that stops the status database and always brings it back.
`docker stop` / `docker start` on the wrapped container, not the testcontainers `stop()`, which
removes the container. The published port is bound explicitly by `import_bff_container`, so the
database comes back at the address the worker is already configured with.
Async, and every blocking call goes through `asyncio.to_thread`, for a reason that cost an
afternoon to find: the Temporal worker runs on this test's event loop. A synchronous
`docker stop` — or a plain `time.sleep` in the test body — starves the worker, so the workflow
makes no progress at all while the database is down and the test ends up asserting that a
*stalled* workflow behaved correctly. Nothing retries on a blocked loop.
The restart is in a `finally` and the engine is disposed on the way in: pooled connections do not
survive the server going away, and a stale one would fail the next test rather than this one.
"""
@contextlib.asynccontextmanager
async def stopped() -> Any:
wrapped = import_bff_container.get_wrapped_container()
await asyncio.to_thread(import_bff_engine.dispose)
await asyncio.to_thread(wrapped.stop, timeout=5)
try:
yield
finally:
await asyncio.to_thread(wrapped.start)
await asyncio.to_thread(_wait_for_postgres, import_bff_engine)
return stopped
def _wait_for_postgres(engine: Engine) -> None:
"""Block until the restarted server accepts a connection again."""
deadline = time.monotonic() + 60
last_error: Exception | None = None
while time.monotonic() < deadline:
try:
with engine.connect() as connection:
connection.execute(text('SELECT 1'))
return
except Exception as error: # noqa: BLE001 - any failure here means "not up yet"
last_error = error
engine.dispose()
time.sleep(0.5)
raise RuntimeError(f'the status database did not come back: {last_error!r}')
# --- running an import --------------------------------------------------------------------------
@@ -834,6 +1046,173 @@ def import_history_events() -> Callable[[WorkflowHandle], Any]:
return fetch
# --- the statement-log window --------------------------------------------------------------------
#
# "The import writes to one table in one database" cannot be asserted from the application side.
# A SQLAlchemy `before_cursor_execute` listener is blind to the psycopg pool `Storage` holds, which
# is exactly the client that would betray the claim; `pg_stat_statements` is aggregated, so it
# cannot be scoped to one import. The server's own log is the only vantage point that sees every
# client, which is why both containers run with `log_statement=all`.
_STATEMENT_MARKERS = ('statement:', 'execute <unnamed>:', 'execute ')
# A new log entry starts with the server's timestamp; anything else continues the previous one.
# Statements are logged across as many lines as they were written on, so without this a multi-line
# `UPDATE ... WHERE id = ...` would be judged on its first line alone.
_LOG_ENTRY_START = re.compile(r'^\d{4}-\d{2}-\d{2} ')
# Statements that name no *user* relation and are therefore not evidence either way.
_RELATIONLESS_PREFIXES = (
'begin',
'commit',
'rollback',
'set ',
'show ',
'discard',
'select 1',
'select current_database',
'select current_schema',
'select pg_backend_pid',
'select version()',
)
# System catalogs. Every driver reads these on connection setup — psycopg2 asks `pg_type` for its
# OID map before the application issues anything — so a statement whose only relations are catalogs
# is connection bookkeeping, not a table the feature chose to touch.
_CATALOG_IDENTIFIERS = (
'pg_catalog',
'pg_type',
'pg_namespace',
'pg_range',
'pg_class',
'pg_attribute',
'pg_proc',
'pg_settings',
'pg_database',
'information_schema',
'typarray',
)
def _container_log_text(container: Any) -> str:
"""Everything the container has written to stdout and stderr so far, as one string."""
streams = container.get_logs()
parts = streams if isinstance(streams, tuple) else (streams,)
return '\n'.join(
part.decode('utf-8', errors='replace') if isinstance(part, bytes) else str(part)
for part in parts
if part
)
def _statements_between(log_text: str, begin: str, end: str) -> list[str]:
"""Every logged statement between the two marker strings, in order.
Anchored on markers rather than on timestamps: the log is shared with every other client, and a
timestamp window would either include their statements or race with the clock.
"""
try:
start_index = log_text.index(begin)
except ValueError as error:
raise AssertionError(f'the opening marker {begin!r} is not in the log') from error
try:
end_index = log_text.index(end, start_index)
except ValueError as error:
raise AssertionError(f'the closing marker {end!r} is not in the log') from error
# Expand to whole lines. Slicing on the marker offsets themselves would cut the marker's own log
# line in half at both ends, and the tail half of the closing one reads like a real statement.
region_start = log_text.find('\n', start_index)
region_start = 0 if region_start == -1 else region_start + 1
region_end = log_text.rfind('\n', start_index, end_index)
region = log_text[region_start:region_end] if region_end > region_start else ''
statements: list[str] = []
collecting = False
for line in region.splitlines():
if _LOG_ENTRY_START.match(line):
collecting = False
for marker in _STATEMENT_MARKERS:
position = line.find(marker)
if position != -1:
statements.append(line[position + len(marker) :].strip())
collecting = True
break
elif collecting and statements:
# A continuation line of the statement above it.
statements[-1] = f'{statements[-1]} {line.strip()}'
return [statement for statement in statements if statement and begin not in statement]
_SINGLE_QUOTED = re.compile(r"'(?:[^']|'')*'")
def strip_literals(statement: str) -> str:
"""The statement with every single-quoted literal removed.
Necessary for any "this relation name does not appear" assertion: the import writes the bundle's
own `model_project` into the JSON detail, so the *value* `factory_project` would otherwise read
as a reference to a `project` table. A relation name appears outside a literal or not at all.
"""
return _SINGLE_QUOTED.sub("''", statement)
def names_no_relation(statement: str) -> bool:
"""True when the statement touches no *user* relation.
Two kinds: transaction and session bookkeeping (`BEGIN`, `SET`, `SELECT 1`), and the system
catalogs every driver reads when it connects. Neither is a table the feature chose to write.
"""
lowered = statement.strip().lower()
if any(lowered.startswith(prefix) for prefix in _RELATIONLESS_PREFIXES):
return True
return any(identifier in lowered for identifier in _CATALOG_IDENTIFIERS)
@pytest.fixture
def import_log_window(
import_bff_engine: Engine,
import_bff_container: PostgresContainer,
postgres_engine: Any,
postgres_container: Any,
) -> Any:
"""A context manager bracketing an import with markers on **both** Postgres servers.
On exit it fills the dict it yielded with `bff` and `predictions` — the statements each server
logged between the markers. A marker is written to both servers so each window is bounded on the
server it is about, rather than one server's window being inferred from the other's clock.
"""
@contextlib.contextmanager
def window() -> Iterator[dict[str, list[str]]]:
token = uuid.uuid4().hex[:12]
begin = f'e2e_window_begin_{token}'
end = f'e2e_window_end_{token}'
captured: dict[str, list[str]] = {}
_emit_marker(import_bff_engine, begin)
_emit_marker(postgres_engine, begin)
try:
yield captured
finally:
_emit_marker(import_bff_engine, end)
_emit_marker(postgres_engine, end)
captured['bff'] = _statements_between(
_container_log_text(import_bff_container), begin, end
)
captured['predictions'] = _statements_between(
_container_log_text(postgres_container), begin, end
)
return window
def _emit_marker(engine: Engine, marker: str) -> None:
"""A statement whose only purpose is to be findable in the server's log."""
with engine.begin() as connection:
connection.execute(text(f"SELECT '{marker}' AS marker")) # noqa: S608 - a generated uuid
# --- reading a failure ---------------------------------------------------------------------------
#
# Temporal wraps what actually went wrong: `WorkflowFailureError` -> `ActivityError` ->

642
e2e/scenarios.md
View File

@@ -4,12 +4,71 @@ This document describes all possible test scenarios for the `predictions_batch`
## Running automated E2E tests (`e2e/`)
### One-time setup: disable the testcontainers reaper
Do this once per machine, before running anything under `e2e/`:
```bash
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
```
**Why.** Before starting any test container, testcontainers starts a reaper container ("Ryuk") whose
job is to clean up containers if the test process dies without running its teardown. Ryuk bind-mounts
the Docker socket. Under **Docker Desktop on Linux** the socket lives at
`~/.docker/desktop/docker.sock`, which Docker Desktop's File Sharing does not expose, so the mount is
denied and the reaper cannot start:
```
docker.errors.APIError: 500 Server Error ... /containers/<id>/start:
Mounts denied: The path /socket_mnt/home/<user>/.docker/desktop/docker.sock
is not shared from the host and is not known to Docker.
```
Every test then errors **at setup** — `ERROR`, not `FAILED`, and never inside an assertion. If you see
that, this is why.
**What you give up.** Nothing in a normal run: every fixture in `conftest.py` and
`import_harness.py` stops its own containers in teardown. Ryuk only covers the case where the process
is killed without teardown (`kill -9`, a crash, a hard stop from the IDE). After one of those, clean
up by hand:
```bash
docker ps -a --filter "label=org.testcontainers=true" -q | xargs -r docker rm -f
```
The setting is read by `testcontainers/core/config.py`, which accepts either `ryuk.disabled` in
`~/.testcontainers.properties` or the `TESTCONTAINERS_RYUK_DISABLED` environment variable. The
properties file is preferred here because it is per-machine and applies to every project and every
shell, so nobody has to remember a prefix on the command line.
*Alternative, if you want to keep the reaper:* point Docker at the classic socket with
`export DOCKER_HOST=unix:///var/run/docker.sock`. That only works if a native Docker Engine is
running alongside Docker Desktop; if the Desktop daemon is the only one, disabling Ryuk is the way.
### Requirements and invocation
- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers.
- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`).
- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio,mongodb]`).
- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`.
- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios.
- **OPC tests (real server)**: `e2e/test_opc_real_server.py` uses an in-process **asyncua** server and real `OpcRepository` (`test_activities_real_opc`). Scenarios 3.1.2, 3.2.2, 3.2.4, and 3.2.5 are covered there. Other E2E modules keep the OPC mock.
- Run only OPC real-server tests: `pytest e2e/test_opc_real_server.py -m "integration and opc"`.
- **Model import tests**: the `import_model` suite is opt-in and has its own harness, containers and
fidelity contract — see [§4.0](#40-running-the-import-e2e-tests). Run it with
`pytest e2e -m import_e2e`.
- **joblib flavor tests**: `e2e/test_predictions_batch_joblib_flavor.py` covers reading a model
through the explicit `joblib` flavor (predict and transform), and
`e2e/test_retrain_joblib_flavor.py` covers writing one back on a retrain — see
[§5](#5-joblib-flavor-scenarios). Neither `joblib.dump` nor `joblib.load` is mocked in either
module: the artifact fixture holds a real `joblib.dump(..., compress=3)` pickle whose class ships
inside the artifact's `code/` directory, so only the flavor's own `sys.path` handling makes it
deserialize. Run them with
`pytest e2e/test_predictions_batch_joblib_flavor.py e2e/test_retrain_joblib_flavor.py -m integration`.
- **Nothing under `e2e/` is covered by `validate`**: that script lints `laborious/` and `tests/`
only, and its test step is `pytest tests/` (`testpaths = ["tests"]`). Lint these files by hand
(`ruff format --check e2e/ && ruff check e2e/`), and remember that `addopts` carries
`--strict-markers`: a test without `@pytest.mark.integration` is not collected, and a run that
collects nothing still exits 0. Check the collected count, not the exit code.
## Workflow Overview
@@ -561,4 +620,583 @@ These paths do **not** rely on Temporal activity retries for export failures: th
- Error notification sent
- Workflow completes
---
---
## 4. Model Import Workflow Scenarios (`import_model`)
The `import_model` workflow turns an encrypted `.sientia` bundle already in MinIO into a registered
MLflow model plus one MongoDB `models` document — or into a durable error on the row the BFF
inserted. It shares no code path with the predictions workflows above and has its own harness.
### 4.0 Running the import E2E tests
- **Marker**: every test below carries `import_e2e`; the real-clock retry tests also carry `slow`.
`testpaths` stays `["tests"]`, so none of this is collected by the unit gate.
- **Invocation**: `pytest e2e -m import_e2e`.
- **Runtime**: Docker for three containers (the shared MinIO, a dedicated Postgres for
`sientia-core-mlops-bff`, MongoDB) **plus** an `mlflow server` subprocess started from this
repository's venv. Dev requirements install `testcontainers[postgres,minio,mongodb]`.
- **Fidelity**: on the import path nothing is mocked — real MinIO, real Postgres, real MongoDB with a
unique index, a real MLflow tracking server with a working registry, real Argon2id and real
`secretstream`. Only `CoreNotificationHandler`'s Mongo client stays mocked, because notifications
are a side channel this suite counts rather than a dependency it tests.
- **Fixtures**: `e2e/import_harness.py`, registered as a pytest plugin from `e2e/conftest.py`. Every
fixture is `import_`-prefixed.
- **Prerequisite**: the one-time reaper setup at the top of this document
([Running automated E2E tests](#one-time-setup-disable-the-testcontainers-reaper)). Without it every
test here errors at setup with a Docker mount error and never reaches an assertion.
### 4.1 Success Scenarios
---
#### Scenario 4.1.1: Producer-Emitted Bundle Imports End to End
**Description**: The committed golden `.sientia` file — bytes the producer's own code wrote — is
imported at the producer's real Argon2id cost (256 MiB, 3 passes)
**Input**:
- `tests/fixtures/bundle/golden_model_v3.sientia` uploaded under `imported_models/`
- One `PENDING` / `run_type='IMPORT'` row carrying the object key and the uploader's digest
- The workflow input: `import_run_id`, `bucket`, `object_key`, `expected_digest`, `password_envelope`
**Expected Behavior**:
- The row is claimed, the seven gates pass, the experiment and a new run are created
- The artifact tree is logged with `prediction_model` and `data_model` at the run's root
- A registry version is created and left unstaged; the model document is written inactive
- One terminal write sets `IMPORT_SUCCESS`
**Assertions**:
- Exactly one run named `import-fixture_model-v3`, status `FINISHED`
- `prediction_model` contains `MLmodel` and `model.pkl`
- Bundle parameters logged, origin logged under `source.*`
- Destination version is `1`, independent of the bundle's `model_version` of `3`
- Model document `id` is the seeded maximum plus one; `active` is `false`
- Row: `IMPORT_SUCCESS`, `error_message` null, both names filled
---
#### Scenario 4.1.2: Factory-Built Bundle Imports End to End
**Description**: The same assertions over a synthesised tree with names unique to the test
**Input**: A bundle from `tests/helpers/bundle_factory.py` at the cheap KDF cost
**Expected Behavior**: As 4.1.1
**Assertions**:
- As 4.1.1, plus `model_config` holds exactly `transform_flavor`, `predict_flavor`,
`retention_minutes` and `target`
- `import.workflow_id` in the row's detail JSON matches the run
---
#### Scenario 4.1.3: The Registry Source Resolves Through the Legacy Loader
**Description**: The `source` the importer writes must be one the *existing* loader can parse
**Input**: A successful import's registered version
**Expected Behavior**: `MLFlowRepository.get_model_run_id` extracts the run id from the source
**Assertions**:
- The source ends in `/prediction_model`
- `get_model_run_id` returns the run the import created
- Both artifact paths are present under that run
- Called with the stage the version is actually in, because an import promotes nothing
---
#### Scenario 4.1.4: Second Import of the Same Model
**Description**: The same model imported twice, with different declared `model_version`s
**Input**: Two bundles declaring `model_version` `2` then `7`, same model and experiment names
**Expected Behavior**: Additive — a new run and a new registry version, nothing promoted, the
existing model document untouched
**Assertions**:
- Two versions (`1`, `2`) of one registered model; destination numbering ignores the bundle's
- Two runs in the same experiment, each logging its own `source.model_version`
- No version in `Production`
- The model document is byte-identical to what the first import wrote
---
#### Scenario 4.1.5: A Model Name Already in the Collection
**Description**: An operator-owned document already carries the bundle's model name
**Input**: A document with `active: True`, `predict_flavor: pyfunc` and a hand-set target
**Expected Behavior**: The import notices the name and leaves the document alone; the import still
succeeds
**Assertions**:
- The collection is unchanged, exactly one document carries that name, no duplicate
- Row: `IMPORT_SUCCESS` — a document that is already there is not a failure
---
### 4.2 Gate Rejection Scenarios
Every case asserts the same five things: terminal status `ORCHESTRATOR_VALIDATION_ERROR`, the
recorded step, the recorded code, `error_message` **equal** to the catalog sentence for that step,
and that no experiment, run, artifact, registry version, model document or extracted file remains.
---
#### Scenario 4.2.1: Object Above the Size Ceiling (gate 1)
**Description**: An object larger than `IMPORT_MAX_OBJECT_BYTES`, rejected from `stat_object` before
its body is fetched
**Input**: A factory bundle against a ceiling narrowed to 1 KiB for this test class
**Expected Behavior**: Rejected at `download`
**Assertions**: The five shared assertions
---
#### Scenario 4.2.2: Malformed Header (gate 2)
**Description**: Five variants — wrong magic, unsupported format version, KDF memory above the
ceiling, KDF passes above the ceiling, a header digest that does not describe the chunk stream
**Expected Behavior**: Each rejected at `download`, sharing gate 1's step
**Assertions**:
- The five shared assertions
- The recorded sentence does **not** mention the password: gates 1 and 2 must not implicate a
password that is fine
---
#### Scenario 4.2.3: Decryption Failures (gate 3)
**Description**: Three variants — a wrong bundle password inside a valid envelope, a tampered chunk
with the header digest repaired, and a stream whose every chunk is framed `MESSAGE`
**Expected Behavior**: Each rejected at `decryption`
**Assertions**:
- The five shared assertions
- The tampered case proves the file is self-consistent and only the AEAD can tell
- The framing case proves the digest cannot see an incomplete stream
---
#### Scenario 4.2.4: Hostile Archive (gate 4)
**Description**: Four variants — an entry escaping the tree, an absolute entry path, a symbolic
link, and a declared expansion over the ceiling
**Expected Behavior**: Each rejected at `archive_inspection`, from the central directory alone
**Assertions**:
- The five shared assertions
- **Nothing extracted**: the worker work directory is empty, not merely free of provisioning
---
#### Scenario 4.2.5: A Member Extraction Cannot Place (gate 5)
**Description**: An archive declaring the same member name twice — passes gate 4, which checks names
and modes but not duplicates
**Expected Behavior**: Rejected at `extraction` on the second member
**Assertions**: The five shared assertions
---
#### Scenario 4.2.6: Incomplete Structure (gate 6)
**Description**: Five variants — no `metadata.json`, metadata missing `run_id`, parameters missing
`target_variable`, no `prediction_model` directory, and a 60-character experiment name
**Expected Behavior**: Each rejected at `structure_validation`
**Assertions**:
- The five shared assertions
- For the over-long name: the row's `experiment_name` is still the inserted placeholder, so no write
reached the `VARCHAR(50)` column, and that column is confirmed still 50 wide
---
#### Scenario 4.2.7: Forbidden File Type (gate 7)
**Description**: An `artifacts/payload.exe` member with a suffix outside the allow-list
**Expected Behavior**: Rejected at `content_policy`
**Assertions**: The five shared assertions
---
### 4.3 Digest Enforcement Scenarios
---
#### Scenario 4.3.1: A Digest Computed Over Other Bytes
**Description**: A perfectly valid bundle whose `expected_digest` belongs to a different file
**Expected Behavior**: Fails at `download`, before any header is parsed
**Assertions**:
- Terminal status, step, code and catalog sentence
- The failure names the digest, not the magic or the version
- `run_name` is still null and `experiment_name` still the placeholder — the bundle was never opened
---
#### Scenario 4.3.2: The Object Is Replaced After Upload
**Description**: The same key, different bytes — also a valid bundle, so nothing but the digest can
tell
**Expected Behavior**: Fails at `download`
**Assertions**: As 4.3.1; only a real object store can show this at all
---
#### Scenario 4.3.3: No Digest at All
**Description**: `expected_digest` absent from the input
**Expected Behavior**: Fails at `received`, before anything is downloaded
**Assertions**: The request-invalid code, and the work directory was never created
---
### 4.4 Provisioning Failure Scenarios
Each asserts the platform terminal status `IMPORT_ERROR` — not the validation status, because a
provisioning failure is not the user's to fix — its own step and code, the workflow identifier in
the detail JSON, and `error_message` **equal** to the catalog sentence.
---
#### Scenario 4.4.1: The Tracking Server Is Stopped
**Description**: The `mlflow server` process is genuinely terminated for the duration
**Expected Behavior**: Fails at `experiment_creation`
**Assertions**:
- The shared assertions
- The names are already on the row: the bundle was opened before MLflow was reached
- Nothing in Mongo; nothing provisioned
---
#### Scenario 4.4.2: Artifact Upload Fails
**Description**: `log_artifacts` is patched — the one step where neither hostile input nor a broken
dependency can isolate the failure
**Expected Behavior**: Fails at `artifact_upload`
**Assertions**:
- The shared assertions
- The experiment exists (the previous step succeeded); no version, no document
---
#### Scenario 4.4.3: Registration Fails
**Description**: `create_model_version` is patched. MLflow rejects only an empty registered-model
name, and one process serves both tracking and the registry, so nothing else isolates this step
**Expected Behavior**: Fails at `registration`
**Assertions**:
- The shared assertions
- The run and both artifact paths survive — "the files were stored, the model was not registered"
---
#### Scenario 4.4.4: The Model Document Is Refused
**Description**: A real unique index on `model_config.target` refuses the activity's own insert.
(A duplicate `id` cannot be constructed: the activity derives `max(int(id)) + 1` from the same
collection, so the id it computes can never already exist)
**Expected Behavior**: Fails at `model_document`
**Assertions**:
- The shared assertions
- One registry version exists — registered, and unusable
---
### 4.5 Status Record Scenarios
---
#### Scenario 4.5.1: The Claim Precedes Every Gate
**Description**: The row is observed from inside `stat_object`, the first call gate 1 makes
**Assertions**: Running status and `import.workflow_id` already present at that moment
---
#### Scenario 4.5.2: The Names Are Visible Mid-Flight
**Description**: The row is observed from inside the first call experiment creation makes, during an
import that then succeeds
**Assertions**: `experiment_name` and `run_name` already populated before anything exists in MLflow
---
#### Scenario 4.5.3: The Inserting Side's Columns Are Never Rewritten
**Description**: After a success and after a failure
**Assertions**: `username`, `request_data`, `bucket_name`, `file_name`, `run_type`, `created_at`
unchanged; `status` and `updated_at` did change
---
#### Scenario 4.5.4: Unclaimable Rows
**Description**: Three variants — an unknown id, a `run_type='TRAIN'` row, and a row already held by
another workflow id
**Expected Behavior**: The Temporal run ends failed with the not-claimable error type, and **nothing
is written to any row**
**Assertions**: The error type from the history, the whole table byte-identical, nothing provisioned
---
#### Scenario 4.5.5: The Row Names a Different Object
**Description**: `file_name` disagrees with the input's `object_key`
**Expected Behavior**: Fails at `received` **with the error recorded on the row** — unlike 4.5.4,
this row was successfully claimed
**Assertions**: The received code and sentence on the row; nothing provisioned
---
#### Scenario 4.5.6: Constraint Conformance
**Description**: The real constraints, exercised directly
**Assertions**:
- Every status a normal import writes is accepted on an `IMPORT` row
- All four terminal values are refused across run types, naming
`chk_experiment_run_terminal_status_run_type`; the shared stage values are accepted on both
- A training row still walks its own vocabulary
- The table has exactly the 13 columns `V1`…`V12` leave (no `ALTER TABLE` by this feature)
---
### 4.6 Absence Scenarios
---
#### Scenario 4.6.1: One Table, One Database
**Description**: A full import bracketed by marker statements on both Postgres servers, read from
`log_statement=all`
**Assertions**:
- Every statement touching a user relation names `experiment_run`
- No `project`, `models`, `import_run`, `model_project` or `model_registry` reference (with quoted
literals stripped, so the bundle's own `model_project` *value* is not mistaken for a table)
- The predictions database logs **no** statement at all, though the import holds working credentials
for it
- A rejected import — which writes more status rows than a successful one — behaves the same
---
### 4.7 Cleanup, Retention and History Scenarios
---
#### Scenario 4.7.1: The Worker Filesystem Is Clean
**Description**: After a success and after rejections at two different steps (decryption and
structure validation, which unwind on different code paths)
**Assertions**: No extracted tree, no downloaded object, no temporary archive
---
#### Scenario 4.7.2: The Retention Rule Is Present
**Assertions**: `sientia-import-bundle-expiry` on the import bucket, enabled, filtered on
`imported_models/`, expiring in 7 days — and the uploaded object still there, because the import
never deletes it
---
#### Scenario 4.7.3: The Record Outlives the Object
**Description**: The uploaded object is deleted, as retention will eventually do
**Assertions**: The row still reads `IMPORT_SUCCESS` with its names
---
#### Scenario 4.7.4: No Plaintext Secret in the Event History
**Assertions**:
- The bundle password, the base64 envelope key and the key's hex form appear nowhere in the history
- `password_envelope` **is** in the history, so the negative assertions above cannot pass vacuously
---
### 4.8 Retry Envelope and Missing-`V12` Scenarios
These run on a **real clock** (`WorkflowEnvironment.start_local()`) with the status retry intervals
shortened; the production values stay asserted by the unit suite. All are marked `slow`.
---
#### Scenario 4.8.1: A Database That Comes Back
**Description**: The status database is stopped for two seconds after provisioning and restarted
inside the retry envelope
**Assertions**: The terminal write lands `IMPORT_SUCCESS` with the correct workflow id
---
#### Scenario 4.8.2: A Broken Progress Update
**Description**: A brief early outage that only the progress hints see
**Assertions**: The import completes with the correct verdict and does not spend the terminal
write's envelope on a hint nobody reads
---
#### Scenario 4.8.3: A Database That Stays Down After a Successful Import
**Description**: The database is taken away as soon as the row is claimed and never comes back
**Expected Behavior**: The model is fully provisioned and the workflow **completes** — reporting
failure would tell somebody to re-import a model that is already on the platform
**Assertions**:
- The run returns its result; the model, version and document all exist
- The failure channel carries `IMPORT_STATUS_NOT_RECORDED`, `step=status_write` and
`verdict_not_recorded=`
- The row is left in the running status
---
#### Scenario 4.8.4: A Failed Import Whose Failure Cannot Be Written Either
**Description**: The import fails at `artifact_upload` and the terminal write is then refused
**Assertions**:
- The reported cause is the artifact-upload failure, not the status write
- The status write reports itself only through the channel
- The row is left running
---
#### Scenario 4.8.5: A Table Without `V12`
**Description**: The evidence for why the migration gates enabling the import worker
**Expected Behavior**: The claim succeeds (`V11` gave the table `run_type`; `PENDING` and the running
status are pre-`V12` values), the model is fully provisioned, and only the terminal value is new — so
the last write of the import is the first thing that fails, non-retryably
**Assertions**:
- The run completes; the model, version and document all exist
- The row is left in the running status, with both names filled
- The channel carries the not-recorded code and the verdict
- Elapsed time shows the refusal was **not** retried
- `pg_constraint` confirms the fixture's `chk_experiment_run_status` really does not allow
`IMPORT_SUCCESS`, so the test cannot pass vacuously
---
## 5. joblib Flavor Scenarios
`joblib` is an explicit model flavor, declared in `model_config.predict_flavor` /
`model_config.transform_flavor` like any other. It reads the artifact directly — download the
`.pkl`, put the artifact's `code/` on `sys.path`, `joblib.load` — and **never** calls an MLflow
flavor loader. There is no fallback in either direction: a failing loader propagates.
Covered by `e2e/test_predictions_batch_joblib_flavor.py` (read) and
`e2e/test_retrain_joblib_flavor.py` (write). Both use the `joblib_flavor_artifacts` fixture, which
builds one artifact directory per model type holding a real compressed `model.pkl` plus a `code/`
directory carrying the module where the model's class lives. The fixture takes that directory off
`sys.path` again before yielding, so the pickle is **unloadable** until the flavor puts it back —
that is what makes these tests prove module resolution rather than assume it. The fixture also
restores `sys.path` in teardown, because the flavor deliberately never removes what it inserted.
### 5.1 Read Scenarios
---
#### Scenario 5.1.1: Predict Through the joblib Flavor
**Description**: `predict_flavor: 'joblib'` on a normal `predictions_batch` run
**Input**:
- `model_config.predict_flavor = 'joblib'`, `transform_flavor = 'sklearn'`
- `client.download_artifacts` serving the real `prediction_model` artifact
**Expected Behavior**:
- The artifact is downloaded once and deserialized with `joblib.load`
- No `models:/<name>/production` URI is built, so no flavor loader ever sees this model
- The prediction the artifact's own object produces is what lands in Postgres
**Assertions**:
- Exactly one prediction row, `prediction_status='Good'`, value matching the artifact's model
- `download_artifacts` called once, with `'prediction_model'`
- No `mlflow.{sklearn,pyfunc,pytorch}.load_model` call names this model
---
#### Scenario 5.1.2: Transform Through the joblib Flavor
**Description**: `transform_flavor: 'joblib'`, which had no E2E coverage before
**Input**:
- `model_config.transform_flavor = 'joblib'`, `predict_flavor = 'sklearn'`
- `client.download_artifacts` serving the real `data_model` artifact
**Expected Behavior**:
- `data_model` is downloaded and deserialized; the prediction model stays on its MLflow flavor
- No artifact URI is resolved for the transform, so `mlflow.get_run` is never called
**Assertions**:
- The prediction lands with the mocked prediction model's value
- `download_artifacts` called once, with `'data_model'`
- `mlflow.get_run` not called; no flavor loader was given a `/data_model` URI
---
#### Scenario 5.1.3: An MLflow Flavor Is Unaffected
**Description**: A model declaring `sklearn` behaves exactly as before
**Assertions**:
- The prediction lands with the mocked model's value
- `download_artifacts` NOT called — the joblib path is never entered as a safety net
---
### 5.2 Error Scenarios
---
#### Scenario 5.2.1: joblib Flavor Failure Has a Single Cause
**Description**: The downloaded artifact has no top-level `.pkl`
**Expected Behavior**: The workflow stops. The reported cause is the missing pickle and nothing
else — with no fallback there is no second exception chained to a first
**Assertions**:
- The run stops (`assert_stop`)
- The log carries `No .pkl file found in prediction_model artifact`
- The log does **not** carry `Original error`, the marker of a chained cause
---
### 5.3 Write Scenarios
---
#### Scenario 5.3.1: Retraining a joblib Model Round-Trips
**Description**: `retrain_model` on a model whose `predict_flavor` and `transform_flavor` are both
`joblib` — write and read in one test
**Input**:
- Training frame indexed by `DATETIME_FORMAT_WITH_TZ` timestamps, carrying the target column
- MLflow's artifact I/O replaced by a fake store that copies into an inspectable run directory
**Expected Behavior**:
- Both models load through the joblib flavor; no flavor loader runs
- `log_model` writes `prediction_model/model.pkl` (a `joblib.dump(..., compress=3)` pickle) and
re-logs `prediction_model/code/`, without an `MLmodel` manifest
- Loading the **written** artifact through the joblib flavor deserializes it and predicts, after
the module the first load imported has been dropped from `sys.modules`
**Assertions**:
- `retrain_model` returns `success: True`
- `download_artifacts` was asked for both `data_model` and `prediction_model`
- `sklearn`/`pyfunc` loaders have zero calls
- The written run has `prediction_model/model.pkl`, `prediction_model/code/<module>.py` and
`data_model/model.pkl`, and **no** `prediction_model/MLmodel`
- The reloaded model's prediction equals the original's, and the written `code/` is on `sys.path`
---

184
e2e/test_import_model_absence.py Normal file
View File

@@ -0,0 +1,184 @@
"""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)]

247
e2e/test_import_model_cleanup.py Normal file
View File

@@ -0,0 +1,247 @@
"""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'
)

4
e2e/test_import_model_gates.py
View File

@@ -400,6 +400,10 @@ class TestGate5Extraction:
and modes, not duplicates — and is refused here on the second member.
"""
# `zipfile` warns while *writing* the duplicate entry. That is the fixture doing exactly what
# this test needs, so it is silenced here rather than left to look like a problem — and scoped to
# this one test, so a duplicate name appearing anywhere else still gets noticed.
@pytest.mark.filterwarnings('ignore:Duplicate name:UserWarning')
async def test_a_duplicate_member_is_rejected_at_extraction(
self,
import_status_table: Any,

108
e2e/test_import_model_provisioning_failures.py
View File

@@ -13,15 +13,19 @@ failure is always the second kind.
| `experiment_creation` | the MLflow server is genuinely stopped | broken dependency |
| `artifact_upload` | targeted patch on `log_artifacts` | patch, justified in place |
| `registration` | targeted patch on `create_model_version` | patch, justified in place |
| `model_document` | a document already holding the id the activity will compute, against the real unique index | broken dependency |
| `model_document` | a real unique index refuses the activity's own insert | broken dependency |
Two patches rather than the one design E7 anticipated. The reason is the same for both and is
recorded in each test: with one server serving both tracking and registry, and MLflow validating
nothing about a registered model name except that it is non-empty (which gate 6 already requires),
there is no way to make the world hostile to *just* one of these two steps. Stopping the server
fails experiment creation first; a malformed name is not malformed to MLflow. The patch is on the
client call the step makes, not on the activity or the workflow, so what is being tested is still
"the workflow attributes this step's failure to this step".
Two departures from what the tasks anticipated, both recorded in the test that found them:
- **Registration needs a patch too**, not just artifact upload. One process serves both tracking and
the registry, so stopping it fails experiment creation three steps earlier; and MLflow validates
nothing about a registered model name except that it is non-empty, which gate 6 already requires.
Both patches are on the single client call the step *is*, never on the activity or the workflow, so
what is under test — the workflow attributing this step's failure to this step — is untouched.
- **The model document's unique-index violation is on `model_config.target`**, not on `id`. The
duplicate-`id` state the task describes cannot be constructed at all; `TestModelDocument`'s
docstring shows why, and the mechanism the task wanted — a real index refusing a real insert — is
kept intact.
"""
from __future__ import annotations
@@ -292,9 +296,26 @@ class TestRegistration:
class TestModelDocument:
"""The last step, failed by a real unique index on a real collection."""
"""The last step, failed by a real unique index refusing the real insert.
async def test_a_duplicate_id_fails_at_model_document(
**Why the index is on `model_config.target` and not on `id`.** Task 7.4 asks for "a document
already holding the id the activity is about to compute". That state cannot be constructed, and
why is worth more than the test would have been: `_next_model_document_id` computes
`max(int(id)) + 1` over `distinct('id')` of the *same* collection, so every value that
stringifies to the colliding id is counted by that maximum. Put a document at `'8'` and the next
id becomes `'9'`. Make the id unparseable and it is ignored by the numbering *and* does not
collide. There is no third case — the numbering cannot collide with itself, and the only real
duplicate-key path on `id` is two imports racing, which this change lists as a non-goal. That is
a good property of the implementation, not a gap.
So the mechanism the task asked for is kept — a real unique index refusing the activity's own
insert — on a key that can actually collide. `model_config.target` is `'SE'` on the seeded
documents and `'SE'` on the one this import writes, so a unique index there makes the insert a
genuine `DuplicateKeyError` from the server. No patch, no mock: the same class of failure, from
the same enforcement mechanism, at the same step.
"""
async def test_a_unique_index_violation_fails_at_model_document(
self,
import_status_table: Any,
import_models_collection: Any,
@@ -308,50 +329,47 @@ class TestModelDocument:
import_envelope: Any,
import_read_status_row: Any,
) -> None:
"""A document already holding the id `max(int(id)) + 1` will compute, under another name.
Under another name on purpose: a document with the *same* name would make the activity skip
the insert entirely, which is a different behaviour (task 8.6) and would not fail. This one
collides only on `id`, so the insert is attempted and the index refuses it.
No patch and no mock: the failure is `pymongo`'s duplicate-key error from a real index,
which is the thing that would happen in production if two imports raced.
"""
collision_id = str(max(int(seeded) for seeded in SEEDED_MODEL_IDS) + 1)
import_models_collection.insert_one(
{'id': collision_id, 'name': 'a_different_model', 'active': True, 'model_config': {}}
)
handle, import_run_id, object_key = await _upload_and_start(
import_bundle_factory=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,
)
with pytest.raises(WorkflowFailureError):
await handle.result()
# All three seeded documents carry target 'SE', so the index has to be buildable first:
# give two of them their own target and leave exactly one holding 'SE'. That one is what the
# import's own document — also target 'SE' — collides with.
for ordinal, seeded_id in enumerate(sorted(SEEDED_MODEL_IDS)[1:], start=1):
import_models_collection.update_one(
{'id': seeded_id}, {'$set': {'model_config.target': f'OTHER_{ordinal}'}}
)
import_models_collection.create_index('model_config.target', unique=True, name='e2e_target')
try:
handle, import_run_id, object_key = await _upload_and_start(
import_bundle_factory=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,
parameters={'target_variable': 'SE'},
)
with pytest.raises(WorkflowFailureError):
await handle.result()
row = import_read_status_row(import_run_id)
finally:
# The collection is dropped and reseeded per test, but dropping the index here keeps the
# failure local even if that ever changes.
import_models_collection.drop_index('e2e_target')
_assert_platform_failure(
row=import_read_status_row(import_run_id),
row=row,
step=ImportStep.MODEL_DOCUMENT,
workflow_id=handle.id,
object_key=object_key,
)
# No document for this model, and the colliding one is untouched.
# No document for this model, and the seeded ones are exactly as they were.
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
blocker = import_models_collection.find_one({'id': collision_id}, {'_id': False})
assert blocker == {
'id': collision_id,
'name': 'a_different_model',
'active': True,
'model_config': {},
}
assert {
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
} == set(SEEDED_MODEL_IDS)
# Everything before it stands: "the model was registered, but the settings ... could not be
# saved" is the catalog sentence, and this is what makes it true rather than merely stored.
# Everything before it stands — registered, and unusable, which is the point of the split
# between this step's sentence and the registration step's.
mlflow_client = _fresh_client(import_mlflow_server)
versions = list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
assert len(versions) == 1

514
e2e/test_import_model_retry_recovery.py Normal file
View File

@@ -0,0 +1,514 @@
"""What happens when the *record* cannot be written — the cases that gate this feature's release.
Everything else in this suite asks "did the import do the right thing". This file asks the harder
question: the import did the right thing, and then could not say so. Design D20 answers it in two
halves, and both are release-critical.
**A database that comes back.** The terminal write retries, so a status database restarted during a
deploy still records the correct verdict. Without this, every rolling restart would leave imports
looking unfinished forever.
**A database that does not come back.** The model is fully provisioned, and the workflow must not
claim otherwise. It **completes**, and the verdict it could not write goes out through the failure
channel — log, metric, notification — under `IMPORT_STATUS_NOT_RECORDED`. An import that worked must
never be reported as an import that failed just because the audit row was unreachable: somebody
would re-import a model that is already there.
That last point was worth discovering the hard way. The first version of this file asserted the run
*failed*; it does not, and it should not. The spec says so in as many words — "the run is not
reported as failed" — and the implementation agrees with the spec.
And the case that is the evidence for a migration in another repository: **a table without `V12`**.
The claim succeeds (`V11` already gave the table `run_type`, and `PENDING` and the running status are
pre-`V12` values), the model is fully provisioned, and only the *terminal* value is new — so the last
write of the import is the first thing that fails. `ImportStatusRejectedError` is non-retryable, so
it lands on the first attempt with no wait, and the row is left running forever. That is what
"apply `V12` before turning `IMPORT_MODEL_ENABLED` on" buys, as a test rather than a warning in a
document.
Everything here runs on a **real clock** (design E9): under time skipping the backoff sleeps are
virtual, so a restarted container would never get wall-clock time to come back and the recovery test
would race by construction. The retry intervals are shortened by `import_status_retry_shortening`;
the production values stay asserted by unit test 12.1 of `import-model-workflow`, because what these
tests are about is the behaviour, not the arithmetic.
"""
from __future__ import annotations
import asyncio
import json
import time
from typing import Any
import mlflow
import pytest
from temporalio.client import WorkflowFailureError
from e2e.import_harness import failure_text, failure_types
pytestmark = [pytest.mark.import_e2e, pytest.mark.slow, pytest.mark.asyncio]
FACTORY_PASSWORD = 'Factory1Model!Test'
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
NOT_RECORDED_CODE = 'IMPORT_STATUS_NOT_RECORDED'
STATUS_WRITE_MARKER = 'step=status_write'
PROVISIONING_STEPS = ('experiment_creation', 'artifact_upload', 'registration', 'model_document')
def _as_dict(value: Any) -> dict[str, Any]:
return json.loads(value) if isinstance(value, str) else (value or {})
def _client(server: dict[str, Any]) -> Any:
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
def _channel_text(mock_logger: Any) -> str:
"""Everything the failure channel said, as one string.
The channel is a log line, a metric and a notification, and the log line is the only sink that
carries the whole classification — code, marker and the verdict it could not record. Reading the
logger is therefore reading the channel, and it is synchronous, which the notification path is
not.
"""
return '\n'.join(str(call) for call in mock_logger.mock_calls)
def _assert_fully_provisioned(
*, server: dict[str, Any], names: dict[str, str], collection: Any
) -> str:
"""Everything the import was supposed to create exists. Returns the run id."""
client = _client(server)
experiment = client.get_experiment_by_name(names['experiment_name'])
assert experiment is not None, 'no experiment: the import never got as far as provisioning'
runs = client.search_runs([experiment.experiment_id])
assert len(runs) == 1, f'expected one run, found {len(runs)}'
assert {a.path for a in client.list_artifacts(runs[0].info.run_id)} >= {
'prediction_model',
'data_model',
}
versions = list(client.search_model_versions(f"name='{names['model_name']}'"))
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
assert collection.find_one({'name': names['model_name']}) is not None, 'no model document'
return runs[0].info.run_id
def _assert_no_provisioning_step_is_blamed(text_of_failure: str) -> None:
"""The import's own steps all succeeded, so none of them may appear as the cause."""
for step in PROVISIONING_STEPS:
assert step not in text_of_failure, f'{step!r} is blamed, but that step succeeded'
async def _wait_for_claim(read_row: Any, import_run_id: int, timeout: float = 30.0) -> None:
"""Block until the workflow has taken ownership of the row.
The earliest deterministic point after which every status write belongs to this import. Used to
place an outage that must cover the whole rest of the run.
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
row = read_row(import_run_id)
if row is not None and row['status'] == RUNNING_STATUS:
return
await asyncio.sleep(0.1)
raise AssertionError('the row was never claimed, so the outage could not be placed')
async def _wait_for_provisioning(
*, server: dict[str, Any], names: dict[str, str], collection: Any, timeout: float = 90.0
) -> None:
"""Block until the model document exists — the observable end of provisioning.
Polling the two destinations rather than patching an activity: "the document exists" *is* the
end of provisioning, so there is nothing to hook and nothing to fake. This is what lets the
outage below be placed exactly on the terminal write.
"""
client = _client(server)
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
document = collection.find_one({'name': names['model_name']})
if document is not None and list(
client.search_model_versions(f"name='{names['model_name']}'")
):
return
await asyncio.sleep(0.2)
raise AssertionError('provisioning never finished, so the outage could not be placed')
class TestTheDatabaseComesBack:
"""Stopped and restarted inside the retry envelope: the verdict still lands."""
async def test_a_restart_inside_the_envelope_still_records_the_verdict(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
) -> None:
"""The outage lands on the terminal write, and the server is back before the retries run out.
The stop is placed after provisioning rather than at the start, so what is being tested is
the *terminal* write's retry envelope specifically. A stop at the start would also exercise
the claim's envelope, which is a different (and already covered) code path.
"""
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_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await _wait_for_provisioning(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
# Down for two seconds — several retry attempts, well inside the shortened envelope.
async with import_bff_stopped():
await asyncio.sleep(2)
result = await handle.result()
assert result['model_name'] == import_names['model_name']
run_id = _assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
assert result['run_id'] == run_id
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS', (
f'the verdict did not land after the restart: the row says {row["status"]!r}'
)
assert row['error_message'] is None
assert _as_dict(row['orchestrator_response_data'])['import']['workflow_id'] == handle.id
async def test_a_broken_progress_update_does_not_delay_the_verdict(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
) -> None:
"""A mid-flight outage that only the progress hints see. The import does not care.
Progress writes are hints nobody reads — the terminal write rewrites status, code, reason and
step from one classification — so they get their own short policy. If they shared the terminal
write's envelope, a database the import does not need would stall a working import for
minutes at every stage update.
The outage is placed early, so it lands on hints rather than on the terminal write, and the
assertion is that the import still finished with the correct verdict.
"""
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_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
# Early and brief: the claim has landed, the gates are running, the terminal write is far off.
await asyncio.sleep(1.0)
async with import_bff_stopped():
await asyncio.sleep(1.5)
started = time.monotonic()
result = await handle.result()
elapsed = time.monotonic() - started
assert result['model_name'] == import_names['model_name']
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS'
assert elapsed < 60, (
f'the import took {elapsed:.1f}s after the outage ended; a lost hint should cost '
f'its own short policy and nothing more'
)
class TestTheDatabaseStaysDown:
"""Past the envelope. The model is real, the record is not, and the run says which is which."""
@pytest.fixture
def import_status_retry_shortening(self) -> dict[str, Any]:
"""Shorter still, because these tests wait for the envelope to *expire*.
Six half-second attempts inside an eight-second envelope: long enough that retries really
happen, short enough that giving up is something a test suite can wait for.
"""
return {
'initial_interval': 0.5,
'backoff_coefficient': 1.0,
'maximum_attempts': 6,
'hint_maximum_attempts': 2,
'envelope_seconds': 8,
'write_timeout_seconds': 3,
}
async def test_a_successful_import_is_not_reported_as_failed(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_stopped: Any,
mock_logger: Any,
) -> None:
"""The most important assertion in this file.
The model is registered, the document is written, the artifacts are stored — and the row
cannot be updated. The workflow **completes**: reporting failure here would tell somebody to
re-import a model that is already on the platform. The verdict goes out through the failure
channel instead, under the not-recorded code, carrying what it could not write.
The outage is placed as soon as the row has been **claimed**, not once provisioning has
finished. Waiting for provisioning looks more precise and is in fact a race: the terminal
write follows the model document by milliseconds, while `docker stop` takes seconds, so the
write lands before the database is actually gone and the test proves nothing. Stopping right
after the claim covers everything that follows — the progress hints (which are allowed to
fail) and the terminal write (which is the subject) — with no timing to get right.
"""
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_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await _wait_for_claim(import_read_status_row, import_run_id)
async with import_bff_stopped():
result = await handle.result()
# The run completed. This is the assertion, not an incidental detail.
assert result['model_name'] == import_names['model_name']
assert result['import_run_id'] == import_run_id
_assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel, (
'the unrecordable verdict never reached the failure channel'
)
assert STATUS_WRITE_MARKER in channel, (
'the channel does not mark this as a status-write failure'
)
assert 'verdict_not_recorded=' in channel, (
'the channel reports the failure but not the verdict it could not write'
)
_assert_no_provisioning_step_is_blamed(
'\n'.join(
line for line in channel.splitlines() if 'model import failed' in line
).replace('step=model_document', '')
)
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS, (
f'the row should still be running, not {row["status"]!r}'
)
async def test_a_failed_import_keeps_its_original_cause(
self,
monkeypatch: pytest.MonkeyPatch,
import_status_table_pre_v12: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_bundle_factory: Any,
import_activities: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
mock_logger: Any,
) -> None:
"""An import that already failed, whose failure then cannot be written down either.
Two failures, and only one of them answers "why did this import not work". The status write
is the second, and it must not displace the first: somebody told "the status could not be
recorded" would go and look at the database, when what actually happened is that their model's
files could not be stored.
**Why a pre-`V12` table instead of stopping the database.** The outage has to cover the
terminal write and nothing earlier, and an import that fails at a gate reaches its terminal
write within a second of starting — there is no window to aim at, and the first version of
this test proved only that a stopped database *after* a completed write changes nothing.
A pre-`V12` table has no timing at all: `IMPORT_ERROR` is a value it has never heard of, so
the terminal write of a *failed* import is refused deterministically while every write before
it succeeds. The failure being tested is the same one either way — the record cannot take the
verdict — and this version cannot race.
"""
client = import_activities.import_mlflow_repository.client
def refuse(*_args: Any, **_kwargs: Any) -> None:
raise OSError('artifact store is unreachable')
monkeypatch.setattr(client, 'log_artifacts', refuse)
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_real_clock(
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) as failure:
await handle.result()
# The reported cause is the step that actually failed.
text_of_failure = failure_text(failure.value)
assert 'artifact store is unreachable' in text_of_failure, (
f'the reported cause lost the original failure: {text_of_failure}'
)
assert 'ImportStatusRejectedError' not in failure_types(failure.value), (
f'the status write displaced the original cause: {failure_types(failure.value)}'
)
assert 'the import log database refused' not in text_of_failure, (
'the run is reported as a status-write failure rather than an artifact-upload failure'
)
# The status write still reported *itself*, through the channel, as the secondary failure.
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel, 'the unrecordable verdict was never reported'
assert STATUS_WRITE_MARKER in channel
# And the row is left running, because the verdict never landed.
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS
class TestMissingV12:
"""A table without `V12`. The evidence for why the migration gates enabling the worker."""
async def test_the_terminal_write_is_refused_on_its_first_attempt(
self,
import_status_table_pre_v12: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_mlflow_server: dict[str, Any],
import_bundle_factory: Any,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute_real_clock: Any,
import_envelope: Any,
import_read_status_row: Any,
import_bff_engine: Any,
mock_logger: Any,
) -> None:
"""Provisioned, then unrecordable — and unrecordable *immediately*.
The shape matters as much as the outcome. `ImportStatusRejectedError` is non-retryable: a
`CHECK` violation is not going to start passing because it was tried again, and waiting the
full 151-second envelope would turn a misconfiguration into a two-and-a-half-minute stall on
every single import. The elapsed time is what proves it — the shortened policy still allows
24 attempts, so a retried failure could not finish this fast.
"""
from sqlalchemy import text as sql_text
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
started = time.monotonic()
handle = await import_execute_real_clock(
import_build_input(
import_run_id=import_run_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
result = await handle.result()
elapsed = time.monotonic() - started
# The run completed: an import that provisioned a model is not a failed import, even when
# the last thing it wanted to do was impossible.
assert result['model_name'] == import_names['model_name']
# The claim and the names write both landed — the pre-V12 table was usable right up to the
# last statement, which is exactly what makes this failure mode easy to miss.
row = import_read_status_row(import_run_id)
assert row is not None
assert row['status'] == RUNNING_STATUS, (
f'the row is {row["status"]!r}; a pre-V12 table cannot hold a terminal import status, '
f'so the row is left looking as though the import never finished'
)
assert row['experiment_name'] == import_names['experiment_name']
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
# The model is fully provisioned. This is the whole problem: the work is done and invisible.
_assert_fully_provisioned(
server=import_mlflow_server,
names=import_names,
collection=import_models_collection,
)
# The verdict reached the failure channel under the not-recorded code.
channel = _channel_text(mock_logger)
assert NOT_RECORDED_CODE in channel
assert STATUS_WRITE_MARKER in channel
assert 'verdict_not_recorded=' in channel
# Non-retryable: refused on the first attempt, with no retry wait.
assert elapsed < 60, (
f'the import took {elapsed:.1f}s; a non-retryable rejection must not be retried'
)
# And the constraint that refused it is the one V12 adds — asserted directly, so the test
# says *why* the write was impossible rather than only that it was.
with import_bff_engine.begin() as connection:
statuses = connection.execute(
sql_text(
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
"WHERE conname = 'chk_experiment_run_status'"
)
).scalar_one()
assert 'IMPORT_SUCCESS' not in statuses, (
'the pre-V12 fixture already allows the import statuses, so this proved nothing'
)

270
e2e/test_import_model_second_import.py Normal file
View File

@@ -0,0 +1,270 @@
"""Importing the same model twice, and what the model document is allowed to say.
Two claims, and they pull in opposite directions, which is why they share a file.
**A second import is additive.** It creates a new run in the *same* experiment and a new registry
version of the *same* registered model, and it promotes nothing. That is the whole of it: no
overwrite, no stage change, no reconciliation.
**A second import touches the document not at all.** `active`, `target` and the flavors of a model
already on the platform are operator-owned — somebody turned that model on, and re-importing it must
not turn it off, retarget it or reset how it is loaded. So the document is written on the first
import and left byte-identical on every one after it.
The document's content is asserted exactly rather than loosely, because it *is* the model listing:
`model_config.transform_flavor`, `model_config.predict_flavor` and `model_config.target` are what
the runtime reads to know how to load the model, and a missing or wrong one makes an imported model
either invisible or unloadable.
"""
from __future__ import annotations
from typing import Any
import mlflow
import pytest
from e2e.import_harness import SEEDED_MODEL_IDS
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
FACTORY_PASSWORD = 'Factory1Model!Test'
CONTRACTED_MODEL_CONFIG_KEYS = {
'transform_flavor',
'predict_flavor',
'retention_minutes',
'target',
}
def _client(server: dict[str, Any]) -> Any:
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
async def _import_once(
*,
payload: bytes,
import_upload_bundle: Any,
import_insert_status_row: Any,
import_build_input: Any,
import_execute: Any,
import_envelope: Any,
) -> dict[str, Any]:
"""One successful import. Raises if it did not succeed, so a caller never asserts on a failure."""
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(FACTORY_PASSWORD),
)
)
result = await handle.result()
return {'result': result, 'import_run_id': import_run_id, 'object_key': object_key}
class TestSecondImport:
"""The same model, imported twice in one test — the only way to observe the second as second."""
async def test_a_second_import_adds_a_version_and_a_run_and_nothing_else(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_bundle_factory: Any,
import_mlflow_server: dict[str, 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:
"""Both bundles declare the same model and experiment, and different `model_version`s.
The second bundle declaring `model_version` '7' while the first declared '2' is the point of
task 8.2: the destination version number is the registry's own count, and the bundle's
number is a record of where the model came from, logged as `source.model_version` and used
for nothing else. If the two were confused, the second version here would be '7'.
"""
first = await _import_once(
payload=import_bundle_factory(metadata={'model_version': '2'}),
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,
)
document_after_first = import_models_collection.find_one(
{'name': import_names['model_name']}, {'_id': False}
)
assert document_after_first is not None, 'the first import wrote no document'
second = await _import_once(
payload=import_bundle_factory(metadata={'model_version': '7'}),
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,
)
model_name = import_names['model_name']
client = _client(import_mlflow_server)
# 8.1 — a new version of the same registered model.
assert first['result']['version'] == '1'
assert second['result']['version'] == '2'
versions = list(client.search_model_versions(f"name='{model_name}'"))
assert len(versions) == 2, f'expected two versions, found {len(versions)}'
assert {v.version for v in versions} == {'1', '2'}
# 8.2 — destination numbering is independent of what the bundles declared.
assert second['result']['version'] != '7'
runs_by_id = {
run.info.run_id: run
for run in client.search_runs(
[client.get_experiment_by_name(import_names['experiment_name']).experiment_id]
)
}
assert runs_by_id[first['result']['run_id']].data.params['source.model_version'] == '2'
assert runs_by_id[second['result']['run_id']].data.params['source.model_version'] == '7'
# 8.1 — a new run, in the *same* experiment, and both runs are still there.
assert first['result']['run_id'] != second['result']['run_id']
assert len(runs_by_id) == 2
assert {run.info.run_name for run in runs_by_id.values()} == {
f'import-{model_name}-v2',
f'import-{model_name}-v7',
}
# 8.1 — nothing promoted, by either import.
assert all(version.current_stage in (None, 'None') for version in versions), [
(v.version, v.current_stage) for v in versions
]
# 8.3 — the document is byte-identical to what the first import wrote.
document_after_second = import_models_collection.find_one(
{'name': model_name}, {'_id': False}
)
assert document_after_second == document_after_first
# And both rows record their own success.
for run in (first, second):
row = import_read_status_row(run['import_run_id'])
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS'
class TestModelDocumentContent:
"""What the written document says, asserted exactly rather than by presence."""
async def test_the_document_continues_the_collection_and_carries_the_contract(
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,
) -> None:
"""`id` is `max(int(id)) + 1` across gapped ids — 1, 2, 7 give 8, not 4 and not "3".
The gaps are what make this worth asserting: a count would give 4, and a lexicographic
maximum would give "8" here but "10" for a collection holding "9" and "30". The seeded ids
are strings, as the baseline migration stores them, so `int()` is doing real work.
"""
await _import_once(
payload=import_bundle_factory(parameters={'target_variable': 'PRODUCT_QUALITY'}),
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,
)
document = import_models_collection.find_one(
{'name': import_names['model_name']}, {'_id': False}
)
assert document is not None
expected_id = str(max(int(seeded) for seeded in SEEDED_MODEL_IDS) + 1)
assert document['id'] == expected_id, (
f'{document["id"]!r} does not continue {SEEDED_MODEL_IDS}'
)
assert document['id'] == '8'
assert document['name'] == import_names['model_name']
assert document['active'] is False
assert set(document['model_config']) == CONTRACTED_MODEL_CONFIG_KEYS, (
f'model_config holds {sorted(document["model_config"])}, not exactly the four '
f'contracted keys'
)
assert document['model_config'] == {
'transform_flavor': 'sklearn',
'predict_flavor': 'joblib',
'retention_minutes': 60,
'target': 'PRODUCT_QUALITY',
}
# The document has exactly these fields and no others — nothing was smuggled in alongside.
assert set(document) == {'id', 'name', 'active', 'model_config'}
async def test_a_name_already_in_the_collection_is_left_untouched(
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:
"""An operator-owned document with this model's name survives the import unchanged.
`active: True` and a hand-set target are exactly what a re-import must not undo: somebody
turned this model on and pointed it at something. The import notices the name and stops.
"""
existing = {
'id': '42',
'name': import_names['model_name'],
'active': True,
'model_config': {
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
'retention_minutes': 5,
'target': 'OPERATOR_CHOSE_THIS',
},
}
import_models_collection.insert_one(dict(existing))
before = list(import_models_collection.find({}, {'_id': False}))
outcome = await _import_once(
payload=import_bundle_factory(parameters={'target_variable': 'SOMETHING_ELSE'}),
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,
)
# The import still succeeded — a document that is already there is not a failure.
row = import_read_status_row(outcome['import_run_id'])
assert row is not None
assert row['status'] == 'IMPORT_SUCCESS'
after = list(import_models_collection.find({}, {'_id': False}))
assert after == before, 'the collection changed'
matching = [d for d in after if d['name'] == import_names['model_name']]
assert len(matching) == 1, 'a duplicate was inserted'
assert matching[0] == existing

596
e2e/test_import_model_status_record.py Normal file
View File

@@ -0,0 +1,596 @@
"""The import log row: who owns it, when, and which columns the import is allowed to touch.
This is the file where the *real* `CHECK` constraints matter. Everything here would pass against a
table without `V12` or without the length and pattern checks, and none of it would mean anything.
Four claims:
- **Ordering.** The claim is the workflow's first act, before any gate runs, and the names are on
the row while the import is still in flight — because a person watching the import list has to see
*which* model is being imported while it is being imported.
- **Ownership.** The columns the inserting side owns — `username`, `request_data`, `bucket_name`,
`file_name`, `run_type`, `created_at` — are never rewritten. The import updates status, the
reason, the two names and the JSON detail, and nothing else.
- **Refusal.** A row this workflow may not own is not written to at all. Not with an error, not with
a status: the loudness goes to the Temporal error, the log and the metric, because a row that
belongs to somebody else is not the place to be loud.
- **Conformance.** Every status value an import writes is one the real constraint accepts, a
training row's vocabulary still works beside it, and the feature added no column.
Two observations here are made from *inside* a running import, by wrapping the first call the step
under observation makes. That is a patch, and it is an observer rather than a failure injector: it
snapshots the row and delegates. Ordering cannot be asserted from a finished run — the terminal
write is precisely what erases the intermediate state — so the assertion has to stand somewhere in
the middle.
"""
from __future__ import annotations
import json
from typing import Any
import pytest
from sqlalchemy import text
from sqlalchemy.exc import IntegrityError
from temporalio.client import WorkflowFailureError
from e2e.import_harness import IMPORT_PREFIX, failure_types
from laborious.activities.model_import_errors import IMPORT_FAILURE_CATALOG
from laborious.utils.bundle.steps import ImportStep
# `asyncio` is applied per class rather than to the whole module: `TestConstraintConformance`
# talks to Postgres directly and its tests are synchronous, and a module-wide `asyncio` mark on
# a sync test is a pytest warning, not a no-op.
pytestmark = [pytest.mark.import_e2e]
FACTORY_PASSWORD = 'Factory1Model!Test'
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
# The columns the inserting side owns. The import may never rewrite one of these.
INSERTER_OWNED_COLUMNS = (
'username',
'request_data',
'bucket_name',
'file_name',
'run_type',
'created_at',
)
# The 13 columns V1…V12 leave. D21 says the feature adds none; this is what checks it.
EXPECTED_COLUMNS = {
'id',
'experiment_name',
'run_name',
'username',
'status',
'error_message',
'created_at',
'updated_at',
'bucket_name',
'file_name',
'request_data',
'orchestrator_response_data',
'run_type',
}
def _as_dict(value: Any) -> dict[str, Any]:
return json.loads(value) if isinstance(value, str) else (value or {})
@pytest.mark.asyncio
class TestOrdering:
"""When each write happens, observed from inside the run rather than after it."""
async def test_the_row_is_claimed_before_the_first_gate_runs(
self,
import_status_table: Any,
import_models_collection: Any,
import_bundle_factory: Any,
import_activities: 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:
"""Snapshot the row from inside `stat_object` — the first call gate 1 makes.
Standing here rather than after the run is not a convenience: the terminal write rewrites
status, reason and detail, so a finished row cannot say what it looked like before the gates.
`stat_object` is patched with a wrapper that records and then delegates — an observer, not a
failure. Everything the import does still really happens.
"""
client = import_activities.import_minio_repository.minio_client
real_stat = client.stat_object
snapshots: list[dict[str, Any] | None] = []
def observing_stat(bucket: str, key: str, *args: Any, **kwargs: Any) -> Any:
snapshots.append(import_read_status_row(observed_run_id[0]))
return real_stat(bucket, key, *args, **kwargs)
observed_run_id: list[int] = [0]
object_key, digest = import_upload_bundle(import_bundle_factory())
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
client.stat_object = observing_stat
try:
handle = await import_execute(
import_build_input(
import_run_id=observed_run_id[0],
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await handle.result()
finally:
client.stat_object = real_stat
assert snapshots, 'gate 1 never ran, so nothing was observed'
claimed = snapshots[0]
assert claimed is not None
assert claimed['status'] == RUNNING_STATUS, (
f'at the first gate the row was {claimed["status"]!r}, not the running status'
)
detail = _as_dict(claimed['orchestrator_response_data'])
assert detail['import']['workflow_id'] == handle.id
assert detail['import']['step'] in (
ImportStep.RECEIVED.value,
ImportStep.DOWNLOAD.value,
), detail
async def test_the_names_are_on_the_row_while_the_import_is_still_running(
self,
import_status_table: Any,
import_models_collection: Any,
import_names: dict[str, str],
import_bundle_factory: Any,
import_activities: 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:
"""Observed at the first call experiment creation makes, so nothing exists in MLflow yet.
This is the stronger form of the same claim task 7.1 makes from a failure: here the import
*succeeds*, and the names were still already written — so their presence is not an artefact
of the run having stopped.
"""
repository = import_activities.import_mlflow_repository
real_get_experiment = repository.get_experiment
snapshots: list[dict[str, Any] | None] = []
observed_run_id: list[int] = [0]
def observing_get_experiment(*args: Any, **kwargs: Any) -> Any:
snapshots.append(import_read_status_row(observed_run_id[0]))
return real_get_experiment(*args, **kwargs)
object_key, digest = import_upload_bundle(import_bundle_factory())
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
repository.get_experiment = observing_get_experiment
try:
handle = await import_execute(
import_build_input(
import_run_id=observed_run_id[0],
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
await handle.result()
finally:
repository.get_experiment = real_get_experiment
assert snapshots, 'experiment creation never ran'
mid_flight = snapshots[0]
assert mid_flight is not None
assert mid_flight['status'] == RUNNING_STATUS
assert mid_flight['experiment_name'] == import_names['experiment_name']
assert mid_flight['run_name'] == f'import-{import_names["model_name"]}-v2'
@pytest.mark.asyncio
class TestOwnership:
"""The inserting side's columns survive an import untouched, whatever the outcome."""
@pytest.mark.parametrize('outcome', ['success', 'failure'])
async def test_the_inserted_columns_are_never_rewritten(
self,
outcome: str,
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,
) -> None:
payload = (
import_bundle_factory()
if outcome == 'success'
else import_bundle_factory(include_metadata=False)
)
object_key, digest = import_upload_bundle(payload)
import_run_id = import_insert_status_row(
object_key=object_key,
expected_digest=digest,
username='owner.of.the.row@example.com',
)
# Snapshotted before the run, not after: `import_execute` waits for the workflow, so a
# snapshot taken from its return value would already be the post-import row and every
# comparison below would be a row against itself.
before = import_read_status_row(import_run_id)
assert before is not None
assert before['status'] == 'PENDING'
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),
)
)
if outcome == 'success':
await handle.result()
else:
with pytest.raises(WorkflowFailureError):
await handle.result()
after = import_read_status_row(import_run_id)
assert after is not None
for column in INSERTER_OWNED_COLUMNS:
assert after[column] == before[column], (
f'{column} was rewritten: {before[column]!r} -> {after[column]!r}'
)
# And the import *did* write what it owns, so this is not passing by doing nothing at all.
assert after['status'] != before['status']
assert after['updated_at'] >= before['updated_at']
@pytest.mark.asyncio
class TestClaimGuards:
"""A row this workflow may not own is not written to. At all."""
async def test_an_unknown_id_fails_the_run_and_writes_nothing(
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_bff_engine: Any,
import_assert_nothing_provisioned: Any,
) -> None:
object_key, digest = import_upload_bundle(import_bundle_factory())
# A real, claimable row exists — so "nothing was written" is a claim about *this* run
# choosing the right row, not about the table being empty.
bystander_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
before = _snapshot_table(import_bff_engine)
handle = await import_execute(
import_build_input(
import_run_id=bystander_id + 10_000,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
with pytest.raises(WorkflowFailureError) as failure:
await handle.result()
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
assert _snapshot_table(import_bff_engine) == before, 'the table changed'
import_assert_nothing_provisioned(
experiment_name=import_names['experiment_name'],
model_name=import_names['model_name'],
object_key=object_key,
)
async def test_a_training_row_cannot_be_claimed_by_an_import(
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_bff_engine: Any,
import_assert_nothing_provisioned: Any,
) -> None:
"""`run_type='TRAIN'` is a different feature's row, and the claim will not take it."""
object_key, digest = import_upload_bundle(import_bundle_factory())
training_id = import_insert_status_row(
object_key=object_key,
expected_digest=digest,
run_type='TRAIN',
experiment_name='a-training-experiment',
)
before = _snapshot_table(import_bff_engine)
handle = await import_execute(
import_build_input(
import_run_id=training_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
with pytest.raises(WorkflowFailureError) as failure:
await handle.result()
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
assert _snapshot_table(import_bff_engine) == before
import_assert_nothing_provisioned(
experiment_name=import_names['experiment_name'],
model_name=import_names['model_name'],
object_key=object_key,
)
async def test_a_row_held_by_another_workflow_cannot_be_taken(
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_bff_engine: Any,
import_assert_nothing_provisioned: Any,
) -> None:
"""Already running under somebody else's workflow id.
The claim is idempotent for *its own* workflow id and refuses every other, which is what
makes a duplicated workflow start safe and a stolen row impossible.
"""
object_key, digest = import_upload_bundle(import_bundle_factory())
held_id = import_insert_status_row(
object_key=object_key,
expected_digest=digest,
status=RUNNING_STATUS,
orchestrator_response_data={
'import': {'workflow_id': 'import-held-by-somebody-else', 'step': 'received'}
},
)
before = _snapshot_table(import_bff_engine)
handle = await import_execute(
import_build_input(
import_run_id=held_id,
object_key=object_key,
expected_digest=digest,
password_envelope=import_envelope(FACTORY_PASSWORD),
)
)
with pytest.raises(WorkflowFailureError) as failure:
await handle.result()
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
assert _snapshot_table(import_bff_engine) == before
import_assert_nothing_provisioned(
experiment_name=import_names['experiment_name'],
model_name=import_names['model_name'],
object_key=object_key,
)
@pytest.mark.asyncio
class TestInputAgainstTheRow:
"""A request that contradicts the row it claimed — recorded *on* that row, because it owns it."""
async def test_a_file_name_disagreeing_with_the_object_key_is_recorded_on_the_row(
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 difference from the claim guards above is where the error goes.
A row that could not be claimed is somebody else's, so nothing is written to it. This row
*was* claimed — the contradiction is only found afterwards — so the failure belongs on it,
where the person who submitted the import will see it.
"""
object_key, digest = import_upload_bundle(import_bundle_factory())
import_run_id = import_insert_status_row(
object_key=f'{IMPORT_PREFIX}a-different-object.sientia',
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(FACTORY_PASSWORD),
)
)
with pytest.raises(WorkflowFailureError):
await handle.result()
row = import_read_status_row(import_run_id)
assert row is not None
code, reason = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
assert row['status'] == VALIDATION_ERROR_STATUS
assert row['error_message'] == reason
detail = _as_dict(row['orchestrator_response_data'])['import']
assert detail['step'] == ImportStep.RECEIVED.value
assert detail['code'] == code.value
# The record still names the object *it* was given, not the one the request named.
assert row['file_name'] == f'{IMPORT_PREFIX}a-different-object.sientia'
import_assert_nothing_provisioned(
experiment_name=import_names['experiment_name'],
model_name=import_names['model_name'],
object_key=object_key,
)
@pytest.mark.asyncio
class TestTrainingRowsAreUntouched:
"""A training row sitting in the same table is never modified by an import."""
async def test_a_training_row_beside_an_import_is_unchanged(
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,
) -> None:
object_key, digest = import_upload_bundle(import_bundle_factory())
training_id = import_insert_status_row(
object_key=object_key,
expected_digest=digest,
run_type='TRAIN',
status='TRAINING_SUCCESS',
experiment_name='someone-elses-training',
username='trainer@example.com',
)
training_before = import_read_status_row(training_id)
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(FACTORY_PASSWORD),
)
)
await handle.result()
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
assert import_read_status_row(training_id) == training_before
class TestConstraintConformance:
"""The real constraints, exercised directly — and the column list D21 promised not to change."""
def test_every_status_an_import_writes_is_accepted_on_an_import_row(
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
) -> None:
"""`PENDING`, the running value, and both import terminals, on `run_type='IMPORT'`."""
row_id = import_insert_status_row(
object_key=f'{IMPORT_PREFIX}conformance.sientia', expected_digest='a' * 64
)
for status in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS, 'IMPORT_SUCCESS'):
_set_status(import_bff_engine, row_id, status)
_set_status(import_bff_engine, row_id, 'IMPORT_ERROR')
def test_the_terminal_values_do_not_cross_run_types(
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
) -> None:
"""`chk_experiment_run_terminal_status_run_type`, in both directions.
This is the constraint `V12` adds, and it is the reason the migration gates enabling the
worker: without it an import could write `TRAINING_SUCCESS` and a training run could write
`IMPORT_ERROR`, and the import list would be unreadable.
"""
import_row = import_insert_status_row(
object_key=f'{IMPORT_PREFIX}crossing.sientia', expected_digest='b' * 64
)
training_row = import_insert_status_row(
object_key=f'{IMPORT_PREFIX}crossing.sientia',
expected_digest='b' * 64,
run_type='TRAIN',
)
for row_id, forbidden in (
(import_row, 'TRAINING_SUCCESS'),
(import_row, 'TRAINING_ERROR'),
(training_row, 'IMPORT_SUCCESS'),
(training_row, 'IMPORT_ERROR'),
):
with pytest.raises(IntegrityError) as refusal:
_set_status(import_bff_engine, row_id, forbidden)
assert 'chk_experiment_run_terminal_status_run_type' in str(refusal.value)
# The shared stage values are deliberately outside the pairing: both run types pass through.
for row_id in (import_row, training_row):
for shared in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS):
_set_status(import_bff_engine, row_id, shared)
def test_a_training_row_still_walks_its_own_vocabulary(
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
) -> None:
training_row = import_insert_status_row(
object_key=f'{IMPORT_PREFIX}training.sientia',
expected_digest='c' * 64,
run_type='TRAIN',
)
for status in ('PENDING', RUNNING_STATUS, 'TRAINING_SUCCESS'):
_set_status(import_bff_engine, training_row, status)
_set_status(import_bff_engine, training_row, 'TRAINING_ERROR')
def test_the_feature_added_no_column(
self, import_status_table: Any, import_bff_engine: Any
) -> None:
"""Design D21: the error code and the workflow id live in the existing JSONB column.
A quiet `ALTER TABLE` would be the easiest way for this feature to break the BFF that owns
the table, and the hardest to notice, so it is asserted rather than trusted.
"""
with import_bff_engine.begin() as connection:
columns = {
row[0]
for row in connection.execute(
text(
'SELECT column_name FROM information_schema.columns '
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
)
)
}
assert columns == EXPECTED_COLUMNS, (
f'unexpected: {sorted(columns - EXPECTED_COLUMNS)}, '
f'missing: {sorted(EXPECTED_COLUMNS - columns)}'
)
def _set_status(engine: Any, row_id: int, status: str) -> None:
"""Write one status onto one row, so the real constraint gets to answer."""
with engine.begin() as connection:
connection.execute(
text('UPDATE public.experiment_run SET status = :status WHERE id = :id'),
{'status': status, 'id': row_id},
)
def _snapshot_table(engine: Any) -> list[dict[str, Any]]:
"""Every row, ordered, as plain values — for "the table is byte-identical" assertions."""
with engine.begin() as connection:
return [
dict(row)
for row in connection.execute(
text('SELECT * FROM public.experiment_run ORDER BY id')
).mappings()
]

112
e2e/test_predictions_batch_joblib_fallback.py → e2e/test_predictions_batch_joblib_flavor.py
View File

@@ -1,9 +1,10 @@
"""
End-to-end tests for the joblib fallback on the MLflow model load paths.
End-to-end tests for the explicit `joblib` flavor on the MLflow model load paths.
`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn
estimator serialized with `joblib.dump(..., compress=3)`, so a successful run
proves the fallback really deserializes a compressed pickle and serves it.
`joblib.load` is NOT mocked here: the artifact fixture holds a real object serialized with
`joblib.dump(..., compress=3)` whose class lives only inside the artifact's `code/` directory.
A passing run therefore proves the flavor decompresses the pickle **and** resolves the module,
without any `mlflow.*.load_model` being called for that model.
"""
import pandas as pd
@@ -12,7 +13,7 @@ from sqlalchemy import text
from temporalio.testing import WorkflowEnvironment
from temporalio.worker import Worker
from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME
from e2e.conftest import JOBLIB_FLAVOR_MODEL_NAME
from e2e.helpers import (
assert_prediction,
assert_stop,
@@ -25,8 +26,15 @@ from laborious.workflows.predictions_batch import PredictionsBatch
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch')
def build_input_data(model_id: int, model_name: str) -> dict:
def build_input_data(
model_id: int,
model_name: str,
transform_flavor: str = 'sklearn',
predict_flavor: str = 'sklearn',
) -> dict:
"""Workflow input for a plain transform + predict run of a given model."""
return {
'schedule_name': 'test-schedule',
@@ -55,15 +63,15 @@ def build_input_data(model_id: int, model_name: str) -> dict:
'prediction_store_policy': 'lts:1',
'model_config': {
'retention_minutes': 0,
'transform_flavor': 'sklearn',
'predict_flavor': 'sklearn',
'transform_flavor': transform_flavor,
'predict_flavor': predict_flavor,
},
'datetime_columns': ['timestamp', 'created_at'],
}
def expected_fallback_prediction(estimator) -> float:
"""Prediction the fixture estimator produces for the mocked transform output."""
def expected_joblib_prediction(model) -> float:
"""Prediction the artifact's model produces for the transform output of this suite."""
features = pd.DataFrame(
{
'feature_1': [0.234],
@@ -71,7 +79,7 @@ def expected_fallback_prediction(estimator) -> float:
'timestamp': ['2024-01-01 12:00:00'],
}
)
return float(estimator.predict(features)[0])
return float(model.predict(features)[0])
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
@@ -98,42 +106,86 @@ def logged_messages(mock_logger, level: str) -> str:
return ' '.join(str(logged.args[0]) for logged in calls)
def loader_uris(patch_mlflow) -> list[str]:
"""Every model URI handed to an MLflow flavor loader during the run."""
return [
str(loaded.args[0])
for flavor in FLAVOR_LOADERS
for loaded in getattr(patch_mlflow, flavor).load_model.call_args_list
]
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_recovered_by_joblib_fallback(
async def test_prediction_workflow_with_joblib_flavor(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
joblib_fallback_artifact,
joblib_flavor_artifacts,
mlflow_client,
mock_logger,
patch_mlflow,
):
"""Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands."""
"""`predict_flavor: 'joblib'` deserializes the real artifact, and the prediction lands."""
client = temporal_test_env.client
model_id = 271
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback')
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor')
)
assert_single_prediction_close(
postgres_engine,
model_id,
expected_fallback_prediction(joblib_fallback_artifact['estimator']),
expected_joblib_prediction(joblib_flavor_artifacts['prediction_model']['model']),
)
assert mlflow_client.download_artifacts.call_count == 1
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning')
# The joblib flavor never builds a `models:/` URI, so no flavor loader saw this model.
assert all(JOBLIB_FLAVOR_MODEL_NAME not in uri for uri in loader_uris(patch_mlflow))
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_without_fallback_is_unchanged(
async def test_prediction_workflow_with_joblib_transform_flavor(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
joblib_flavor_artifacts,
mlflow_client,
patch_mlflow,
mock_mlflow_models,
):
"""`transform_flavor: 'joblib'` downloads and deserializes `data_model` on its own."""
client = temporal_test_env.client
model_id = 274
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, transform_flavor='joblib')
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-transform')
)
assert_prediction(postgres_engine, model_id, prediction=0.5)
mock_mlflow_models['predict_model'].predict.assert_called()
assert mlflow_client.download_artifacts.call_count == 1
assert mlflow_client.download_artifacts.call_args.args[1] == 'data_model'
# The joblib flavor resolves no artifact URI, so `get_run` is never reached and no flavor
# loader was given a `data_model` URI.
patch_mlflow.get_run.assert_not_called()
assert all(not uri.endswith('/data_model') for uri in loader_uris(patch_mlflow))
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
@@ -141,14 +193,14 @@ async def test_prediction_workflow_without_fallback_is_unchanged(
mlflow_client,
mock_mlflow_models,
):
"""A model whose flavor loader works keeps the current result and downloads nothing."""
"""A model on an MLflow flavor keeps the current result and downloads nothing."""
client = temporal_test_env.client
model_id = 272
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, 'test_model')
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback')
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-sklearn-flavor')
)
assert_prediction(postgres_engine, model_id, prediction=0.5)
@@ -158,29 +210,31 @@ async def test_prediction_workflow_without_fallback_is_unchanged(
@pytest.mark.asyncio
@pytest.mark.integration
async def test_prediction_workflow_reports_both_causes_when_fallback_fails(
async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause(
temporal_test_env: WorkflowEnvironment,
temporal_worker: Worker,
test_activities: Activities,
postgres_engine,
joblib_fallback_artifact_without_pickle,
joblib_flavor_artifact_without_pickle,
mlflow_client,
mock_logger,
):
"""Both attempts fail: workflow stops and the logged traceback carries both causes."""
mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle
"""The joblib flavor failing stops the workflow with one cause, not a chain of two."""
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: (
joblib_flavor_artifact_without_pickle
)
client = temporal_test_env.client
model_id = 273
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME)
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
await start_and_await_workflow(
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed')
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor-failed')
)
assert_stop(postgres_engine, model_id)
logged = logged_messages(mock_logger, 'debug')
assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged
assert 'No .pkl file found in prediction_model artifact' in logged
assert 'Original error' not in logged

126
e2e/test_retrain_joblib_flavor.py Normal file
View File

@@ -0,0 +1,126 @@
"""
End-to-end test for retraining a model served through the explicit `joblib` flavor.
Retraining is not part of the `predictions_batch` workflow, so it gets its own module. The test
is a write→read round trip in one run: `retrain_model` loads the production artifact through the
joblib flavor, `log_model` writes the retrained model back as `model.pkl` + `code/`, and the
written artifact is then loaded again through the same flavor. Neither `joblib.dump` nor
`joblib.load` is mocked; only MLflow's own I/O is.
"""
import importlib
import shutil
import sys
from pathlib import Path
import pandas as pd
import pytest
from e2e.conftest import JOBLIB_FLAVOR_CODE_MODULE, JOBLIB_FLAVOR_MODEL_NAME
from laborious.activities.activities import Activities
RETRAIN_INDEX = [
'2024-01-01 12:00:00+0000',
'2024-01-01 12:01:00+0000',
'2024-01-01 12:02:00+0000',
]
METADATA = {
'model_id': 'test_model',
'model_name': JOBLIB_FLAVOR_MODEL_NAME,
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
}
def build_retrain_data() -> pd.DataFrame:
"""Training frame indexed by timestamp strings, carrying the target column."""
return pd.DataFrame(
{
'sensor_1': [60.0, 61.0, 62.0],
'sensor_2': [78.2, 78.4, 78.6],
'target': [1.0, 2.0, 3.0],
},
index=RETRAIN_INDEX,
)
class FakeArtifactStore:
"""Minimal stand-in for the run's artifact tree, so the written layout is inspectable."""
def __init__(self, root: Path):
self.root = root
def log_artifact(self, local_path, artifact_path=None):
if artifact_path is None:
# Retraining also logs the CSVs at the run root; irrelevant to this test.
return
target = self.root / artifact_path
target.mkdir(parents=True, exist_ok=True)
shutil.copy(local_path, target / Path(local_path).name)
def log_artifacts(self, local_dir, artifact_path=None):
target = self.root / artifact_path if artifact_path else self.root
shutil.copytree(local_dir, target, dirs_exist_ok=True)
@pytest.mark.asyncio
@pytest.mark.integration
async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
test_activities: Activities,
joblib_flavor_artifacts,
mlflow_client,
patch_mlflow,
tmp_path,
):
"""A joblib retrain writes `model.pkl` + `code/`, and the joblib flavor reads it back."""
repository = test_activities.model_monitoring_repository
store = FakeArtifactStore(tmp_path / 'run-retrained')
patch_mlflow.log_artifact.side_effect = store.log_artifact
patch_mlflow.log_artifacts.side_effect = store.log_artifacts
patch_mlflow.get_run.return_value.data.params = {}
patch_mlflow.start_run.return_value.__enter__.return_value.info.run_id = 'run-retrained'
output = await repository.retrain_model(
data=build_retrain_data(),
model_name=JOBLIB_FLAVOR_MODEL_NAME,
model_config={
'target': 'target',
'transform_flavor': 'joblib',
'predict_flavor': 'joblib',
},
metadata=METADATA,
)
assert output['success'] is True, output.get('traceback')
# Both models were loaded through the joblib flavor, not through a flavor loader.
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
assert sorted(downloaded) == ['data_model', 'prediction_model']
assert patch_mlflow.sklearn.load_model.call_count == 0
assert patch_mlflow.pyfunc.load_model.call_count == 0
written = store.root / 'prediction_model'
assert (written / 'model.pkl').is_file()
assert (written / 'code' / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').is_file()
assert not (written / 'MLmodel').exists()
assert (store.root / 'data_model' / 'model.pkl').is_file()
# Drop the module the first load put in place: the reload must resolve it from the `code/`
# the retrain wrote, not from what is already imported.
for code_dir in (artifact['code_dir'] for artifact in joblib_flavor_artifacts.values()):
if code_dir in sys.path:
sys.path.remove(code_dir)
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
importlib.invalidate_caches()
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)
reloaded = await repository.load_predict_model(JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib')
features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)
assert reloaded.predict(features) == pytest.approx(expected)
assert str(written / 'code') in sys.path