ray.air.integrations.wandb.setup_wandb#

ray.air.integrations.wandb.setup_wandb(config: Dict | None = None, api_key: str | None = None, api_key_file: str | None = None, rank_zero_only: bool = True, **kwargs) → wandb.wandb_run.Run | wandb.sdk.lib.disabled.RunDisabled[source]#

Set up a Weights & Biases session.

This function can be used to initialize a Weights & Biases session in a (distributed) training or tuning run.

By default, the run ID is the trial ID, the run name is the trial name, and the run group is the experiment name. These settings can be overwritten by passing the respective arguments as kwargs, which will be passed to wandb.init().

In distributed training with Ray Train, only the zero-rank worker will initialize wandb. All other workers will return a disabled run object, so that logging is not duplicated in a distributed run. This can be disabled by passing rank_zero_only=False, which will then initialize wandb in every training worker.

The config argument will be passed to Weights and Biases and will be logged as the run configuration.

If no API key or key file are passed, wandb will try to authenticate using locally stored credentials, created for instance by running wandb login.

Keyword arguments passed to setup_wandb() will be passed to wandb.init() and take precedence over any potential default settings.

Parameters:

  • config – Configuration dict to be logged to Weights and Biases. Can contain arguments for wandb.init() as well as authentication information.

  • api_key – API key to use for authentication with Weights and Biases.

  • api_key_file – File pointing to API key for with Weights and Biases.

  • rank_zero_only – If True, will return an initialized session only for the rank 0 worker in distributed training. If False, will initialize a session for all workers.

  • kwargs – Passed to wandb.init().

Example

from ray.air.integrations.wandb import setup_wandb

def training_loop(config):
    wandb = setup_wandb(config)
    # ...
    wandb.log({"loss": 0.123})

PublicAPI (alpha): This API is in alpha and may change before becoming stable.