Passing data into a Tune run through search spaces#

The first way of passing inputs into Trainables is thesearch space (it may also be calledparameter space orconfig). In the Trainable itself, it maps to theconfig dict passed in as an argument to the function. You define the search space using theparam_space argument of theTuner. The search space is a dict and may be composed ofdistributions, which will sample a different value for each Trial, or of constant values. The search space may be composed of nested dictionaries, and those in turn can have distributions as well.

For example, passing in a large pandas DataFrame or an unserializable model object as a value in the search space will lead to unwanted behavior. At best it will cause large slowdowns and disk space usage as Trial metadata saved to disk will also contain this data. At worst, an exception will be raised, as the data cannot be sent over to the Trial workers. For more details, seeHow can I avoid bottlenecks?.

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

In our example, we want to tune the two model hyperparameters. We also want to set the number of epochs, so that we can easily tweak it later. For the hyperparameters, we will use thetune.uniform distribution. We will also modify thetraining_function to obtain those values from theconfig dictionary.

deftraining_function(config):# For now, we have nothing here.data=Nonemodel={"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".forepochinrange(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,},)

Usingtune.with_parameters access data in Tune runs#

If we have large objects that are constant across Trials, we can use thetune.with_parameters utility to pass them into the Trainable directly. The objects will be stored in theRay object store so that each Trial worker may access them to obtain a local copy to use in its process.

Note that the serialization (once) and deserialization (for each Trial) of large objects may incur a performance overhead.

In our example, we will pass thedata DataFrame usingtune.with_parameters. In order to do that, we need to modify our function signature to includedata as an argument.

deftraining_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".forepochinrange(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 step is to wrap thetraining_function usingtune.with_parameters before passing it into theTuner. Every keyword argument of thetune.with_parameters call will be mapped to the keyword arguments 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),)

Loading data in a Tune Trainable#

You can also load data directly in Trainable from e.g. cloud storage, shared file storage such as NFS, or from the local disk of the Trainable worker.

A common use-case is to load the dataset from S3 or any other cloud storage with pandas, arrow or any other framework.

The working directory of the Trainable worker will be automatically changed to the corresponding Trial directory. For more details, seeHow do I access relative filepaths in my Tune training function?.

Our tuning run can now be run, though we will not yet obtain any meaningful outputs back.

results=tuner.fit()

Reporting metrics with Tune#

Metrics are values passed through themetrics argument in atune.report call. Metrics can be used by TuneSearch Algorithms andSchedulers to direct the search. After the tuning run is complete, you cananalyze the results, which include the reported metrics.

In our example, we want to maximize themetric. We will report it each epoch to Tune, and set themetric andmode arguments intune.TuneConfig to let Tune know that it should use it as the optimization objective.

fromrayimporttraindeftraining_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".forepochinrange(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"),)

Logging metrics with Tune callbacks#

Every metric logged usingtune.report can be accessed during the tuning run through TuneCallbacks. Ray Tune providesseveral built-in integrations with popular frameworks, such as MLFlow, Weights & Biases, CometML and more. You can also use theCallback API to create your own callbacks.

Callbacks are passed in thecallback argument of theTuner’sRunConfig.

In our example, we’ll use the MLFlow callback to track the progress of our tuning run and the changing value of themetric (requiresmlflow to be installed).

importray.tunefromray.tune.logger.mlflowimportMLflowLoggerCallbackdeftraining_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".forepochinrange(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")]),)

Getting data out of Tune using checkpoints & other artifacts#

Aside from metrics, you may want to save the state of your trained model and any other artifacts to allow resumption from training failure and further inspection and usage. Those cannot be saved as metrics, as they are often far too large and may not be easily serializable. Finally, they should be persisted on disk or cloud storage to allow access after the Tune run is interrupted or terminated.

Ray Train provides aCheckpoint API for that purpose.Checkpoint objects can be created from various sources (dictionaries, directories, cloud storage).

In Ray Tune,Checkpoints are created by the user in their Trainable functions and reported using the optionalcheckpoint argument oftune.report.Checkpoints can contain arbitrary data and can be freely passed around the Ray cluster. After a tuning run is over,Checkpoints can beobtained from the results.

Ray Tune can be configured toautomatically sync checkpoints to cloud storage, keep only a certain number of checkpoints to save space (withray.tune.CheckpointConfig) and more.

In our example, we want to be able to resume the training from the latest checkpoint, and to save thetrained_model in a checkpoint every iteration. To accomplish this, we will use thesession andCheckpoint APIs.

importosimportpickleimporttempfilefromrayimporttunedeftraining_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=0ifcheckpoint:withcheckpoint.as_directory()ascheckpoint_dir:withopen(os.path.join(checkpoint_dir,"model.pkl"),"w")asf:checkpoint_dict=pickle.load(f)start_epoch=checkpoint_dict["epoch"]+1model=checkpoint_dict["state"]# Simulate training & evaluation - we obtain back a "metric" and a "trained_model".forepochinrange(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.withtempfile.TemporaryDirectory()astemp_checkpoint_dir:withopen(os.path.join(temp_checkpoint_dir,"model.pkl"),"w")asf: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 all of those changes implemented, we can now run our tuning and obtain 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).

Checkpoints, metrics, and the log directory for each trial can be accessed through theResultGrid output of a Tune experiment. For more information on how to interact with the returnedResultGrid, seeAnalyzing Tune Experiment Results.

How do I access Tune results after I am finished?#

After you have finished running the Python session, you can still access the results and checkpoints. By default, Tune will save the experiment results to the~/ray_results local directory. You can configure Tune to persist results in the cloud as well. SeeHow to Configure Persistent Storage in Ray Tune for more information on how to configure storage options for persisting experiment results.

You can restore the Tune experiment by callingTuner.restore(path_or_cloud_uri,trainable), wherepath_or_cloud_uri points to a location either on the filesystem or cloud where the experiment was saved to. After theTuner has been restored, you can access the results and checkpoints by callingTuner.get_results() to receive theResultGrid object, and then proceeding as outlined in the previous section.