Track models

What is model tracking?

MLflow allows to serialize and deserialize models to a common format, track those models in MLflow Tracking and manage them using MLflow Model Registry. Many popular Machine / Deep Learning frameworks have built-in support through what MLflow calls flavors. Even if there is no flavor for your framework of choice, it is easy to create your own flavor and integrate it with MLflow.

How to track models using MLflow in Kedro project?

kedro-mlflow introduces two new DataSet types that can be used in DataCatalog called MlflowModelTrackingDataset and MlflowModelLocalFileSystemDataset. The two have very similar API, except that:

• the MlflowModelTrackingDataset is used to load from and save to from the mlflow artifact store. It can load from any given model_uri, including registered model, local models, models tied to a run…

• the MlflowModelLocalFileSystemDataset is used to load from and save to a given local path. It uses the standard filepath argument in the constructor of Kedro DataSets. Note that it does not log in mlflow.

Note: If you use MlflowModelTrackingDataset, it will be saved as a “LoggedModel” object and will be linked to your current run. However, you will need to specify the run id to predict with (since it is not persisted locally, it will not pick the latest model by default). You may prefer to combine MlflowModelLocalFileSystemDataset and MlflowArtifactDataset to make persist it both locally and remotely, see further.

Suppose you would like to register a scikit-learn model of your DataCatalog in mlflow, you can use the following yaml API:

my_sklearn_model:
    type: kedro_mlflow.io.models.MlflowModelTrackingDataset
    flavor: mlflow.sklearn

More informations on available parameters are available in the dedicated section.

You are now able to use my_sklearn_model in your nodes. Since this model is registered in mlflow, you can also leverage the mlflow model serving abilities or predicting on batch abilities, as well as the mlflow models registry to manage the lifecycle of this model.

Frequently asked questions

How is it working under the hood?

For MlflowModelTrackingDataset

During save, a model object from node output is logged to mlflow using log_model function of the specified flavor. It is logged as a LoggedModel mlflow object, and eventually (this will actually always happen if you are using the kedro run command, because the MlflowHook creates a new run before each pipeline run) linked to your currently active mlflow run. If the model_uri is specified, the saving operation will fail. Consequently it will never be possible to save in a specific mlflow run_id if you launch a pipeline with the kedro run command because the MlflowHook creates a new run before each pipeline run.

During load, the model is retrieved from the model_uri if specified, else the previously saved is reloaded.

For MlflowModelLocalFileSystemDataset

During save, a model object from node output is saved locally under specified filepath using save_model function of the specified flavor.

When model is loaded, the latest version stored locally is read using load_model function of the specified flavor. You can also load a model from a specific kedro run by specifying the version argument to the constructor.

How can I track a custom MLflow model flavor?

To track a custom MLflow model flavor you need to set the flavor parameter to import the module of your custom flavor and to specify a pyfunc workflow which can be set either to python_model or loader_module. The former is the more high level and user friendly and is recommend by mlflow while the latter offers more control. We haven’t tested the integration in kedro-mlflow of this second workflow extensively, and it should be used with caution.

my_custom_model:
    type: kedro_mlflow.io.models.MlflowModelTrackingDataset
    flavor: my_package.custom_mlflow_flavor
    pyfunc_workflow: python_model # or loader_module

How can I save model locally and log it in MLflow in one step?

How can I save model locally and log it in MLflow in one step?

If you want to save your model both locally and remotely within the same run, you can leverage MlflowArtifactDataset:

sklearn_model:
    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
    dataset:
        type: kedro_mlflow.io.models.MlflowModelLocalFileSystemDataset
        flavor: mlflow.sklearn
        filepath: data/06_models/sklearn_model

This might be useful if you want to always read the lastest model saved locally and log it to MLflow each time the new model is being trained for tracking purpose.