Airflow Task Grouping과 Dynamic Dags 알아보기 - TIL230621

📚 KDT WEEK 12 DAY 3 TIL

• Task Grouping
• Dynamic Dags

 

---

 

📦 Task Grouping

 

Task 수가 많은 DAG라면 Task들을 성격에 따라 관리할 수 있음

Airflow 2.0에서부터는 SubDAG 대신 Task Grouping 사용

 

ex) start > 파일 1, 2, 3 다운로드 Task > 파일 1, 2, 3 Process Task > end 구조로 DAG를 작성했다.

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
    # Task Group 안에 Task Group 넣기
    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')

	# Task Group으로 순서 작성 가능
    start >> section_1 >> section_2 >> end

 

코드를 보면, Task Group안에서 Task들의 순서를 작성할 수도 있다.

 

이 DAG의 구조를 Airflow 웹 UI에서 확인해보자.

 

 

Task Grouping 1

 

파란 박스를 클릭해보면 안에 들어있는 Task들을 확인할 수 있다.

 

Task Grouping 2

 

---

 

☸️ Dynamic Dags

Template 형태(비슷한 형태)로 Dag를 많이 만들어내는 경우, 코드 재사용을 최대화할 수 있는 방법

 

Jinja 템플릿과 YAML 기반으로 DAG를 동적으로 만들기

Jinja가 DAG 자체 템플릿을 디자인하고, YAML을 통해 템플릿에 파라미터를 제공한다.

 

---

 

ex) 주식 심볼에 따라 간단하게 DAG 만들기

 

1. templated_dag.jinja2

 

jinja2 포맷으로 DAG 형태의 템플릿을 만들었다.

{{ }}의 형식으로 YAML 파일안의 값을 불러오게 된다.

 

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: # yml의 catchup 값을 가져오되 없으면 True로 넣기

    @task
    def extract(symbol):
        return symbol

    @task
    def process(symbol):
        return symbol

    @task
    def store(symbol):
        return symbol

    store(process(extract("{{ symbol }}")))

 

2. config_appl.yml

 

jinja2 템플릿안에 들어갈 값을 yml 파일에 각각 정의한다.

여기에 정의된 key : value 형태를 통해 key의 값을 jinja2 템플릿에서 불러오게 된다.

거기에 해당하는 값이 채워지는 형식으로, 아래의 generator.py를 통해 DAG 파이썬 파일로 만들어진다.

 

dag_id: 'APPL'
schedule: '@daily'
catchup: False
symbol: 'APPL'

 

3. generator.py

 

jinja2 템플릿에 yml 파일의 값을 채워넣고, 파이썬 파일(get_price_APPL.py)로 만들어지는 파이썬 파일이다.

 

from jinja2 import Environment, FileSystemLoader
import yaml
import os

file_dir = os.path.dirname(os.path.abspath(__file__)) # 현재 파일의 절대경로

# 생성했던 jinja2 템플릿 가져오기
env = Environment(loader=FileSystemLoader(file_dir)) 
template = env.get_template('templated_dag.jinja2')

# 현재 파일과 동일한 위치에서 yml 파일을 찾아 특정 위치에 DAG 파이썬 파일 생성
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))

 

generator.py는 airflow 폴더의 루트 위치에서 실행되는 것을 기준으로 작성되었다.

그리고 generator.py와 yml파일, jinja2 템플릿은 같은 폴더에 있어야 한다.

 

generator.py를 루트위치에서 실행시키면, get_price_APPL.py 파일이 정상적으로 생성된다.

 

결과 ) get_price_APPL.py

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

with DAG(dag_id="get_price_APPL",
    start_date=datetime(2023, 6, 15),
    schedule='@daily',
    catchup=True) as dag:

    @task
    def extract(symbol):
        return symbol

    @task
    def process(symbol):
        return symbol

    @task
    def store(symbol):
        return symbol

    store(process(extract("APPL")))