> ## Documentation Index
> Fetch the complete documentation index at: https://wb-21fd5541-locadex-parallel-t9n-main-cs60c8p4o6ik99tylxgp3.mintlify.site/llms.txt
> Use this file to discover all available pages before exploring further.

# Add W&B to a Python library

> Best practices for integrating W&B into your Python library for experiment tracking, system monitoring, and model management.

This guide describes patterns for integrating W\&B into a Python library, framework, or SDK. It covers dependency management, authentication, optional logging, run initialization, artifacts, hyperparameter tuning, and distributed execution.

Use this guide when your integration spans reusable library code rather than a single training script or notebook. For an introduction to W\&B, see [Experiment Tracking](/models/track).

## Design the integration

Before adding W\&B to your library, decide:

* Whether wandb is a required or optional dependency.
* Whether users can run your library without saving or uploading W\&B data.
* Which configuration values, metrics, and artifacts your library logs.
* Whether to support hyperparameter tuning with W\&B Sweeps.
* Whether to share artifacts, such as models and datasets, through your organization's W\&B Registry.
* How your library handles distributed training and multiple processes.

These decisions determine how your library imports W\&B, manages run lifecycles, and behaves when W\&B is unavailable.

## Decide how to install W\&B

Choose whether to install W\&B automatically with your library or expose it as an optional feature.

### Require W\&B as a dependency

If W\&B is central to your library, add `wandb` to its dependencies.

For a `requirements.txt` file:

```txt title="requirements.txt" type="text" theme={null}
torch
wandb
```

For a `pyproject.toml` file:

```toml title="pyproject.toml" type="toml" theme={null}
[project]
name = "my_awesome_lib"
version = "0.1.0"
dependencies = [
    "torch",
    "wandb",
]
```

Consider specifying a compatible version range based on the W\&B features your integration uses.

### Make W\&B an optional dependency

If W\&B is an optional feature, enable installation of your library without `wandb`.

Declare W\&B as an optional dependency in `pyproject.toml`:

```toml title="pyproject.toml" theme={null}
[project]
name = "my_awesome_lib"
version = "0.1.0"
dependencies = [
    "torch",
]

[project.optional-dependencies]
wandb = [
    "wandb",
]
```

Then install the integration with:

```bash theme={null}
pip install "my_awesome_lib[wandb]"
```

## Authenticate users

W\&B uses API keys to authenticate users and machines. Before you can log runs from your library, you must generate an API key and make it available to the `wandb` client.

W\&B supports authentication through the CLI, environment variables, and `wandb.login()`.

### Create an API key

