Problem defining workflow dependency in airflow 2.0 taskflow API

[HashSalmon]

Let's say I have the follow dummy DAG defined as below:

@dag(default_args=default_args,
     schedule_interval=None,
     start_date=days_ago(2))
def airflow_taskflow_api_dag():
    cur_day = '2020-01-01'

    @task()
    def A(current_day: str):
        return current_day + ' 10'

    @task(multiple_outputs=True)
    def B(current_date_time: str):
        timestamp = calendar.timegm(
            datetime.strptime(
                current_date_time, '%Y-%m-%d %H').timetuple())

        another_current_date_time_format = '2020-01-01T10:00:00'

        return {
            "timestamp": timestamp,
            "new_cur_date_time": another_current_date_time_format
        }

    @task()
    def C(livy_batch_id, start_time: str):
        logging.info("Executing task C {}".format(start_time))
       return "c_id"

    @task()
    def D(upstream_id, start_time: str):
        logging.info("Executing task D {}".format(start_time))

    # workflow dependencies
    get_current_time = A(cur_day)
    new_date_time = B(get_current_time)

    livy_task = LivyOperator(
        task_id='livy_task',
        file="dummy_file",
        class_name="dummy_class",
        args=[
            new_date_time['timestamp']
        ],
        driver_memory='1G',
        driver_cores='4',
        executor_memory='1G',
        executor_cores='4',
        num_executors='4',
        queue='default_queue',
        livy_conn_id='livy_conn_id'
    )

    c_task = C(livy_task.output, new_date_time['new_cur_date_time'])

    d_task = D(c_task, new_date_time['new_cur_date_time'])

airflow_taskflow_api_dag = airflow_taskflow_api_dag()

If I wrote my DAG like this (passing two params from upstream tasks), the dependency would be like:

But I want to define the dependency with: A >> B >> livy_task >> C >> D

Is there a way to do this by using taskflow api? It seems like if I pass the task B's output to livy_task, C and D, the three tasks will be run in parallel.

Thanks!

[turbaszek]

This should help:

livy_task >> c_task.operator >> d_task.operator

[HashSalmon]

Hey, thanks for the help. I tried to define the dependency as the way you mentioned. However, it still runs c_task and d_task as another parallel branch. Please see the image below.

I think the problem is the return value new_date_time['new_cur_date_time'] from B task is passed into c_task and d_task. In my use case, this is how I want to use the xcom variables. B task computes all of the xcom variables for the downstream tasks. Is there a way to fix this ? Thanks

[turbaszek]

Have you checked graph view instead of tree view? This is how your DAG is rendering using

airflow dags show airflow_taskflow_api_dag --imgcat

And in graph view in webserver:

[HashSalmon]

I checked the graph view. It was like this

But still, this is not what I wanted to achieve. I want the task executed as a sequential. A -> B -> livy_task -> C -> D
Then I want to generate the xcom variables in B task and pass the xcom vars to livy_task, C and D.

In the current airflow 1.0, I can do it. I definitely want to use the taskflow API to achieve the same functionality because it simplifies the xcom operations.

Thanks!

[turbaszek]

@HashSalmon it may not be obvious but the tasks will be executed in the following order A -> B -> livy_task -> C -> D. Task to be executed has to have all tasks succeeded in upstream.

1. Execute A
2. Execute B
3. Execute livy, C, D
4. C, D does not meet requirement of succeded update, C waits for livy task, D waits for C
5. livy succeded, C is being executed, D waits for C
6. C succeeded, execute D

Or am I wrong? Have you run the DAG?

On graphs with task flow api, you will see not only order relation but also data relation, thus the arrows between B->C, B->D

[HashSalmon]

Hey, thanks for your help. I think I know where the problem are. There are two parts I get confused. The first part is the flow chart in Graph View. I thought it was like airflow 1.0 which the graph view represents the task flow. But it is not what I thought. I think it is just purely a dependency chart. The second part is my fake livy operator. If I use my own livy operator which extends from BaseOperator, the dependency will break. If I use the official LivyOperator/EmailOperator provided under airflow.operators, the workflow will work as I expected.

Thanks anyway !