Get data in and out of Tune

You often need to pass data into Ray Tune trainables and get data out of them. Data going in includes datasets, models, and other large parameters. Data coming out includes metrics, checkpoints, and other artifacts. This guide explores different ways to do both and when to use each one.

Start by defining a simple function trainable. This guide extends the function step by step.

import random
import time
import pandas as pd


def training_function(config):
    # For now, we have nothing here.
    data = None
    model = {"hyperparameter_a": None, "hyperparameter_b": None}
    epochs = 0

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()
        trained_model = {"state": model, "epoch": epoch}

As inputs, training_function requires a pandas DataFrame, a model with some hyperparameters, and the number of epochs to train the model for. The model’s hyperparameters affect the returned metric, and each epoch, or iteration of training, changes the trained_model state.

Use the Tuner API to run hyperparameter optimization.

from ray.tune import Tuner
from ray import tune

tuner = Tuner(training_function, tune_config=tune.TuneConfig(num_samples=4))

Get data into Tune

First, provide the inputs for the trainable. Inputs fall broadly into two categories: variables and constants.

Variables are the parameters you want to tune. They’re different for every trial. For example, variables might be the learning rate and batch size for a neural network, the number of trees and the maximum depth for a random forest, or the data partition if you’re using Tune as an execution engine for batch training.

Constants are the parameters that are the same for every trial. Examples include the number of epochs, model hyperparameters you want to set but not tune, and the dataset. Constants are often large, such as the dataset or the model.

Warning

Objects from the outer scope of training_function are also automatically serialized and sent to trial actors, which might lead to unintended behavior. For example, a global lock doesn’t work because each actor operates on its own copy, and you might hit general serialization errors. As a best practice, don’t refer to any objects from the outer scope in training_function.

Pass data into a Tune run through search spaces

Note

In short, use the param_space argument to specify small, serializable constants and variables.

The first way to pass inputs into trainables is the search space, also called the parameter space or config. In the trainable, the search space maps to the config dictionary passed to the function as an argument. You define the search space with the param_space argument of the Tuner. The search space is a dictionary whose values can be constants or distributions, which sample a different value for each trial. The search space can also contain nested dictionaries, which can contain distributions as well.

Warning

Tune saves each value in the search space directly in the trial metadata, so every value in the search space must be serializable and take up a small amount of memory.

For example, passing a large pandas DataFrame or an unserializable model object as a value in the search space leads to unwanted behavior. At best, it causes large slowdowns and disk space usage, because the trial metadata saved to disk also contains this data. At worst, an exception is raised because the data can’t be sent to the trial workers. For more details, see How can I avoid bottlenecks?.

Instead, use strings or other identifiers as your values, and initialize or load the objects directly inside your trainable based on those values.

Note

You can use Ray Data datasets directly as values in the search space.

In this example, you tune the two model hyperparameters and set the number of epochs so you can adjust it later. Use the tune.uniform distribution for the hyperparameters, and modify training_function to read those values from the config dictionary.

def training_function(config):
    # For now, we have nothing here.
    data = None

    model = {
        "hyperparameter_a": config["hyperparameter_a"],
        "hyperparameter_b": config["hyperparameter_b"],
    }
    epochs = config["epochs"]

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()
        trained_model = {"state": model, "epoch": epoch}


tuner = Tuner(
    training_function,
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
)

Use tune.with_parameters to access data in Tune runs

Note

In short, use the tune.with_parameters utility function to specify large constant parameters.

If you have large objects that are constant across trials, you can use the tune.with_parameters utility to pass them directly into the trainable. The objects are stored in the Ray object store, so each trial worker can access them to get a local copy to use in its process.

Note

Objects put into the Ray object store must be serializable.

Large objects are serialized once and deserialized for each trial, which might add performance overhead.

In this example, pass the data DataFrame with tune.with_parameters. To do that, add data as an argument to the function signature.

def training_function(config, data):
    model = {
        "hyperparameter_a": config["hyperparameter_a"],
        "hyperparameter_b": config["hyperparameter_b"],
    }
    epochs = config["epochs"]

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()
        trained_model = {"state": model, "epoch": epoch}


tuner = Tuner(
    training_function,
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
)

Next, wrap training_function with tune.with_parameters before passing it to the Tuner. Each keyword argument of the tune.with_parameters call maps to a keyword argument in the trainable signature.

data = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})

tuner = Tuner(
    tune.with_parameters(training_function, data=data),
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
    tune_config=tune.TuneConfig(num_samples=4),
)

Load data in a Tune trainable

