Contents

External library integrations (tune.integration)¶

Keras (tune.integration.keras)¶

ray.tune.integration.keras.TuneReportCallback¶

alias of <MagicMock spec=’str’ id=’140359995641424’>

ray.tune.integration.keras.TuneReportCheckpointCallback¶

alias of <MagicMock spec=’str’ id=’140359995622160’>

MLflow (tune.integration.mlflow)¶

See also here.

class ray.air.callbacks.mlflow.MLflowLoggerCallback(tracking_uri: Optional[str] = None, registry_uri: Optional[str] = None, experiment_name: Optional[str] = None, tags: Optional[Dict] = None, save_artifact: bool = False)[source]¶

MLflow Logger to automatically log Tune results and config to MLflow.

MLflow (https://mlflow.org) Tracking is an open source library for recording and querying experiments. This Ray Tune LoggerCallback sends information (config parameters, training results & metrics, and artifacts) to MLflow for automatic experiment tracking.

Parameters

  • tracking_uri – The tracking URI for where to manage experiments and runs. This can either be a local file path or a remote server. This arg gets passed directly to mlflow initialization. When using Tune in a multi-node setting, make sure to set this to a remote server and not a local file path.

  • registry_uri – The registry URI that gets passed directly to mlflow initialization.

  • experiment_name – The experiment name to use for this Tune run. If the experiment with the name already exists with MLflow, it will be reused. If not, a new experiment will be created with that name.

  • tags – An optional dictionary of string keys and values to set as tags on the run

  • save_artifact – If set to True, automatically save the entire contents of the Tune local_dir as an artifact to the corresponding run in MlFlow.

Example:

from ray.air.callbacks.mlflow import MLflowLoggerCallback

tags = { "user_name" : "John",
         "git_commit_hash" : "abc123"}

tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
    },
    callbacks=[MLflowLoggerCallback(
        experiment_name="experiment1",
        tags=tags,
        save_artifact=True)])
ray.tune.integration.mlflow.mlflow_mixin(func: Callable)[source]¶

MLflow (https://mlflow.org) Tracking is an open source library for recording and querying experiments. This Ray Tune Trainable mixin helps initialize the MLflow API for use with the Trainable class or the @mlflow_mixin function API. This mixin automatically configures MLflow and creates a run in the same process as each Tune trial. You can then use the mlflow API inside the your training function and it will automatically get reported to the correct run.

For basic usage, just prepend your training function with the @mlflow_mixin decorator:

from ray.tune.integration.mlflow import mlflow_mixin

@mlflow_mixin
def train_fn(config):
    ...
    mlflow.log_metric(...)

You can also use MlFlow’s autologging feature if using a training framework like Pytorch Lightning, XGBoost, etc. More information can be found here (https://mlflow.org/docs/latest/tracking.html#automatic-logging).

from ray.tune.integration.mlflow import mlflow_mixin

@mlflow_mixin
def train_fn(config):
    mlflow.autolog()
    xgboost_results = xgb.train(config, ...)

The MlFlow configuration is done by passing a mlflow key to the config parameter of tune.run() (see example below).

The content of the mlflow config entry is used to configure MlFlow. Here are the keys you can pass in to this config entry:

Parameters

  • tracking_uri – The tracking URI for MLflow tracking. If using Tune in a multi-node setting, make sure to use a remote server for tracking.

  • experiment_id – The id of an already created MLflow experiment. All logs from all trials in tune.run will be reported to this experiment. If this is not provided or the experiment with this id does not exist, you must provide an``experiment_name``. This parameter takes precedence over experiment_name.

  • experiment_name – The name of an already existing MLflow experiment. All logs from all trials in tune.run will be reported to this experiment. If this is not provided, you must provide a valid experiment_id.

  • token – A token to use for HTTP authentication when logging to a remote tracking server. This is useful when you want to log to a Databricks server, for example. This value will be used to set the MLFLOW_TRACKING_TOKEN environment variable on all the remote training processes.

Example:

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

import mlflow

# Create the MlFlow expriment.
mlflow.create_experiment("my_experiment")

@mlflow_mixin
def train_fn(config):
    for i in range(10):
        loss = config["a"] + config["b"]
        mlflow.log_metric(key="loss", value=loss)
    tune.report(loss=loss, done=True)

tune.run(
    train_fn,
    config={
        # define search space here
        "a": tune.choice([1, 2, 3]),
        "b": tune.choice([4, 5, 6]),
        # mlflow configuration
        "mlflow": {
            "experiment_name": "my_experiment",
            "tracking_uri": mlflow.get_tracking_uri()
        }
    })

MXNet (tune.integration.mxnet)¶

class ray.tune.integration.mxnet.TuneReportCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None)[source]¶

MXNet to Ray Tune reporting callback

Reports metrics to Ray Tune.

This has to be passed to MXNet as the eval_end_callback.

Parameters

metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to MXNet, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to MXNet.

Example:

from ray.tune.integration.mxnet import TuneReportCallback

# mlp_model is a MXNet model
mlp_model.fit(
    train_iter,
    # ...
    eval_metric="acc",
    eval_end_callback=TuneReportCallback({
        "mean_accuracy": "accuracy"
    }))
class ray.tune.integration.mxnet.TuneCheckpointCallback(filename: str = 'checkpoint', frequency: int = 1)[source]¶

MXNet checkpoint callback

Saves checkpoints after each epoch.

This has to be passed to the epoch_end_callback of the MXNet model.

Checkpoint are currently not registered if no tune.report() call is made afterwards. You have to use this in conjunction with the TuneReportCallback to work!

Parameters

  • filename – Filename of the checkpoint within the checkpoint directory. Defaults to “checkpoint”.

  • frequency – Integer indicating how often checkpoints should be saved.

Example:

from ray.tune.integration.mxnet import TuneReportCallback,             TuneCheckpointCallback

# mlp_model is a MXNet model
mlp_model.fit(
    train_iter,
    # ...
    eval_metric="acc",
    eval_end_callback=TuneReportCallback({
        "mean_accuracy": "accuracy"
    }),
    epoch_end_callback=TuneCheckpointCallback(
        filename="mxnet_cp",
        frequency=3
    ))

PyTorch Lightning (tune.integration.pytorch_lightning)¶

class ray.tune.integration.pytorch_lightning.TuneReportCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, on: Union[str, List[str]] = 'validation_end')[source]¶

