Skip to main content

Weights and Biases

Weights and Biases is a suite of tools for experiment tracking for Machine Learning practitioners and the integration with Anyscale is at the code-level. Although wandb code can be called directly from your Anyscale cluster, the W&B integration with Ray is treated as a first class integration in Anyscale with additional features.

Ray currently provides two lightweight integrations with W&B:

  • WandbLoggerCallback is used to automatically log metrics reported to Ray Tune or Train to W&B
  • setup_wandb initializes a W&B session in a Ray Tune or Train method so wandb.log() can be used as normal to log metrics

To simplify the experience of using the Ray integration with W&B in Anyscale, the following features are added to the Anyscale integration with W&B. These will all automatically apply if either WandbLoggerCallback or setup_wandb is used in Ray application code running in a workspace, Anyscale production job, or Ray job running on Anyscale.

  • Default values for the W&B project and group populated based on the type of Anyscale execution (Ray job, production job, workspace) that is used, if these fields are not already populated through code or environment variables. This allows for an easy conceptual mapping and organization of W&B runs in context of the execution in Anyscale.
  • Link from the Anyscale execution detail page to W&B resource for easier navigation between Anyscale execution details and W&B metrics. Integration in Anyscale UI
  • Link back to the Anyscale execution detail page from the W&B run config. This is especially useful to view logs related to your W&B run in a single file/viewer, even if the execution may have occurred over multiple nodes across machines. Integration in W&B UI
  • If the Anyscale cloud has access to read from the cloud provider's secrets manager, support fetching the W&B API key from the secrets manager based on the user and cloud of the code execution.

Example with Ray Job run on Anyscale​

import ray
import os
from ray import tune
from ray.air.integrations.wandb import WandbLoggerCallback
import wandb

def train_fn(config):
    for i in range(10):
        loss = config["a"] + config["b"]
    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]),
    },
    callbacks=[WandbLoggerCallback()]
)

In an Anyscale cluster/workspace that has wandb installed, run wandb login and paste in your API key. Then submit the above code as a Ray job with RAY_ADDRESS=anyscale://my_cluster ray job submit --working-dir . -- python wandb_demo.py

Methods of specifying API key​

The W&B API key can be specified through the following methods for code running in workspaces, Anyscale jobs, and Ray jobs running on Anyscale:

  • Running wandb login on the cluster before executing code. This only works if WandbLoggerCallback is being used.
  • Specifying WANDB_API_KEY in the environment where code is being executed.
  • Providing API key as argument to WandbLoggerCallback or setup_wandb
  • Providing API key file as argument to WandbLoggerCallback or setup_wandb. If using setup_wandb, make sure the API key file is synced to all worker nodes, for example, through runtime environments.

API key in Secrets Manager​

A more secure approach to specifying the API key is to store it in the secrets manager of the cloud provider associated with your Anyscale cloud following our naming convention:

  • For AWS clouds: anyscale_{anyscale_cloud_id}/{anyscale_user_id}/wandb_api_key
  • For GCP clouds: anyscale_{anyscale_cloud_id}-{anyscale_user_id}-wandb_api_key

The Anyscale Cloud ID can be found at https://console.anyscale.com/clouds and the current user ID can be found at https://console.anyscale.com/api/v2/userinfo/.

Alternatively, WANDB_API_KEY_NAME can be specified as an environment variable to use a different naming convention. If the API key is not provided through any other method above and the cluster has access to the secrets manager, the W&B integration on Anyscale will automatically fetch the appropriate secret for the current user and cloud per our naming convention or the value in WANDB_API_KEY_NAME.

Please work with your Anyscale Admin and Solution Architect to set up the correct policies to allow your cluster to access the secrets manager.

Environment Dependencies: Any version of Ray compatible with the open source integration and any Anyscale version can be used if running the open source integration on an Anyscale cluster. However, to get the additional features specific to the Anyscale integration, please use nightly Ray and anyscale>=0.5.74 in your Anyscale cluster. The correct dependencies are installed in the Anyscale Ray nightly base images, for example anyscale/ray:nightly-py310.