MLflow

Logging and visualizing experiments with MLflow.

The MLflow Experiment Tracker is an Experiment Tracker flavor provided with the MLflow ZenML integration that uses the MLflow tracking service to log and visualize information from your pipeline steps (e.g. models, parameters, metrics).

When would you want to use it?

MLflow Tracking is a very popular tool that you would normally use in the iterative ML experimentation phase to track and visualize experiment results. That doesn't mean that it cannot be repurposed to track and visualize the results produced by your automated pipeline runs, as you make the transition toward a more production-oriented workflow.

You should use the MLflow Experiment Tracker:

• if you have already been using MLflow to track experiment results for your project and would like to continue doing so as you are incorporating MLOps workflows and best practices in your project through ZenML.
• if you are looking for a more visually interactive way of navigating the results produced from your ZenML pipeline runs (e.g. models, metrics, datasets)
• if you or your team already have a shared MLflow Tracking service deployed somewhere on-premise or in the cloud, and you would like to connect ZenML to it to share the artifacts and metrics logged by your pipelines

You should consider one of the other Experiment Tracker flavors if you have never worked with MLflow before and would rather use another experiment tracking tool that you are more familiar with.

How do you configure it?

The MLflow Experiment Tracker flavor is provided by the MLflow ZenML integration, you need to install it on your local machine to be able to register an MLflow Experiment Tracker and add it to your stack:

zenml integration install mlflow -y

The MLflow Experiment Tracker can be configured to accommodate the following MLflow deployment scenarios:

• Localhost (default) and Local Tracking with Local Database: This scenario requires that you use a local Artifact Store alongside the MLflow Experiment Tracker in your ZenML stack. The local Artifact Store comes with limitations regarding what other types of components you can use in the same stack. This scenario should only be used to run ZenML locally and is not suitable for collaborative and production settings. No parameters need to be supplied when configuring the MLflow Experiment Tracker, e.g:

# Register the MLflow experiment tracker
zenml experiment-tracker register mlflow_experiment_tracker --flavor=mlflow

# Register and set a stack with the new experiment tracker
zenml stack register custom_stack -e mlflow_experiment_tracker ... --set

If you use the default local MLflow configuration without setting a tracking_uri, ZenML stores new MLflow tracking metadata in a SQLite database at <LOCAL_ARTIFACT_STORE>/mlflow.db and stores MLflow artifacts in <LOCAL_ARTIFACT_STORE>/mlflow_artifacts.

Earlier ZenML versions used MLflow's file store at <LOCAL_ARTIFACT_STORE>/mlruns for this default local setup. Existing local runs in that directory are not migrated automatically and won't appear in the new default SQLite-backed MLflow UI. The files remain on disk and can still be inspected separately with MLflow's file-store backend.

• Remote Experiment Tracking with MLflow Tracking Server: This scenario assumes that you have already deployed an MLflow Tracking Server enabled with proxied artifact storage access. There is no restriction regarding what other types of components it can be combined with. This option requires authentication-related parameters to be configured for the MLflow Experiment Tracker.

Due to a critical severity vulnerability found in older versions of MLflow, we recommend using MLflow version 2.2.1 or higher. ZenML supports both MLflow 2.x and 3.x versions.

• Databricks scenario: This scenario assumes that you have a Databricks workspace, and you want to use the managed MLflow Tracking server it provides. This option requires authentication-related parameters to be configured for the MLflow Experiment Tracker.

Authentication Methods

Remote MLflow tracking servers require a tracking_uri and exactly one authentication method. Set tracking_uri to the URL of a remote MLflow Tracking Server, or to "databricks" when using the MLflow Tracking Server managed by Databricks. You can also set tracking_insecure_tls=True to skip SSL certificate verification when connecting to the tracking server.

The supported authentication methods are:

• Username/password: set both tracking_username and tracking_password. Use this for tracking servers that accept basic authentication.
• Token: set tracking_token. Use this for tracking servers that accept bearer token authentication.
• Databricks OAuth M2M: set tracking_uri="databricks", databricks_host, databricks_client_id, and databricks_client_secret. Use this for Databricks-managed MLflow with a service principal, as described in the section below.

For Databricks-managed MLflow, databricks_host must be the Databricks workspace URL. The databricks_client_id and databricks_client_secret fields are only used for Databricks OAuth M2M authentication.

Databricks OAuth M2M

Databricks OAuth machine-to-machine (M2M) authentication authorizes unattended workloads, such as pipelines and automation scripts, to access Databricks resources as a service principal. See the Databricks documentation on authorizing service principal access with OAuth for the full setup flow.

For ZenML, the Databricks OAuth M2M values map to the MLflow experiment tracker configuration as follows:

• Databricks workspace URL: databricks_host
• Service principal application ID / client ID: databricks_client_id
• Service principal OAuth secret: databricks_client_secret

To get these values, create or select a service principal in your Databricks workspace, generate an OAuth secret for it, and copy the displayed client ID and secret value. Make sure the service principal has access to the workspace and to the MLflow experiment path you want to use.

Basic Authentication

This option configures the credentials for the MLflow tracking service directly as stack component attributes.

This is not recommended for production settings as the credentials won't be stored securely and will be clearly visible in the stack configuration.

# Register the MLflow experiment tracker
zenml experiment-tracker register mlflow_experiment_tracker --flavor=mlflow \ 
    --tracking_uri=<URI> --tracking_token=<token>

