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?

score 1 · Accepted Answer · answered Aug 17 '22 at 19:00

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()

answered Aug 17 '22 at 19:00

Hussein Awala

4,285
2
9
23

Thank you so much Hussein! Let me try that out... – Diana Vazquez Romo Aug 17 '22 at 19:05
A tiny followup: What if `task_b` returns a value that I want to feed to a third task, say `task_c`? If I write: `task_a() >> task_c(task_b())`. Would following dependence be respected `task_a >> task_b >> task_c`? – Diana Vazquez Romo Aug 17 '22 at 19:25
I don't think so, it will be `A >> C` and `B >> C` but A and B can be run in parallel, the best way is storing the task instances in variables, and use these variables to define the dependencies and provide the arguments – Hussein Awala Aug 17 '22 at 19:28
Yeah that's what I thought... Oh okay so in your above example, `A = task_a()`, A stores an instance of a task, NOT the return value of the function? – Diana Vazquez Romo Aug 17 '22 at 19:32
1

It doesn't store the value, but when you pass it as an argument to another task, airflow will replace this instance by the result of the task at run_time, this is one of the decorator features. – Hussein Awala Aug 17 '22 at 19:41
Oh wow! that's so cool, thanks Hussein! – Diana Vazquez Romo Aug 17 '22 at 21:20
Would this be the right way to express it. Again `task_a` no return value, `task_b` return value that `task_c` takes in. `` A = task_a(); B = task_b(); task_c(B); A >> B `` Think dependency here would be A >> B >> task_c() – Diana Vazquez Romo Aug 17 '22 at 21:24
1

Yes, this should do the job! – Hussein Awala Aug 17 '22 at 21:29
Thanks Hussein you saved me! Trying to get my DAG to work here < 3 – Diana Vazquez Romo Aug 17 '22 at 21:42

Explicitly set dependency TaskFlowAPI

1 Answers1