Control Flow

Branches

The pipeline structure can be dynamically altered based on the output of other tasks. Create the following two files in the main project directory:

control_flow.pysubmit.sh

from sdag import Elif, Else, If, pipeline


@pipeline
def control_flow():
    with If(condition1()):
        option1()
    with Elif(condition2()):
        option2()
    with Else():
        option3()


@control_flow.task("submit.sh")
def condition1() -> bool:
    return False


@control_flow.task("submit.sh")
def condition2() -> bool:
    return True


@control_flow.task("submit.sh")
def option1():
    print("Option 1 selected")


@control_flow.task("submit.sh")
def option2():
    print("Option 2 selected")


@control_flow.task("submit.sh")
def option3():
    print("Option 3 selected")

sdag-execute

condition1 and condition2 are tasks that must return a Boolean value. Run:

sdag run control_flow --local

to execute the example. The expected flow is:

1. condition1 is executed and returns False.
2. Because the output is False, option1 is skipped and condition2 is evaluated.
3. The output of condition2 is True, so option2 is executed while option3 is skipped.

OneOf

One might expect the following snippet to work as expected:

@pipeline
def dag():
    with If(condition1()):
        t = option1()
    with Else():
        t = option2()
    final_task(t)

Unfortunately, it's not! t in the example will always refer to option2. The oneof node is inspired by Kubeflow and it can be used to reference tasks running on separate branches. Add the following snippet to the same control_flow.py module:

from sdag import oneof


@pipeline
def oneof_example():
    with If(condition()):
        t1 = first_alternative()
    with Else():
        t2 = second_alternative()
    t = oneof(t1, t2)
    final_task(t)


@oneof_example.task("submit.sh")
def condition() -> bool:
    return False


@oneof_example.task("submit.sh")
def first_alternative():
    print("alternative 1 selected")


@oneof_example.task("submit.sh")
def second_alternative():
    print("alternative 2 selected")


@oneof_example.task("submit.sh")
def final_task():
    print("final task")

Here, final_task will only run if at least one between first_alternative and second_alternative executes successfully. You can run the example through:

sdag run oneof_example --local

You can also use oneof when multiple concurrent tasks executed. It will become the first task of the pool that completes the execution without errors.

@pipeline
def dag():
    t_fail = will_fail()
    t_slow = slow_task()
    t_fast = fast_task()

    # Likely equal to t_fast
    first_to_complete = oneof(t_fail, t_slow, t_fast)

Info

To run these examples with Slurm, turn submit.sh into a sbatch script as explained in the Hello World section and execute pipelines without the --local option.