PyTorch Lightning to Ray Tune reporting callback

Reports metrics to Ray Tune.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to PyTorch Lightning, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to PyTorch Lightning.

  • on – When to trigger checkpoint creations. Must be one of the PyTorch Lightning event hooks (less the on_), e.g. “batch_start”, or “train_end”. Defaults to “validation_end”.

Example:

import pytorch_lightning as pl
from ray.tune.integration.pytorch_lightning import TuneReportCallback

# Report loss and accuracy to Tune after each validation epoch:
trainer = pl.Trainer(callbacks=[TuneReportCallback(
        ["val_loss", "val_acc"], on="validation_end")])

# Same as above, but report as `loss` and `mean_accuracy`:
trainer = pl.Trainer(callbacks=[TuneReportCallback(
        {"loss": "val_loss", "mean_accuracy": "val_acc"},
        on="validation_end")])
class ray.tune.integration.pytorch_lightning.TuneReportCheckpointCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, filename: str = 'checkpoint', on: Union[str, List[str]] = 'validation_end')[source]¶

PyTorch Lightning report and checkpoint callback

Saves checkpoints after each validation step. Also reports metrics to Tune, which is needed for checkpoint registration.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to PyTorch Lightning, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to PyTorch Lightning.

  • filename – Filename of the checkpoint within the checkpoint directory. Defaults to “checkpoint”.

  • on – When to trigger checkpoint creations. Must be one of the PyTorch Lightning event hooks (less the on_), e.g. “batch_start”, or “train_end”. Defaults to “validation_end”.

Example:

import pytorch_lightning as pl
from ray.tune.integration.pytorch_lightning import (
    TuneReportCheckpointCallback)

# Save checkpoint after each training batch and after each
# validation epoch.
trainer = pl.Trainer(callbacks=[TuneReportCheckpointCallback(
    metrics={"loss": "val_loss", "mean_accuracy": "val_acc"},
    filename="trainer.ckpt", on="validation_end")])

Weights and Biases (tune.integration.wandb)¶

See also here.

class ray.air.callbacks.wandb.WandbLoggerCallback(project: str, group: Optional[str] = None, api_key_file: Optional[str] = None, api_key: Optional[str] = None, excludes: Optional[List[str]] = None, log_config: bool = False, save_checkpoints: bool = False, **kwargs)[source]¶

