[Airflow] Task-Group & Dynamic Dags

태스크 그룹

태스크 그룹핑의 필요성

•

태스크 수가 많은 DAG의 경우, 태스크들을 성격에 따라 관리하고 싶은 요구가 존재합니다.

◦

SubDAG가 사용되던 중, Airflow 2.0에서는 Task Grouping이 더 많이 사용되고 있습니다.

▪

SubDAG는 비슷한 일을 하는 태스크들을 SubDAG라는 Child DAG로 만들어서 관리합니다.

▪

태스크 그룹으로 묶이는 것들은 비슷한 성격의 일을 수행하는 DAG일 가능성이 높다

•

다수의 파일을 처리하는 DAG의 경우,

◦

파일 다운로드 태스크들과 파일 체크 태스크와 데이터 처리 태스크들로 구성됩니다.

•

소스 코드

◦

Learn Task Groups

from airflow.utils.task_group import TaskGroup

	start = EmptyOperator(task_id="start")

	with TaskGroup("Download", tooltip="Tasks for downloading data") as section_1:
	
		task_1 = EmptyOperator(task_id="task_1")
		task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
		task_3 = EmptyOperator(task_id="task_3")

	task_1 >> [task_2, task_3]
start >> section_1
Python
복사

▪

TaskGroup 안에 TaskGroup nesting 가능

▪

TaskGroup도 태스크처럼 실행 순서 정의 가능

▪

tooltip으로 TaskGroup 네이밍이 가능해진다.

▪

/Learn_TaskGroups.py

from airflow.models.dag import DAG
from airflow.operators.empty import EmptyOperator
from airflow.operators.bash import BashOperator
from airflow.utils.task_group import TaskGroup
import pendulum

with DAG(dag_id="Learn_Task_Group", start_date=pendulum.today('UTC').add(days=-2), tags=["example"]) as dag:
    start = EmptyOperator(task_id="start")

    # Task Group #1
    with TaskGroup("Download", tooltip="Tasks for downloading data") as section_1:
        task_1 = EmptyOperator(task_id="task_1")
        task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
        task_3 = EmptyOperator(task_id="task_3")

        task_1 >> [task_2, task_3]

    # Task Group #2
    with TaskGroup("Process", tooltip="Tasks for processing data") as section_2:
        task_1 = EmptyOperator(task_id="task_1")

        with TaskGroup("inner_section_2", tooltip="Tasks for inner_section2") as inner_section_2:
            task_2 = BashOperator(task_id="task_2", bash_command='echo 1')
            task_3 = EmptyOperator(task_id="task_3")
            task_4 = EmptyOperator(task_id="task_4")

            [task_2, task_3] >> task_4

    end = EmptyOperator(task_id='end')

    start >> section_1 >> section_2 >> end
Python
복사

Dynamic Dag란 무엇인가?

Dynamic Dags를 사용해서 코드 재사용을 최대화할 수 있다! Jinja template와 YAML을 사용해야 함

•

Dynamic Dag란 무엇인가?

◦

템플릿과 YAML을 기반으로 DAG를 동적으로 만들어보자

▪

Jinja를 기반으로 DAG 자체의 템플릿을 디자인하고 YAML을 통해 앞서 만든 템플릿에 파라미터를 제공

◦

이를 통해 비슷한 DAG를 계속해서 매뉴얼하게 개발하는 것을 방지

◦

DAG를 계속해서 만드는 것과 한 DAG안에서 태스크를 늘리는 것 사이의 밸런스 필요

▪

오너가 다르거나 태스크의 수가 너무 커지는 경우 DAG를 복제해나가는 것이 더 좋음

•

Dynamic Dag의 기본적인 아이디어

•

바뀌어야 하는 파라미터만 jinja 템플릿(templated_dag.jinja2)을 통해 바꿈

•

generator.py를 통해서 템플릿을 통해 Dag 파일을 생성해준다.

•

간단한 예제

◦

템플릿을 통한 최종 DAG 파일 생성

▪

$ python3 dags/dynamic_dags/generator.py

◦

이는 dags 폴더에 yml 파일의 수 만큼의 DAG 코드를 생성해줌

▪

generator 실행을 언제할지 결정이 필요

▪

generator를 통해 아래 get_price_XXXX.py DAG를 생성

▪

generator.py

from jinja2 import Environment, FileSystemLoader
import yaml
import os

file_dir = os.path.dirname(os.path.abspath(__file__))
env = Environment(loader=FileSystemLoader(file_dir))
template = env.get_template('templated_dag.jinja2')

for f in os.listdir(file_dir):
    if f.endswith(".yml"):
        with open(f"{file_dir}/{f}", "r") as cf:
            config = yaml.safe_load(cf)
            with open(f"dags/get_price_{config['dag_id']}.py", "w") as f:
                f.write(template.render(config))
Python
복사

▪

templated_dag.jinja2

from airflow import DAG
from airflow.decorators import task
from datetime import datetime

with DAG(dag_id="get_price_{{ dag_id }}",
    start_date=datetime(2023, 6, 15),
    schedule='{{ schedule }}',
    catchup={{ catchup or True }}) as dag:

    @task
    def extract(symbol):
        return symbol

    @task
    def process(symbol):
        return symbol

    @task
    def store(symbol):
        return symbol

    store(process(extract("{{ symbol }}")))
Python
복사

▪

config_appl.yml

dag_id: 'APPL'
schedule: '@daily'
catchup: False
symbol: 'APPL'
Python
복사

▪

config_goog.yml

dag_id: 'GOOG'
schedule: '@weekly'
symbol: 'GOOG'
Python
복사

catch up이 없으면 True로 jinja에서 세팅해준다.