Notebook

Using MLflow with Tune¶

(tune-mlflow-ref)=

:::{warning} If you are using these MLflow integrations with {ref}ray-client, it is recommended that you setup a remote Mlflow tracking server instead of one that is backed by the local filesystem. :::

MLflow is an open source platform to manage the ML lifecycle, including experimentation, reproducibility, deployment, and a central model registry. It currently offers four components, including MLflow Tracking to record and query experiments, including code, data, config, and results.

{image}
:align: center
:alt: MLflow
:height: 80px
:target: https://www.mlflow.org/

Ray Tune currently offers two lightweight integrations for MLflow Tracking. One is the {ref}MLflowLoggerCallback <tune-mlflow-logger>, which automatically logs metrics reported to Tune to the MLflow Tracking API.

The other one is the {ref}@mlflow_mixin <tune-mlflow-mixin> decorator, which can be used with the function API. It automatically initializes the MLflow API with Tune's training information and creates a run for each Tune trial. Then within your training function, you can just use the MLflow like you would normally do, e.g. using mlflow.log_metrics() or even mlflow.autolog() to log to your training process.

{contents}
:backlinks: none
:local: true

Running an MLflow Example¶

In the following example we're going to use both of the above methods, namely the MLflowLoggerCallback and the mlflow_mixin decorator to log metrics. Let's start with a few crucial imports:

In [ ]:

import os
import tempfile
import time

import mlflow

from ray import tune
from ray.tune.integration.mlflow import MLflowLoggerCallback, mlflow_mixin

Next, let's define an easy objective function (a Tune Trainable) that iteratively computes steps and evaluates intermediate scores that we report to Tune.

In [ ]:

def evaluation_fn(step, width, height):
    return (0.1 + width * step / 100) ** (-1) + height * 0.1


def easy_objective(config):
    width, height = config["width"], config["height"]

    for step in range(config.get("steps", 100)):
        # Iterative training function - can be any arbitrary training procedure
        intermediate_score = evaluation_fn(step, width, height)
        # Feed the score back to Tune.
        tune.report(iterations=step, mean_loss=intermediate_score)
        time.sleep(0.1)

Given an MLFlow tracking URI, you can now simply use the MLflowLoggerCallback as a callback argument to your tune.run() call:

In [ ]:

def tune_function(mlflow_tracking_uri, finish_fast=False):
    tune.run(
        easy_objective,
        name="mlflow",
        num_samples=5,
        callbacks=[
            MLflowLoggerCallback(
                tracking_uri=mlflow_tracking_uri,
                experiment_name="example",
                save_artifact=True,
            )
        ],
        config={
            "width": tune.randint(10, 100),
            "height": tune.randint(0, 100),
            "steps": 5 if finish_fast else 100,
        },
    )

To use the mlflow_mixin decorator, you can simply decorate the objective function from earlier. Note that we also use mlflow.log_metrics(...) to log metrics to MLflow. Otherwise, the decorated version of our objective is identical to its original.

In [ ]:

@mlflow_mixin
def decorated_easy_objective(config):
    # Hyperparameters
    width, height = config["width"], config["height"]

    for step in range(config.get("steps", 100)):
        # Iterative training function - can be any arbitrary training procedure
        intermediate_score = evaluation_fn(step, width, height)
        # Log the metrics to mlflow
        mlflow.log_metrics(dict(mean_loss=intermediate_score), step=step)
        # Feed the score back to Tune.
        tune.report(iterations=step, mean_loss=intermediate_score)
        time.sleep(0.1)

With this new objective function ready, you can now create a Tune run with it as follows:

In [ ]:

def tune_decorated(mlflow_tracking_uri, finish_fast=False):
    # Set the experiment, or create a new one if does not exist yet.
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment(experiment_name="mixin_example")
    tune.run(
        decorated_easy_objective,
        name="mlflow",
        num_samples=5,
        config={
            "width": tune.randint(10, 100),
            "height": tune.randint(0, 100),
            "steps": 5 if finish_fast else 100,
            "mlflow": {
                "experiment_name": "mixin_example",
                "tracking_uri": mlflow.get_tracking_uri(),
            },
        },
    )

If you hapen to have an MLFlow tracking URI, you can set it below in the mlflow_tracking_uri variable and set smoke_test=False. Otherwise, you can just run a quick test of the tune_function and tune_decorated functions without using MLflow.

In [ ]:

smoke_test = True

if smoke_test:
    mlflow_tracking_uri = os.path.join(tempfile.gettempdir(), "mlruns")
else:
    mlflow_tracking_uri = "<MLFLOW_TRACKING_URI>"

tune_function(mlflow_tracking_uri, finish_fast=smoke_test)
if not smoke_test:
    df = mlflow.search_runs(
        [mlflow.get_experiment_by_name("example").experiment_id]
    )
    print(df)

tune_decorated(mlflow_tracking_uri, finish_fast=smoke_test)
if not smoke_test:
    df = mlflow.search_runs(
        [mlflow.get_experiment_by_name("mixin_example").experiment_id]
    )
    print(df)

This completes our Tune and MLflow walk-through. In the following sections you can find more details on the API of the Tune-MLflow integration.

MLflow AutoLogging¶

You can also check out {doc}here </tune/examples/includes/mlflow_ptl_example> for an example on how you can leverage MLflow auto-logging, in this case with Pytorch Lightning

MLflow Logger API¶

(tune-mlflow-logger)=

{eval-rst}
.. autoclass:: ray.tune.integration.mlflow.MLflowLoggerCallback
   :noindex:

MLflow Mixin API¶

(tune-mlflow-mixin)=

{eval-rst}
.. autofunction:: ray.tune.integration.mlflow.mlflow_mixin
   :noindex:

More MLflow Examples¶

{doc}/tune/examples/includes/mlflow_ptl_example: Example for using MLflow and Pytorch Lightning with Ray Tune.