Databricks

Integrate W&B with Databricks for experiment tracking, metric logging, and model management on Spark clusters.

W&B integrates with Databricks by customizing the W&B Jupyter notebook experience in the Databricks environment. This page shows you how to install and authenticate W&B on a Databricks cluster so that you can track experiments and log metrics from notebooks running on Spark.

Configure Databricks

To use W&B from a Databricks notebook, you must install the wandb package on the cluster and configure authentication so your notebooks can log to W&B.

1. Install wandb in the cluster

   In your cluster configuration, choose your cluster, then click Libraries > Install New > PyPI, and add the package wandb.

2. Set up authentication

   To authenticate your W&B account, add a Databricks secret that your notebooks can query at runtime. This avoids hard-coding your API key in notebooks.

   # install databricks cli
   pip install databricks-cli
   
   # Generate a token from databricks UI
   databricks configure --token
   
   # Create a scope with one of the two commands (depending if you have security features enabled on databricks):
   # with security add-on
   databricks secrets create-scope --scope wandb
   # without security add-on
   databricks secrets create-scope --scope wandb --initial-manage-principal users
   
   # Create an API key at https://wandb.ai/settings
   databricks secrets put --scope wandb --key api_key
   

Examples

The following examples show how to use the preceding secret to log in and begin logging from a Databricks notebook.

Basic example

import os
import wandb

api_key = dbutils.secrets.get("wandb", "api_key")
wandb.login(key=api_key)

with wandb.init() as run:
    run.log({"foo": 1})

Sweeps

Notebooks that use wandb.sweep() or wandb.agent() must set the entity and project as environment variables:

import os

os.environ["WANDB_ENTITY"] = "my-entity"
os.environ["WANDB_PROJECT"] = "my-project-that-exists"