Metaflow

Integrate W&B with Metaflow to track experiments and manage ML workflows with automatic metric and artifact logging.

Overview

Metaflow is a framework created by Netflix for creating and running ML workflows.

This integration lets you apply decorators to Metaflow steps and flows to automatically log parameters and artifacts to W&B, so you can track experiments and inspect lineage across the workflows you build with Metaflow without writing custom logging code:

• Decorating a step turns logging off or on for certain types within that step.
• Decorating the flow turns logging off or on for every step in the flow.

Quickstart

The following sections walk you through authenticating with W&B, installing the required libraries, and adding the wandb_log decorator to your Metaflow steps and flows.

Sign up and create an API key

An API key authenticates your machine to W&B. You can generate an API key from your user profile.

For a more streamlined approach, create an API key by going directly to User Settings. Copy the newly created API key immediately and save it in a secure location such as a password manager.

To find your API key in the W&B app:

1. Click your user profile icon in the upper right corner.
2. Select User Settings, then scroll to the API Keys section.

Install the wandb library and log in

To install the wandb library locally and log in:

For wandb version 0.19.8 or below, install fastcore version 1.8.0 or below (fastcore<1.8.0) instead of plum-dispatch.

Command line

1. Set the WANDB_API_KEY environment variable to your API key.

   export WANDB_API_KEY=[YOUR-API-KEY]
   

2. Install the wandb library and log in.

   pip install -Uqqq metaflow "plum-dispatch<3.0.0" wandb
   
   wandb login
   

Python

pip install -Uqqq metaflow "plum-dispatch<3.0.0" wandb

import wandb
wandb.login()

Python notebook

!pip install -Uqqq metaflow "plum-dispatch<3.0.0" wandb

import wandb
wandb.login()

Decorate your flows and steps

Step

Decorating a step turns logging off or on for certain types within that step.

In this example, the integration logs all datasets and models in start:

from wandb.integration.metaflow import wandb_log

class WandbExampleFlow(FlowSpec):
    @wandb_log(datasets=True, models=True, settings=wandb.Settings(...))
    @step
    def start(self):
        self.raw_df = pd.read_csv(...).    # pd.DataFrame -> upload as dataset
        self.model_file = torch.load(...)  # nn.Module    -> upload as model
        self.next(self.transform)

Flow

Decorating a flow is equivalent to decorating all the constituent steps with a default.

In this case, all steps in WandbExampleFlow default to logging datasets and models by default, equivalent to decorating each step with @wandb_log(datasets=True, models=True):

from wandb.integration.metaflow import wandb_log

@wandb_log(datasets=True, models=True)  # decorate all @step 
class WandbExampleFlow(FlowSpec):
    @step
    def start(self):
        self.raw_df = pd.read_csv(...).    # pd.DataFrame -> upload as dataset
        self.model_file = torch.load(...)  # nn.Module    -> upload as model
        self.next(self.transform)

Flow and steps

Decorating the flow is equivalent to decorating all steps with a default. That means if you later decorate a Step with another @wandb_log, it overrides the flow-level decoration.

In this example:

• start and mid log both datasets and models.
• end logs neither datasets nor models.

from wandb.integration.metaflow import wandb_log

@wandb_log(datasets=True, models=True)  # same as decorating start and mid
class WandbExampleFlow(FlowSpec):
  # this step will log datasets and models
  @step
  def start(self):
    self.raw_df = pd.read_csv(...).    # pd.DataFrame -> upload as dataset
    self.model_file = torch.load(...)  # nn.Module    -> upload as model
    self.next(self.mid)

  # this step will also log datasets and models
  @step
  def mid(self):
    self.raw_df = pd.read_csv(...).    # pd.DataFrame -> upload as dataset
    self.model_file = torch.load(...)  # nn.Module    -> upload as model
    self.next(self.end)

  # this step is overwritten and will NOT log datasets OR models
  @wandb_log(datasets=False, models=False)
  @step
  def end(self):
    self.raw_df = pd.read_csv(...).    
    self.model_file = torch.load(...)

Access your data programmatically

Once your flows and steps are decorated, runs send parameters and artifacts to W&B each time the flow executes. You can access the captured information in three ways: inside the original Python process being logged using the wandb client library, with the web app UI, or programmatically using the Public API. Parameters are saved to the W&B config and can be found in the Overview tab. datasets, models, and others are saved to W&B Artifacts and can be found in the Artifacts tab. Base python types are saved to the W&B summary dict and can be found in the Overview tab. See the guide to the Public API for details on using the API to get this information programmatically from outside.

Quick reference

Data	Client library	UI
Parameter(...)	wandb.Run.config	Overview tab, Config
datasets, models, others	wandb.Run.use_artifact("{var_name}:latest")	Artifacts tab
Base Python types (dict, list, str, etc.)	wandb.Run.summary	Overview tab, Summary

wandb_log kwargs

kwarg	Options
datasets	True: Log instance variables that are a dataset False
models	True: Log instance variables that are a model False
others	True: Log anything else that is serializable as a pickle False
settings	wandb.Settings(…): Specify your own wandb settings for this step or flow None: Equivalent to passing wandb.Settings() By default, if: settings.run_group is None, it’s set to {flow_name}/{run_id} settings.run_job_type is None, it’s set to {run_job_type}/{step_name}

Frequently asked questions

The following sections answer common questions about logging behavior, supported data types, and artifact lineage.

What exactly do you log

wandb_log only logs instance variables. Local variables are never logged. This is useful to avoid logging unnecessary data.

Which data types get logged

W&B supports these types:

Logging setting	Type
default (always on)	dict, list, set, str, int, float, bool
datasets	pd.DataFrame pathlib.Path
models	nn.Module sklearn.base.BaseEstimator
others	Anything that is pickle-able and JSON serializable

Configure logging behavior

Kind of variable	Behavior	Example	Data type
Instance	Auto-logged	self.accuracy	float
Instance	Logged if datasets=True	self.df	pd.DataFrame
Instance	Not logged if datasets=False	self.df	pd.DataFrame
Local	Never logged	accuracy	float
Local	Never logged	df	pd.DataFrame

Artifact lineage tracking

If you have an artifact that is an output of step A and an input to step B, W&B automatically constructs the lineage directed acyclic graph (DAG) for you.

For an example of this behavior, see this Metaflow integration example notebook and its corresponding Artifacts page.