This is the Flyte SDK (v2) version of this plugin. The Flytekit version is available as flytekitplugins-wandb.
Weights & Biases
flyteplugins-wandb
This plugin provides integration between Flyte and Weights & Biases (W&B) for experiment tracking, including support for distributed training with PyTorch Elastic.
pip install flyteplugins-wandbQuick Start(example, may need adjustment)
See full examplespip install flyteplugins-wandb
from flytekit import task, workflow
from flyteplugins.wandb import wandb_init, get_wandb_run, wandb_config, wandb_sweep
@task
def my_task() -> None:
wandb_init(...)
@workflow
def my_workflow() -> None:
my_task()Available Imports (9)
Task for Weights & Biases.
from flyteplugins.wandb import wandb_init
Get the current wandb run if within a `@wandb_init` decorated task or trace.
from flyteplugins.wandb import get_wandb_run
Configuration type for Weights & Biases.
from flyteplugins.wandb import wandb_config
Task for Weights & Biases.
from flyteplugins.wandb import wandb_sweep
Get the current wandb `sweep_id` if within a `@wandb_sweep` decorated task.
from flyteplugins.wandb import get_wandb_sweep_id
Get wandb config from current Flyte context.
from flyteplugins.wandb import get_wandb_context
Configuration type for Weights & Biases.
from flyteplugins.wandb import Elastic
Generates a Weights & Biases run link.
extends dataclass — configuration or data structure for plugin setup
from flyteplugins.wandb import Wandb
Generates a Weights & Biases Sweep link.
extends dataclass — configuration or data structure for plugin setup
from flyteplugins.wandb import WandbSweep
Dependencies
Related Plugins
Weights & Biases
The Weights and Biases MLOps platform helps AI developers streamline their ML workflow from end-to-end. This plugin
Comet ML
Comet’s machine learning platform integrates with your existing infrastructure and tools so you can manage, visualize, and optimize models—from training runs to production monitoring. This plugin integrates Flyte with Comet.ml by configuring links between the two platforms.
MLflow
MLflow enables us to log parameters, code, and results in machine learning experiments and compare them using an interactive UI.
Neptune
Neptune is the MLOps stack component for experiment tracking. It offers a single place to log, compare, store, and collaborate on experiments and models. This plugin integrates Flyte with Neptune by configuring links between the two platforms.
