Return list of tasks from function that should be run in sequence in Airflow

Question

I want to return 2 or more tasks from a function that should be run in sequence in the spot they're inserted in the dependencies, see below.

t1 = PythonOperator()

def generate_tasks():
    t2 = PythonOperator()
    t3 = PythonOperator()
    return magic(t2, t3) # magic needed here (preferably)

t1 >> generate_tasks() # otherwise here
# desired result: t1 >> t2 >> t3

Is this doable? As I understand it Airflow 2.0 seems to achieve this with a TaskGroup, but we're on Google's Composer, and 2.0 won't be available for a while.

Best workaround I've found:

t1 = PythonOperator()

def generate_tasks():
    t2 = PythonOperator()
    t3 = PythonOperator()
    return [t2, t3]

tasks = generate_tasks()
t1 >> tasks[0] >> tasks[1]

But I'd really like that to be abstracted away, as it more or less defeats the purpose of having multiple operators returned from a single function. We want it to be a single unit as far as the end user knows, even though it can be composed of 2 or more tasks.

How to do it with TaskGroup in Airflow 2.0:

class Encryptor:
    def encrypt_and_archive(self):
        with TaskGroup("archive_and_encrypt") as section_1:
            encrypt = DummyOperator(task_id="encrypt")
            archive = BashOperator(task_id="archive", bash_command='echo 1')
            encrypt >> archive
        return section_1

with DAG(dag_id="example_return_task_group", start_date=days_ago(2), tags=["example"]) as dag:
    start = DummyOperator(task_id="start")
    encrypt_and_archive = Encryptor().encrypt_and_archive()
    end = DummyOperator(task_id='end')

             #  single variable, containing two tasks
    start >> encrypt_and_archive >> end

Which creates the following graph:

Is something similar remotely doable before 2.0?

user_14624409 · Answer

You didn't explain what magic(t2, t3) is. TaskGroup is strictly UI feature it doesn't effect on the DAG logic. According to your description it seems that you are looking for a specific logic (otherwise what is magic?).

I believe this is what you are after:

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2021, 1, 24),
}
def generate_tasks():
    operator_list =[]
    for i in range(5): # Replace to generate the logic you wish to dynamically create tasks
        op = DummyOperator(task_id=f"t{str(i)}_task", dag=dag)
        if i>0:
            operator_list[i - 1] >> op
        operator_list.append(op)
    return operator_list

with DAG(
    dag_id='loop',
    default_args=default_args,
    schedule_interval=None,
) as dag:
    start_op = DummyOperator(task_id='start_task')
    end_op = DummyOperator(task_id='end_task')
    tasks = generate_tasks()
    start_op >> tasks[0]
    tasks[-1] >> end_op

You can replace the DummyOperator with any operator you'd like.

Return list of tasks from function that should be run in sequence in Airflow

1 Answers