Files
sientia-dataops-laborious_t…/docs/model-flavors.md
vitor-aignosi 25626b1da3
Some checks failed
Quality gate / quality-gate (push) Has been cancelled
Add docs
2026-08-27 14:39:04 -03:00

6.6 KiB

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.

{
  "model_config": {
    "transform_flavor": "sklearn",
    "predict_flavor": "joblib",
    "retention_minutes": 60,
    "target": "SE"
  }
}

predict_flavor drives load_predict_model (artifact prediction_model); transform_flavor drives load_transform_model (artifact data_model). Both default to sklearn when the key is absent.

Accepted values

Flavor Read Write (retrain) How it loads
sklearn yes yes mlflow.sklearn.load_model(<uri>)
pyfunc yes yes mlflow.pyfunc.load_model(<uri>); the wrapper path also unwraps ._model_impl.python_model
pytorch yes yes mlflow.pytorch.load_model(<uri>)
joblib yes yes downloads the artifact and joblib.loads its pickle — no MLflow flavor loader involved

Anything else raises ValueError with INVALID_FLAVOR_MESSAGE, before any artifact is resolved or downloaded. The same message and the same four values govern reading and writing.

There is no fallback between flavors. A failing loader propagates: one attempt, one cause. A model configured with the wrong flavor fails visibly instead of being rescued silently.

The joblib flavor

Some models the pipeline produces do not load through an MLflow flavor loader even though the artifact is intact: the .pkl is in the compressed form joblib.dump writes, which MLflow's pickle.load does not read. Declaring joblib is how such a model is served.

Reading. Resolve the Production run, client.download_artifacts(run_id, <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.

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.

sys.path is process-global. Two different models that embed a package with the same name (utils/, typically) and are served by the same worker resolve to whichever entered first — the second silently runs the first one's code. This is pre-existing behaviour of MLflow's own pyfunc path with the same artifact layout; the mitigation, if it ever bites, is one worker per model.

Writing / retraining. joblib is a write flavor too, so a joblib model retrains. log_model writes joblib.dump(model, 'model.pkl', compress=3) under the artifact path and re-logs the source artifact's code/ to <artifact>/code, so the retrained model stays deserializable. Registration and promotion to Production are unchanged.

The written artifact has no MLmodel manifest. Only this runtime's joblib flavor reads it: mlflow.sklearn.load_model and models:/<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.

Switching a model to joblib

The runtime does not migrate documents. Editing the models collection is an operator step, and the order matters:

  1. Inventory first, before the deploy. Collect the models that fail in their configured flavor loader — run predict/transform and note which ones raise. There is no automatic discovery afterwards: once deployed, such a model simply fails, and the failure names the flavor loader, not the fix.
  2. Edit the document. Set model_config.predict_flavor: 'joblib' and/or model_config.transform_flavor: 'joblib' for each model in the inventory. Nothing else changes; models already on a working flavor need no edit and behave identically.
  3. Imported models arrive ready. The import_model workflow writes predict_flavor: 'joblib' itself — see model-import.md § 12.
  4. Check the first retrain. In MLflow, confirm the new run carries prediction_model/model.pkl and prediction_model/code/ when the source model had a code/ directory, and that the promoted version loads on the next predict.
  5. Rolling back is not just a revert. Reverting the code leaves any artifact a joblib retrain already wrote in MLflow: that Production version stops being loadable, so the rollback also means promoting the previous version back in the registry, and setting the documents' flavors back — manual, in reverse order.

Known gap: how the run id is resolved

get_model_run_id resolves the Production run with source.split('/')[2] on the registered version's source. This has not been confirmed against the real tracking server. No artifact scheme reproduced locally puts the run id in that position (a file store yields '', S3 yields the bucket, mlflow-artifacts:/ yields the experiment id), yet predict works in production — so the real source has a shape none of the local probes produced.

This is not specific to joblib: every flavor that resolves a run goes through the same line, and the joblib read path depends on it. Confirm it with a read-only query against the tracking server before relying on it, ideally before the next deploy.