Conditional and dynamic workflows

Conditional logic in flytekit allows workflows to branch based on the results of previous tasks or workflow inputs. Unlike standard Python if statements, which execute immediately, flytekit's conditional constructs are compiled into a static graph (an IfElseBlock) that the Flyte engine evaluates at runtime.

Defining Conditional Branches

To create a conditional branch, use the conditional function from flytekit.core.condition. This function returns a ConditionalSection that provides a fluent API for building if, elif, and else branches.

Every conditional block must:

1. Start with an .if_() call.
2. Define the branch action using .then().
3. Optionally include one or more .elif_() branches.
4. Always terminate with an .else_() branch or a .fail() call.

The entire conditional block acts as an expression that returns the output of the branch that eventually executes.

from flytekit import task, workflow, conditional

@task
def double(n: float) -> float:
    return n * 2.0

@task
def square(n: float) -> float:
    return n * n

@workflow
def conditional_wf(my_input: float) -> float:
    return (
        conditional("fractions")
        .if_((my_input > 0.1) & (my_input < 1.0))
        .then(double(n=my_input))
        .elif_(my_input >= 1.0)
        .then(square(n=my_input))
        .else_()
        .then(my_input)
    )

Supported Expressions

Flytekit does not support standard Python truthiness for workflow variables (which are Promise objects during compilation). Instead, you must use comparison operators or specific methods on the Promise class in flytekit.core.promise.

• Comparisons: ==, !=, <, <=, >, >=
• Conjunctions: Use & for AND and | for OR. Standard Python and/or keywords will not work.
• Boolean Methods: Use .is_true(), .is_false(), or .is_none() for explicit checks.

Note: Unary expressions like if_(my_bool_promise) are explicitly rejected by the Case class constructor to prevent common errors where users expect Python's native evaluation.

Handling Failures

You can use the .fail() method to terminate a workflow execution if a specific branch is reached. This is useful for validating inputs or handling unsupported states within a workflow.

@workflow
def validated_wf(my_input: float) -> float:
    return (
        conditional("check-input")
        .if_(my_input < 0.0)
        .fail("Input must be non-negative")
        .else_()
        .then(square(n=my_input))
    )

Nested Conditionals

Conditionals can be nested by passing another conditional block into the .then() method of a branch.

@workflow
def nested_wf(my_input: float) -> float:
    return (
        conditional("outer")
        .if_(my_input > 0.0)
        .then(
            conditional("inner")
            .if_(my_input < 10.0)
            .then(double(n=my_input))
            .else_()
            .then(square(n=my_input))
        )
        .else_()
        .fail("Negative input")
    )

Internal Implementation and Execution

Flytekit handles conditionals differently depending on the execution context:

Compilation Mode

When the workflow is compiled (e.g., during registration), ConditionalSection captures the structure of the branches. The end_branch method in ConditionalSection transforms the collected Case objects into a BranchNode. This node contains an IfElseBlock which is part of the Flyte workflow specification.

Local Execution

During local execution (e.g., running the workflow function directly in Python), flytekit uses LocalExecutedConditionalSection.

1. When start_branch is called, it evaluates the ComparisonExpression or ConjunctionExpression using the eval() method.
2. If the expression evaluates to True, it calls ctx.execution_state.take_branch().
3. Only the code inside the .then() of the matching branch is executed.
4. If a branch is skipped, flytekit uses SkippedConditionalSection to ensure that tasks within that branch are not invoked.

Output Consistency

The compute_output_vars method in ConditionalSection ensures that all branches in a conditional block return a consistent set of outputs. If branches return different types or different numbers of values, the compiler will attempt to find the intersection of output variables. If no common outputs exist, it returns a VoidPromise.