# You can also register it like this:
# zenml experiment-tracker register mlflow_experiment_tracker --flavor=mlflow \ 
#    --tracking_uri=<URI> --tracking_username=<USERNAME> --tracking_password=<PASSWORD>

# Register and set a stack with the new experiment tracker
zenml stack register custom_stack -e mlflow_experiment_tracker ... --set

Databricks OAuth M2M

This option configures authentication to a Databricks-managed MLflow tracking server with a Databricks-managed service principal and OAuth M2M credentials.

# Create a secret with the Databricks service principal OAuth credentials
zenml secret create databricks_oauth_secret \
    --client_id=<CLIENT_ID> \
    --client_secret=<CLIENT_SECRET>

# Register the MLflow experiment tracker
zenml experiment-tracker register mlflow_databricks --flavor=mlflow \
    --tracking_uri=databricks \
    --databricks_host=<DATABRICKS_WORKSPACE_URL> \
    --databricks_client_id={{databricks_oauth_secret.client_id}} \
    --databricks_client_secret={{databricks_oauth_secret.client_secret}}

# Register and set a stack with the new experiment tracker
zenml stack register custom_stack -e mlflow_databricks ... --set

ZenML Secret (Recommended)

This method requires you to configure a ZenML secret to store the MLflow tracking service credentials securely.

You can create the secret using the zenml secret create command:

# Create a secret called `mlflow_secret` with key-value pairs for the
# username and password to authenticate with the MLflow tracking server
zenml secret create mlflow_secret \
    --username=<USERNAME> \
    --password=<PASSWORD>

Once the secret is created, you can use it to configure the MLflow Experiment Tracker:

# Reference the username and password in our experiment tracker component
zenml experiment-tracker register mlflow \
    --flavor=mlflow \
    --tracking_username={{mlflow_secret.username}} \
    --tracking_password={{mlflow_secret.password}} \
    ...

PowerShell Terminal Note

When using the zenml experiment-tracker register command in PowerShell, referencing secrets using the {{secret_name.key}} syntax without quotes can cause the following error:

zenml.exe : The command parameter was already specified.

This is a quirk of how PowerShell interprets braces in command-line arguments.

To resolve this, enclose the secret references in double quotes:

--tracking_username="{{mlflow_secret.username}}" --tracking_password="{{mlflow_secret.password}}"

Read more about ZenML Secrets in the ZenML documentation.

For more, up-to-date information on the MLflow Experiment Tracker implementation and its configuration, you can have a look at the SDK docs .

How do you use it?

To be able to log information from a ZenML pipeline step using the MLflow Experiment Tracker component in the active stack, you need to enable an experiment tracker using the @step decorator. Then use MLflow's logging or auto-logging capabilities as you would normally do, e.g.:

import mlflow
import numpy as np
import tensorflow as tf
from zenml import step


@step(experiment_tracker="<MLFLOW_TRACKER_STACK_COMPONENT_NAME>")
def tf_trainer(
    x_train: np.ndarray,
    y_train: np.ndarray,
) -> tf.keras.Model:
    """Train a neural net from scratch to recognize MNIST digits return our
    model or the learner"""

    # compile model

    mlflow.tensorflow.autolog()

    # train model

    # log additional information to MLflow explicitly if needed

    mlflow.log_param(...)
    mlflow.log_metric(...)
    mlflow.log_artifact(...)

    return model

Instead of hardcoding an experiment tracker name, you can also use the Client to dynamically use the experiment tracker of your active stack:

from zenml.client import Client

experiment_tracker = Client().active_stack.experiment_tracker

@step(experiment_tracker=experiment_tracker.name)
def tf_trainer(...):
    ...

MLflow UI

MLflow comes with its own UI that you can use to find further details about your tracked experiments.

You can find the URL of the MLflow experiment linked to a specific ZenML run via the metadata of the step in which the experiment tracker was used:

from zenml.client import Client

last_run = client.get_pipeline("<PIPELINE_NAME>").last_run
trainer_step = last_run.steps["<STEP_NAME>"]
tracking_url = trainer_step.run_metadata["experiment_tracker_url"].value
print(tracking_url)

This will be the URL of the corresponding experiment in your deployed MLflow instance, or the local MLflow tracking URI if you are using ZenML's default local MLflow configuration.

If you are using ZenML's default local MLflow configuration, the MLflow tracking metadata is stored in a SQLite database inside your active local artifact store, and MLflow artifacts are stored next to it. You can use the mlflow ui command to start MLflow at localhost:5000 where you can then explore the UI in your browser.

You can find the active artifact store path with zenml artifact-store describe.

mlflow ui \
  --backend-store-uri sqlite:///<PATH_TO_LOCAL_ARTIFACT_STORE>/mlflow.db \
  --default-artifact-root file:<PATH_TO_LOCAL_ARTIFACT_STORE>/mlflow_artifacts

Additional configuration

For additional configuration of the MLflow experiment tracker, you can pass MLFlowExperimentTrackerSettings to create nested runs or add additional tags to your MLflow runs:

import mlflow
from zenml.integrations.mlflow.flavors.mlflow_experiment_tracker_flavor import MLFlowExperimentTrackerSettings

mlflow_settings = MLFlowExperimentTrackerSettings(
    nested=True,
    tags={"key": "value"}
)


@step(
    experiment_tracker="<MLFLOW_TRACKER_STACK_COMPONENT_NAME>",
    settings={
        "experiment_tracker": mlflow_settings
    }
)
def step_one(
    data: np.ndarray,
) -> np.ndarray:
    ...

Check out the SDK docs for a full list of available attributes and this docs page for more information on how to specify settings.