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

14
.env.example
View File

@@ -55,14 +55,12 @@ IMPORT_BUNDLE_PREFIX="imported_models/"
# MongoDB collection holding the model listing the import writes its document into.
IMPORT_MODELS_COLLECTION="models"
# The import log lives in the Spring Boot BFF database, a second Postgres
# connection independent of POSTGRES_* above. The schema (public) and the table
# (experiment_run) are constants of the code, not configuration.
IMPORT_STATUS_DB_HOST="localhost"
IMPORT_STATUS_DB_PORT="5432"
# The import log lives in the Spring Boot BFF database on the POSTGRES_* server
# above: same host, port, user and password, another database name. Only that
# name is configuration here; the schema (public) and the table (experiment_run)
# are constants of the code. Set IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} only
# if the log really is on a different server than POSTGRES_* — they override it.
IMPORT_STATUS_DB_NAME="sientia-core-mlops-bff"
IMPORT_STATUS_DB_USER="mlops_bff_user"
IMPORT_STATUS_DB_PASSWORD="bff_password"
# --- scripts/import_model_manual_run.py only ---------------------------------
# Inputs for the manual run that stands in for the frontend. The script writes
@@ -70,6 +68,6 @@ IMPORT_STATUS_DB_PASSWORD="bff_password"
IMPORT_SCRIPT_BUNDLE_PATH="tests/fixtures/bundle/golden_model_v3.sientia"
IMPORT_SCRIPT_BUNDLE_PASSWORD="the-password-the-bundle-was-written-with"
IMPORT_SCRIPT_USERNAME="manual.run@example.com"
# Cell 2's guard: restate IMPORT_STATUS_DB_HOST here (or in the cell) after reading what cell 1
# Cell 2's guard: restate POSTGRES_HOST here (or in the cell) after reading what cell 1
# printed. Left empty on purpose — the script refuses to write until it matches.
IMPORT_SCRIPT_CONFIRM_TARGET=""

74
README.md
View File

