metaflow_train.py

import os

from metaflow import FlowSpec, current, step

def select_model(a, b):
    if a > b:
        return a
    return b

class TrainingFlow(FlowSpec):
    """
    The flow performs the following steps:
    1) Loads the new data from NewDataFlow and initializes a new model.
    2) In parallel branches:
       - A) Fits the model with approach A.
       - B) Fits the model with approach B.
    3) Evaluates both models and exports the best.
    """

    @step
    def start(self):
        """
        Use the Metaflow client to retrieve the latest successful run from our
        NewDataFlow, then initialize new model.
        """

        self.next(self.train_approach_a, self.train_approach_b)

    @step
    def train_approach_a(self):
        """
        Train model A
        """

        self.model = 42 # Store resulting trained model
        self.next(self.join)

    @step
    def train_approach_b(self):
        """
        Train model B
        """

        self.model = 43 # Store resulting trained model
        self.next(self.join)

    @step
    def join(self, inputs):
        """
        Pick the best of both models
        """

        self.model = select_model(inputs.train_approach_a.model,
                                  inputs.train_approach_b.model) # Evaluate best model

        self.model_artifact = '/'.join([current.pathspec, 'model'])
        self.next(self.end)

    @step
    def end(self):
        print("All done!")

if __name__ == '__main__':
    TrainingFlow()