<Note>
  For a more streamlined approach, go to [User Settings](https://wandb.ai/settings) and create an API key. Copy the API key immediately and save it in a secure location such as a password manager.
</Note>

### Authenticate from the command line

From your terminal, run:

```bash theme={null}
wandb login
```

W\&B prompts you to enter an API key.

Alternatively, set the `WANDB_API_KEY` environment variable:

```bash theme={null}
export WANDB_API_KEY="<api_key>"
```

For more information, see [Environment variables](/models/track/environment-variables).

### Authenticate from Python

In an interactive Python environment or notebook, call:

```python theme={null}
import wandb

wandb.login()
```

Avoid calling [`wandb.login()`](/models/ref/python/functions/login) automatically from reusable library code because it can interrupt noninteractive workflows.

## Make W\&B logging optional

Choose whether W\&B saves or uploads run data.

| Mode       | Behavior                                                                                                                   |
| ---------- | -------------------------------------------------------------------------------------------------------------------------- |
| `disabled` | Deactivates W\&B logging. W\&B does not save or upload run data. Calls to W\&B methods have no effect.                     |
| `offline`  | Saves run data locally without uploading it. You can upload the run later with [`wandb sync`](/models/ref/cli/wandb-sync). |

Use `disabled` when you do not want to save run data. Use `offline` when they want to save data locally for later upload.

### Use `disabled` mode

Pass `mode="disabled"` to `wandb.init()`:

```python theme={null}
with wandb.init(mode="disabled") as run:
    # Training logic.
```

Or set `WANDB_MODE` in the environment:

```bash theme={null}
export WANDB_MODE=disabled
```

You can also set the mode with the W\&B CLI:

```bash theme={null}
wandb disabled
```

### Use `offline` mode

Pass `mode="offline"` to `wandb.init()`:

```python theme={null}
with wandb.init(mode="offline") as run:
    # Training logic.
```

Or set `WANDB_MODE` in the environment:

```bash theme={null}
export WANDB_MODE=offline
```

In Python, set the environment variable before calling `wandb.init()`:

```python theme={null}
import os

os.environ["WANDB_MODE"] = "offline"
```

You can also set the mode with the W\&B CLI:

```bash theme={null}
wandb offline
```

Upload an offline run later with:

```bash theme={null}
wandb sync <run_directory>
```

## Initialize a run

After authenticating, initialize a [run](/models/runs) to log metrics, [configuration](/models/track/config) values, and [artifacts](/models/artifacts) from your library.

Call [`wandb.init()`](/models/ref/python/functions/init) and specify the project and team entity. If you omit the project, W\&B stores the run in the default `"uncategorized"` project.

Use `wandb.init()` as a context manager around the training loop. When the block exits, W\&B finishes the run and processes pending data before the process ends.

For example, suppose your library contains the following training loop:

```python icon="python" title="Python" theme={null}
import random # For simulating data

def model(training_data: int) -> int:
    """Model simulation for demonstration purposes."""
    return training_data * 2 + random.randint(-1, 1)  

# Simulate weights and noise
weights = random.random() # Initialize random weights
noise = random.random() / 5  # Small random noise to simulate noise

for epoch in range(epochs):
    xb = weights + noise  # Simulated input training data
    yb = weights + noise * 2  # Simulated target output (double the input noise)
    
    y_pred = model(xb)  # Model prediction
    loss = (yb - y_pred) ** 2  # Mean Squared Error loss

    print(f"epoch={epoch}, loss={loss}")
```

To integrate W\&B, initialize a run with a context manager, pass configuration values to `wandb.init()`, and log metrics with `wandb.Run.log()`:

```python icon="python" title="Python" highlight={2,13-16,19,30-33} theme={null}
import random # For simulating data
import wandb

def model(training_data: int) -> int:
    """Model simulation for demonstration purposes."""
    return training_data * 2 + random.randint(-1, 1)  

# Simulate weights and noise
weights = random.random() # Initialize random weights
noise = random.random() / 5  # Small random noise to simulate noise

# Hyperparameters and configuration
config = {
    "epochs": 10,  # Number of epochs to train
    "learning_rate": 0.01,  # Learning rate for the optimizer
}

# Use context manager to initialize and close W&B runs
with wandb.init(entity="your-entity", project="your-project-name", config=config) as run:    
    # Simulate training loop
    for epoch in range(config["epochs"]):
        xb = weights + noise  # Simulated input training data
        yb = weights + noise * 2  # Simulated target output (double the input noise)
        
        y_pred = model(xb)  # Model prediction
        loss = (yb - y_pred) ** 2  # Mean Squared Error loss

        print(f"epoch={epoch}, loss={loss}")
        # Log epoch and loss to W&B
        run.log({
            "epoch": epoch,
            "loss": loss,
        })
```

<Tip>
  **When to call `wandb.init()`**

  Call `wandb.init()` before the work you want W\&B to monitor. Use it as a context manager around the entire training loop so the run captures relevant standard output, standard error, and error messages for debugging.
</Tip>

The [next section](#log-configuration-and-metrics) describes in detail how to log configuration values and metrics from your library.

## Log configuration and metrics

Log configuration values and metrics to W\&B so you, and members of your team, can [compare](/models/runs/compare-runs), [filter](/models/runs/filter-runs#example-filter-run-configuration-values-with-contains), [group](/models/runs/grouping), and [reproduce](/models/track/reproduce_experiments) experiments.

### Log configuration values

Pass a [configuration](/models/track/config) dictionary to `wandb.init()` to record hyperparameters and other metadata.

Use descriptive keys and JSON-serializable values:

```python theme={null}
config = {
    "batch_size": 32,
    "learning_rate": 0.001,
    "optimizer": "adam",
    "model": {
        "type": "resnet",
        "depth": 50,
    },
    "dataset": {
        "name": "CIFAR-10",
        "num_classes": 10,
    },
}
```

Pass the dictionary to `wandb.init()`:

```python theme={null}
with wandb.init(entity="your-entity", project="your-project-name", config=config) as run:
    # Training logic.
```

Some values might not be available when the run starts. Add them later with `wandb.Run.config.update()`:

```python theme={null}
run.config.update(
    {
        "model_parameters": 3500,
    }
)
```

For more information, see [Configure experiments](/models/track/config).

### Log metrics

Log metrics such as loss or accuracy during training. Create a dictionary where each key is the name of a metric and the value is the metric value. Pass this dictionary to [`wandb.Run.log()`](/models/ref/python/experiments/run#method-run-log).

The following code snippet logs training and validation metrics to W\&B:

```python icon="python" title="Python" theme={null}
import wandb

with wandb.init(entity="your-entity", project="your-project-name") as run:
    metrics = {
        "train/loss": 0.4,
        "train/learning_rate": 0.4,
        "val/loss": 0.5, 
        "val/accuracy": 0.7
    }
    run.log(metrics)
```

Use prefixes such as `train/` and `val/` to group related metrics.

For guidance on logging supported data types, automatically tracked data, and best practices, see [Log metrics and data](/models/track/log).

To track models and datasets, see the [Track models and datasets with artifacts](#track-models-and-datasets-with-artifacts) section.

### Define a custom log axis

By default, W\&B plots logged metrics against an automatically incremented step. Each call to `wandb.Run.log()` advances the step.

Use `wandb.Run.define_metric()` to plot a metric against another value, such as an epoch or global step.

The following example defines `x_axis_squared` as the x-axis for `validation_loss`. For each loop iteration, `x_axis_squared` is the square of the index `i`, and `validation_loss` is a randomly generated value:

```python title="Python" icon="python" theme={null}
import wandb
import random

with wandb.init() as run:
    run.define_metric(step_metric = "x_axis_squared", name = "validation_loss")

    for i in range(10):
        log_dict = {
            "x_axis_squared": i**2,
            "validation_loss": random.random(),
        }
        run.log(log_dict)
```

For more information, see [Define a custom log axis](/models/track/log/customize-logging-axes).

## Track models and datasets with artifacts

In addition to metrics, you can persist the models and datasets your library produces or consumes so that you, and members of your team, can reproduce and compare runs.

Use [W\&B Artifacts](/models/artifacts) to version models, datasets, and other files that your library produces or consumes.

Before adding artifact support, decide:

* Which files to log.
* Whether artifact logging is optional.
* How frequently to log checkpoints.
* How to name artifacts and aliases.
* Which artifacts represent run inputs and outputs.

### Log model checkpoints

Log model checkpoints as artifacts so you can recover, version, and share trained weights. Include the run ID in the artifact name to associate each checkpoint with its source run.

The following example logs a checkpoint every 10 epochs. It creates an artifact whose name includes the run ID, adds model weights from a local directory, and logs the artifact with a [custom alias](/models/artifacts/create-a-custom-alias).

```python title="Python" icon="python" theme={null}
import wandb
with wandb.init(entity="your-entity", project="your-project-name") as run:

    # Training data and model training logic here

    if epoch % 10 == 0: # Log model checkpoint every 10 epochs

        metadata = {"eval/accuracy": 0.8, "train/steps": 800} 

        artifact = wandb.Artifact(
                        name=f"model-{run.id}", 
                        metadata=metadata, 
                        type="model"
                        )
        artifact.add_dir("./models/output_model") # local directory where the model weights are stored

        run.log_artifact(artifact, aliases=f"epoch_{epoch}")
```

### Track run inputs

Use [`wandb.Run.use_artifact()`](/models/ref/python/experiments/run#method-run-use_artifact) when a run consumes an artifact, such as a dataset or model checkpoint. W\&B records the artifact as an input to the run.

Specify the name of the artifact and an optional alias to reference a specific version of that artifact. The name of the artifact is in the format `artifact_name:version` or `artifact_name:alias`.

```python title="Python" icon="python" theme={null}
import wandb

# Initialize a run
with wandb.init(entity="your-entity", project="your-project") as run:
  # Get artifact, mark it as a dependency
  artifact = run.use_artifact(artifact_or_name="artifact_name:alias_or_version")
```

### Track run outputs

Use [`wandb.Run.log_artifact()`](/models/ref/python/experiments/run#method-run-log_artifact) to log an artifact as an output of a run.

1. Create an artifact with `wandb.Artifact()`.
2. Add one or more files to the artifact.
3. Log the artifact with `wandb.Run.log_artifact()`.

```python theme={null}
import wandb

# Initialize a run
with wandb.init(entity="your-entity", project="project-name") as run:
  
  # Create an artifact
  artifact = wandb.Artifact(name = "artifact-name", type = "artifact-type")
  artifact.add_file(local_path = "path/to/file", name="optional-filename") # Add a file to the artifact

  # Log the artifact as an output of the run
  run.log_artifact(artifact_or_path = artifact)
```

For more information about creating, managing, and using artifacts, see [Artifacts](/models/artifacts).

### Download artifacts

Use the [W\&B Public API](/models/ref/python/public-api) to download an artifact without creating a run or recording a run input relationship.

```python theme={null}
import wandb

api = wandb.Api()

artifact = api.artifact(
    "<entity>/<project>/<artifact_name>:<alias_or_version>"
)

local_path = artifact.download()
```

For information, see [Download and use artifacts](/models/registry/download_use_artifact).

## Link artifacts to the W\&B Registry

Use [W\&B Registry](/models/registry) to share and manage artifact versions across teams.

Before you link an artifact version, decide:

* Which artifact versions to share.
* Which [collections](/models/registry/create_collection) to link them to.
* Which [registry](/models/registry) to publish them in.

<Note>
  The destination registry must already exist.
</Note>

The following example logs an artifact and links it to a collection in a registry:

```python theme={null}
import wandb
import random

# Specify the name of the collection and registry
# you want to publish the artifact to
COLLECTION_NAME = "collection-name"
REGISTRY_NAME = "registry-name"

# Initialize a W&B Run to track the artifact
with wandb.init(project="project-name") as run:
    # Create a simulated model file so that you can log it
    with open("my_model.txt", "w") as f:
        f.write("Model: " + str(random.random()))

    # Log the artifact to W&B
    logged_artifact = run.log_artifact(
        artifact_or_path="./my_model.txt", 
        name="artifact-name",
        type="artifact-type" # Specifies artifact type
    )

    # Link the artifact to the registry
    run.link_artifact(
        artifact=logged_artifact, 
        target_path=f"wandb-registry-{REGISTRY_NAME}/{COLLECTION_NAME}"
    )
```

For more information, see [Link an artifact version to a collection](/models/registry/link_version).

## Tune hyperparameters

If your library supports hyperparameter tuning, integrate [W\&B Sweeps](/models/sweeps) to run grid, random, or Bayesian searches.

<Tip>
  W\&B recommends that you keep the sweep configuration separate from the underlying training logic. The training function should read its parameters from `wandb.Run.config` rather than depending directly on sweep-specific logic.
</Tip>

For more information, see [Sweeps](/models/sweeps).

## Support distributed training

If your library supports multiple processes or machines, define which processes create runs and log data.

Common approaches include:

* Log only from the main process. This approach avoids duplicate metrics and artifacts.
* Create one run for each process and group the runs with a shared `group` value.

For more information, see [Log distributed training experiments](/models/track/log/distributed-training).