Weights and biases (https://www.wandb.ai/) is a tool for experiment tracking, model optimization, and dataset versioning. This Ray Tune LoggerCallback sends metrics to Wandb for automatic tracking and visualization.

Parameters

  • project – Name of the Wandb project. Mandatory.

  • group – Name of the Wandb group. Defaults to the trainable name.

  • api_key_file – Path to file containing the Wandb API KEY. This file only needs to be present on the node running the Tune script if using the WandbLogger.

  • api_key – Wandb API Key. Alternative to setting api_key_file.

  • excludes – List of metrics that should be excluded from the log.

  • log_config – Boolean indicating if the config parameter of the results dict should be logged. This makes sense if parameters will change during training, e.g. with PopulationBasedTraining. Defaults to False.

  • save_checkpoints – If True, model checkpoints will be saved to Wandb as artifacts. Defaults to False.

  • **kwargs – The keyword arguments will be pased to wandb.init().

Wandb’s group, run_id and run_name are automatically selected by Tune, but can be overwritten by filling out the respective configuration values.

Please see here for all other valid configuration settings: https://docs.wandb.ai/library/init

Example:

from ray.tune.logger import DEFAULT_LOGGERS
from ray.air.callbacks.wandb import WandbLoggerCallback
tune.run(
    train_fn,
    config={
        # define search space here
        "parameter_1": tune.choice([1, 2, 3]),
        "parameter_2": tune.choice([4, 5, 6]),
    },
    callbacks=[WandbLoggerCallback(
        project="Optimization_Project",
        api_key_file="/path/to/file",
        log_config=True)])
ray.tune.integration.wandb.wandb_mixin(func: Callable)[source]¶

Weights and biases (https://www.wandb.ai/) is a tool for experiment tracking, model optimization, and dataset versioning. This Ray Tune Trainable mixin helps initializing the Wandb API for use with the Trainable class or with @wandb_mixin for the function API.

For basic usage, just prepend your training function with the @wandb_mixin decorator:

from ray.tune.integration.wandb import wandb_mixin

@wandb_mixin
def train_fn(config):
    wandb.log()

Wandb configuration is done by passing a wandb key to the config parameter of tune.run() (see example below).

The content of the wandb config entry is passed to wandb.init() as keyword arguments. The exception are the following settings, which are used to configure the WandbTrainableMixin itself:

Parameters

  • api_key_file – Path to file containing the Wandb API KEY. This file must be on all nodes if using the wandb_mixin.

  • api_key – Wandb API Key. Alternative to setting api_key_file.

Wandb’s group, run_id and run_name are automatically selected by Tune, but can be overwritten by filling out the respective configuration values.

Please see here for all other valid configuration settings: https://docs.wandb.ai/library/init

Example:

from ray import tune
from ray.tune.integration.wandb import wandb_mixin

@wandb_mixin
def train_fn(config):
    for i in range(10):
        loss = self.config["a"] + self.config["b"]
        wandb.log({"loss": loss})
    tune.report(loss=loss, done=True)

tune.run(
    train_fn,
    config={
        # define search space here
        "a": tune.choice([1, 2, 3]),
        "b": tune.choice([4, 5, 6]),
        # wandb configuration
        "wandb": {
            "project": "Optimization_Project",
            "api_key_file": "/path/to/file"
        }
    })

XGBoost (tune.integration.xgboost)¶

class ray.tune.integration.xgboost.TuneReportCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, results_postprocessing_fn: Optional[Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]] = None)[source]¶

XGBoost to Ray Tune reporting callback

Reports metrics to Ray Tune.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to XGBoost, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to XGBoost. If this is None, all metrics will be reported to Tune under their default names as obtained from XGBoost.

  • results_postprocessing_fn – An optional Callable that takes in the dict that will be reported to Tune (after it has been flattened) and returns a modified dict that will be reported instead. Can be used to eg. average results across CV fold when using xgboost.cv.

Example:

import xgboost
from ray.tune.integration.xgboost import TuneReportCallback

config = {
    # ...
    "eval_metric": ["auc", "logloss"]
}

# Report only log loss to Tune after each validation epoch:
bst = xgb.train(
    config,
    train_set,
    evals=[(test_set, "eval")],
    verbose_eval=False,
    callbacks=[TuneReportCallback({"loss": "eval-logloss"})])
