Ray Tune

Integrate W&B with Ray Tune to track hyperparameter tuning trials, log metrics, and compare experiment results.

This page describes how to use W&B with Ray Tune so you can track hyperparameter tuning trials, log metrics, and compare experiment results across runs. W&B offers two lightweight integrations with Ray, and you can choose the one that best fits your training workflow:

• The WandbLoggerCallback function automatically logs metrics reported to Tune to the W&B API.
• The setup_wandb() function, which you can use with the function API, automatically initializes the W&B API with Tune’s training information. You can use the W&B API as usual, such as by calling run.log() to log your training process.

Configure the integration

This section describes how to configure the WandbLoggerCallback, which is the most direct way to send Tune trial metrics to W&B.

from ray.air.integrations.wandb import WandbLoggerCallback

To configure W&B, pass a wandb key to the config parameter of tune.run(). See the example for usage.

The integration passes the content of the wandb config entry to wandb.init() as keyword arguments. The exceptions are the settings that configure the WandbLoggerCallback itself.

Parameters

The WandbLoggerCallback accepts the following parameters:

• project (str): Name of the W&B project. Required.
• api_key_file (str): Path to file containing the W&B API key.
• api_key (str): W&B API key. Alternative to setting api_key_file.
• excludes (list): List of metrics to exclude from the log.
• log_config (bool): Whether to log the config parameter of the results dictionary. Defaults to False.
• upload_checkpoints (bool): If True, uploads model checkpoints as artifacts. Defaults to False.

Example

from ray import tune, train
from ray.air.integrations.wandb import WandbLoggerCallback


def train_fc(config):
    for i in range(10):
        train.report({"mean_accuracy": (i + config["alpha"]) / 10})


tuner = tune.Tuner(
    train_fc,
    param_space={
        "alpha": tune.grid_search([0.1, 0.2, 0.3]),
        "beta": tune.uniform(0.5, 1.0),
    },
    run_config=train.RunConfig(
        callbacks=[
            WandbLoggerCallback(
                project="[YOUR-PROJECT]", api_key="[YOUR-API-KEY]", log_config=True
            )
        ]
    ),
)

results = tuner.fit()

setup_wandb

Use setup_wandb() when you want direct control over W&B logging from inside your training function, for example, to call run.log() with custom metrics alongside Tune’s reporting.

from ray.air.integrations.wandb import setup_wandb

This utility function helps initialize W&B for use with Ray Tune. For basic usage, call setup_wandb() in your training function:

from ray.air.integrations.wandb import setup_wandb


def train_fn(config):
    # Initialize wandb
    wandb = setup_wandb(config)
    run = wandb.init(
        project=config["wandb"]["project"],
        api_key_file=config["wandb"]["api_key_file"],
    )

    for i in range(10):
        loss = config["a"] + config["b"]
        run.log({"loss": loss})
        tune.report(loss=loss)
    run.finish()


tuner = tune.Tuner(
    train_fn,
    param_space={
        # 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"},
    },
)
results = tuner.fit()

Example code

For end-to-end references, see the following examples that show how the integration works:

• Try the integration in Colab: A demo to try the integration.
• View the example dashboard: View the dashboard generated from the example.