From e95e65b72b5e6e5c1755b109de85d0babb70c01e Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 26 Aug 2026 15:45:33 -0300 Subject: [PATCH] QTZPOC-20 QTZPOC-20: Implement explicit joblib model flavor and enhance model import workflow with extensive E2E tests. This commit introduces a first-class joblib model flavor for direct loading and saving of Python pickle artifacts, including custom code within the artifact's code/ directory, bypassing MLflow's flavor loaders. The model import workflow's status database configuration is now streamlined to default to the POSTGRES_* server, with explicit IMPORT_STATUS_DB_* overrides. Comprehensive end-to-end tests have been added for the model import workflow, covering success paths, gate rejections, digest enforcement, provisioning failures, status record integrity, cleanup, and retry/recovery scenarios using a real-clock environment. New E2 --- .env.example | 14 +- README.md | 74 +- docs/model-flavors.md | 112 +++ docs/model-import.md | 59 +- e2e/conftest.py | 182 +++-- e2e/import_harness.py | 401 ++++++++++- e2e/scenarios.md | 642 ++++++++++++++++- e2e/test_import_model_absence.py | 184 +++++ e2e/test_import_model_cleanup.py | 247 +++++++ e2e/test_import_model_gates.py | 4 + ...test_import_model_provisioning_failures.py | 108 +-- e2e/test_import_model_retry_recovery.py | 514 ++++++++++++++ e2e/test_import_model_second_import.py | 270 ++++++++ e2e/test_import_model_status_record.py | 596 ++++++++++++++++ ...> test_predictions_batch_joblib_flavor.py} | 112 ++- e2e/test_retrain_joblib_flavor.py | 126 ++++ laborious/activities/activities.py | 4 +- laborious/activities/model_import.py | 6 +- laborious/utils/connectors_config.py | 33 +- .../utils/repository/model_repository.py | 379 +++++----- laborious/worker/worker.py | 4 +- pyproject.toml | 11 + requirements-dev.txt | 5 +- scripts/encode.sh | 20 - scripts/encrypt.py | 113 --- scripts/import_model_manual_run.py | 19 +- scripts/inter_arrival.py | 27 - scripts/model_convert.ipynb | 107 --- .../utils/repository/test_model_repository.py | 647 ++++++++++++------ .../laborious/utils/test_connectors_config.py | 50 +- tests/scripts/__init__.py | 0 tests/scripts/test_import_model_manual_run.py | 217 ------ 32 files changed, 4247 insertions(+), 1040 deletions(-) create mode 100644 docs/model-flavors.md create mode 100644 e2e/test_import_model_absence.py create mode 100644 e2e/test_import_model_cleanup.py create mode 100644 e2e/test_import_model_retry_recovery.py create mode 100644 e2e/test_import_model_second_import.py create mode 100644 e2e/test_import_model_status_record.py rename e2e/{test_predictions_batch_joblib_fallback.py => test_predictions_batch_joblib_flavor.py} (56%) create mode 100644 e2e/test_retrain_joblib_flavor.py delete mode 100755 scripts/encode.sh delete mode 100644 scripts/encrypt.py delete mode 100644 scripts/inter_arrival.py delete mode 100644 scripts/model_convert.ipynb delete mode 100644 tests/scripts/__init__.py delete mode 100644 tests/scripts/test_import_model_manual_run.py diff --git a/.env.example b/.env.example index d13ef47..0f631da 100644 --- a/.env.example +++ b/.env.example @@ -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="" \ No newline at end of file diff --git a/README.md b/README.md index c818329..48c75d6 100644 --- a/README.md +++ b/README.md @@ -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//start: +Mounts denied: The path /socket_mnt/home//.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). diff --git a/docs/model-flavors.md b/docs/model-flavors.md new file mode 100644 index 0000000..a8db934 --- /dev/null +++ b/docs/model-flavors.md @@ -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()` | +| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model()`; the wrapper path also unwraps `._model_impl.python_model` | +| `pytorch` | yes | yes | `mlflow.pytorch.load_model()` | +| `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, )`, 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://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 `/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 `/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://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. diff --git a/docs/model-import.md b/docs/model-import.md index 125f36a..51ca951 100644 --- a/docs/model-import.md +++ b/docs/model-import.md @@ -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). diff --git a/e2e/conftest.py b/e2e/conftest.py index 878a114..51af260 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -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): diff --git a/e2e/import_harness.py b/e2e/import_harness.py index 4589317..517f802 100644 --- a/e2e/import_harness.py +++ b/e2e/import_harness.py @@ -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 :', '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` -> diff --git a/e2e/scenarios.md b/e2e/scenarios.md index 8841424..71abd5c 100644 --- a/e2e/scenarios.md +++ b/e2e/scenarios.md @@ -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//start: +Mounts denied: The path /socket_mnt/home//.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 ---- \ No newline at end of file +--- +## 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://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/.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` + +--- diff --git a/e2e/test_import_model_absence.py b/e2e/test_import_model_absence.py new file mode 100644 index 0000000..98f164f --- /dev/null +++ b/e2e/test_import_model_absence.py @@ -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)] diff --git a/e2e/test_import_model_cleanup.py b/e2e/test_import_model_cleanup.py new file mode 100644 index 0000000..9d94509 --- /dev/null +++ b/e2e/test_import_model_cleanup.py @@ -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' + ) diff --git a/e2e/test_import_model_gates.py b/e2e/test_import_model_gates.py index 593c7fc..be9ab0e 100644 --- a/e2e/test_import_model_gates.py +++ b/e2e/test_import_model_gates.py @@ -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, diff --git a/e2e/test_import_model_provisioning_failures.py b/e2e/test_import_model_provisioning_failures.py index a0f15e0..1062db1 100644 --- a/e2e/test_import_model_provisioning_failures.py +++ b/e2e/test_import_model_provisioning_failures.py @@ -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 diff --git a/e2e/test_import_model_retry_recovery.py b/e2e/test_import_model_retry_recovery.py new file mode 100644 index 0000000..03db085 --- /dev/null +++ b/e2e/test_import_model_retry_recovery.py @@ -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' + ) diff --git a/e2e/test_import_model_second_import.py b/e2e/test_import_model_second_import.py new file mode 100644 index 0000000..78ffddf --- /dev/null +++ b/e2e/test_import_model_second_import.py @@ -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 diff --git a/e2e/test_import_model_status_record.py b/e2e/test_import_model_status_record.py new file mode 100644 index 0000000..5ab5a5d --- /dev/null +++ b/e2e/test_import_model_status_record.py @@ -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() + ] diff --git a/e2e/test_predictions_batch_joblib_fallback.py b/e2e/test_predictions_batch_joblib_flavor.py similarity index 56% rename from e2e/test_predictions_batch_joblib_fallback.py rename to e2e/test_predictions_batch_joblib_flavor.py index a807c3e..87b6f9d 100644 --- a/e2e/test_predictions_batch_joblib_fallback.py +++ b/e2e/test_predictions_batch_joblib_flavor.py @@ -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 diff --git a/e2e/test_retrain_joblib_flavor.py b/e2e/test_retrain_joblib_flavor.py new file mode 100644 index 0000000..dabda57 --- /dev/null +++ b/e2e/test_retrain_joblib_flavor.py @@ -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 diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index e62e0d1..a06eb16 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -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 diff --git a/laborious/activities/model_import.py b/laborious/activities/model_import.py index 0af7a88..886ee73 100644 --- a/laborious/activities/model_import.py +++ b/laborious/activities/model_import.py @@ -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. diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index aa2a019..4b31496 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -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', } diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index e89b685..500622f 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -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://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 + `/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: `/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) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index f10a6cb..c3b3014 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 06f349f..2f2ef0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/requirements-dev.txt b/requirements-dev.txt index a2b8893..57224a1 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -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 diff --git a/scripts/encode.sh b/scripts/encode.sh deleted file mode 100755 index 73d0534..0000000 --- a/scripts/encode.sh +++ /dev/null @@ -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}" diff --git a/scripts/encrypt.py b/scripts/encrypt.py deleted file mode 100644 index 14eb3da..0000000 --- a/scripts/encrypt.py +++ /dev/null @@ -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() diff --git a/scripts/import_model_manual_run.py b/scripts/import_model_manual_run.py index 2c393f7..b260313 100644 --- a/scripts/import_model_manual_run.py +++ b/scripts/import_model_manual_run.py @@ -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.' ) diff --git a/scripts/inter_arrival.py b/scripts/inter_arrival.py deleted file mode 100644 index f3e2cef..0000000 --- a/scripts/inter_arrival.py +++ /dev/null @@ -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() -# %% diff --git a/scripts/model_convert.ipynb b/scripts/model_convert.ipynb deleted file mode 100644 index d04f63e..0000000 --- a/scripts/model_convert.ipynb +++ /dev/null @@ -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 -} diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 4a89d69..2c449de 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -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) diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 8dd82f5..08639c1 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -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') diff --git a/tests/scripts/__init__.py b/tests/scripts/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/scripts/test_import_model_manual_run.py b/tests/scripts/test_import_model_manual_run.py deleted file mode 100644 index 2d61784..0000000 --- a/tests/scripts/test_import_model_manual_run.py +++ /dev/null @@ -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)}'