class ray.tune.integration.xgboost.TuneReportCheckpointCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, filename: str = 'checkpoint', frequency: int = 5, results_postprocessing_fn: Optional[Callable[[Dict[str, Union[float, List[float]]]], float]] = None)[source]¶

XGBoost report and checkpoint callback

Saves checkpoints after each validation step. Also reports metrics to Tune, which is needed for checkpoint registration.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to XGBoost, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to XGBoost.

  • filename – Filename of the checkpoint within the checkpoint directory. Defaults to “checkpoint”. If this is None, all metrics will be reported to Tune under their default names as obtained from XGBoost.

  • frequency – How often to save checkpoints. Per default, a checkpoint is saved every five iterations.

  • results_postprocessing_fn – An optional Callable that takes in the dict that will be reported to Tune (after it has been flattened) and returns a modified dict that will be reported instead. Can be used to eg. average results across CV fold when using xgboost.cv.

Example:

import xgboost
from ray.tune.integration.xgboost import TuneReportCheckpointCallback

config = {
    # ...
    "eval_metric": ["auc", "logloss"]
}

# Report only log loss to Tune after each validation epoch.
# Save model as `xgboost.mdl`.
bst = xgb.train(
    config,
    train_set,
    evals=[(test_set, "eval")],
    verbose_eval=False,
    callbacks=[TuneReportCheckpointCallback(
        {"loss": "eval-logloss"}, "xgboost.mdl)])

LightGBM (tune.integration.lightgbm)¶

class ray.tune.integration.lightgbm.TuneReportCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, results_postprocessing_fn: Optional[Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]] = None)[source]¶

Create a callback that reports metrics to Ray Tune.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to LightGBM, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to LightGBM. If this is None, all metrics will be reported to Tune under their default names as obtained from LightGBM.

  • results_postprocessing_fn – An optional Callable that takes in the dict that will be reported to Tune (after it has been flattened) and returns a modified dict that will be reported instead.

Example:

import lightgbm
from ray.tune.integration.lightgbm import TuneReportCallback

config = {
    # ...
    "metric": ["binary_logloss", "binary_error"],
}

# Report only log loss to Tune after each validation epoch:
bst = lightgbm.train(
    config,
    train_set,
    valid_sets=[test_set],
    valid_names=["eval"],
    verbose_eval=False,
    callbacks=[TuneReportCallback({"loss": "eval-binary_logloss"})])
class ray.tune.integration.lightgbm.TuneReportCheckpointCallback(metrics: Optional[Union[str, List[str], Dict[str, str]]] = None, filename: str = 'checkpoint', frequency: int = 5, results_postprocessing_fn: Optional[Callable[[Dict[str, Union[float, List[float]]]], Dict[str, float]]] = None)[source]¶

Creates a callback that reports metrics and checkpoints model.

Saves checkpoints after each validation step. Also reports metrics to Tune, which is needed for checkpoint registration.

Parameters

  • metrics – Metrics to report to Tune. If this is a list, each item describes the metric key reported to LightGBM, and it will reported under the same name to Tune. If this is a dict, each key will be the name reported to Tune and the respective value will be the metric key reported to LightGBM.

  • filename – Filename of the checkpoint within the checkpoint directory. Defaults to “checkpoint”. If this is None, all metrics will be reported to Tune under their default names as obtained from LightGBM.

  • frequency – How often to save checkpoints. Per default, a checkpoint is saved every five iterations.

  • results_postprocessing_fn – An optional Callable that takes in the dict that will be reported to Tune (after it has been flattened) and returns a modified dict that will be reported instead.

Example:

import lightgbm
from ray.tune.integration.lightgbm import (
    TuneReportCheckpointCallback
)

config = {
    # ...
    "metric": ["binary_logloss", "binary_error"],
}

# Report only log loss to Tune after each validation epoch.
# Save model as `lightgbm.mdl`.
bst = lightgbm.train(
    config,
    train_set,
    valid_sets=[test_set],
    valid_names=["eval"],
    verbose_eval=False,
    callbacks=[TuneReportCheckpointCallback(
        {"loss": "eval-binary_logloss"}, "lightgbm.mdl)])

previous

Scikit-Learn API (tune.sklearn)

next

Tune Internals