QTZPOC-20

QTZPOC-20: Unify artifact code directory handling for all model flavors and standardize joblib pickle path

Implemented a robust mechanism to ensure an artifact's code/ directory is prepended to sys.path before loading for sklearn, pyfunc, and joblib flavors. sklearn and pyfunc now resolve MLflow URIs to local directories first to facilitate this sys.path modification. Standardized the joblib flavor to explicitly read and write model.pkl within the artifact directory. Updated documentation and added comprehensive tests to reflect these changes.

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, `<artifact>/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)**.
 

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(<uri>)` |
-| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; 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(<dir>)` |
+| `pyfunc` | yes | yes | resolves the URI to a local directory, then `mlflow.pyfunc.load_model(<dir>)`; the wrapper path also unwraps `._model_impl.python_model` |
+| `joblib` | yes | yes | downloads the same artifact and `joblib.load`s its `<artifact>/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, <artifact>)`, take the
-**first top-level `.pkl` in name order**, and `joblib.load` it. The listing does not recurse, so a
-nested wrapper artifact (`artifacts/*.pkl`) is never deserialized in place of the model. No
-`MLmodel` is parsed and no `models:/<name>/production` URI is built.
+```
+<run>/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 `<artifact>/code` to `sys.path` when that directory exists — the directory itself, matching
-MLflow's own convention where `code/utils/…` imports as the package `utils`. The entry is not
-duplicated and is **not** removed after the load, because the deserialized object may import
-lazily. A missing `code/` is not an error, and no extra download is performed: `download_artifacts`
-already brings `code/` along.
+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:/<name>/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 `<artifact>/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 `<artifact>/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'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:
+
+- `prepend_artifact_code_dir` inserts `<artifact>/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.
+
+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=<artifact>, code_path=…)`. Two read
+paths, and the second is **orthogonal to the flavor**:
+
+- Plain: `mlflow.pyfunc.load_model(<uri>)`, called as `predict(context, model_input)`.
+- Wrapper (`download_model(load_wrapper=True)`): downloads to `./tmp/artifacts/<model>/<artifact>`
+  via `dowload_artifacts`, loads as pyfunc, then unwraps `._model_impl.python_model`. Inside such an
+  artifact the real pickle is **nested** — `<artifact>/artifacts/stacking_model.pkl` for predict,
+  `<artifact>/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, <artifact>)`, then
+  `joblib.load(<artifact>/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, '<artifact>/model.pkl', compress=3)`, plus a re-log of the source
+  artifact's `code/` to `<artifact>/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:/<name>/production` both fail on it. A model retrained in
-joblib **stays** joblib — `predict_flavor` needs no edit after a retrain, and no consumer outside
-this runtime (the MLflow UI, MLflow serving, another service) recognizes that artifact as a model.
-
-**What `joblib` is for.** A pickled object exposing `predict(data)` — and `fit(data)` if it is
-retrained. It is **not** for a bare pyfunc `PythonModel` wrapper: `get_prediction_data` and
-`get_cached_operation` call `predict(data)` for every non-pyfunc flavor, so a wrapper expecting
-`predict(context, model_input)` raises `TypeError` at predict and retrain time.
-
-**Observability.** A joblib load is an ordinary load: `MODEL_READ_LAG` / `MODEL_READ_COUNT` on
-success and `MODEL_READ_ERROR_COUNT` on failure, under the base `operation_type`
-(`load_predict_model` / `load_transform_model`). A joblib write emits the usual `MODEL_WRITE_*`
-under `log_model`. No suffixed operation type, no new metric, no new label.
+  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.

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/<relpath>` 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/<relpath>`, 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.

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(

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 (`<run>/artifacts/prediction_model/`, `<run>/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 — `<run>/artifacts/prediction_model` or
+        `<run>/artifacts/data_model`, the very artifacts the other flavors
+        resolve — never the run's `artifacts/` root. So the pickle sought is
+        `<local_dir>/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 `<local_dir>/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 `<run>/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:/<name>/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 — `<local_dir>/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 `<run>/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
+        `<run>/artifacts/<artifact_name>/` 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)
 

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(