You can also load data directly in the trainable, for example from cloud storage, from shared file storage such as NFS, or from the local disk of the trainable worker.

Warning

When loading from disk, make sure that all nodes in your cluster have access to the file you’re loading.

A common use case is loading the dataset from S3 or another cloud storage service with pandas, Arrow, or another framework.

The working directory of the trainable worker is automatically set to the corresponding trial directory. For more details, see How do I access relative file paths in my Tune training function?.

You can now run tuning, though you won’t get any meaningful outputs back yet.

results = tuner.fit()

Get data out of Tune

You can now run tuning with the training_function trainable. Next, report metrics to Tune that it can use to guide the optimization. You also want to checkpoint your trained models so you can resume training after an interruption and use the models for prediction later.

Use the ray.tune.report API to get data out of the trainable workers. You can call it multiple times in the function trainable. Each call corresponds to one iteration of training, such as an epoch, a step, or a tree.

Report metrics with Tune

Metrics are values passed through the metrics argument in a tune.report call. Tune search algorithms and schedulers can use metrics to direct the search. After the tuning run completes, you can analyze the results, which include the reported metrics.

Note

As with search space values, Tune saves each value reported as a metric directly in the trial metadata, so every value reported as a metric must be serializable and take up a small amount of memory.

Note

Tune automatically includes some metrics, such as the training iteration and timestamp. See the full list of automatically included metrics.

In this example, you want to maximize metric. Report it to Tune each epoch, and set the metric and mode arguments in tune.TuneConfig to tell Tune to use it as the optimization objective.

from ray import tune


def training_function(config, data):
    model = {
        "hyperparameter_a": config["hyperparameter_a"],
        "hyperparameter_b": config["hyperparameter_b"],
    }
    epochs = config["epochs"]

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()
        trained_model = {"state": model, "epoch": epoch}
        tune.report(metrics={"metric": metric})


tuner = Tuner(
    tune.with_parameters(training_function, data=data),
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
    tune_config=tune.TuneConfig(num_samples=4, metric="metric", mode="max"),
)

Log metrics with Tune callbacks

During the tuning run, Tune callbacks can access every metric logged with tune.report. Ray Tune provides built-in integrations with popular frameworks, such as MLflow, Weights & Biases, and CometML. You can also create your own callbacks with the Callback API.

Pass callbacks in the callbacks argument of the RunConfig that you pass to the Tuner.

In this example, use the MLflow callback to track the progress of the tuning run and the changing value of metric. To use this callback, install mlflow.

import ray.tune
from ray.tune.logger.mlflow import MLflowLoggerCallback


def training_function(config, data):
    model = {
        "hyperparameter_a": config["hyperparameter_a"],
        "hyperparameter_b": config["hyperparameter_b"],
    }
    epochs = config["epochs"]

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()
        trained_model = {"state": model, "epoch": epoch}
        tune.report(metrics={"metric": metric})


tuner = tune.Tuner(
    tune.with_parameters(training_function, data=data),
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
    tune_config=tune.TuneConfig(num_samples=4, metric="metric", mode="max"),
    run_config=tune.RunConfig(
        callbacks=[MLflowLoggerCallback(experiment_name="example")]
    ),
)

Get data out of Tune with checkpoints and other artifacts

Besides metrics, you might want to save the state of your trained model and other artifacts, so you can resume after a training failure and inspect or use them later. You can’t save these as metrics, because they’re often far too large and might be hard to serialize. Persist them on disk or in cloud storage so you can access them after the Tune run is interrupted or terminated.

Ray Tune provides a Checkpoint API for that purpose. You can create Checkpoint objects from various sources, such as local directories and cloud storage.

In Ray Tune, you create Checkpoint objects in your function trainable and report them with the optional checkpoint argument of tune.report. A checkpoint can contain arbitrary data, and you can pass it freely around the Ray cluster. After a tuning run finishes, you can get the checkpoints from the results.

You can also configure how Ray Tune handles checkpoints. For example, you can configure Ray Tune to automatically sync checkpoints to cloud storage, or use ray.tune.CheckpointConfig to keep only a certain number of checkpoints to save space.

Note

Tune checkpoints the experiment state itself separately. For more details, see Appendix: Types of data stored by Tune.

In this example, you want to resume training from the latest checkpoint and save the trained_model in a checkpoint every iteration. To do this, use tune.get_checkpoint() and the Checkpoint API.

import os
import pickle
import tempfile

from ray import tune

