Skip to main content

Parallel Branches in Workflows

In this episode, you will build a workflow to train multiple models in parallel. This is done using the branching pattern of Metaflow. Specifically, you will see how to define steps so that Metaflow knows to execute them in parallel on multiple CPU cores or cloud instances. To show this you will combine the previous two flows into and train the RandomForestClassifier and XGBoostClassifer in parallel.

1Write Your First Branching Flow

The flow has the following structure:

  • Parameter values are defined at the beginning of the class.
  • The start step loads and splits a dataset to be used in downstream tasks.
    • Notice that this step calls two downstream steps in, self.train_xgb). This is called branching. This means the train_rf and train_xgb steps will be run in parallel.
  • The train_rf step fits a sklearn.ensemble.RandomForestClassifier for the classification task using cross-validation.
  • The train_xgb step fits a xgboost.XGBClassifier for the classification task using cross-validation.
  • The score step evaluates each classifier on a held-out dataset for testing.
    • This step is referred to as a join step.
    • It takes in an extra argument that contains the results of the tasks that call
  • The end step prints the accuracy scores for each classifier.
from metaflow import FlowSpec, step, Parameter

class ParallelTreesFlow(FlowSpec):

max_depth = Parameter("max_depth", default=None)
random_state = Parameter("seed", default=21)
n_estimators = Parameter("n-est", default=10)
min_samples_split = Parameter("min-samples", default=2)
eval_metric = Parameter("eval-metric", default='mlogloss')
k_fold = Parameter("k", default=5)

def start(self):
from sklearn import datasets
self.iris = datasets.load_iris()
self.X = self.iris['data']
self.y = self.iris['target'], self.train_xgb)

def train_rf(self):
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
self.clf = RandomForestClassifier(
self.model_name = "Random Forest"
self.scores = cross_val_score(
self.clf, self.X, self.y, cv=self.k_fold)

def train_xgb(self):
from xgboost import XGBClassifier
from sklearn.model_selection import cross_val_score
self.clf = XGBClassifier(
self.model_name = "XGBoost"
self.scores = cross_val_score(
self.clf, self.X, self.y, cv=self.k_fold)

def score(self, modeling_tasks):
import numpy as np
self.scores = [
for model in modeling_tasks

def end(self):
self.experiment_results = []
for name, mean, std in self.scores:
msg = "{} Model Accuracy: {} \u00B1 {}%"
print(msg.format(name, round(mean, 3), round(std, 3)))

if __name__ == "__main__":

2Run the Flow

python run
     Workflow starting (run-id 1666720729727507):
[1666720729727507/start/1 (pid 52718)] Task is starting.
[1666720729727507/start/1 (pid 52718)] Task finished successfully.
[1666720729727507/train_rf/2 (pid 52725)] Task is starting.
[1666720729727507/train_xgb/3 (pid 52726)] Task is starting.
[1666720729727507/train_rf/2 (pid 52725)] Task finished successfully.
[1666720729727507/train_xgb/3 (pid 52726)] Task finished successfully.
[1666720729727507/score/4 (pid 52731)] Task is starting.
[1666720729727507/score/4 (pid 52731)] Task finished successfully.
[1666720729727507/end/5 (pid 52734)] Task is starting.
[1666720729727507/end/5 (pid 52734)] Random Forest Model Accuracy: 0.96 ± 0.025%
[1666720729727507/end/5 (pid 52734)] XGBoost Model Accuracy: 0.967 ± 0.021%
[1666720729727507/end/5 (pid 52734)] Task finished successfully.

In this episode, you trained two models in parallel using multiple CPU cores. In the next episode, you will transition from authoring and running flows to focusing on how to analyze the results produced by flows. See you there!