MLflow

Learn how to use the MLflow Experiment Tracker with ZenML.

The ZenML MLflow Experiment Tracker integration and stack component allows you to log and visualize information from your pipeline steps using MLflow, without having to write extra MLflow code.

Prerequisites

To use the MLflow Experiment Tracker, you'll need:

ZenML mlflow integration installed (zenml integration install mlflow -y)

An MLflow deployment, either local (scenario 1) or remote with proxied artifact storage (scenario 5)

Configuring the Experiment Tracker

There are two main MLflow deployment scenarios:

Local (scenario 1): Use a local artifact store, only suitable for running ZenML locally. No extra configuration needed.

zenml experiment-tracker register mlflow_experiment_tracker --flavor=mlflow

zenml stack register custom_stack -e mlflow_experiment_tracker ... --set

Remote with proxied artifact storage (scenario 5): Can be used with any stack components. Requires authentication configuration.

For remote, you'll need to configure authentication using one of:

Basic authentication (not recommended for production)

ZenML secrets (recommended)

To use ZenML secrets:

zenml secret create mlflow_secret \

--username=<USERNAME> \

--password=<PASSWORD>

zenml experiment-tracker register mlflow \

--flavor=mlflow \

--tracking_username={{mlflow_secret.username}} \

--tracking_password={{mlflow_secret.password}} \

...

Using the Experiment Tracker

To log information with MLflow in a pipeline step:

Enable the experiment tracker using the @step decorator

Use MLflow's logging or auto-logging capabilities as usual

import mlflow

@step(experiment_tracker="<MLFLOW_TRACKER_STACK_COMPONENT_NAME>")

def train_step(...):

mlflow.tensorflow.autolog()

mlflow.log_param(...)

mlflow.log_metric(...)

mlflow.log_artifact(...)

...

Viewing Results

You can find the URL to the MLflow experiment for a ZenML run:

last_run = client.get_pipeline("<PIPELINE_NAME>").last_run