def training_function(config, data):
    model = {
        "hyperparameter_a": config["hyperparameter_a"],
        "hyperparameter_b": config["hyperparameter_b"],
    }
    epochs = config["epochs"]

    # Load the checkpoint, if there is any.
    checkpoint = tune.get_checkpoint()
    start_epoch = 0
    if checkpoint:
        with checkpoint.as_directory() as checkpoint_dir:
            with open(os.path.join(checkpoint_dir, "model.pkl"), "rb") as f:
                checkpoint_dict = pickle.load(f)
        start_epoch = checkpoint_dict["epoch"] + 1
        model = checkpoint_dict["state"]

    # Simulate training & evaluation - we obtain back a "metric" and a "trained_model".
    for epoch in range(start_epoch, epochs):
        # Simulate doing something expensive.
        time.sleep(1)
        metric = (0.1 + model["hyperparameter_a"] * epoch / 100) ** (
            -1
        ) + model["hyperparameter_b"] * 0.1 * data["A"].sum()

        checkpoint_dict = {"state": model, "epoch": epoch}

        # Create the checkpoint.
        with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
            with open(os.path.join(temp_checkpoint_dir, "model.pkl"), "wb") as f:
                pickle.dump(checkpoint_dict, f)
            tune.report(
                {"metric": metric},
                checkpoint=tune.Checkpoint.from_directory(temp_checkpoint_dir),
            )


tuner = tune.Tuner(
    tune.with_parameters(training_function, data=data),
    param_space={
        "hyperparameter_a": tune.uniform(0, 20),
        "hyperparameter_b": tune.uniform(-100, 100),
        "epochs": 10,
    },
    tune_config=tune.TuneConfig(num_samples=4, metric="metric", mode="max"),
    run_config=tune.RunConfig(
        callbacks=[MLflowLoggerCallback(experiment_name="example")]
    ),
)

With these changes in place, you can run tuning and get meaningful metrics and artifacts.

results = tuner.fit()
results.get_dataframe()

2022-11-30 17:40:28,839 INFO tune.py:762 -- Total run time: 15.79 seconds (15.65 seconds for the tuning loop).

	metric	time_this_iter_s	should_checkpoint	done	timesteps_total	episodes_total	training_iteration	trial_id	experiment_id	date	...	hostname	node_ip	time_since_restore	timesteps_since_restore	iterations_since_restore	warmup_time	config/epochs	config/hyperparameter_a	config/hyperparameter_b	logdir
0	-58.399962	1.015951	True	False	NaN	NaN	10	0b239_00000	acf38c19d59c4cf2ad7955807657b6ea	2022-11-30_17-40-26	...	ip-172-31-43-110	172.31.43.110	10.282120	0	10	0.003541	10	18.065981	-98.298928	/home/ubuntu/ray_results/training_function_202...
1	-24.461518	1.030420	True	False	NaN	NaN	10	0b239_00001	5ca9e03d7cca46a7852cd501bc3f7b38	2022-11-30_17-40-28	...	ip-172-31-43-110	172.31.43.110	10.362581	0	10	0.004031	10	1.544918	-47.741455	/home/ubuntu/ray_results/training_function_202...
2	18.510299	1.034228	True	False	NaN	NaN	10	0b239_00002	aa38dd786c714486a8d69fa5b372df48	2022-11-30_17-40-28	...	ip-172-31-43-110	172.31.43.110	10.333781	0	10	0.005286	10	8.129285	28.846415	/home/ubuntu/ray_results/training_function_202...
3	-16.138780	1.020072	True	False	NaN	NaN	10	0b239_00003	5b401e15ab614332b631d552603a8d77	2022-11-30_17-40-28	...	ip-172-31-43-110	172.31.43.110	10.242707	0	10	0.003809	10	17.982020	-27.867871	/home/ubuntu/ray_results/training_function_202...

4 rows × 23 columns

You can access the checkpoints, metrics, and log directory for each trial through the ResultGrid output of a Tune experiment. For more information on working with the returned ResultGrid, see Analyzing Tune Experiment Results.

Access Tune results after tuning finishes

After your Python session ends, you can still access the results and checkpoints. By default, Tune saves the experiment results to the local ~/ray_results directory. You can also configure Tune to persist results in the cloud. For more information on configuring storage for experiment results, see Configure persistent storage in Ray Tune.

To restore the Tune experiment, call Tuner.restore(path_or_cloud_uri, trainable), where path_or_cloud_uri points to the filesystem or cloud location where the experiment was saved. After you restore the Tuner, call Tuner.get_results() to get the ResultGrid object, then access the results and checkpoints as described in the previous section.