Conditional branch documentation + example usage

From the flowspec docstring:

- Conditional branch:
    self.next(self.if_true, self.if_false, condition='boolean_variable')
  In this situation, both `if_true` and `if_false` are methods in the current class
  decorated with the `@step` decorator and `boolean_variable` is a variable name
  in the current class that evaluates to True or False. The `if_true` step will be
  executed if thecondition variable evaluates to True and the `if_false` step will
  be executed otherwise

It’d be great to have this mentioned on Metaflow’s website with an example.

On a related note, this capability makes it easier to introduce cycles into the DAG and while the documentation mentions:

Metaflow infers a directed (typically acyclic) graph based on the transitions between step functions.

The acyclicity seems to be checked by the linter, however that piece of validation is not disabled using the --no-pylint flag. It’s alluded to the possibility of a graph with cycles, but it doesn’t seem possible to do for now.

Example:

from metaflow import FlowSpec, step


class TestFlowConditional(FlowSpec):
    """
    A toy flow to mimic a hyperparameter tuning strategy.
    
    The flow performs the following steps:
    1) Load data.
    2) Generate hyperparameter candidates.
    3) Fan-out training over hyperparameter candidates to evaluate using foreach.
    4) Join results.
    5) Conditionally stop at max iterations or keep evaluating.
    """

    @step
    def start(self):
        # self.data = ...
        self._iteration = 0
        self._max_iteration = 15
        self._num_candidates = 3
        self.results = []
        self.next(self.generate_candidates)

    @step
    def generate_candidates(self):
        candidates = []
        for _ in range(self._num_candidates):
            candidate = {
                "hyperparameters": { 
                    # ... 
                },
            }
            candidates.append(candidate)
        self.candidates = candidates
        self._iteration += len(candidates)
        self.next(self.train, foreach='candidates')

    @step
    def train(self):
        hyperparams = self.input['hyperparameters']
        
        # ...
        
        self.next(self.join)

    @step
    def join(self, inputs):
        """
        Combine results for hyperparameter candidates.
        """
        # ...
        self.next(self.should_stop)

    @step
    def should_stop(self):
        """
        Conditional branch to end when max iterations is reached, otherwise evaluate more candidates.
        """
        self._done = self._iteration < self._max_iteration
        self.next(self.end, self.generate_candidates, condition='_done')

    @step
    def end(self):
        print("Finished.")

if __name__ == '__main__':
    TestFlowConditional()

python TestMetaflowConditional.py --no-pylint run

Results in:

Metaflow 2.0.1 executing TestFlowConditional for user:russell
Validating your flow...
    Validity checker found an issue on line 62:
    There is a loop in your flow: generate_candidates->train->join->should_stop->generate_candidates. Break the loop by fixing self.next() transitions.

Keep up the great work and I’ve been enjoying Metaflow so far!

[tuulos]

@mikejmills curious, what’s your use case? I listed some commonly used workarounds here https://gitter.im/metaflow_org/community?at=6079df40a2ac0d38e7be28fc

[sgaseretto]

Are there any plans on adding conditional branching in the near future? I think this is a key feature to implement workflows like simple automatic retraining (or continuous training), where you compare if a new model performs better than an old one, and if so, the new model replaces the old one