diff --git a/README.md b/README.md index d889abd..f61454e 100644 --- a/README.md +++ b/README.md @@ -1003,10 +1003,7 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an | `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_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 | +| `IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD}` | Override, only when the import log is not on the `POSTGRES_*` server | the `POSTGRES_*` value | 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 | @@ -1113,9 +1110,12 @@ Legacy MinIO object layout (relative key): `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). +`joblib` reads the artifact directly — download, `/model.pkl`, `joblib.load` — and writes on +retrain (`model.pkl` + `code/`, no `MLmodel`, so only this runtime reads it back). + +Every flavor puts the artifact's own `code/` on `sys.path` before loading, **when that directory +exists**: MLflow's loaders only do it when the `MLmodel` manifest declares `code`, so `sklearn` and +`pyfunc` resolve the URI to a local directory first and load from there. Full reference, including how to switch an existing model over: **[docs/model-flavors.md](docs/model-flavors.md)**. diff --git a/docs/model-flavors.md b/docs/model-flavors.md index 82c0f60..74c0d72 100644 --- a/docs/model-flavors.md +++ b/docs/model-flavors.md @@ -22,9 +22,9 @@ reads the artifact's `MLmodel` manifest to guess it. | 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` | -| `joblib` | yes | yes | downloads the artifact and `joblib.load`s its pickle — **no MLflow flavor loader involved** | +| `sklearn` | yes | yes | resolves the URI to a local directory, then `mlflow.sklearn.load_model()` | +| `pyfunc` | yes | yes | resolves the URI to a local directory, then `mlflow.pyfunc.load_model()`; the wrapper path also unwraps `._model_impl.python_model` | +| `joblib` | yes | yes | downloads the same artifact and `joblib.load`s its `/model.pkl` — **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 three values govern reading and writing. @@ -35,80 +35,108 @@ not installed in the runtime, so `mlflow.pytorch.load_model` could never have lo **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 +## Same artifact, three readers -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. +The flavor changes the reader, never the path. All three resolve the same two artifact directories of +the Production run: -**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. +``` +/artifacts/ +├── data_model/ ← transform_flavor loads this one +│ ├── MLmodel ← parsed by sklearn/pyfunc; ignored by joblib +│ ├── model.pkl ← what joblib.load reads +│ └── code/ ← put on sys.path by every flavor, when present +├── prediction_model/ ← predict_flavor loads this one +└── model_card/, *.csv +``` -**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. +The run's `artifacts/` root holds **no** `.pkl` — every model pickle lives one level down, inside its +artifact directory. Predict resolves that directory through the registry URI +(`models://production`, which MLflow points at the same `prediction_model`); transform builds +the run artifact URI directly with `get_model_uri`, since the run is already resolved. Observability +is the same for the three: `MODEL_READ_*` on load, `MODEL_WRITE_*` on `log_model`, under the base +`operation_type` — no per-flavor metric, operation type or label. -> `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. +## `code/` on `sys.path`, for every flavor -**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. +An artifact may ship the modules its own classes live in, under `/code/`. Without that +directory on `sys.path`, unpickling raises `ModuleNotFoundError` — the artifact is intact, the class +is simply not importable. -**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. +MLflow's loaders do add it, but **from the manifest**: `_add_code_from_conf_to_system_path` acts only +when the `MLmodel` flavor config declares a `code` entry. An artifact carrying `code/` on disk +without that entry — code logged with `log_artifacts`, a re-log that dropped the config, anything a +`joblib` retrain wrote — loads with no module path at all. So this runtime reads the **directory**, +which no manifest can misdescribe: -**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. +- `prepend_artifact_code_dir` inserts `/code` at the front of `sys.path`, **only when that + directory exists**, and only once (an entry already present is not duplicated). +- Because the insert must happen *before* the loader unpickles, `sklearn` and `pyfunc` resolve their + URI to a local directory first (`mlflow.artifacts.download_artifacts`) and the loader is handed + that directory. Same bytes, same download count — MLflow would have downloaded it anyway. +- The wrapper path (`download_model(load_wrapper=True)`) already had its local directory and adds + the entry itself; `joblib` adds it inside `_load_model_pickle`. +- When the manifest *does* declare `code`, MLflow inserts the same directory a second time. Harmless. -**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. +The entry is never removed: a loaded model may import lazily, long after the load returned. + +> `sys.path` is process-global: two models embedding a package with the same name (`utils/`) on the +> same worker resolve to whichever entered first. Mitigation, if it bites: one worker per model. + +## `sklearn` + +Default. A standard MLflow sklearn artifact: `mlflow.sklearn.log_model` writes `MLmodel` + +`model.pkl`, `mlflow.sklearn.load_model` reads it, and every outside consumer (MLflow UI, MLflow +serving) recognizes it as a model. + +## `pyfunc` + +For models logged as `mlflow.pyfunc`, including the wrapper models this pipeline produces. Retrain +writes through the wrapper's own `store_model(artifact_path=, code_path=…)`. Two read +paths, and the second is **orthogonal to the flavor**: + +- Plain: `mlflow.pyfunc.load_model()`, called as `predict(context, model_input)`. +- Wrapper (`download_model(load_wrapper=True)`): downloads to `./tmp/artifacts//` + via `dowload_artifacts`, loads as pyfunc, then unwraps `._model_impl.python_model`. Inside such an + artifact the real pickle is **nested** — `/artifacts/stacking_model.pkl` for predict, + `/artifacts/training_transformer.pkl` for transform (`PREDICTION_COMPRESSED_PATH` / + `TRANSFORMED_COMPRESSED_PATH`). That nesting is why the `joblib` listing does not recurse. + +## `joblib` + +For models whose `.pkl` is in the compressed form `joblib.dump` writes, which MLflow's `pickle.load` +does not read — the artifact is intact, but no flavor loader opens it. + +- **Read**: `client.download_artifacts(run_id, )`, then + `joblib.load(/model.pkl)`. `model.pkl` is `MODEL_PICKLE_NAME`, the same constant + `log_model` writes; another `.pkl` directly inside the artifact is a fallback, in name order. +- **Write**: `joblib.dump(model, '/model.pkl', compress=3)`, plus a re-log of the source + artifact's `code/` to `/code`. +- **`code/` goes on `sys.path`** before the load, when present — see the section above; the joblib + flavor is where that started, and it is now shared by all three. +- **No `MLmodel` is written**, so only this runtime reads the artifact back — + `mlflow.sklearn.load_model` and `models://production` both fail on it. A model retrained in + joblib stays joblib. +- **The pickled object must expose `predict(data)`**, plus `fit(data)` if retrained. A pyfunc + `PythonModel` wrapper expecting `predict(context, model_input)` raises `TypeError`. ## Switching a model to `joblib` -The runtime does not migrate documents. Editing the `models` collection is an operator step, and the -order matters: +The runtime does not migrate documents; editing `models` is an operator step. -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. +1. **Inventory before the deploy** — run predict/transform and note which models raise. There is no + discovery afterwards: the failure names the loader, not the fix. +2. **Edit the document** — `predict_flavor` and/or `transform_flavor` to `'joblib'`. Imported models + arrive ready: `import_model` writes it itself, see [`model-import.md`](model-import.md) § 12. +3. **Check the first retrain** — the new run must carry `prediction_model/model.pkl` and, when the + source had one, `prediction_model/code/`. +4. **Rollback is not just a revert** — artifacts a joblib retrain already wrote stop being loadable, + so it also means promoting the previous version back and reverting the documents' flavors. ## 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. +`get_model_run_id` takes `source.split('/')[2]` of the registered version's `source`. **Not confirmed +against the real tracking server**: no scheme reproduced locally puts the run id there (a file store +yields `''`, S3 the bucket, `mlflow-artifacts:/` the experiment id), yet predict works in production. +Every flavor that resolves a run goes through this line — confirm it with a **read-only** query +before relying on it. diff --git a/docs/sientia-bundle-format.md b/docs/sientia-bundle-format.md index 79dbc73..209cb82 100644 --- a/docs/sientia-bundle-format.md +++ b/docs/sientia-bundle-format.md @@ -1,111 +1,48 @@ # The `.sientia` bundle format -A `.sientia` file is one exported MLflow model: its whole artifact tree, its run parameters and six -fields of origin metadata, zipped and then encrypted under a password the exporter chose. It is -written by the Streamlit platform (`sientia-projects-templates`) and read by this repository's -importer. Two programs in two repositories agree on these bytes, so the constants below are -duplicated by necessity — the drift is held in place by tests on both sides, not by a shared library. +A `.sientia` file is one exported MLflow model: its whole artifact tree, its run parameters and its +origin metadata, zipped and then encrypted under a password the exporter chose. -This document covers both directions: **§1–§3 how a bundle is generated**, so the producer can be -audited against it, and **§4–§7 how it is read**, so a rejection can be explained. +**It is written by the Streamlit platform, `sientia-projects-templates`** (`app/src/operations/ +model_export/`). That application owns the export — its screen, its flow, its code — and documents +it. This page is the short version: the parameters both sides have to agree on, and what this +repository's reader accepts. -Normative implementations: +## How it is generated -| Side | Where | +The producer downloads the run's artifact tree whole, copies `run.data.params` **verbatim** (with +`target_variable` among them, or the import fails at gate 6), writes `metadata.json` first into a +`ZIP_DEFLATED` archive with the artifacts following as `artifacts/` and no directory +entries, encrypts that zip, and offers it as `{model_name}_v{model_version}.sientia`. + +`metadata.json` has exactly two keys: `parameters` — the run's parameters verbatim — and `metadata`, +six fields: `model_name`, `experiment_name`, `run_id`, `model_version`, `model_project`, +`export_timestamp` (UTC ISO 8601). + +## Encryption parameters + +The reader derives the key with the values the header carries, but refuses anything above the +producer's own cost — a header declaring more is a resource-exhaustion attempt, refused before any +memory is allocated. + +| | Value | |---|---| -| Writer | `sientia-projects-templates`, `app/src/operations/model_export/{crypto,bundle,exporter}.py` | -| Reader | this repository, `laborious/utils/bundle/{format,reader}.py` | -| Test writer | this repository, `tests/helpers/bundle_factory.py` — mirrors the producer field for field, for building hostile variants | -| Golden fixture | `tests/fixtures/bundle/` — a real producer artefact, with its password and producer commit recorded in that directory's `README.md` | +| KDF | Argon2id (`crypto_pwhash_ALG_ARGON2ID13`) | +| memlimit | `268435456` (256 MiB) — a **ceiling** on the reader's side | +| opslimit | `3` — likewise a ceiling | +| parallelism | `1` (reserved, not configurable) | +| salt | 16 bytes, fresh per bundle | +| key | 32 bytes | +| AEAD | XChaCha20-Poly1305 `secretstream`, 64 KiB plaintext chunks, last tagged `FINAL` | +| digest | SHA-256 of the **pushed chunks only** — on disk, `sha256(file[164:])`, past the 140-byte header and the 24-byte stream header. Hashing from 140 rejects every valid bundle | +| password | normalised `NFC`, encoded UTF-8 — no trimming, no case folding | ---- +The export-side password rule (12+ characters, mixed case, digit, symbol) is enforced there only: by +the time a bundle exists, the password is whatever it was, and the reader cannot re-check it. -## 1. Generation, step by step +## Shape of the file -Entry point: `exporter.export_model_bundle(*, client, run_id, model_name, experiment_name, -model_version, model_project, password) -> bytes`. - -1. **Password strength** — `crypto.validate_password_strength`: at least 12 characters, one - uppercase, one lowercase, one digit, one symbol. **Export-side only.** The importer cannot - re-check it and does not try: by the time a bundle exists, the password is whatever it was. -2. **Exportability** — `bundle.validate_run_is_exportable`: the run's top-level artifacts must - include **both** `data_model` and `prediction_model`, checked with `client.list_artifacts(run_id)` - before anything is downloaded. -3. **Whole-tree download** — `client.download_artifacts(run_id, '', staging)` into a - `tempfile.mkdtemp(prefix='sientia_export_')` directory. The entire artifact tree, not a subset. -4. **Run parameters** — `dict(run.data.params)`, copied **verbatim**. Whatever the run logged is - what the bundle carries, including `target_variable` (see §3). -5. **Metadata document** — `bundle.build_metadata_json`, exactly two top-level keys: - - `parameters`: the run parameters, verbatim; - - `metadata`: `model_name`, `experiment_name`, `run_id`, `model_version`, `model_project` and - `export_timestamp = datetime.now(UTC).isoformat()`. Six fields, no others. -6. **Zip** — `bundle.build_zip`, `zipfile.ZIP_DEFLATED`. `metadata.json` is written **first**, with - `json.dumps(...)` and no indent. Every file under the staging directory follows as - `artifacts/`, with `\` replaced by `/`. **No directory entries are stored**, so a reader - must derive directories from file paths. -7. **Encryption** — `crypto.encrypt_archive`: a fresh 16-byte `os.urandom` salt, Argon2id key - derivation, then a libsodium `secretstream` push in 64 KiB plaintext chunks, the last chunk tagged - `FINAL` and the rest `MESSAGE`. -8. **Download name** — `{model_name}_v{model_version}.sientia`. The staging directory is removed in a - `finally`, and the password is popped from the Streamlit session state. - -Producer pins: `pynacl==1.5.0`, `mlflow==2.10.1`. - -## 2. Constants - -Copied from `crypto.py` unless noted. The reader's names are the ones in -`laborious/utils/bundle/format.py`. - -| Producer | Value | Reader | -|---|---|---| -| `MAGIC` | `b'SIENTIA1'` | `MAGIC` | -| `_HEADER_STRUCT` | `struct.Struct('>8sBBB16sQQB32s64s')` | `HEADER_STRUCT` | -| `HEADER_STRUCT_SIZE` | `140` | `HEADER_SIZE` | -| `FORMAT_VERSION` | `1` | `FORMAT_VERSION` | -| `AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM` | `1` | same name | -| `KDF_ID_ARGON2ID` | `1` | same name | -| `SALT_SIZE` | `16` | `SALT_SIZE` | -| `KDF_MEMLIMIT_BYTES` | `268435456` (256 MiB) | `MAX_KDF_MEMLIMIT` — a **ceiling** on the reader's side | -| `KDF_OPSLIMIT` | `3` | `MAX_KDF_OPSLIMIT` — likewise a ceiling | -| `KDF_PARALLELISM_RESERVED` | `1` | `PARALLELISM` | -| `DIGEST_SIZE` | `32` | `DIGEST_SIZE` | -| `RESERVED_SIGNATURE_SIZE` | `64` | `RESERVED_SIGNATURE_SIZE` | -| `CHUNK_SIZE` | `65536` (64 KiB) | `CHUNK_SIZE` | -| KDF algorithm | `crypto_pwhash_ALG_ARGON2ID13` | `KDF_ALG` | -| Key length | `crypto_secretstream_xchacha20poly1305_KEYBYTES` (32) | `KEY_SIZE` | - -Fixed by libsodium rather than by the producer, but needed to frame the stream: - -- `crypto_secretstream_xchacha20poly1305_HEADERBYTES == 24` → `STREAM_HEADER_SIZE` -- `crypto_secretstream_xchacha20poly1305_ABYTES == 17` → `ABYTES` -- pre-upload prefix `140 + 24 + (65536 + 17) = 65717` → `PREFIX_SIZE`, the smallest number of leading - bytes from which a header and one full chunk can be checked - -**Password normalisation.** `_derive_key` applies `unicodedata.normalize('NFC', password)` and -encodes UTF-8 — no trimming, no case folding, no padding. The reader normalises identically, so a -password typed in decomposed form opens a bundle written from the composed form. - -## 3. `target_variable`, and where it does not come from - -Gate 6 requires `parameters.target_variable`: it is the only source of `model_config.target` in the -model document, and every workflow that serves a model indexes that key directly. On the producer -side: - -| Export path | Logs `target_variable`? | Evidence | -|---|---|---| -| Time-series training template | yes | `app/pages/template_01.py` L1212-1214 | -| Model manager, "Save Experiment" | yes, if the session carried it | `app/pages/model_manager.py` L791-796 re-logs `st.session_state.parameters` verbatim | -| Model ensemble | same condition | `app/pages/model_ensemble.py` L695-700 | -| Pipeline / timeseries template | **no** | `app/pages/timeseries_template.py` L1761-1762, L1811-1812 log `run_metadata` only | -| AutoML | **no** | `automl_logging_operations.py` L454-491; the target appears as `"target"` inside the model card JSON, never as a run parameter | - -So a gate-6 rejection naming `target_variable` means the origin run came from the pipeline or AutoML -path. The remedy is to re-export from a run trained by the time-series template, which is what the -importer's `IMPORT_BUNDLE_INCOMPLETE` sentence already asks for. - ---- - -## 4. The file on disk +`HEADER_STRUCT` is `>8sBBB16sQQB32s64s`. Big-endian, no alignment padding, so the offsets are exact: ``` offset width field constraint the reader enforces @@ -125,72 +62,39 @@ offset width field constraint the reader enforces 164 ... chunk stream 64 KiB plaintext chunks, each +17 bytes of AEAD tag ``` -All integers are big-endian; the format string starts with `>`, so there is no alignment padding and -the offsets above are exact. +`format_version` is checked for equality, not for a range: an unknown version is refused at gate 2 +rather than parsed optimistically. `reserved_signature` must be all zero today, so a signed bundle +will be a new format version and not a silent change of meaning. -**The digest boundary.** The producer computes `hashlib.sha256(ciphertext)` where `ciphertext` is the -concatenation of the *pushed chunks only*, and then writes `header + stream_header + ciphertext`. On -disk, therefore, the digest covers everything after offset **164**, not after 140: +`laborious/utils/bundle/format.py` is the normative layout on this side, and the golden fixture is +what keeps it and the producer from drifting apart. -> `ciphertext_digest == sha256(file[HEADER_SIZE + STREAM_HEADER_SIZE:])` +## How it is read -A reader that hashes from 140 will reject every valid bundle. The fixture test that recomputes both -values and asserts the first matches while the second does not is what pins this, and is the reason -the golden fixture exists. +`BundleReader.open()` runs seven ordered gates, each mapped to an import step and to one message a +person can read: -**Version policy.** `format_version` is checked for equality, not for a range: a version this reader -does not know is refused at gate 2 rather than parsed optimistically. `reserved_signature` is -reserved for a future signature and must be all zero today, so a signed bundle is a new format -version and not a silent change of meaning. - -## 5. The seven gates - -`BundleReader.open()` runs them in this order, in one function, so the order cannot be changed by a -caller. Each raises `BundleGateError(gate=n)`, which the activity layer maps to an import step. - -| Gate | Step recorded | What it checks | Written to disk | -|---|---|---|---| -| 1 | `download` | object size ceiling, key prefix, `.sientia` suffix — from `stat_object`, before the body is fetched; then the SHA-256 of the downloaded bytes against the digest the uploader computed | the downloaded object only | -| 2 | `download` | the whole header table of §4, including the chunk-stream digest | nothing | -| 3 | `decryption` | Argon2id derivation with the header's parameters, then `secretstream` pull. The AEAD is the integrity check: a wrong password and an altered file fail identically, and the message says so | the decrypted zip, in the work directory | -| 4 | `archive_inspection` | central directory only: entry count, total uncompressed size, compression ratio, absolute names, `..` components, symlink and non-regular modes | nothing | -| 5 | `extraction` | extraction into a **fresh** directory (a pre-existing one is a failure), every member's `realpath` contained in it, no overwrite of a pre-existing file; a rejection removes its own partial tree | the extracted tree | -| 6 | `structure_validation` | `metadata.json` + `artifacts/` layout, exactly two top-level keys, all six metadata fields present and non-blank, `experiment_name` between 3 and 50 characters, `parameters.target_variable` present and non-blank, and `prediction_model`/`data_model` each holding `MLmodel` and `model.pkl` | nothing | -| 7 | `content_policy` | every file is either a known name (`MLmodel`, `metadata.json`, `model.pkl`, `conda.yaml`, `python_env.yaml`, `requirements.txt`, `model_card.json`, `model_card.svg`) or a known suffix (`.csv .json .yaml .yml .txt .pkl .svg .md`) | nothing | - -Gates 1 and 2 share the `download` step deliberately: an oversized object, a wrong magic, an -unsupported version and a digest that disagrees all say *these are not the bytes that were sent*, and -none of them may borrow the `decryption` sentence, which implicates the password. - -The `experiment_name` bound in gate 6 is not cosmetic: the import record's `experiment_name` column is -`VARCHAR(50)`, and the name is written onto the record before provisioning starts. Enforcing the -column's limit at the gate is what keeps it from surfacing three activities later as a constraint -violation. - -## 6. Ceilings and where they come from - -Gate 1 and gate 4 read their limits from `build_import_config()`, so a deployment can tighten them: - -| Variable | Default | Gate | +| Gate | Step | Refuses | |---|---|---| -| `IMPORT_MAX_OBJECT_BYTES` | 1 GiB | 1 | -| `IMPORT_MAX_ARCHIVE_ENTRIES` | 5000 | 4 | -| `IMPORT_MAX_UNCOMPRESSED_BYTES` | 4 GiB | 4 | -| `IMPORT_MAX_COMPRESSION_RATIO` | 200 | 4 | -| `IMPORT_BUNDLE_PREFIX` | `imported_models/` | 1 | +| 1 | `download` | object too large, wrong prefix or suffix, digest ≠ the one the uploader sent | +| 2 | `download` | header that is not this format or not this version | +| 3 | `decryption` | wrong password, or altered bytes — the AEAD cannot tell them apart, and the message says so | +| 4 | `archive_inspection` | zip bomb, absolute or `..` entries, symlinks — read from the central directory, before extracting | +| 5 | `extraction` | anything that escapes the fresh work directory, or overwrites a file | +| 6 | `structure_validation` | missing `metadata.json` fields, missing `target_variable`, missing `prediction_model`/`data_model` | +| 7 | `content_policy` | a file that is not a known model file by name or suffix | -The KDF ceilings are **not** configurable: they are the producer's own cost parameters, and a header -declaring more is a resource-exhaustion attempt, refused before any memory is allocated. +Gates 1 and 2 share the `download` step on purpose: neither may borrow the `decryption` sentence, +which implicates the password. A gate-6 rejection naming `target_variable` means the origin run never +logged it — the fix is a re-export, on the producer's side. -## 7. Fixtures +Error codes, sentences and the ceilings the gates read from configuration: +[`model-import.md`](model-import.md). -- `tests/fixtures/bundle/golden_model_v3.sientia` — emitted by the producer's own code in its venv, - with the password and producer commit recorded in `tests/fixtures/bundle/README.md`. Regenerate it - with `tests/fixtures/bundle/regenerate.py`. It pays the producer's real Argon2id cost (256 MiB, - 3 passes), so it is used where the real thing matters and not in every test. -- `tests/helpers/bundle_factory.py` — the test-only writer. Same header fields at the same offsets, - but 8 MiB and 1 pass, and with a knob for every hostile variant: wrong magic, unknown version, KDF - parameters above the ceiling, digest mismatch, tampered ciphertext, truncated final chunk, non-zero - reserved signature, zip bomb, `..` entry, absolute entry, symlink, missing `data_model`, missing - `MLmodel`, missing `model.pkl`, malformed metadata, missing `target_variable`, extra `.sh` file. A - test asserts its header bytes are structurally identical to the golden fixture's. +## Fixtures + +- `tests/fixtures/bundle/` — a real bundle emitted by the producer, with its password and producer + commit recorded in that directory's `README.md`. It pays the producer's real Argon2id cost, so it + is used where the real bytes matter and not in every test. +- `tests/helpers/bundle_factory.py` — the test-only writer: same header at the same offsets, cheap + KDF parameters, and a knob for every hostile variant the gates above refuse. diff --git a/e2e/conftest.py b/e2e/conftest.py index 29076b3..bf4310f 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -546,6 +546,15 @@ def patch_mlflow(mock_mlflow_models, joblib_flavor_artifacts): mock_mlflow.pyfunc = MagicMock() mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model) + # The sklearn and pyfunc paths resolve the URI to a local directory before loading, so that the + # artifact's `code/` can go on `sys.path` whatever the `MLmodel` manifest says. Echoing the URI + # keeps the loader stubs above — which branch on the URI text — working unchanged, and points + # the code-directory lookup at a path that does not exist. + mock_mlflow.artifacts = MagicMock() + mock_mlflow.artifacts.download_artifacts = MagicMock( + side_effect=lambda artifact_uri: str(artifact_uri) + ) + # Mock other mlflow methods that might be called mock_mlflow.set_tracking_uri = MagicMock() mock_mlflow.get_run = MagicMock( diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 54e3391..ea13d07 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -45,6 +45,10 @@ ARTIFACTS_PATH = './tmp/artifacts' TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' +# Pickle file name the `joblib` flavor writes and looks for first, inside the model artifact's own +# directory (`/artifacts/prediction_model/`, `/artifacts/data_model/`). +MODEL_PICKLE_NAME = 'model.pkl' + INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'." @@ -358,21 +362,31 @@ class MLFlowRepository(SientiaMonitoring): """ Locate the pickle file of a downloaded model artifact. - 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. + Primary path of the `joblib` flavor. `local_dir` is the directory of the + model artifact itself — `/artifacts/prediction_model` or + `/artifacts/data_model`, the very artifacts the other flavors + resolve — never the run's `artifacts/` root. So the pickle sought is + `/model.pkl`, the name `log_model` writes. + + `MODEL_PICKLE_NAME` wins whenever it is present, so write and read agree + by construction. Any other `.pkl` directly inside the artifact is taken + as a fallback, in name order, for artifacts built by hand. The listing + does not recurse: a pyfunc wrapper nested inside the artifact + (`prediction_model/artifacts/stacking_model.pkl`) 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. + local_dir (str): Local directory of the downloaded model artifact. artifact_name (str): Artifact name, e.g. 'prediction_model'. Returns: - str: Path to the first top-level `.pkl` file, in name order. + str: Path to `/model.pkl` when it exists, otherwise to + the first other `.pkl` directly inside `local_dir`, in name + order. Raises: - ValueError: If the top level of `local_dir` has no `.pkl` file. + ValueError: If `local_dir` holds no `.pkl` file directly. """ entries = sorted(listdir(local_dir)) pkl_files = [ @@ -383,11 +397,97 @@ class MLFlowRepository(SientiaMonitoring): if not pkl_files: raise ValueError( - f'No .pkl file found in {artifact_name} artifact. ' - f'Top-level entries in {local_dir}: {entries}' + f'No .pkl file found in {artifact_name} artifact. Entries in {local_dir}: {entries}' ) - return path.join(local_dir, pkl_files[0]) + chosen = MODEL_PICKLE_NAME if MODEL_PICKLE_NAME in pkl_files else pkl_files[0] + + return path.join(local_dir, chosen) + + def prepend_artifact_code_dir( + self, + local_dir: str | None, + flavor: str, + metadata: dict[str, Any] | None = None, + ) -> str | None: + """ + Put a downloaded artifact's own `code` directory on `sys.path`. + + Every flavor goes through here, not only `joblib`. MLflow's loaders + extend `sys.path` from the **manifest**: `_add_code_from_conf_to_system_path` + acts only when the `MLmodel` flavor config declares a `code` entry. An + artifact that carries `code/` on disk without that entry — anything + whose code was logged with `log_artifacts`, anything re-logged by a + path that dropped the config, anything a `joblib` retrain wrote — loads + with no module path and raises `ModuleNotFoundError` at unpickle time, + with the artifact intact on disk the whole while. Reading the directory + itself depends on no manifest, so it works for all three flavors. + + Prepending before the loader runs is deliberate: the loader is what + unpickles, so the path has to be in place first. When the manifest + *does* declare `code`, MLflow inserts the same directory again and the + duplicate entry is harmless. + + The entry is never removed. A loaded model may import lazily, long after + the load returned. + + Args: + local_dir (str | None): Local directory of the downloaded model + artifact, e.g. the download of `/artifacts/prediction_model`. + flavor (str): Flavor doing the load, for the log line only. + metadata (dict[str, Any] | None): Metadata used for structured + logging. + + Returns: + str | None: The directory added, `None` when the artifact carries no + `code` directory or it was already on `sys.path`. + """ + code_dir = self.resolve_artifact_code_dir(local_dir) + + if not code_dir: + return None + + if code_dir in sys.path: + self.debug(f'{flavor} flavor: {code_dir} already on sys.path', metadata) + return None + + sys.path.insert(0, code_dir) + self.debug(f'{flavor} flavor: added {code_dir} to sys.path', metadata) + + return code_dir + + def _localize_model_uri( + self, model_uri: str, flavor: str, metadata: dict[str, Any] | None = None + ) -> str: + """ + Resolve a model URI to a local directory and put its `code` on `sys.path`. + + The MLflow loaders download the artifact themselves and take a local + path just as readily as a URI, so resolving here costs no extra + download — it only moves the download one step earlier, which is what + makes the artifact's `code` directory visible before the loader + unpickles anything. + + Args: + model_uri (str): `models://production` or a run artifact URI. + flavor (str): Flavor doing the load, for the log line only. + metadata (dict[str, Any] | None): Metadata used for structured + logging. + + Returns: + str: Local directory the model artifact was downloaded to, to be + handed to the flavor loader in place of the URI. + + Raises: + Exception: Any error raised while downloading the artifact. + """ + local_path = mlflow.artifacts.download_artifacts(artifact_uri=model_uri) + + self.debug(f'{flavor} flavor: {model_uri} downloaded to {local_path}', metadata) + + self.prepend_artifact_code_dir(local_path, flavor, metadata) + + return local_path def _load_model_pickle( self, @@ -404,13 +504,14 @@ class MLFlowRepository(SientiaMonitoring): `._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. + The artifact's own `code` directory — `/code` — is prepended + to `sys.path` before deserializing, when it exists, so that classes + whose modules ship inside the artifact resolve at unpickle time. That + step is shared with every other flavor: see `prepend_artifact_code_dir`. Args: - local_dir (str): Local directory holding the downloaded artifact. + local_dir (str): Local directory of the downloaded model artifact, + e.g. the download of `/artifacts/prediction_model`. artifact_name (str): Artifact name, e.g. 'prediction_model'. metadata (dict[str, Any] | None): Metadata used for structured logging, so the chosen pickle path correlates with the other @@ -420,15 +521,12 @@ class MLFlowRepository(SientiaMonitoring): Any: The object deserialized from the artifact's pickle file. Raises: - ValueError: If the artifact directory has no top-level `.pkl` file. + ValueError: If the artifact directory holds no `.pkl` file directly. Exception: Any error raised by `joblib.load` itself. """ pkl_path = self._find_model_pickle(local_dir, artifact_name) - 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.prepend_artifact_code_dir(local_dir, 'joblib', metadata) self.debug(f'joblib flavor: loading {artifact_name} from {pkl_path}', metadata) @@ -448,8 +546,10 @@ class MLFlowRepository(SientiaMonitoring): `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. + downloading — is deliberately not reused. `artifact_name` is the same + model artifact the other flavors resolve, so the download yields + `/artifacts//` and the pickle read from it is that + directory's `model.pkl`. The deserialized object is returned as is. 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 — @@ -470,7 +570,7 @@ class MLFlowRepository(SientiaMonitoring): artifact was downloaded to. Raises: - ValueError: If the downloaded artifact has no top-level `.pkl` file. + ValueError: If the downloaded artifact holds no `.pkl` file directly. Exception: Any error raised while resolving the Production run, downloading the artifact, or deserializing the pickle. """ @@ -559,10 +659,13 @@ 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. + - The artifact's `code` directory, when present, is put on + `sys.path` before the load, whatever the flavor. `sklearn` and + `pyfunc` are therefore resolved to a local directory first and the + loader is handed that directory instead of the URI. - 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. + `joblib.load`, which needs no `MLmodel` manifest. """ if flavor not in ('pyfunc', 'sklearn', 'joblib'): raise ValueError(INVALID_FLAVOR_MESSAGE) @@ -582,10 +685,12 @@ class MLFlowRepository(SientiaMonitoring): core_labels = self.get_core_labels(metadata, operation_type='load_predict_model') start_time = time.time() try: + local_path = self._localize_model_uri(model_uri, flavor, metadata) + if flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(model_uri) + model = mlflow.pyfunc.load_model(local_path) else: - model = mlflow.sklearn.load_model(model_uri) + model = mlflow.sklearn.load_model(local_path) except Exception as e: await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) raise e @@ -617,11 +722,14 @@ class MLFlowRepository(SientiaMonitoring): model cannot be loaded. Notes: + - The artifact's `code` directory, when present, is put on + `sys.path` before the load, whatever the flavor. `sklearn` and + `pyfunc` are therefore resolved to a local directory first and the + loader is handed that directory instead of the URI. - 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. + which needs no `MLmodel` manifest. """ if flavor not in ('sklearn', 'pyfunc', 'joblib'): @@ -646,10 +754,12 @@ class MLFlowRepository(SientiaMonitoring): core_labels = self.get_core_labels(metadata, operation_type='load_transform_model') start_time = time.time() try: + local_path = self._localize_model_uri(model_uri, flavor, metadata) + if flavor == 'sklearn': - model = mlflow.sklearn.load_model(model_uri) + model = mlflow.sklearn.load_model(local_path) else: - model = mlflow.pyfunc.load_model(model_uri) + model = mlflow.pyfunc.load_model(local_path) except Exception as e: await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) raise e @@ -684,6 +794,10 @@ class MLFlowRepository(SientiaMonitoring): 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. + - The wrapper branch downloads before it loads, so it puts the + artifact's `code` directory on `sys.path` itself. The two plain + branches delegate, and the code directory is handled inside + `load_predict_model` / `load_transform_model`. """ self.info( @@ -706,6 +820,8 @@ class MLFlowRepository(SientiaMonitoring): f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' ) + self.prepend_artifact_code_dir(artifact_path, 'pyfunc', metadata) + raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model @@ -1255,7 +1371,7 @@ class MLFlowRepository(SientiaMonitoring): self.debug('Model uploaded successfully', metadata) elif flavor == 'joblib': with tempfile.TemporaryDirectory() as tmp_dir: - pkl_path = path.join(tmp_dir, 'model.pkl') + pkl_path = path.join(tmp_dir, MODEL_PICKLE_NAME) joblib.dump(model, pkl_path, compress=3) mlflow.log_artifact(pkl_path, artifact_path=model_type) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 02e8cd5..99f6194 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -17,6 +17,7 @@ from laborious import metrics from laborious.utils.repository import model_repository from laborious.utils.repository.model_repository import ( INVALID_FLAVOR_MESSAGE, + MODEL_PICKLE_NAME, MLFlowRepository, force_memory_release, ) @@ -68,6 +69,11 @@ def mlflow_repository(): @pytest.fixture def mlflow(): with patch('laborious.utils.repository.model_repository.mlflow') as mlflow: + # Every URI load resolves the artifact to a local directory before calling the flavor + # loader, and `os.path` refuses a MagicMock. Echoing the URI back keeps the loader's + # argument readable and points the code-directory lookup at a path that does not exist, + # so no test touches `sys.path` unless it builds a real artifact directory. + mlflow.artifacts.download_artifacts.side_effect = lambda artifact_uri: str(artifact_uri) yield mlflow @@ -467,6 +473,11 @@ async def test_load_predict_model_error(mlflow_repository): mlflow_repository.observe_lag.assert_not_called() +# A real string, because the transform load now resolves its URI to a local directory before +# calling the loader, and that resolution goes through `os.path`. +TRANSFORM_MODEL_URI = 'runs:/run-123/data_model' + + def validate_common_load_transform_model_mocks(mlflow_repository, model_name): mlflow_repository.get_model_run_id.assert_called_once_with( model_name=model_name, stage='Production' @@ -479,14 +490,14 @@ def validate_common_load_transform_model_mocks(mlflow_repository, model_name): @pytest.mark.asyncio async def test_load_transform_model_sklearn(mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock() - mlflow_repository.get_model_uri = MagicMock() + mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') assert result == mlflow.sklearn.load_model.return_value - mlflow.sklearn.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) + mlflow.sklearn.load_model.assert_called_once_with(TRANSFORM_MODEL_URI) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) mlflow_repository.emit_metric.assert_called_once_with( @@ -497,14 +508,14 @@ async def test_load_transform_model_sklearn(mlflow, mlflow_repository): @pytest.mark.asyncio async def test_load_transform_model_pyfunc(mlflow, mlflow_repository): mlflow_repository.get_model_run_id = MagicMock() - mlflow_repository.get_model_uri = MagicMock() + mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) result = await mlflow_repository.load_transform_model('test_model', {}, 'pyfunc') validate_common_load_transform_model_mocks(mlflow_repository, 'test_model') assert result == mlflow.pyfunc.load_model.return_value - mlflow.pyfunc.load_model.assert_called_once_with(mlflow_repository.get_model_uri.return_value) + mlflow.pyfunc.load_model.assert_called_once_with(TRANSFORM_MODEL_URI) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY) mlflow_repository.emit_metric.assert_called_once_with( @@ -543,8 +554,8 @@ async def test_download_model_invalid_model_type(mlflow_repository): @pytest.mark.parametrize( 'model_type', [('predict', 'prediction_model'), ('transform', 'data_model')] ) -async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type): - mlflow_repository.dowload_artifacts = AsyncMock() +async def test_download_model_load_wrapper(mlflow, mlflow_repository, model_type, tmp_path): + mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) result = await mlflow_repository.download_model('test_model', {}, model_type[0], 'pyfunc', True) @@ -1740,14 +1751,14 @@ FLAVOR_LOADERS = ('sklearn', 'pyfunc') @pytest.fixture def preserved_sys_path(): - """Restore `sys.path`: the joblib flavor inserts entries and never removes them.""" + """Restore `sys.path`: every flavor inserts artifact code entries and never removes them.""" original = list(sys.path) yield sys.path[:] = original def build_artifact_dir(tmp_path, *filenames): - """Create a fake downloaded artifact directory with the given top-level files.""" + """Create a fake download of one model artifact directory, holding the given files.""" for filename in filenames: (tmp_path / filename).write_bytes(b'pickle-bytes') return str(tmp_path) @@ -1979,7 +1990,12 @@ async def test_no_load_emits_a_fallback_operation_type( @pytest.mark.parametrize( ('filenames', 'expected'), [ + # The canonical name `log_model` writes, alongside a real MLflow manifest. (('MLmodel', 'model.pkl'), 'model.pkl'), + # `model.pkl` wins over any other candidate, whatever the name order. + (('a_model.pkl', 'model.pkl'), 'model.pkl'), + (('model.pkl', 'z_model.pkl'), 'model.pkl'), + # Fallback for hand-built artifacts: no `model.pkl`, so name order decides. (('MLmodel', 'python_model.pkl'), 'python_model.pkl'), (('b_model.pkl', 'a_model.pkl'), 'a_model.pkl'), ], @@ -1993,6 +2009,7 @@ def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, exp def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path): + """A pyfunc wrapper's own `artifacts/*.pkl` is never taken for the model.""" nested = tmp_path / 'artifacts' nested.mkdir() (nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes') @@ -2003,6 +2020,15 @@ def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path): assert 'No .pkl file found' in str(exc_info.value) +def test_find_model_pickle_reads_the_name_log_model_writes(mlflow_repository, tmp_path): + """Read and write agree: both go through `MODEL_PICKLE_NAME`.""" + local_dir = build_artifact_dir(tmp_path, MODEL_PICKLE_NAME) + + result = mlflow_repository._find_model_pickle(local_dir, 'prediction_model') + + assert result == str(tmp_path / MODEL_PICKLE_NAME) + + 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') @@ -2099,6 +2125,125 @@ def test_load_model_pickle_keeps_code_dir_after_load( assert code_dir in sys.path +""" +Tests for the artifact `code/` directory on every flavor's load path +""" + + +def observe_sys_path(recorder): + """A loader stand-in that records `sys.path` as it was when the load was called.""" + + def loader(_path): + recorder.extend(sys.path) + return MagicMock() + + return loader + + +def test_prepend_artifact_code_dir_inserts_at_the_front( + mlflow_repository, tmp_path, preserved_sys_path +): + code_dir = build_code_dir(tmp_path) + + result = mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn') + + assert result == code_dir + assert sys.path[0] == code_dir + + +def test_prepend_artifact_code_dir_without_code_dir_leaves_sys_path_alone( + mlflow_repository, tmp_path, preserved_sys_path +): + before = list(sys.path) + + assert mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'pyfunc') is None + assert sys.path == before + + +def test_prepend_artifact_code_dir_does_not_duplicate( + mlflow_repository, tmp_path, preserved_sys_path +): + code_dir = build_code_dir(tmp_path) + + mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn') + + assert mlflow_repository.prepend_artifact_code_dir(str(tmp_path), 'sklearn') is None + assert sys.path.count(code_dir) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('flavor', FLAVOR_LOADERS) +async def test_load_predict_model_puts_code_dir_on_sys_path_before_the_loader( + flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path +): + """MLflow extends `sys.path` only when `MLmodel` declares `code`; the directory decides here.""" + code_dir = build_code_dir(tmp_path) + local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl') + mlflow.artifacts.download_artifacts.side_effect = None + mlflow.artifacts.download_artifacts.return_value = local_dir + sys_path_at_load = [] + getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load) + + await mlflow_repository.load_predict_model('test_model', {}, flavor) + + mlflow.artifacts.download_artifacts.assert_called_once_with( + artifact_uri='models:/test_model/production' + ) + getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir) + assert sys_path_at_load[0] == code_dir + + +@pytest.mark.asyncio +@pytest.mark.parametrize('flavor', FLAVOR_LOADERS) +async def test_load_transform_model_puts_code_dir_on_sys_path_before_the_loader( + flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path +): + code_dir = build_code_dir(tmp_path) + local_dir = build_artifact_dir(tmp_path, 'MLmodel', 'model.pkl') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) + mlflow.artifacts.download_artifacts.side_effect = None + mlflow.artifacts.download_artifacts.return_value = local_dir + sys_path_at_load = [] + getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load) + + await mlflow_repository.load_transform_model('test_model', {}, flavor) + + mlflow.artifacts.download_artifacts.assert_called_once_with(artifact_uri=TRANSFORM_MODEL_URI) + getattr(mlflow, flavor).load_model.assert_called_once_with(local_dir) + assert sys_path_at_load[0] == code_dir + + +@pytest.mark.asyncio +@pytest.mark.parametrize('flavor', FLAVOR_LOADERS) +async def test_load_predict_model_without_code_dir_leaves_sys_path_alone( + flavor, mlflow, mlflow_repository, tmp_path, preserved_sys_path +): + mlflow.artifacts.download_artifacts.side_effect = None + mlflow.artifacts.download_artifacts.return_value = build_artifact_dir(tmp_path, 'MLmodel') + before = list(sys.path) + + await mlflow_repository.load_predict_model('test_model', {}, flavor) + + assert sys.path == before + + +@pytest.mark.asyncio +async def test_download_model_load_wrapper_puts_code_dir_on_sys_path( + mlflow, mlflow_repository, tmp_path, preserved_sys_path +): + """The wrapper branch downloads itself, so it is the one that has to add the directory.""" + code_dir = build_code_dir(tmp_path) + mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) + sys_path_at_load = [] + mlflow.pyfunc.load_model.side_effect = observe_sys_path(sys_path_at_load) + + await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True) + + assert sys_path_at_load[0] == code_dir + assert sys.path[0] == code_dir + + @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_log_model_joblib_writes_compressed_pickle_and_code(