Explicitly set dependency TaskFlowAPI

Question

Suppose I have a dag as follows:

def my_dag():

    @task
    def task_a():
        # Some code here
        return

    @task
    def task_b():
        # Some code here
        return

    task_a()
    task_b()

my_dag_instance = my_dag()

I would like to set task_b downstream from task_a. That is, I want task_b to execute only if task_a has executed successfully. How can I do so? Note that task_a does not return anything.

I tried task_a >> task_b right before the last line of my code (my_dag_instance = my_dag()) yet got an error along the lines of " TaskDecorator >> TaskDecorator is invalid operation". Any advice?

Answer 1

Indeed, you can do that, inside the dag context, and with task instance instead of task method reference:

def my_dag():

    @task
    def task_a():
        # Some code here
        return

    @task
    def task_b():
        # Some code here
        return

    task_a() >> task_b()
    # or
    # A = task_a()
    # B = task_b()
    # A >> B

my_dag_instance = my_dag()