@@ -45,6 +45,7 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
- [Testing](#testing)
- [Test Structure](#test-structure)
- [Test Execution](#test-execution)
- [End-to-End Tests](#end-to-end-tests-e2e)
- [Monitoring and Metrics](#monitoring-and-metrics)
- [Application Health Metrics](#application-health-metrics)
- [Prediction Operation Metrics](#prediction-operation-metrics)
@@ -881,6 +882,57 @@ pytest tests/laborious/activities/test_gates.py
pytest tests/laborious/workflows/test_predictions_batch.py
```
### End-to-End Tests (`e2e/`)
The `e2e/` suite runs against real dependencies in Docker containers via
[testcontainers](https://testcontainers.com/): PostgreSQL, MinIO, MongoDB, and — for the model-import
suite — an `mlflow server` subprocess started from this repository's venv. It lives outside
`testpaths`, so it is never collected by the unit gate.
**One-time setup, required before running anything under `e2e/`:**
```bash
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
```
Without it, **every** e2e test fails at setup with a Docker mount error and never reaches an
assertion:
```
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.
```
The cause is not the tests. Before starting any test container, testcontainers starts a reaper
container ("Ryuk") that bind-mounts the Docker socket so it can clean up if the test process dies
without teardown. Under Docker Desktop on Linux that socket is `~/.docker/desktop/docker.sock`, a path
Docker Desktop's File Sharing does not expose, so the reaper's own container cannot start. Note the
symptom is `ERROR`, not `FAILED` — the failure is in fixture setup.
Disabling it costs nothing in a normal run: every fixture stops its own containers in teardown. Ryuk
only matters when the process is killed outright (`kill -9`, a crash, a hard stop from the IDE). After
one of those, clean up with:
```bash
docker ps -a --filter "label=org.testcontainers=true" -q | xargs -r docker rm -f
```
**Running it:**
```bash
source ./venv/bin/activate
# the predictions and OPC suites
pytest e2e -m "not import_e2e"
# the .sientia model-import suite (opt-in; see e2e/scenarios.md § 4)
pytest e2e -m import_e2e
```
Full scenario catalogue and the import suite's fidelity contract:
[`e2e/scenarios.md`](e2e/scenarios.md).
## 📊 Monitoring and Metrics
The Laborious system exposes comprehensive Prometheus metrics for operational visibility and performance monitoring:
@@ -950,11 +1002,11 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | Gate 4 uncompressed size ceiling | `4294967296` | No |
| `IMPORT_MAX_COMPRESSION_RATIO` | Gate 4 compression ratio ceiling | `200` | No |
| `IMPORT_MODELS_COLLECTION` | MongoDB collection holding the model listing | `models` | No |
| `IMPORT_STATUS_DB_HOST` | Import log database host (the BFF's) | `localhost` | If importing |
| `IMPORT_STATUS_DB_PORT` | Import log database port | `5432` | No |
| `IMPORT_STATUS_DB_NAME` | Import log database name | `sientia-core-mlops-bff` | No |
| `IMPORT_STATUS_DB_USER` | Import log database user | `mlops_bff_user` | If importing |
| `IMPORT_STATUS_DB_PASSWORD` | Import log database password | _(empty)_ | If importing |
| `IMPORT_STATUS_DB_NAME` | Import log database name (the BFF's, on the `POSTGRES_*` server) | `sientia-core-mlops-bff` | If importing |
| `IMPORT_STATUS_DB_HOST` | Override: import log host, when it is not the `POSTGRES_*` server | `POSTGRES_HOST` | No |
| `IMPORT_STATUS_DB_PORT` | Override: import log port | `POSTGRES_PORT` | No |
| `IMPORT_STATUS_DB_USER` | Override: import log user | `POSTGRES_USER` | No |
| `IMPORT_STATUS_DB_PASSWORD` | Override: import log password | `POSTGRES_PASSWORD` | No |
| `OPC_CONFIG` | OPC server configuration (JSON) | `{}` | No |
| `OPC_ID` | OPC server identifier | `1` | No |
| `OPC_URL` | OPC server URL | `opc.tcp://localhost:4840` | No |
@@ -1055,6 +1107,18 @@ Legacy MinIO object layout (relative key):
`training_datasets/{model_name}/{object_prefix}_{timestamp}.parquet` where `object_prefix` is sanitized
(slashes replaced by underscores) to keep a stable model-level directory.
## Model Flavors
`model_config.predict_flavor` / `model_config.transform_flavor` accept `sklearn`, `pyfunc`,
`pytorch` and `joblib`. The flavor is configuration, never discovered from the artifact, and there is
no fallback between flavors: a failing loader propagates.
`joblib` reads the artifact directly — download, first top-level `.pkl`, `joblib.load` — with the
artifact's `code/` put on `sys.path` first, and it writes on retrain (`model.pkl` + `code/`, no
`MLmodel`, so only this runtime reads it back).
Full reference, including how to switch an existing model over: **[docs/model-flavors.md](docs/model-flavors.md)**.
## OPC UA Communication
Full reference: **[docs/opc-communication.md](docs/opc-communication.md)** (connection lifecycle, Tier-1 `Bad*` reconnect, connection lock / session readiness, metrics, PostgreSQL confidence **12** vs **14**, tests).

112
docs/model-flavors.md Normal file
View File

@@ -0,0 +1,112 @@
# Model flavors
The *flavor* is how this runtime is told to deserialize a model. It is **configuration, never
discovery**: the value comes from the model document in the `models` collection and the code never
reads the artifact's `MLmodel` manifest to guess it.
```json
{
"model_config": {
"transform_flavor": "sklearn",
"predict_flavor": "joblib",
"retention_minutes": 60,
"target": "SE"
}
}
```
`predict_flavor` drives `load_predict_model` (artifact `prediction_model`); `transform_flavor` drives
`load_transform_model` (artifact `data_model`). Both default to `sklearn` when the key is absent.
## Accepted values
| Flavor | Read | Write (retrain) | How it loads |
|---|---|---|---|
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` |
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; the wrapper path also unwraps `._model_impl.python_model` |
| `pytorch` | yes | yes | `mlflow.pytorch.load_model(<uri>)` |
| `joblib` | yes | yes | downloads the artifact and `joblib.load`s its pickle — **no MLflow flavor loader involved** |
Anything else raises `ValueError` with `INVALID_FLAVOR_MESSAGE`, **before** any artifact is resolved
or downloaded. The same message and the same four values govern reading and writing.
**There is no fallback between flavors.** A failing loader propagates: one attempt, one cause. A
model configured with the wrong flavor fails visibly instead of being rescued silently.
## The `joblib` flavor
Some models the pipeline produces do not load through an MLflow flavor loader even though the
artifact is intact: the `.pkl` is in the compressed form `joblib.dump` writes, which MLflow's
`pickle.load` does not read. Declaring `joblib` is how such a model is served.
**Reading.** Resolve the Production run, `client.download_artifacts(run_id, <artifact>)`, take the
**first top-level `.pkl` in name order**, and `joblib.load` it. The listing does not recurse, so a
nested wrapper artifact (`artifacts/*.pkl`) is never deserialized in place of the model. No
`MLmodel` is parsed and no `models:/<name>/production` URI is built.
**The artifact's `code/` goes on `sys.path` first.** A model whose class lives inside the artifact
would otherwise raise `ModuleNotFoundError` at unpickle time. Before `joblib.load`, the flavor
prepends `<artifact>/code` to `sys.path` when that directory exists — the directory itself, matching
MLflow's own convention where `code/utils/…` imports as the package `utils`. The entry is not
duplicated and is **not** removed after the load, because the deserialized object may import
lazily. A missing `code/` is not an error, and no extra download is performed: `download_artifacts`
already brings `code/` along.
> `sys.path` is process-global. Two different models that embed a package with the same name
> (`utils/`, typically) and are served by the same worker resolve to whichever entered first — the
> second silently runs the first one's code. This is pre-existing behaviour of MLflow's own pyfunc
> path with the same artifact layout; the mitigation, if it ever bites, is one worker per model.
**Writing / retraining.** `joblib` is a write flavor too, so a joblib model retrains. `log_model`
writes `joblib.dump(model, 'model.pkl', compress=3)` under the artifact path and re-logs the source
artifact's `code/` to `<artifact>/code`, so the retrained model stays deserializable. Registration
and promotion to `Production` are unchanged.
**The written artifact has no `MLmodel` manifest.** Only this runtime's `joblib` flavor reads it:
`mlflow.sklearn.load_model` and `models:/<name>/production` both fail on it. A model retrained in
joblib **stays** joblib — `predict_flavor` needs no edit after a retrain, and no consumer outside
this runtime (the MLflow UI, MLflow serving, another service) recognizes that artifact as a model.
**What `joblib` is for.** A pickled object exposing `predict(data)` — and `fit(data)` if it is
retrained. It is **not** for a bare pyfunc `PythonModel` wrapper: `get_prediction_data` and
`get_cached_operation` call `predict(data)` for every non-pyfunc flavor, so a wrapper expecting
`predict(context, model_input)` raises `TypeError` at predict and retrain time.
**Observability.** A joblib load is an ordinary load: `MODEL_READ_LAG` / `MODEL_READ_COUNT` on
success and `MODEL_READ_ERROR_COUNT` on failure, under the base `operation_type`
(`load_predict_model` / `load_transform_model`). A joblib write emits the usual `MODEL_WRITE_*`
under `log_model`. No suffixed operation type, no new metric, no new label.
## Switching a model to `joblib`
The runtime does not migrate documents. Editing the `models` collection is an operator step, and the
order matters:
1. **Inventory first, before the deploy.** Collect the models that fail in their configured flavor
loader — run predict/transform and note which ones raise. There is no automatic discovery
afterwards: once deployed, such a model simply fails, and the failure names the flavor loader,
not the fix.
2. **Edit the document.** Set `model_config.predict_flavor: 'joblib'` and/or
`model_config.transform_flavor: 'joblib'` for each model in the inventory. Nothing else changes;
models already on a working flavor need no edit and behave identically.
3. **Imported models arrive ready.** The `import_model` workflow writes
`predict_flavor: 'joblib'` itself — see [`model-import.md`](model-import.md) § 12.
4. **Check the first retrain.** In MLflow, confirm the new run carries
`prediction_model/model.pkl` **and** `prediction_model/code/` when the source model had a
`code/` directory, and that the promoted version loads on the next predict.
5. **Rolling back is not just a revert.** Reverting the code leaves any artifact a joblib retrain
already wrote in MLflow: that Production version stops being loadable, so the rollback also means
promoting the previous version back in the registry, and setting the documents' flavors back —
manual, in reverse order.
## Known gap: how the run id is resolved
`get_model_run_id` resolves the Production run with `source.split('/')[2]` on the registered
version's `source`. **This has not been confirmed against the real tracking server.** No artifact
scheme reproduced locally puts the run id in that position (a file store yields `''`, S3 yields the
bucket, `mlflow-artifacts:/` yields the experiment id), yet predict works in production — so the
real `source` has a shape none of the local probes produced.
This is not specific to `joblib`: every flavor that resolves a run goes through the same line, and
the joblib read path depends on it. Confirm it with a **read-only** query against the tracking
server before relying on it, ideally before the next deploy.

59
docs/model-import.md
View File

@@ -88,14 +88,18 @@ is at fault.
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | 4 GiB | gate 4 |
| `IMPORT_MAX_COMPRESSION_RATIO` | `200` | gate 4 |
| `IMPORT_MODELS_COLLECTION` | `models` | the MongoDB collection that *is* the model listing |
| `IMPORT_STATUS_DB_HOST` | `localhost` | the BFF database holding the import log |
| `IMPORT_STATUS_DB_PORT` | `5432` | |
| `IMPORT_STATUS_DB_NAME` | `sientia-core-mlops-bff` | |
| `IMPORT_STATUS_DB_USER` | `mlops_bff_user` | |
| `IMPORT_STATUS_DB_PASSWORD` | empty | |
| `IMPORT_STATUS_DB_NAME` | `sientia-core-mlops-bff` | the BFF database holding the import log |
| `IMPORT_STATUS_DB_HOST` | `POSTGRES_HOST` | override: the log is on another server |
| `IMPORT_STATUS_DB_PORT` | `POSTGRES_PORT` | override |
| `IMPORT_STATUS_DB_USER` | `POSTGRES_USER` | override |
| `IMPORT_STATUS_DB_PASSWORD` | `POSTGRES_PASSWORD` | override |
`IMPORT_STATUS_DB_*` is a **second** Postgres connection, independent of the one Laborious already
holds for predictions. The import writes to that database and no other, and to one table in it.
The import log is a **second connection but the same server**: host, port, user and password come
from `POSTGRES_*`, and only the database name is import configuration, because the log lives in the
BFF's database rather than in the one Laborious holds for predictions. The import writes to that
database and no other, and to one table in it. The four `IMPORT_STATUS_DB_*` connection variables
are overrides for a deployment whose log genuinely sits elsewhere — the e2e suite, which runs the
two databases as two containers, is exactly that case.
A missing or malformed `IMPORT_PASSWORD_KEY` is an import failure at the decryption step, not a boot
failure: the worker starts either way.
@@ -324,10 +328,11 @@ A `COMPLETED` import is **not** a servable model. Stated plainly:
So an operator must promote the version to `Production` and set `active: true`. This workflow never
promotes and never activates — that runbook is QTZPOC-21's.
One more thing to know before serving an imported model: `predict_flavor: 'joblib'` is a **read-only**
flavor. It is what the artifact actually is, and it is enough to load and predict, but nothing in this
platform writes that flavor — so an imported model cannot be retrained until a write flavor is
configured for it.
One more thing to know before serving an imported model: `predict_flavor: 'joblib'` is what the
artifact actually is, and it loads, predicts **and** retrains — `log_model` writes that flavor too,
as `model.pkl` + `code/`. The artifact it writes carries no `MLmodel` manifest, so an imported model
stays readable only by this runtime, before and after a retrain. Full contract:
[`model-flavors.md`](model-flavors.md).
## 14. Retry, and what does not retry
@@ -347,3 +352,35 @@ change.
There is no rollback, no compensation, no resume and no reconciliation. A failed import stays failed;
trying again is a new upload and a new row.
## 15. How this contract is verified
Two suites, and they prove different things.
**Unit** — `pytest tests/`, part of the standard gate (`validate --project-name=laborious`). Covers
the bundle reader against committed golden bytes, every gate, the failure catalog, the status SQL and
the workflow's ordering, all against fakes.
**End to end** — `pytest e2e -m import_e2e`, opt-in and outside `testpaths`. Runs the real workflow on
a real worker against real MinIO, a real Postgres holding this document's `public.experiment_run`
(with and without `V12`), a real MongoDB with the unique index on `models.id`, and a real
`mlflow server`. Nothing on the import path is mocked. This is what proves the claims a fake cannot:
that the `CHECK` constraints of § 6 accept every status written, that the registry source of § 11
resolves through the existing loader, that the `V12` failure shape of § 7 is what it says, and that no
plaintext password reaches the Temporal event history.
Requires Docker plus one one-time setup per machine:
```bash
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
```
Without it every e2e test fails in fixture setup with a Docker mount error — see
[`README.md` § End-to-End Tests](../README.md#end-to-end-tests-e2e) for why and what it costs.
Scenario catalogue: [`e2e/scenarios.md`](../e2e/scenarios.md) § 4.
Requirement-to-test mapping: `openspec/changes/import-model-e2e-suite/traceability.md`.
Doing one import by hand, standing in for the frontend that does not exist yet:
[`scripts/import_model_manual_run.py`](../scripts/import_model_manual_run.py) — see
[`README.md` § Running One Model Import by Hand](../README.md#running-one-model-import-by-hand).

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

4
laborious/activities/activities.py
View File

@@ -82,8 +82,8 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
notification handler keeps its own client either way.
import_config: `build_import_config()`; read from the environment when absent
import_status_config: `build_import_status_config()`; read from the environment when
absent. This is the BFF database holding the import log, a second Postgres
connection that is never the one Storage holds.
absent. This is the BFF database holding the import log: another database on the
`POSTGRES_*` server, and a second connection that is never the one Storage holds.
Raises:
Exception: If any parent class initialization fails

6
laborious/activities/model_import.py
View File

@@ -288,8 +288,10 @@ class ModelImport(SientiaMonitoring):
"""
Args:
status_config: `build_import_status_config()` — the BFF database holding the import log.
This is a **second** Postgres connection, independent of the one Laborious already
holds, which is what makes the transactional independence of design D7 literal.
This is a **second** Postgres connection — another database, on the `POSTGRES_*`
server unless the IMPORT_STATUS_DB_* overrides move it — and its separateness from
the one Laborious already holds is what makes design D7's transactional
independence literal.
import_config: `build_import_config()` — bucket, prefix, ceilings, work directory,
envelope key and the models collection name.
minio_repository: the shared MinIO repository, for the uploaded bundle.

33
laborious/utils/connectors_config.py
View File

@@ -136,27 +136,36 @@ def build_import_config() -> dict[str, Any]:
def build_import_status_config() -> dict[str, Any]:
"""
Build the connection to the import log database — a **second** Postgres connection.
Build the connection to the import log database — the POSTGRES_* server, another database.
The import log is `public.experiment_run` in the database owned by the Spring Boot BFF
(`sientia-core-mlops-bff`), not the `sientia` database Laborious itself uses. The table and
schema are constants of the code, never inputs; only the connection is configuration.
(`sientia-core-mlops-bff`), which sits on the same Postgres server as the `sientia` database
Laborious itself uses. Only the database name differs, so only the database name has a variable
of its own; host, port, user and password come from POSTGRES_*. The table and schema are
constants of the code, never inputs.
The four IMPORT_STATUS_DB_* connection variables remain as overrides, for a deployment where the
log genuinely lives on another server — the e2e suite is exactly that, two containers.
Environment Variables:
IMPORT_STATUS_DB_HOST: host (default: localhost)
IMPORT_STATUS_DB_PORT: port (default: 5432)
IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff)
IMPORT_STATUS_DB_USER: user (default: mlops_bff_user)
IMPORT_STATUS_DB_PASSWORD: password (default: empty)
POSTGRES_HOST: server host, unless IMPORT_STATUS_DB_HOST overrides it (default: localhost)
POSTGRES_PORT: server port, unless IMPORT_STATUS_DB_PORT overrides it (default: 5432)
POSTGRES_USER: user, unless IMPORT_STATUS_DB_USER overrides it (default: sientia)
POSTGRES_PASSWORD: password, unless IMPORT_STATUS_DB_PASSWORD overrides it (default: empty)
Returns:
dict: import status database configuration dictionary
"""
# Written as `override or fallback or default` rather than
# `getenv(override) or getenv(fallback, default)`: inside an `or`, the two-argument `getenv`
# widens to `str | None` for mypy, so the `int(...)` below fails to type-check. Putting the
# default last as its own term keeps every value a plain `str` and reads in precedence order.
return {
'host': getenv('IMPORT_STATUS_DB_HOST', 'localhost'),
'port': int(getenv('IMPORT_STATUS_DB_PORT', '5432')),
'dbname': getenv('IMPORT_STATUS_DB_NAME', 'sientia-core-mlops-bff'),
'user': getenv('IMPORT_STATUS_DB_USER', 'mlops_bff_user'),
'password': getenv('IMPORT_STATUS_DB_PASSWORD', ''),
'host': getenv('IMPORT_STATUS_DB_HOST') or getenv('POSTGRES_HOST') or 'localhost',
'port': int(getenv('IMPORT_STATUS_DB_PORT') or getenv('POSTGRES_PORT') or '5432'),
'dbname': getenv('IMPORT_STATUS_DB_NAME') or 'sientia-core-mlops-bff',
'user': getenv('IMPORT_STATUS_DB_USER') or getenv('POSTGRES_USER') or 'sientia',
'password': getenv('IMPORT_STATUS_DB_PASSWORD') or getenv('POSTGRES_PASSWORD') or '',
'schema': 'public',
}

379
laborious/utils/repository/model_repository.py
View File

@@ -16,11 +16,11 @@ Capabilities:
import ctypes
import gc
import inspect
import sys
import tempfile
import threading
import time
import traceback
from collections.abc import Callable
from datetime import datetime, timedelta
from io import StringIO
from os import environ, listdir, makedirs, path
@@ -45,7 +45,7 @@ ARTIFACTS_PATH = './tmp/artifacts'
TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch' or 'joblib'."
def force_memory_release(logger: Logger):
@@ -354,23 +354,19 @@ class MLFlowRepository(SientiaMonitoring):
self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata)
return dataframe
def _find_model_pickle(
self, local_dir: str, artifact_name: str, flavor_error: Exception
) -> str:
def _find_model_pickle(self, local_dir: str, artifact_name: str) -> str:
"""
Locate the pickle file of a downloaded model artifact.
Fallback path only: this is reached after the MLflow flavor loader has
already failed. The target artifact is always a single `.pkl` at the top
level of the artifact directory, so the listing does not recurse —
nested wrapper artifacts must never be deserialized in place of the
model. No other extension is considered and `MLmodel` is not parsed.
Primary path of the `joblib` flavor. The target artifact is always a
single `.pkl` at the top level of the artifact directory, so the listing
does not recurse — nested wrapper artifacts must never be deserialized
in place of the model. No other extension is considered and `MLmodel` is
not parsed.
Args:
local_dir (str): Local directory holding the downloaded artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader,
reported in the message so the original cause is not lost.
Returns:
str: Path to the first top-level `.pkl` file, in name order.
@@ -388,7 +384,6 @@ class MLFlowRepository(SientiaMonitoring):
if not pkl_files:
raise ValueError(
f'No .pkl file found in {artifact_name} artifact. '
f'Original error: {flavor_error}. '
f'Top-level entries in {local_dir}: {entries}'
)
@@ -398,24 +393,28 @@ class MLFlowRepository(SientiaMonitoring):
self,
local_dir: str,
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any] | None = None,
) -> Any:
"""
Deserialize a downloaded model artifact with `joblib`.
Fallback path only. `joblib.load` reads the compressed `.pkl` produced
by `joblib.dump`, which is what the MLflow flavor loader fails to read.
The deserialized object is returned as is: no `._model_impl.python_model`
unwrap is applied, because the pickled object already is the model.
Primary path of the `joblib` flavor. `joblib.load` reads the compressed
`.pkl` produced by `joblib.dump`, which is what the MLflow flavor
loaders do not read. The deserialized object is returned as is: no
`._model_impl.python_model` unwrap is applied, because the pickled
object already is the model.
The artifact's top-level `code` directory is prepended to `sys.path`
before deserializing, when it exists, so that classes whose modules ship
inside the artifact resolve at unpickle time. The entry is not removed
afterwards: the deserialized object may import lazily.
Args:
local_dir (str): Local directory holding the downloaded artifact.
artifact_name (str): Artifact name, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any] | None): Metadata used for structured
logging, so the chosen pickle path correlates with the other
fallback log lines.
joblib log lines.
Returns:
Any: The object deserialized from the artifact's pickle file.
@@ -424,9 +423,14 @@ class MLFlowRepository(SientiaMonitoring):
ValueError: If the artifact directory has no top-level `.pkl` file.
Exception: Any error raised by `joblib.load` itself.
"""
pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error)
pkl_path = self._find_model_pickle(local_dir, artifact_name)
self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata)
code_dir = self.resolve_artifact_code_dir(local_dir)
if code_dir and code_dir not in sys.path:
sys.path.insert(0, code_dir)
self.debug(f'joblib flavor: added {code_dir} to sys.path', metadata)
self.debug(f'joblib flavor: loading {artifact_name} from {pkl_path}', metadata)
return joblib.load(pkl_path)
@@ -434,37 +438,36 @@ class MLFlowRepository(SientiaMonitoring):
self,
model_name: str,
artifact_name: str,
flavor_error: Exception,
metadata: dict[str, Any],
run_id: str | None = None,
) -> Any:
) -> tuple[Any, str]:
"""
Download a model artifact and deserialize it with `joblib`.
Fallback path only. The download runs without `dst_path`, so MLflow
writes to a per-call temporary directory: the shared `ARTIFACTS_PATH`
tree is never written to and `dowload_artifacts` — which removes its
destination before downloading — is deliberately not reused. The target
artifact is always a `.pkl` and the deserialized object is returned as
is.
Primary path of the `joblib` flavor. The download runs without
`dst_path`, so MLflow writes to a per-call temporary directory: the
shared `ARTIFACTS_PATH` tree is never written to and
`dowload_artifacts` — which removes its destination before
downloading — is deliberately not reused. The target artifact is always
a `.pkl` and the deserialized object is returned as is.
Callers that already resolved a run to build the flavor URI must pass it
as `run_id`: resolving the Production run again here would open a window
where a promotion between the two lookups makes the fallback serve a
different version than the one the flavor attempt tried and logged. Only
callers whose flavor attempt used a symbolic URI — `models:/<name>/production`,
resolved by MLflow itself — leave it as `None`.
Callers that already resolved a run must pass it as `run_id`, so that
the run is resolved exactly once per load. Callers with no run at hand —
`load_predict_model`, whose flavor URI is symbolic — leave it as `None`.
The local directory is returned alongside the model because the retrain
path needs it to re-log the artifact's `code` directory.
Args:
model_name (str): Registered model name.
artifact_name (str): Artifact to download, e.g. 'prediction_model'.
flavor_error (Exception): Error raised by the MLflow flavor loader.
metadata (dict[str, Any]): Metadata used for structured logging.
run_id (str | None): Run already resolved by the caller. When `None`,
the Production run is resolved here.
Returns:
Any: The object deserialized from the artifact's pickle file.
tuple[Any, str]: The deserialized object and the local directory the
artifact was downloaded to.
Raises:
ValueError: If the downloaded artifact has no top-level `.pkl` file.
@@ -475,86 +478,68 @@ class MLFlowRepository(SientiaMonitoring):
run_id = self.get_model_run_id(model_name=model_name, stage='Production')
self.info(
f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}',
f'joblib flavor: downloading {artifact_name} of {model_name} from run {run_id}',
metadata,
)
local_dir = self.client.download_artifacts(run_id, artifact_name)
return self._load_model_pickle(local_dir, artifact_name, flavor_error, metadata)
model = self._load_model_pickle(local_dir, artifact_name, metadata)
async def _load_with_joblib_fallback(
return model, local_dir
async def _load_joblib_artifact(
self,
primary_loader: Callable[[], Any],
fallback_loader: Callable[[Exception], Any],
failure_description: str,
model_name: str,
artifact_name: str,
operation_type: str,
metadata: dict[str, Any],
fallback_labels: dict[str, Any],
core_labels: dict[str, Any] | None = None,
) -> Any:
run_id: str | None = None,
) -> tuple[Any, str]:
"""
Run a model load with the joblib fallback and its observability.
Load a model artifact through the `joblib` flavor, with observability.
Single owner of the fallback flow shared by every load path: warn on the
flavor failure without emitting an error metric, try the fallback, and
emit lag/count under the fallback labels — or the error count chained to
the flavor error when the fallback also fails. Fixes to the flow belong
here, not in the call sites.
Both loaders are called with no observability concern of their own:
`fallback_loader` receives the flavor error and may be sync or async,
which is what lets the download-based paths and the already-downloaded
wrapper path share this method.
Single owner of the joblib read path: no MLflow flavor loader is called
here, neither before nor after the pickle load. Metrics are emitted with
the **base** `operation_type` the caller passes — a joblib load is a
load, so no suffixed operation type exists.
Args:
primary_loader (Callable[[], Any]): Flavor load attempt. Any
exception it raises triggers the fallback.
fallback_loader (Callable[[Exception], Any]): joblib fallback,
receiving the flavor error. Its result is awaited when
awaitable.
failure_description (str): What failed, e.g. 'Flavor loader sklearn
failed for prediction model x'. The fallback and original error
are appended to it.
model_name (str): Registered model name.
artifact_name (str): Artifact to download, e.g. 'prediction_model'.
operation_type (str): Base metric operation type of the calling
method, e.g. 'load_predict_model'.
metadata (dict[str, Any]): Metadata used for structured logging.
fallback_labels (dict[str, Any]): Metric labels for the fallback
path, already carrying its `operation_type`.
core_labels (dict[str, Any] | None): Metric labels for a successful
flavor load. `None` emits no metric on that path.
run_id (str | None): Run already resolved by the caller. When `None`,
the Production run is resolved during the download.
Returns:
Any: The model loaded by the flavor loader or by the fallback.
tuple[Any, str]: The deserialized model and the local directory the
artifact was downloaded to.
Raises:
Exception: The fallback error, chained from the flavor error, when
both attempts fail.
Exception: Any error raised while resolving the run, downloading the
artifact or deserializing the pickle. It propagates as is: there
is no second attempt.
"""
core_labels = self.get_core_labels(metadata, operation_type=operation_type)
start_time = time.time()
try:
model = primary_loader()
except Exception as flavor_error:
self.warning(
f'{failure_description}, trying joblib fallback. Original error: {flavor_error}',
metadata,
model, local_dir = await self._download_and_load_pickle(
model_name=model_name,
artifact_name=artifact_name,
metadata=metadata,
run_id=run_id,
)
try:
model = fallback_loader(flavor_error)
if inspect.isawaitable(model):
model = await model
except Exception as fallback_error:
await self.emit_metric(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels
)
raise fallback_error from flavor_error
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels)
return model
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
if core_labels is not None:
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
return model, local_dir
async def load_predict_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -565,7 +550,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_name (str): The name of the model to download from the registry.
metadata (dict[str, Any]): Metadata used for structured logging.
flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch')
flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -574,37 +559,42 @@ class MLFlowRepository(SientiaMonitoring):
Notes:
- The model is fetched from the "Production" stage of the MLflow Model Registry.
- Warnings during the model loading process are suppressed.
- With `flavor='joblib'` no MLflow flavor loader is called: the
`prediction_model` artifact is downloaded and deserialized with
`joblib.load`, which needs no `MLmodel` manifest. The artifact's
`code` directory, when present, is put on `sys.path` first.
"""
if flavor not in ('pyfunc', 'sklearn', 'pytorch'):
if flavor not in ('pyfunc', 'sklearn', 'pytorch', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
if flavor == 'joblib':
model, _local_dir = await self._load_joblib_artifact(
model_name=model_name,
artifact_name='prediction_model',
operation_type='load_predict_model',
metadata=metadata,
)
return model
model_uri = f'models:/{model_name}/production'
self.info(f'Loading prediction model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
fallback_labels = self.get_core_labels(
metadata, operation_type='load_predict_model_joblib_fallback'
)
def load_flavor() -> Any:
start_time = time.time()
try:
if flavor == 'pyfunc':
return mlflow.pyfunc.load_model(model_uri)
if flavor == 'sklearn':
return mlflow.sklearn.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
model = mlflow.pyfunc.load_model(model_uri)
elif flavor == 'sklearn':
model = mlflow.sklearn.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name, 'prediction_model', flavor_error, metadata
),
failure_description=(
f'Flavor loader {flavor} failed for prediction model {model_name}'
),
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
async def load_transform_model(
self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
@@ -618,7 +608,7 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_name (str): The name of the model to download.
metadata (dict[str, Any]): Metadata used for structured logging.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns:
@@ -629,45 +619,48 @@ class MLFlowRepository(SientiaMonitoring):
model cannot be loaded.
Notes:
- The Production run is resolved once and reused by the joblib
fallback, so a promotion between both steps cannot make the
fallback serve a different version than the flavor attempt.
- With `flavor='joblib'` no MLflow flavor loader is called and no
model URI is built: the `data_model` artifact of the resolved
Production run is downloaded and deserialized with `joblib.load`,
which needs no `MLmodel` manifest. The artifact's `code`
directory, when present, is put on `sys.path` first.
"""
if flavor not in ('sklearn', 'pyfunc', 'pytorch'):
if flavor not in ('sklearn', 'pyfunc', 'pytorch', 'joblib'):
raise ValueError(INVALID_FLAVOR_MESSAGE)
latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
if flavor == 'joblib':
model, _local_dir = await self._load_joblib_artifact(
model_name=model_name,
artifact_name='data_model',
operation_type='load_transform_model',
metadata=metadata,
run_id=latest_production_id,
)
return model
model_uri = self.get_model_uri(latest_production_id, prediction=False)
self.info(f'Loading data model {model_name} from {model_uri}')
core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
fallback_labels = self.get_core_labels(
metadata, operation_type='load_transform_model_joblib_fallback'
)
def load_flavor() -> Any:
start_time = time.time()
try:
if flavor == 'sklearn':
return mlflow.sklearn.load_model(model_uri)
if flavor == 'pyfunc':
return mlflow.pyfunc.load_model(model_uri)
return mlflow.pytorch.load_model(model_uri)
model = mlflow.sklearn.load_model(model_uri)
elif flavor == 'pyfunc':
model = mlflow.pyfunc.load_model(model_uri)
else:
model = mlflow.pytorch.load_model(model_uri)
except Exception as e:
await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
raise e
return await self._load_with_joblib_fallback(
primary_loader=load_flavor,
fallback_loader=lambda flavor_error: self._download_and_load_pickle(
model_name,
'data_model',
flavor_error,
metadata,
run_id=latest_production_id,
),
failure_description=f'Flavor loader {flavor} failed for data model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
core_labels=core_labels,
)
await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
return model
async def download_model(
self,
@@ -684,11 +677,17 @@ class MLFlowRepository(SientiaMonitoring):
model_name (str): Name of the model to download
metadata (dict[str, Any]): Metadata used for structured logging.
model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib')
load_wrapper (bool): Whether to load wrapper
Returns:
tuple[Any, str | None]: Model object and optional artifact path.
Notes:
- `flavor='joblib'` returns the local directory the artifact was
downloaded to as the artifact path, because the retrain path needs
it to re-log the artifact's `code` directory. Every other flavor
returns `None` unless the wrapper branch was taken.
"""
self.info(
@@ -711,27 +710,24 @@ class MLFlowRepository(SientiaMonitoring):
f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
)
fallback_labels = self.get_core_labels(
metadata, operation_type='download_model_wrapper_joblib_fallback'
)
def load_wrapper_flavor() -> Any:
raw_model = mlflow.pyfunc.load_model(artifact_path)
return raw_model._model_impl.python_model
model = await self._load_with_joblib_fallback(
primary_loader=load_wrapper_flavor,
fallback_loader=lambda flavor_error: self._load_model_pickle(
artifact_path, target, flavor_error, metadata
),
failure_description=f'Wrapper loader failed for {model_type} model {model_name}',
metadata=metadata,
fallback_labels=fallback_labels,
)
raw_model = mlflow.pyfunc.load_model(artifact_path)
model = raw_model._model_impl.python_model
self.debug(
f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata
)
elif flavor == 'joblib':
target = 'prediction_model' if model_type == 'predict' else 'data_model'
operation_type = (
'load_predict_model' if model_type == 'predict' else 'load_transform_model'
)
return await self._load_joblib_artifact(
model_name=model_name,
artifact_name=target,
operation_type=operation_type,
metadata=metadata,
)
else:
if model_type == 'predict':
model = await self.load_predict_model(model_name, metadata, flavor)
@@ -1165,6 +1161,30 @@ class MLFlowRepository(SientiaMonitoring):
}
return retrain_data
def resolve_artifact_code_dir(self, artifact_path: str | None) -> str | None:
"""Resolve the top-level `code` directory of a downloaded artifact.
Sibling of `resolve_code_paths`: both look at the same
`<artifact_path>/code` root, but this one returns the directory itself,
which is what `sys.path` and `mlflow.log_artifacts` take, while
`resolve_code_paths` returns the entries inside it, which is what
`mlflow.pyfunc`'s `code_path=` takes.
Args:
artifact_path (str | None): Local directory of the downloaded model
artifact.
Returns:
str | None: `<artifact_path>/code` when it is a directory, `None`
otherwise — including when `artifact_path` itself is `None`.
"""
if not artifact_path:
return None
code_dir = path.join(artifact_path, 'code')
return code_dir if path.isdir(code_dir) else None
def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]:
"""Resolve the code dependencies to forward when re-logging a pyfunc model.
@@ -1187,11 +1207,12 @@ class MLFlowRepository(SientiaMonitoring):
self.warning('No artifact path available, logging model without code paths', metadata)
return []
code_dir = path.join(artifact_path, 'code')
code_dir = self.resolve_artifact_code_dir(artifact_path)
if not path.isdir(code_dir):
if not code_dir:
self.warning(
f'No code directory found at {code_dir}, logging model without code paths',
f'No code directory found at {path.join(artifact_path, "code")}, '
'logging model without code paths',
metadata,
)
return []
@@ -1207,9 +1228,16 @@ class MLFlowRepository(SientiaMonitoring):
Args:
model_data (dict): Model holder with keys 'model' and optional 'artifact_path'.
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch', 'joblib').
model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'.
metadata (dict): Metadata for structured logging.
Notes:
- `flavor='joblib'` writes a `joblib`-compressed `model.pkl` under
`model_type` and re-logs the source artifact's `code` directory,
so the written model stays deserializable. No `mlflow.*.log_model`
runs and no `MLmodel` manifest is written, which means only this
runtime's `joblib` flavor can read that artifact back.
"""
model = model_data['model']
@@ -1231,6 +1259,19 @@ class MLFlowRepository(SientiaMonitoring):
self.debug('Model uploaded successfully', metadata)
elif flavor == 'pytorch':
mlflow.pytorch.log_model(model, model_type)
elif flavor == 'joblib':
with tempfile.TemporaryDirectory() as tmp_dir:
pkl_path = path.join(tmp_dir, 'model.pkl')
joblib.dump(model, pkl_path, compress=3)
mlflow.log_artifact(pkl_path, artifact_path=model_type)
code_dir = self.resolve_artifact_code_dir(model_data.get('artifact_path'))
if code_dir:
self.debug(f'Re-logging code directory {code_dir}', metadata)
mlflow.log_artifacts(code_dir, artifact_path=f'{model_type}/code')
else:
self.debug(f'No code directory to re-log for {model_type}', metadata)
else:
raise ValueError(INVALID_FLAVOR_MESSAGE)

4
laborious/worker/worker.py
View File

@@ -33,8 +33,8 @@ Environment Variables:
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: laborious)
- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_*: import configuration, see
``laborious.utils.connectors_config``
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_NAME: import configuration, see
``laborious.utils.connectors_config`` (the import log's server is POSTGRES_*)
"""
from temporalio import client, workflow

11
pyproject.toml
View File

@@ -133,6 +133,17 @@ addopts = [
# GC'd. Harmless mock artifacts turned into a hard ERROR by the plugin itself.
"-p", "no:unraisableexception",
]
filterwarnings = [
# `kafka-python`, pulled in transitively by `sientia_do`, reads its protocol schemas through the
# legacy `importlib.resources` API that Python 3.11 deprecated. It emits ~109 warnings per run,
# twice each (once at the call site, once inside `importlib`), which is enough noise to bury a
# warning worth reading. Nothing here can fix it — the change belongs in `kafka-python`.
#
# Matched on the message rather than by module so the filter cannot quietly swallow an unrelated
# deprecation: any *other* DeprecationWarning, including one from our own code, still shows up.
"ignore:read_text is deprecated:DeprecationWarning",
"ignore:open_text is deprecated:DeprecationWarning",
]
markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",

5
requirements-dev.txt
View File

@@ -13,7 +13,10 @@ types-requests>=2.31.0 # Type stubs for requests
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
testcontainers[postgres,minio,mongodb] # PostgreSQL, MinIO and MongoDB containers for E2E tests
testcontainers[postgres,minio,mongodb]>=4.15.0 # PostgreSQL, MinIO and MongoDB containers for E2E tests
# Floor is the `testcontainers.community.*` namespace: the flat
# `testcontainers.minio` / `.postgres` / `.mongodb` paths the e2e suite used
# to import are deprecated shims that warn on import.
# Development Tools
ipython>=8.12.0 # Enhanced Python shell

20
scripts/encode.sh
View File

@@ -1,20 +0,0 @@
#!/bin/bash
# Resolve paths from the script's own location so this works whether it is run
# from the repository root or from inside scripts/.
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
source "${REPO_ROOT}/venv/bin/activate"
pip install pathspec
pip install pyyaml
echo "
.git" >> "${REPO_ROOT}/.gitignore"
python "${SCRIPT_DIR}/encrypt.py" "${REPO_ROOT}" "${REPO_ROOT}/code" --ignore "${REPO_ROOT}/.gitignore" --chunk-size 100000
sed -i '/.git/d' "${REPO_ROOT}/.gitignore"
xdg-open "${REPO_ROOT}"

113
scripts/encrypt.py
View File

@@ -1,113 +0,0 @@
import argparse
import os
from typing import Any
import yaml # type: ignore
from pathspec import PathSpec
r"""
Usage:
python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
"""
def load_ignore_patterns(ignore_file, include_library):
# Ensure the .gitignore file exists
if not os.path.exists(ignore_file):
raise FileNotFoundError(f'Ignore file not found at {ignore_file}')
# Load and parse the .gitignore patterns
with open(ignore_file) as file:
patterns = file.readlines()
if not include_library:
patterns.append('**/deploy/library/')
spec = PathSpec.from_lines('gitwildmatch', patterns)
return spec
def is_ignored(file_path, spec):
"""Check if a file should be ignored based on the ignore patterns."""
return spec.match_file(file_path) if spec else False
def encode_file_tree_to_yaml(directory, ignore_file, include_library):
"""Encode the file tree into a single YAML file."""
ignore_patterns = load_ignore_patterns(ignore_file, include_library) if ignore_file else None
file_tree: dict[str, Any] = {}
for root, dirs, files in os.walk(directory):
# Skip ignored directories
dirs[:] = [d for d in dirs if not is_ignored(os.path.join(root, d), ignore_patterns)]
for file in files:
file_path = os.path.join(root, file)
# Skip ignored files
if is_ignored(file_path, ignore_patterns):
continue
# Read file content
try:
with open(file_path, encoding='utf-8') as f:
content = f.read()
except Exception as e:
print(f'Error reading file {file_path}: {e}')
raise
# Create nested dictionary structure
path_parts = os.path.relpath(file_path, directory).split(os.sep)
current_level = file_tree
# all except the last part (the file name)
for part in path_parts[:-1]:
current_level = current_level.setdefault(part, {})
# Add the file and its content
current_level[path_parts[-1]] = content
return yaml.dump(file_tree, default_flow_style=False)
def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
"""Chunk the YAML content and write it to the output file."""
chunks = (
[yaml_content]
if chunk_size is None
else [yaml_content[i : i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
)
for i, chunk in enumerate(chunks):
chunk_file = f'{output_file}_{i}.yaml'
# Write the file tree to the output YAML file
with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
yaml_file.write(chunk)
def main():
parser = argparse.ArgumentParser(description='Encrypts file tree to yaml file')
parser.add_argument('input_directory', help='Directory to encode')
parser.add_argument('output_yaml_file', help='Output YAML file')
parser.add_argument('--ignore', default=None, help='Path to the ignore file')
parser.add_argument(
'--chunk-size', type=int, default=None, help='Chunk size for the output YAML file'
)
parser.add_argument(
'--library', type=bool, default=False, help='Incude the library in the output YAML file'
)
# Parse arguments
args = parser.parse_args()
# Example usage
directory_to_encode = args.input_directory
ignore_file_path = args.ignore
output_yaml_file = args.output_yaml_file
include_library = args.library
content = encode_file_tree_to_yaml(directory_to_encode, ignore_file_path, include_library)
chunk_and_write_file_tree_to_yaml(content, output_yaml_file, args.chunk_size)
if __name__ == '__main__':
main()

19
scripts/import_model_manual_run.py
View File

@@ -77,13 +77,18 @@ def _line(label: str, value: object) -> str:
# back to localhost, which is right for a service that must boot; it is wrong for a script that
# writes, because a silent fallback is a write to somewhere nobody read out loud.
#
# The import log database: a second Postgres connection, independent of the POSTGRES_* one above it
# in .env. The schema and table are constants of the worker's code, never configuration.
STATUS_DB_HOST = env_required('IMPORT_STATUS_DB_HOST')
STATUS_DB_PORT = int(env_required('IMPORT_STATUS_DB_PORT'))
# The import log database: the POSTGRES_* server from .env, with only the database name of its own,
# because the log lives in the Spring Boot BFF's database rather than in `sientia`. The schema and
# table are constants of the worker's code, never configuration.
#
# `build_import_status_config` also honours IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} as overrides,
# for a deployment whose log sits on another server. This script does not: it is a development tool
# pointed at a real environment, and the override exists for the e2e suite's two containers.
STATUS_DB_HOST = env_required('POSTGRES_HOST')
STATUS_DB_PORT = int(env_required('POSTGRES_PORT'))
STATUS_DB_NAME = env_required('IMPORT_STATUS_DB_NAME')
STATUS_DB_USER = env_required('IMPORT_STATUS_DB_USER')
STATUS_DB_PASSWORD = os.getenv('IMPORT_STATUS_DB_PASSWORD', '')
STATUS_DB_USER = env_required('POSTGRES_USER')
STATUS_DB_PASSWORD = os.getenv('POSTGRES_PASSWORD', '')
STATUS_TABLE = 'public.experiment_run'
# Where the frontend uploads the bundle. The prefix keeps the worker's own default, which is a
@@ -145,7 +150,7 @@ CONFIRM_TARGET = env('IMPORT_SCRIPT_CONFIRM_TARGET')
if CONFIRM_TARGET != STATUS_DB_HOST:
raise RuntimeError(
'refusing to write: the confirmation does not match the resolved status database host. '
f'confirmation={CONFIRM_TARGET!r}, resolved IMPORT_STATUS_DB_HOST={STATUS_DB_HOST!r}. '
f'confirmation={CONFIRM_TARGET!r}, resolved POSTGRES_HOST={STATUS_DB_HOST!r}. '
'Read the targets cell 1 printed, then restate the host here.'
)

27
scripts/inter_arrival.py
View File

@@ -1,27 +0,0 @@
# %%
# Load logs.txt
with open('logs.txt') as file:
lines = file.readlines()
# %%
import re
# Grep "inter-arrival_s=number" with regex
intervals = []
for line in lines:
match = re.search(r'inter-arrival_s=([0-9.]+)', line)
if match:
intervals.append(float(match.group(1)))
# %%
print(intervals)
# %%
import matplotlib.pyplot as plt
plt.plot(intervals)
plt.ylabel('Inter-arrival time (s)')
plt.xlabel('Sample')
plt.title('Inter-arrival time distribution')
plt.show()
# %%

107
scripts/model_convert.ipynb
View File

@@ -1,107 +0,0 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 23,
"id": "e838ff21",
"metadata": {},
"outputs": [],
"source": [
"import csv\n",
"\n",
"\n",
"def csv_to_tag_lists(csv_path: str) -> dict:\n",
" read_tags = []\n",
" write_tags = []\n",
"\n",
" def to_float(val):\n",
" try:\n",
" return float(str(val).strip())\n",
" except Exception:\n",
" return None\n",
"\n",
" with open(csv_path, newline='', encoding='utf-8') as f:\n",
" reader = csv.DictReader(f)\n",
" for row in reader:\n",
" # Basic normalization\n",
" op = (row.get('operation') or '').strip()\n",
"\n",
" if op == 'READ':\n",
" # Build common tag payload with required mappings\n",
" tag = {\n",
" 'server_id': '1',\n",
" 'tag_address': row.get('opc_tag'),\n",
" 'tag_name': row.get('name'),\n",
" 'data_range': [to_float(row.get('min_value')), to_float(row.get('max_value'))],\n",
" 'aggr_func': row.get('aggregation_func').lower(),\n",
" # keep other fields with their original names\n",
" 'frequency': row.get('frequency'),\n",
" 'local': row.get('local'),\n",
" 'area': row.get('area'),\n",
" 'description': row.get('description'),\n",
" }\n",
"\n",
" read_tags.append(tag)\n",
"\n",
" else:\n",
" tag = {\n",
" 'server_id': '1',\n",
" 'addr': row.get('opc_tag'),\n",
" 'tag_name': row.get('name'),\n",
" 'local': row.get('local'),\n",
" 'area': row.get('area'),\n",
" 'description': row.get('description'),\n",
" }\n",
"\n",
" if op == 'WRITE_PREDICTION':\n",
" tag['type'] = 'prediction'\n",
" write_tags.append(tag)\n",
" elif op == 'WRITE_CONFIDENCE':\n",
" tag['type'] = 'confidence'\n",
" write_tags.append(tag)\n",
" # ignore any other operation values silently\n",
"\n",
" return {'read_tags': read_tags, 'write_tags': write_tags}"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4621cd43",
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"\n",
"file_names = ['Courier - Página1.csv']\n",
"\n",
"for file_name in file_names:\n",
" write_file = file_name.replace('.csv', '.json')\n",
"\n",
" with open(write_file, 'w', encoding='utf-8') as f:\n",
" json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

647
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -1,4 +1,8 @@
import ast
import importlib
import sys
from datetime import UTC, datetime
from pathlib import Path
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
import joblib as joblib_lib
@@ -10,6 +14,7 @@ from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sklearn.linear_model import LinearRegression
from laborious import metrics
from laborious.utils.repository import model_repository
from laborious.utils.repository.model_repository import (
INVALID_FLAVOR_MESSAGE,
MLFlowRepository,
@@ -465,7 +470,7 @@ async def test_load_predict_model_pytorch(mlflow, mlflow_repository):
async def test_load_predict_model_error(mlflow_repository):
with pytest.raises(ValueError) as e:
await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
@@ -542,7 +547,7 @@ async def test_load_transform_model_error(mlflow_repository):
with pytest.raises(ValueError) as e:
await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
@@ -1124,7 +1129,7 @@ async def test_log_model_error(mlflow_repository):
await mlflow_repository.log_model(
model_data, 'invalid', 'prediction_model', metadata['metadata']
)
assert str(e) == "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'."
assert str(e) == INVALID_FLAVOR_MESSAGE
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
)
@@ -1768,10 +1773,19 @@ async def test_predict_success_ndarray(mlflow_repository):
"""
Tests for the joblib fallback on the model load paths
Tests for the joblib flavor on the model load and write paths
"""
JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib'
FLAVOR_LOADERS = ('sklearn', 'pyfunc', 'pytorch')
@pytest.fixture
def preserved_sys_path():
"""Restore `sys.path`: the joblib flavor inserts entries and never removes them."""
original = list(sys.path)
yield
sys.path[:] = original
def build_artifact_dir(tmp_path, *filenames):
@@ -1781,6 +1795,21 @@ def build_artifact_dir(tmp_path, *filenames):
return str(tmp_path)
def build_code_dir(tmp_path, module_name='probe_flavor_model'):
"""Add a `code/` directory carrying the model's class module to an artifact dir."""
code_dir = tmp_path / 'code'
code_dir.mkdir()
(code_dir / f'{module_name}.py').write_text(
'class ProbeModel:\n'
' def __init__(self, offset):\n'
' self.offset = offset\n'
'\n'
' def predict(self, data):\n'
' return [value + self.offset for value in data]\n'
)
return str(code_dir)
def emitted_tags(emit_metric_mock, metric_object):
"""Collect the `tags` kwarg of every emit_metric call for a given metric."""
return [
@@ -1790,86 +1819,22 @@ def emitted_tags(emit_metric_mock, metric_object):
]
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_flavor_success_skips_fallback(
joblib_mock, mlflow, mlflow_repository
):
result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
assert result == mlflow.sklearn.load_model.return_value
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
def assert_no_flavor_loader_called(mlflow):
for flavor in FLAVOR_LOADERS:
getattr(mlflow, flavor).load_model.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_flavor_success_skips_fallback(
joblib_mock, mlflow, mlflow_repository
):
mlflow_repository.get_model_run_id = MagicMock()
mlflow_repository.get_model_uri = MagicMock()
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
assert result == mlflow.sklearn.load_model.return_value
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_success_skips_fallback(
joblib_mock, mlflow, mlflow_repository
):
mlflow_repository.dowload_artifacts = AsyncMock()
result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
assert result == (
mlflow.pyfunc.load_model.return_value._model_impl.python_model,
mlflow_repository.dowload_artifacts.return_value,
)
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
mlflow_repository.emit_metric.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository):
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
mlflow_repository.emit_metric.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_fallback_recovers_load(
async def test_load_predict_model_joblib_flavor_loads_pickle_directly(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'MLmodel', 'model.pkl'
)
result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
result = await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
assert result is joblib_mock.load.return_value
mlflow_repository.get_model_run_id.assert_called_once_with(
@@ -1879,46 +1844,43 @@ async def test_load_predict_model_fallback_recovers_load(
'run-123', 'prediction_model'
)
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_fallback_recovers_load(
async def test_load_transform_model_joblib_flavor_loads_pickle_directly(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.get_model_uri = MagicMock()
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'MLmodel', 'model.pkl'
)
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
result = await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
assert result is joblib_mock.load.return_value
mlflow_repository.get_model_run_id.assert_called_with(
model_name='test_model', stage='Production'
)
mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model')
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
mlflow_repository.get_model_uri.assert_not_called()
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_fallback_reuses_resolved_run(
async def test_load_transform_model_joblib_resolves_run_once(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
"""A promotion between both steps must not change the run served."""
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
"""A promotion mid-load must not change the run served: one resolution per load."""
mlflow_repository.get_model_run_id = MagicMock(side_effect=['run-A', 'run-B'])
mlflow_repository.get_model_uri = MagicMock()
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'MLmodel', 'model.pkl'
tmp_path, 'model.pkl'
)
result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
assert result is joblib_mock.load.return_value
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='test_model', stage='Production'
)
@@ -1927,57 +1889,91 @@ async def test_load_transform_model_fallback_reuses_resolved_run(
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_fallback_reuses_downloaded_dir(
joblib_mock, mlflow, mlflow_repository, tmp_path
@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch'])
async def test_load_predict_model_flavor_failure_propagates(
joblib_mock, mlflow, mlflow_repository, flavor
):
artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl')
mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir)
mlflow.pyfunc.load_model.side_effect = Exception('flavor boom')
loader_error = mlflow_lib.exceptions.MlflowException('flavor boom')
getattr(mlflow, flavor).load_model.side_effect = loader_error
model, returned_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'pyfunc', True
)
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.load_predict_model('test_model', {}, flavor)
mlflow_repository.dowload_artifacts.assert_called_once_with(
'test_model', {}, 'prediction_model'
)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
joblib_mock.load.assert_called_once_with(str(tmp_path / 'python_model.pkl'))
assert model is joblib_mock.load.return_value
assert returned_path == artifact_dir
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'load_predict_model'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_both_attempts_fail_chains_causes(
@pytest.mark.parametrize('flavor', ['sklearn', 'pyfunc', 'pytorch'])
async def test_load_transform_model_flavor_failure_propagates(
joblib_mock, mlflow, mlflow_repository, flavor
):
mlflow_repository.get_model_run_id = MagicMock()
mlflow_repository.get_model_uri = MagicMock()
loader_error = mlflow_lib.exceptions.MlflowException('flavor boom')
getattr(mlflow, flavor).load_model.side_effect = loader_error
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.load_transform_model('test_model', {}, flavor)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'load_transform_model'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_failure_propagates(joblib_mock, mlflow, mlflow_repository):
mlflow_repository.dowload_artifacts = AsyncMock()
loader_error = mlflow_lib.exceptions.MlflowException('wrapper boom')
mlflow.pyfunc.load_model.side_effect = loader_error
with pytest.raises(mlflow_lib.exceptions.MlflowException) as exc_info:
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
assert exc_info.value is loader_error
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_joblib_success_metrics(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
flavor_error = Exception('flavor boom')
mlflow.sklearn.load_model.side_effect = flavor_error
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'model.pkl'
)
joblib_mock.load.side_effect = ValueError('joblib boom')
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
assert str(exc_info.value) == 'joblib boom'
assert exc_info.value.__cause__ is flavor_error
read_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT)
assert len(read_tags) == 1
assert read_tags[0]['operation_type'] == 'load_predict_model'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == []
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
mlflow_repository.observe_lag.assert_called_once()
lag_call = mlflow_repository.observe_lag.call_args
assert lag_call.args[1] is metrics.MODEL_READ_LAG
assert lag_call.args[2]['operation_type'] == 'load_predict_model'
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_both_attempts_fail_emits_error_metric(
async def test_load_transform_model_joblib_failure_metrics(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.get_model_uri = MagicMock()
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
@@ -1985,65 +1981,41 @@ async def test_load_transform_model_both_attempts_fail_emits_error_metric(
)
joblib_mock.load.side_effect = ValueError('joblib boom')
with pytest.raises(ValueError):
await mlflow_repository.load_transform_model('test_model', {}, 'sklearn')
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'load_transform_model_joblib_fallback'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_wrapper_both_attempts_fail_chains_causes(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
flavor_error = Exception('flavor boom')
artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl')
mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir)
mlflow.pyfunc.load_model.side_effect = flavor_error
joblib_mock.load.side_effect = ValueError('joblib boom')
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)
await mlflow_repository.load_transform_model('test_model', {}, 'joblib')
assert str(exc_info.value) == 'joblib boom'
assert exc_info.value.__cause__ is flavor_error
assert exc_info.value.__cause__ is None
error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT)
assert len(error_tags) == 1
assert error_tags[0]['operation_type'] == 'download_model_wrapper_joblib_fallback'
assert error_tags[0]['operation_type'] == 'load_transform_model'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == []
mlflow_repository.observe_lag.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_fallback_success_metrics(
async def test_no_load_emits_a_fallback_operation_type(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'model.pkl'
)
mlflow_repository.warning = MagicMock()
await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')
read_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT)
assert len(read_tags) == 1
assert read_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback'
assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == []
emitted = [
emitted.kwargs['tags']['operation_type']
for emitted in mlflow_repository.emit_metric.call_args_list
]
observed = [
lag.args[2]['operation_type'] for lag in mlflow_repository.observe_lag.call_args_list
]
mlflow_repository.observe_lag.assert_called_once()
lag_call = mlflow_repository.observe_lag.call_args
assert lag_call.args[1] is metrics.MODEL_READ_LAG
assert lag_call.args[2]['operation_type'] == 'load_predict_model_joblib_fallback'
mlflow_repository.warning.assert_called_once()
assert 'flavor boom' in mlflow_repository.warning.call_args.args[0]
assert emitted
assert all(not operation.endswith('_joblib_fallback') for operation in emitted + observed)
@pytest.mark.parametrize(
@@ -2057,9 +2029,7 @@ async def test_load_predict_model_fallback_success_metrics(
def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, expected):
local_dir = build_artifact_dir(tmp_path, *filenames)
result = mlflow_repository._find_model_pickle(
local_dir, 'prediction_model', Exception('flavor boom')
)
result = mlflow_repository._find_model_pickle(local_dir, 'prediction_model')
assert result == str(tmp_path / expected)
@@ -2070,16 +2040,14 @@ def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path):
(nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes')
with pytest.raises(ValueError) as exc_info:
mlflow_repository._find_model_pickle(
str(tmp_path), 'prediction_model', Exception('flavor boom')
)
mlflow_repository._find_model_pickle(str(tmp_path), 'prediction_model')
assert 'No .pkl file found' in str(exc_info.value)
def test_find_model_pickle_empty_directory(mlflow_repository, tmp_path):
with pytest.raises(ValueError) as exc_info:
mlflow_repository._find_model_pickle(str(tmp_path), 'data_model', Exception('flavor boom'))
mlflow_repository._find_model_pickle(str(tmp_path), 'data_model')
assert 'data_model' in str(exc_info.value)
@@ -2088,80 +2056,373 @@ def test_find_model_pickle_error_message_carries_context(mlflow_repository, tmp_
local_dir = build_artifact_dir(tmp_path, 'MLmodel')
with pytest.raises(ValueError) as exc_info:
mlflow_repository._find_model_pickle(
local_dir, 'prediction_model', Exception('flavor boom')
)
mlflow_repository._find_model_pickle(local_dir, 'prediction_model')
message = str(exc_info.value)
assert 'prediction_model' in message
assert 'flavor boom' in message
assert 'MLmodel' in message
assert 'Original error' not in message
def test_find_model_pickle_ignores_directory_named_pkl(mlflow_repository, tmp_path):
(tmp_path / 'algo.pkl').mkdir()
with pytest.raises(ValueError) as exc_info:
mlflow_repository._find_model_pickle(
str(tmp_path), 'prediction_model', Exception('flavor boom')
)
mlflow_repository._find_model_pickle(str(tmp_path), 'prediction_model')
assert 'algo.pkl' in str(exc_info.value)
def test_resolve_artifact_code_dir_present(tmp_path, mlflow_repository):
code_dir = build_code_dir(tmp_path)
assert mlflow_repository.resolve_artifact_code_dir(str(tmp_path)) == code_dir
def test_resolve_artifact_code_dir_absent(tmp_path, mlflow_repository):
assert mlflow_repository.resolve_artifact_code_dir(str(tmp_path)) is None
def test_resolve_artifact_code_dir_without_artifact_path(mlflow_repository):
assert mlflow_repository.resolve_artifact_code_dir(None) is None
@patch(JOBLIB_PATH)
def test_load_model_pickle_adds_code_dir_before_load(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
sys_path_at_load = []
joblib_mock.load.side_effect = lambda _path: sys_path_at_load.extend(sys.path)
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert code_dir in sys_path_at_load
assert sys.path[0] == code_dir
@patch(JOBLIB_PATH)
def test_load_model_pickle_without_code_dir_leaves_sys_path_alone(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
before = list(sys.path)
mlflow_repository._load_model_pickle(local_dir, 'data_model')
joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl'))
assert sys.path == before
@patch(JOBLIB_PATH)
def test_load_model_pickle_adds_code_dir_only_once(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
assert sys.path.count(code_dir) == 1
@patch(JOBLIB_PATH)
def test_load_model_pickle_keeps_code_dir_after_load(
joblib_mock, mlflow_repository, tmp_path, preserved_sys_path
):
code_dir = build_code_dir(tmp_path)
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository._load_model_pickle(local_dir, 'prediction_model')
assert code_dir in sys.path
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_get_model_caches_model_loaded_by_fallback(
async def test_log_model_joblib_writes_compressed_pickle_and_code(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
code_dir = build_code_dir(tmp_path)
model_data = {'model': MagicMock(), 'artifact_path': str(tmp_path)}
await mlflow_repository.log_model(
model_data, 'joblib', 'prediction_model', metadata['metadata']
)
joblib_mock.dump.assert_called_once()
dumped_model, dumped_path = joblib_mock.dump.call_args.args
assert dumped_model is model_data['model']
assert dumped_path.endswith('/model.pkl')
assert joblib_mock.dump.call_args.kwargs == {'compress': 3}
mlflow.log_artifact.assert_called_once_with(dumped_path, artifact_path='prediction_model')
mlflow.log_artifacts.assert_called_once_with(code_dir, artifact_path='prediction_model/code')
for flavor in FLAVOR_LOADERS:
getattr(mlflow, flavor).log_model.assert_not_called()
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY)
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_COUNT, tags=ANY
)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_log_model_joblib_without_code_dir_logs_only_pickle(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
model_data = {'model': MagicMock(), 'artifact_path': str(tmp_path)}
await mlflow_repository.log_model(model_data, 'joblib', 'data_model', metadata['metadata'])
mlflow.log_artifact.assert_called_once_with(ANY, artifact_path='data_model')
mlflow.log_artifacts.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_log_model_sklearn_does_not_dump_joblib(joblib_mock, mlflow, mlflow_repository):
model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
await mlflow_repository.log_model(
model_data, 'sklearn', 'prediction_model', metadata['metadata']
)
mlflow.sklearn.log_model.assert_called_once_with(model_data['model'], 'prediction_model')
joblib_mock.dump.assert_not_called()
mlflow.log_artifact.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_joblib_returns_download_dir(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
local_dir = build_artifact_dir(tmp_path, 'model.pkl')
mlflow_repository.client.download_artifacts.return_value = local_dir
mlflow_repository.dowload_artifacts = AsyncMock()
model, artifact_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'joblib'
)
assert model is joblib_mock.load.return_value
assert artifact_path == local_dir
mlflow_repository.dowload_artifacts.assert_not_called()
assert_no_flavor_loader_called(mlflow)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_sklearn_returns_no_artifact_path(
joblib_mock, mlflow, mlflow_repository
):
model, artifact_path = await mlflow_repository.download_model(
'test_model', {}, 'predict', 'sklearn'
)
assert model is mlflow.sklearn.load_model.return_value
assert artifact_path is None
joblib_mock.load.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_download_model_transform_joblib_downloads_data_model(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'model.pkl'
)
first = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn')
await mlflow_repository.download_model('test_model', {}, 'transform', 'joblib')
mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model')
@patch('laborious.utils.repository.model_repository.force_memory_release')
@patch('laborious.utils.repository.model_repository.rmtree')
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_retrain_model_joblib_reaches_joblib_write_branch(
joblib_mock, _rmtree, _force_memory_release, mlflow, mlflow_repository, tmp_path
):
code_dir = build_code_dir(tmp_path)
model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'joblib'}
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.get_model_params = MagicMock(return_value={})
mlflow_repository.get_experiment = MagicMock()
mlflow_repository.get_next_run_name = MagicMock()
mlflow_repository.fit_models = AsyncMock(
return_value={
'prediction_model': {'model': MagicMock(), 'artifact_path': str(tmp_path)},
'data_model': {'model': MagicMock(), 'artifact_path': None},
'prediction_data': MagicMock(spec=DataFrame),
}
)
output = await mlflow_repository.retrain_model(
MagicMock(), 'test_model', model_config, metadata['metadata']
)
assert output['success'] is True
mlflow_repository.fit_models.assert_called_once()
assert mlflow_repository.fit_models.call_args.kwargs['predict_flavor'] == 'joblib'
joblib_mock.dump.assert_called_once()
assert joblib_mock.dump.call_args.kwargs == {'compress': 3}
mlflow.log_artifact.assert_any_call(ANY, artifact_path='prediction_model')
mlflow.log_artifacts.assert_called_once_with(code_dir, artifact_path='prediction_model/code')
mlflow.sklearn.log_model.assert_called_once_with(ANY, 'data_model')
def test_module_defines_no_joblib_fallback_routine():
"""Guard against reintroducing 'try the flavor loader, then try pickle'."""
source = Path(model_repository.__file__).read_text()
assert 'joblib_fallback' not in source
# The fallback threaded the flavor loader's exception into the pickle helpers as
# `flavor_error`; nothing in the module may take that parameter again.
assert 'flavor_error' not in source
pickle_helpers = {
'_find_model_pickle',
'_load_model_pickle',
'_download_and_load_pickle',
'_load_joblib_artifact',
}
for handler in ast.walk(ast.parse(source)):
if not isinstance(handler, ast.ExceptHandler):
continue
for node in ast.walk(handler):
if not isinstance(node, ast.Call):
continue
func = node.func
if isinstance(func, ast.Attribute):
assert func.attr not in pickle_helpers, (
f'{func.attr} is called from an except handler'
)
assert not (isinstance(func.value, ast.Name) and func.value.id == 'joblib'), (
'joblib is called from an except handler'
)
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_predict_model_invalid_flavor_before_any_io(joblib_mock, mlflow_repository):
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.load_predict_model('test_model', {}, 'invalid')
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
mlflow_repository.emit_metric.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_load_transform_model_invalid_flavor_before_any_io(joblib_mock, mlflow_repository):
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.load_transform_model('test_model', {}, 'invalid')
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.load.assert_not_called()
mlflow_repository.client.download_artifacts.assert_not_called()
mlflow_repository.emit_metric.assert_not_called()
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_log_model_invalid_flavor(joblib_mock, mlflow, mlflow_repository):
model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
with pytest.raises(ValueError) as exc_info:
await mlflow_repository.log_model(
model_data, 'invalid', 'prediction_model', metadata['metadata']
)
assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE
joblib_mock.dump.assert_not_called()
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
)
mlflow_repository.observe_lag.assert_not_called()
def test_invalid_flavor_message_lists_every_accepted_flavor():
for flavor in ('sklearn', 'pyfunc', 'pytorch', 'joblib'):
assert f"'{flavor}'" in INVALID_FLAVOR_MESSAGE
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_get_model_caches_joblib_loaded_model(
joblib_mock, mlflow, mlflow_repository, tmp_path
):
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'model.pkl'
)
first = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'joblib')
second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'joblib')
assert first is joblib_mock.load.return_value
assert second is first
assert mlflow.sklearn.load_model.call_count == 1
assert joblib_mock.load.call_count == 1
@patch(JOBLIB_PATH)
@pytest.mark.asyncio
async def test_predict_reports_both_causes_when_fallback_fails(
joblib_mock, mlflow, mlflow_repository, tmp_path
async def test_joblib_flavor_real_round_trip_with_code_dir(
mlflow_repository, tmp_path, preserved_sys_path
):
mlflow.sklearn.load_model.side_effect = Exception('flavor boom')
"""Real joblib, no mock: proves decompression and module resolution together."""
code_dir = build_code_dir(tmp_path)
sys.path.insert(0, code_dir)
try:
module = importlib.import_module('probe_flavor_model')
joblib_lib.dump(module.ProbeModel(2), tmp_path / 'model.pkl', compress=3)
finally:
sys.path.remove(code_dir)
del sys.modules['probe_flavor_model']
importlib.invalidate_caches()
with pytest.raises(ModuleNotFoundError):
importlib.import_module('probe_flavor_model')
mlflow_repository.get_model_run_id = MagicMock(return_value='run-123')
mlflow_repository.client.download_artifacts.return_value = build_artifact_dir(
tmp_path, 'model.pkl'
)
joblib_mock.load.side_effect = ValueError('joblib boom')
mlflow_repository.client.download_artifacts.return_value = str(tmp_path)
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
try:
loaded = await mlflow_repository.load_predict_model('test_model', {}, 'joblib')
output = await mlflow_repository.predict(
'test_model',
data,
{'retention_minutes': 0, 'predict_flavor': 'sklearn'},
metadata['metadata'],
)
assert output['success'] is False
assert 'joblib boom' in output['content']['traceback']
assert 'flavor boom' in output['content']['traceback']
assert loaded.predict([1.0, 2.0]) == [3.0, 4.0]
assert code_dir in sys.path
finally:
sys.modules.pop('probe_flavor_model', None)
importlib.invalidate_caches()
def test_load_model_pickle_real_joblib_round_trip(mlflow_repository, tmp_path):
def test_log_model_joblib_real_round_trip(tmp_path, mlflow, mlflow_repository):
"""The write flavor dumps what the read flavor loads: same dialect by construction."""
estimator = LinearRegression().fit(np.array([[0.0], [1.0], [2.0]]), np.array([0.0, 1.0, 2.0]))
joblib_lib.dump(estimator, tmp_path / 'model.pkl', compress=3)
written = tmp_path / 'model.pkl'
joblib_lib.dump(estimator, written, compress=3)
loaded = mlflow_repository._load_model_pickle(
str(tmp_path), 'prediction_model', Exception('flavor boom')
)
loaded = mlflow_repository._load_model_pickle(str(tmp_path), 'prediction_model')
assert loaded is not estimator
assert loaded.predict(np.array([[3.0]]))[0] == pytest.approx(3.0)

50
tests/laborious/utils/test_connectors_config.py
View File

@@ -29,11 +29,20 @@ IMPORT_ENV_VARS = (
'IMPORT_STATUS_DB_PASSWORD',
)
# The import log connection falls back to the POSTGRES_* server, so a developer's own POSTGRES_*
# would otherwise decide what `build_import_status_config` returns under test.
POSTGRES_ENV_VARS = (
'POSTGRES_HOST',
'POSTGRES_PORT',
'POSTGRES_USER',
'POSTGRES_PASSWORD',
)
@pytest.fixture
def clean_import_env(monkeypatch):
"""Import configuration is read from a known-empty environment."""
for name in IMPORT_ENV_VARS:
for name in IMPORT_ENV_VARS + POSTGRES_ENV_VARS:
monkeypatch.delenv(name, raising=False)
return monkeypatch
@@ -208,12 +217,49 @@ def test_build_import_status_config_with_defaults(clean_import_env):
'host': 'localhost',
'port': 5432,
'dbname': 'sientia-core-mlops-bff',
'user': 'mlops_bff_user',
'user': 'sientia',
'password': '',
'schema': 'public',
}
def test_build_import_status_config_uses_the_postgres_server(clean_import_env):
"""The log is another database on the POSTGRES_* server: only its name is import configuration."""
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
clean_import_env.setenv('POSTGRES_PORT', '5433')
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
clean_import_env.setenv('IMPORT_STATUS_DB_NAME', 'bff-db')
assert build_import_status_config() == {
'host': 'paradedb-rw',
'port': 5433,
'dbname': 'bff-db',
'user': 'sientia-user',
'password': 'sientia-password',
'schema': 'public',
}
def test_build_import_status_config_overrides_beat_the_postgres_server(clean_import_env):
"""A log on a server of its own: the four IMPORT_STATUS_DB_* variables win over POSTGRES_*."""
clean_import_env.setenv('POSTGRES_HOST', 'paradedb-rw')
clean_import_env.setenv('POSTGRES_PORT', '5433')
clean_import_env.setenv('POSTGRES_USER', 'sientia-user')
clean_import_env.setenv('POSTGRES_PASSWORD', 'sientia-password')
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')
clean_import_env.setenv('IMPORT_STATUS_DB_USER', 'bff-user')
clean_import_env.setenv('IMPORT_STATUS_DB_PASSWORD', 'bff-password')
config = build_import_status_config()
assert config['host'] == 'bff-postgres'
assert config['port'] == 6543
assert config['user'] == 'bff-user'
assert config['password'] == 'bff-password'
def test_build_import_status_config_reads_the_environment(clean_import_env):
clean_import_env.setenv('IMPORT_STATUS_DB_HOST', 'bff-postgres')
clean_import_env.setenv('IMPORT_STATUS_DB_PORT', '6543')

0
tests/scripts/__init__.py
View File

217
tests/scripts/test_import_model_manual_run.py
View File

@@ -1,217 +0,0 @@
"""The anti-rot guard for `scripts/import_model_manual_run.py`.
The script itself is only ever smoke-run by a person (task 13.6 of `import-model-e2e-suite`), so
nothing else would notice the three things most likely to be quietly broken in it:
1. the `# %%` cell layout — the whole contract is "cell N does step N", and a lost or reordered
marker turns a ten-step procedure into an unreadable one;
2. the boundary of design E15 — the script stands in for the frontend, so borrowing
`laborious.activities` or `laborious.workflows` would make it exercise the implementation instead
of the contract, and that is exactly the shortcut somebody reaches for when a cell needs one more
value;
3. the configuration rule — no host, user, password or key literal in the file, every target from
the environment, and every variable it reads documented in `.env.example`.
The wire shapes the script produces (the digest, the envelope, the object key, the row) are asserted
by the e2e suite against the real workflow, so they are deliberately not re-asserted here. This file
is about the script staying the *kind of thing* it is.
"""
from __future__ import annotations
import ast
import re
from pathlib import Path
import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
SCRIPT_PATH = REPO_ROOT / 'scripts' / 'import_model_manual_run.py'
ENV_EXAMPLE_PATH = REPO_ROOT / '.env.example'
CELL_MARKER = '# %%'
EXPECTED_CELLS = 10
FORBIDDEN_IMPORT_ROOTS = ('laborious.activities', 'laborious.workflows')
ALLOWED_LABORIOUS_IMPORT = 'laborious.utils.bundle.format'
# Substrings that would mean a target or a credential was written into the file instead of read
# from the environment. `mongodb://` and `postgresql` are schemes, not hosts, so they are not here.
FORBIDDEN_LITERAL_SUBSTRINGS = (
'localhost',
'127.0.0.1',
'0.0.0.0', # noqa: S104 - a denylist entry, not an interface this test binds
'::1',
'.svc.cluster.local',
'minioadmin',
'mlops_bff_user',
)
# A dotted host with a port (`db.example.com:5432`), an IPv4 address, or inline URL credentials.
HOSTLIKE_PATTERNS = (
re.compile(r'[A-Za-z0-9][A-Za-z0-9.-]*\.[A-Za-z]{2,}:\d{2,5}'),
re.compile(r'\b\d{1,3}(?:\.\d{1,3}){3}\b'),
re.compile(r'//[^/\s]+:[^/\s]+@'),
)
@pytest.fixture(scope='module')
def source() -> str:
"""The script's source text."""
return SCRIPT_PATH.read_text(encoding='utf-8')
@pytest.fixture(scope='module')
def tree(source: str) -> ast.Module:
"""The script parsed. A syntax error here is itself the finding — the module must be runnable."""
return ast.parse(source, filename=str(SCRIPT_PATH))
def _string_literals(tree: ast.Module) -> list[str]:
"""Every string constant in the module, comments and docstrings excluded."""
docstrings = set()
for node in ast.walk(tree):
if isinstance(node, ast.Module | ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef):
body = getattr(node, 'body', [])
if (
body
and isinstance(body[0], ast.Expr)
and isinstance(body[0].value, ast.Constant)
and isinstance(body[0].value.value, str)
):
docstrings.add(id(body[0].value))
return [
node.value
for node in ast.walk(tree)
if isinstance(node, ast.Constant)
and isinstance(node.value, str)
and id(node) not in docstrings
]
def _env_names(tree: ast.Module) -> set[str]:
"""Every environment variable the script reads, by name."""
names: set[str] = set()
readers = {'env', 'env_required', 'getenv'}
for node in ast.walk(tree):
if not isinstance(node, ast.Call) or not node.args:
continue
target = node.func
called = (
target.attr
if isinstance(target, ast.Attribute)
else target.id
if isinstance(target, ast.Name)
else ''
)
first = node.args[0]
if called in readers and isinstance(first, ast.Constant) and isinstance(first.value, str):
names.add(first.value)
return names
class TestCellLayout:
"""The `# %%` contract of design E15."""
def test_script_exists(self) -> None:
assert SCRIPT_PATH.is_file(), f'{SCRIPT_PATH} is missing'
def test_every_cell_marker_is_present(self, source: str) -> None:
markers = [line for line in source.splitlines() if line.strip() == CELL_MARKER]
assert len(markers) == EXPECTED_CELLS, (
f'expected {EXPECTED_CELLS} `{CELL_MARKER}` cell markers, found {len(markers)}'
)
def test_cells_are_numbered_in_order(self, source: str) -> None:
"""Each cell opens with `# Cell N — ...`, and N counts 1..10 without a gap."""
lines = source.splitlines()
numbers: list[int] = []
for index, line in enumerate(lines):
if line.strip() != CELL_MARKER:
continue
following = lines[index + 1] if index + 1 < len(lines) else ''
match = re.match(r'#\s*Cell\s+(\d+)\b', following.strip())
assert match, f'the line after the marker on line {index + 1} does not name its cell'
numbers.append(int(match.group(1)))
assert numbers == list(range(1, EXPECTED_CELLS + 1))
def test_the_marker_precedes_the_first_statement(self, source: str) -> None:
"""The module docstring comes first, then cell 1 — nothing executable before a marker."""
body = source.split('"""', 2)[-1]
assert body.lstrip().startswith(CELL_MARKER)
class TestBoundary:
"""Design E15: the script stands in for the frontend, so it borrows no consumer code."""
def test_no_import_from_activities_or_workflows(self, tree: ast.Module) -> None:
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
assert not alias.name.startswith(FORBIDDEN_IMPORT_ROOTS), (
f'the script imports {alias.name}'
)
elif isinstance(node, ast.ImportFrom) and node.module:
assert not node.module.startswith(FORBIDDEN_IMPORT_ROOTS), (
f'the script imports from {node.module}'
)
def test_the_only_laborious_import_is_the_format_helper(self, tree: ast.Module) -> None:
modules: set[str] = set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
modules.update(a.name for a in node.names if a.name.startswith('laborious'))
elif (
isinstance(node, ast.ImportFrom)
and node.module
and node.module.startswith('laborious')
):
modules.add(node.module)
assert modules <= {ALLOWED_LABORIOUS_IMPORT}, (
f'unexpected laborious imports: {sorted(modules - {ALLOWED_LABORIOUS_IMPORT})}'
)
class TestConfiguration:
"""No credential and no target literal in the file; every variable documented."""
@pytest.mark.parametrize('forbidden', FORBIDDEN_LITERAL_SUBSTRINGS)
def test_no_forbidden_substring(self, tree: ast.Module, forbidden: str) -> None:
offenders = [text for text in _string_literals(tree) if forbidden in text]
assert not offenders, f'{forbidden!r} appears in a string literal: {offenders}'
def test_no_hostlike_literal(self, tree: ast.Module) -> None:
for text in _string_literals(tree):
for pattern in HOSTLIKE_PATTERNS:
assert not pattern.search(text), (
f'{text!r} looks like a host, an address or inline credentials'
)
def test_every_variable_it_reads_is_documented(self, tree: ast.Module) -> None:
documented = {
line.split('=', 1)[0].strip()
for line in ENV_EXAMPLE_PATH.read_text(encoding='utf-8').splitlines()
if '=' in line and not line.lstrip().startswith('#')
}
undocumented = sorted(
_env_names(tree) - documented - {'MLFLOW_TRACKING_USERNAME', 'MLFLOW_TRACKING_PASSWORD'}
)
assert not undocumented, f'.env.example does not document: {undocumented}'
def test_it_reads_the_variables_the_contract_names(self, tree: ast.Module) -> None:
names = _env_names(tree)
required = {
'IMPORT_PASSWORD_KEY',
'IMPORT_BUNDLE_BUCKET',
'IMPORT_BUNDLE_PREFIX',
'IMPORT_STATUS_DB_HOST',
'IMPORT_STATUS_DB_PORT',
'IMPORT_STATUS_DB_NAME',
'IMPORT_STATUS_DB_USER',
'IMPORT_STATUS_DB_PASSWORD',
'IMPORT_SCRIPT_BUNDLE_PATH',
'IMPORT_SCRIPT_BUNDLE_PASSWORD',
'IMPORT_SCRIPT_USERNAME',
'RUNTIME',
}
assert required <= names, f'the script no longer reads: {sorted(required - names)}'