# Weights & Biases

[Weights & Biases](https://wandb.ai) (W&B) is a platform for tracking machine learning experiments, visualizing metrics and optimizing hyperparameters. This plugin integrates W&B with Flyte, enabling you to:

- Automatically initialize W&B runs in your tasks without boilerplate
- Link directly from the Flyte UI to your W&B runs and sweeps
- Share W&B runs across parent and child tasks
- Track distributed training jobs across multiple GPUs and nodes
- Run hyperparameter sweeps with parallel agents

## Installation

```bash
pip install flyteplugins-wandb
```

You also need a W&B API key. Store it as a Flyte secret so your tasks can authenticate with W&B.

## Quick start

Here's a minimal example that logs metrics to W&B from a Flyte task:

```
import flyte

from flyteplugins.wandb import get_wandb_run, wandb_config, wandb_init

env = flyte.TaskEnvironment(
    name="wandb-example",
    image=flyte.Image.from_debian_base(name="wandb-example").with_pip_packages(
        "flyteplugins-wandb"
    ),
    secrets=[flyte.Secret(key="wandb_api_key", as_env_var="WANDB_API_KEY")],
)

@wandb_init
@env.task
async def train_model() -> str:
    wandb_run = get_wandb_run()

    # Your training code here
    for epoch in range(10):
        loss = 1.0 / (epoch + 1)
        wandb_run.log({"epoch": epoch, "loss": loss})

    return "Training complete"

if __name__ == "__main__":
    flyte.init_from_config()

    r = flyte.with_runcontext(
        custom_context=wandb_config(
            project="my-project",
            entity="my-team",
        ),
    ).run(train_model)

    print(f"run url: {r.url}")
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/integrations/flyte-plugins/wandb/quick_start.py*

This example demonstrates the core pattern:

1. **Define a task environment** with the plugin installed and your W&B API key as a secret
2. **Decorate your task** with `@wandb_init` (must be the outermost decorator, above `@env.task`)
3. **Access the run** with `get_wandb_run()` to log metrics
4. **Provide configuration** via `wandb_config()` when running the task

The plugin handles calling `wandb.init()` and `wandb.finish()` for you, and automatically adds a link to the W&B run in the Flyte UI.

![UI](https://raw.githubusercontent.com/unionai/unionai-docs-static/main/images/integrations/wandb/ui.png)

## What's next

This integration guide is split into focused sections, depending on how you want to use Weights & Biases with Flyte:

- **[Experiments](https://www.union.ai/docs/v2/union/integrations/wandb/experiments/page.md)**: Create and manage W&B runs from Flyte tasks.
- **[Distributed training](https://www.union.ai/docs/v2/union/integrations/wandb/distributed_training/page.md)**: Track experiments across multi-GPU and multi-node training jobs.
- **[Sweeps](https://www.union.ai/docs/v2/union/integrations/wandb/sweeps/page.md)**: Run hyperparameter searches and manage sweep execution from Flyte tasks.
- **[Downloading logs](https://www.union.ai/docs/v2/union/integrations/wandb/downloading_logs/page.md)**: Download logs and execution metadata from Weights & Biases.
- **[Constraints and best practices](https://www.union.ai/docs/v2/union/integrations/wandb/constraints_and_best_practices/page.md)**: Learn about limitations, edge cases and recommended patterns.
- **[Manual integration](https://www.union.ai/docs/v2/union/integrations/wandb/manual/page.md)**: Use Weights & Biases directly in Flyte tasks without decorators or helpers.

> **📝 Note**
>
> We've included additional examples developed while testing edge cases of the plugin [here](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/wandb/examples).

## Subpages

- [Experiments](https://www.union.ai/docs/v2/union/integrations/wandb/experiments/page.md)
  - Basic usage
  - Accessing the run object
  - Parent-child task relationships
  - Run modes
  - Using `run_mode="new"` for independent runs
  - Using `run_mode="shared"` for explicit sharing
  - Configuration with `wandb_config`
  - Workflow-level configuration
  - Overriding configuration for child tasks
  - Using traces with W&B runs
- [Distributed training](https://www.union.ai/docs/v2/union/integrations/wandb/distributed_training/page.md)
  - Quick start
  - Run modes in distributed training
  - Single-node behavior
  - Multi-node behavior
  - Choosing run mode and rank scope
  - Single-node multi-GPU
  - Basic example with `auto` mode
  - Using `shared` mode for per-rank metrics
  - Using `new` mode for per-rank runs
  - Multi-node training with `Elastic`
  - Global scope (default): Single run across all nodes
  - Worker scope: One run per node
  - Shared mode: All ranks log to the same run
  - New mode: Separate run per rank
  - How it works
  - Run ID patterns
- [Sweeps](https://www.union.ai/docs/v2/union/integrations/wandb/sweeps/page.md)
  - Creating a sweep
  - Running parallel agents
  - Writing objective functions
- [Downloading logs](https://www.union.ai/docs/v2/union/integrations/wandb/downloading_logs/page.md)
  - Automatic download
  - Accessing run directories during execution
- [Constraints and best practices](https://www.union.ai/docs/v2/union/integrations/wandb/constraints_and_best_practices/page.md)
  - Decorator ordering
  - Traces cannot use decorators
  - Maximum sweep agents
  - Configuration priority
  - Run ID generation
  - Sync delay for local files
  - Shared run mode requirements
  - Objective functions for sweeps
  - Error handling
- [Manual integration](https://www.union.ai/docs/v2/union/integrations/wandb/manual/page.md)
  - Using the Wandb link class
  - With a custom run ID
  - Adding links at runtime with override
  - Using the `WandbSweep` link class

---
**Source**: https://github.com/unionai/unionai-docs/blob/main/content/integrations/wandb/_index.md
**HTML**: https://www.union.ai/docs/v2/union/integrations/wandb/