Files
sientia-dataops-laborious-l…/docs/model-flavors.md
vitor-aignosi c1a44a9d7a 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.
2026-08-28 15:10:46 -03:00

143 lines
7.9 KiB
Markdown

# Model flavors
The *flavor* is how this runtime is told to deserialize a model. It is **configuration, never
discovery**: the value comes from the model document in the `models` collection and the code never
reads the artifact's `MLmodel` manifest to guess it.
```json
{
"model_config": {
"transform_flavor": "sklearn",
"predict_flavor": "joblib",
"retention_minutes": 60,
"target": "SE"
}
}
```
`predict_flavor` drives `load_predict_model` (artifact `prediction_model`); `transform_flavor` drives
`load_transform_model` (artifact `data_model`). Both default to `sklearn` when the key is absent.
## Accepted values
| Flavor | Read | Write (retrain) | How it loads |
|---|---|---|---|
| `sklearn` | yes | yes | 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.
`pytorch` was accepted until this change and is now rejected like any other unknown value: `torch` is
not installed in the runtime, so `mlflow.pytorch.load_model` could never have loaded anything there.
**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.
## Same artifact, three readers
The flavor changes the reader, never the path. All three resolve the same two artifact directories of
the Production run:
```
<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 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.
## `code/` on `sys.path`, for every flavor
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.
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.
- **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 `models` is an operator step.
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` 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.