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.
This commit is contained in:
@@ -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.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.
|
||||
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 `<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.
|
||||
|
||||
**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=<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 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.
|
||||
|
||||
Reference in New Issue
Block a user