TaskFlow API로 파이써닉한 DAG 작성하기

TaskFlow API로 파이써닉한 DAG 작성하기 (Pythonic Dags with the TaskFlow API)

첫 튜토리얼에서는 BashOperator 같은 전통적인 Operators로 첫 Airflow DAG을 만들었어요. 이 튜토리얼에서는 Airflow 2.0에서 도입된 TaskFlow API로 워크플로우를 작성하는 더 현대적이고 파이써닉한 방법을 살펴봅니다.

출처: 문서

본문

TaskFlow API는 코드를 더 간단하고 깔끔하며 유지보수하기 쉽게 만들기 위해 설계되었어요. 일반 Python 함수를 작성하고 데코레이트하면, Airflow가 나머지 — task 생성, 의존성 배선, task 간 데이터 전달 —를 처리해요.

이 튜토리얼에서는 TaskFlow API로 간단한 ETL 파이프라인 — Extract → Transform → Load —을 만들 거예요. 시작해 볼까요!

큰 그림: TaskFlow 파이프라인

TaskFlow를 사용한 전체 파이프라인은 다음과 같아요. 다소 낯설어 보여도 걱정하지 마세요 — 단계별로 쪼개서 설명할게요.

airflow/example_dags/tutorial_taskflow_api.py[source]

import json

import pendulum

from airflow.sdk import dag, task
@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
)
def tutorial_taskflow_api():
    """
    ### TaskFlow API Tutorial Documentation
    This is a simple data pipeline example which demonstrates the use of
    the TaskFlow API using three simple tasks for Extract, Transform, and Load.
    Documentation that goes along with the Airflow TaskFlow API tutorial is
    located
    [here](https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html)
    """
    @task()
    def extract():
        """
        #### Extract task
        A simple Extract task to get data ready for the rest of the data
        pipeline. In this case, getting data is simulated by reading from a
        hardcoded JSON string.
        """
        data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'

        order_data_dict = json.loads(data_string)
        return order_data_dict
    @task(multiple_outputs=True)
    def transform(order_data_dict: dict):
        """
        #### Transform task
        A simple Transform task which takes in the collection of order data and
        computes the total order value.
        """
        total_order_value = 0

        for value in order_data_dict.values():
            total_order_value += value

        return {"total_order_value": total_order_value}
    @task()
    def load(total_order_value: float):
        """
        #### Load task
        A simple Load task which takes in the result of the Transform task and
        instead of saving it to end user review, just prints it out.
        """

        print(f"Total order value is: {total_order_value:.2f}")
    order_data = extract()
    order_summary = transform(order_data)
    load(order_summary["total_order_value"])
tutorial_taskflow_api()

1단계: DAG 정의하기

이전과 마찬가지로 DAG은 Airflow가 로드하고 파싱하는 Python 스크립트예요. 하지만 이번에는 @dag 데코레이터로 정의해요.

airflow/example_dags/tutorial_taskflow_api.py[source]

@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
)
def tutorial_taskflow_api():
    """
    ### TaskFlow API Tutorial Documentation
    This is a simple data pipeline example which demonstrates the use of
    the TaskFlow API using three simple tasks for Extract, Transform, and Load.
    Documentation that goes along with the Airflow TaskFlow API tutorial is
    located
    [here](https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html)
    """

이 DAG이 Airflow에서 발견되도록 하려면 @dag로 데코레이트된 Python 함수를 호출할 수 있어요:

airflow/example_dags/tutorial_taskflow_api.py[source]

tutorial_taskflow_api()

버전 2.4 변경: @dag 데코레이터를 사용하거나 with 블록에서 DAG을 정의한다면, 더 이상 전역 변수에 할당할 필요가 없어요. Airflow가 자동으로 찾아요.

Airflow UI에서 DAG을 시각화할 수 있어요! DAG이 로드되면 Graph View로 이동해 task들이 어떻게 연결되는지 볼 수 있어요.

2단계: `@task`로 Task 작성하기

TaskFlow에서 각 task는 일반 Python 함수일 뿐이에요. @task 데코레이터를 사용해 Airflow가 스케줄하고 실행할 수 있는 task로 만들 수 있어요. extract task는 다음과 같아요:

airflow/example_dags/tutorial_taskflow_api.py[source]

@task()
def extract():
    """
    #### Extract task
    A simple Extract task to get data ready for the rest of the data
    pipeline. In this case, getting data is simulated by reading from a
    hardcoded JSON string.
    """
    data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'

    order_data_dict = json.loads(data_string)
    return order_data_dict

함수의 반환 값이 다운스트림 task에 전달돼요 — 수동으로 XComs를 사용할 필요가 없어요. 내부적으로 TaskFlow는 XComs를 사용해 데이터 전달을 자동으로 관리하며, 이전 방식의 수동 XCom 관리 복잡성을 추상화해요. transform과 load task도 같은 패턴으로 정의할 거예요.

위의 @task(multiple_outputs=True) 사용에 주목하세요 — 이것은 함수가 개별 XCom으로 분할되어야 하는 값의 딕셔너리를 반환한다는 것을 Airflow에 알려줘요. 반환된 딕셔너리의 각 키가 자체 XCom 항목이 되어, 다운스트림 task에서 특정 값들을 참조하기 쉽게 해줘요. multiple_outputs=True를 생략하면 전체 딕셔너리가 단일 XCom으로 저장되고 전체로만 접근해야 해요.

3단계: 흐름 구축하기

task들이 정의되면, Python 함수처럼 호출해 파이프라인을 만들 수 있어요. Airflow는 이 함수 호출을 사용해 task 의존성을 설정하고 데이터 전달을 관리해요.

airflow/example_dags/tutorial_taskflow_api.py[source]

order_data = extract()
order_summary = transform(order_data)
load(order_summary["total_order_value"])

이게 전부예요! Airflow는 이 코드만으로도 파이프라인을 스케줄하고 오케스트레이션하는 방법을 알아요.

DAG 실행하기

DAG을 활성화하고 트리거하려면:

Airflow UI로 이동해요.
목록에서 DAG을 찾아 토글을 클릭해 활성화해요.
"Trigger Dag" 버튼을 클릭해 수동으로 트리거하거나, 스케줄에 따라 실행되기를 기다릴 수 있어요.

뒤에서 무슨 일이 일어나나요?

Airflow 1.x를 사용했다면 아마 마법처럼 느껴질 거예요. 내부에서 무슨 일이 일어나는지 비교해 볼게요.

"구식 방식": 수동 배선과 XComs

TaskFlow API 이전에는 PythonOperator 같은 Operators를 사용하고 XComs로 task 사이에 수동으로 데이터를 전달해야 했어요.

전통적 접근 방식으로 같은 DAG이 어떻게 보였을지:

import json
import pendulum
from airflow.sdk import DAG
from airflow.providers.standard.operators.python import PythonOperator


def extract():
    # Old way: simulate extracting data from a JSON string
    data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'
    return json.loads(data_string)


def transform(ti):
    # Old way: manually pull from XCom
    order_data_dict = ti.xcom_pull(task_ids="extract")
    total_order_value = sum(order_data_dict.values())
    return {"total_order_value": total_order_value}


def load(ti):
    # Old way: manually pull from XCom
    total = ti.xcom_pull(task_ids="transform")["total_order_value"]
    print(f"Total order value is: {total:.2f}")


with DAG(
    dag_id="legacy_etl_pipeline",
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
) as dag:
    extract_task = PythonOperator(task_id="extract", python_callable=extract)
    transform_task = PythonOperator(task_id="transform", python_callable=transform)
    load_task = PythonOperator(task_id="load", python_callable=load)

    extract_task >> transform_task >> load_task

Note

이 버전은 TaskFlow API 예시와 같은 결과를 만들지만, XComs와 task 의존성을 명시적으로 관리해야 해요.

TaskFlow 방식

TaskFlow를 사용하면 이 모든 것이 자동으로 처리돼요.

airflow/example_dags/tutorial_taskflow_api.py[source]

import json

import pendulum

from airflow.sdk import dag, task
@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
)
def tutorial_taskflow_api():
    """
    ### TaskFlow API Tutorial Documentation
    This is a simple data pipeline example which demonstrates the use of
    the TaskFlow API using three simple tasks for Extract, Transform, and Load.
    Documentation that goes along with the Airflow TaskFlow API tutorial is
    located
    [here](https://airflow.apache.org/docs/apache-airflow/stable/tutorial_taskflow_api.html)
    """
    @task()
    def extract():
        """
        #### Extract task
        A simple Extract task to get data ready for the rest of the data
        pipeline. In this case, getting data is simulated by reading from a
        hardcoded JSON string.
        """
        data_string = '{"1001": 301.27, "1002": 433.21, "1003": 502.22}'

        order_data_dict = json.loads(data_string)
        return order_data_dict
    @task(multiple_outputs=True)
    def transform(order_data_dict: dict):
        """
        #### Transform task
        A simple Transform task which takes in the collection of order data and
        computes the total order value.
        """
        total_order_value = 0

        for value in order_data_dict.values():
            total_order_value += value

        return {"total_order_value": total_order_value}
    @task()
    def load(total_order_value: float):
        """
        #### Load task
        A simple Load task which takes in the result of the Transform task and
        instead of saving it to end user review, just prints it out.
        """

        print(f"Total order value is: {total_order_value:.2f}")
    order_data = extract()
    order_summary = transform(order_data)
    load(order_summary["total_order_value"])
tutorial_taskflow_api()

Airflow는 여전히 XComs를 사용하고 의존성 그래프를 만든다는 점을 기억하세요 — 단지 추상화되어 비즈니스 로직에 집중할 수 있을 뿐이에요.

XComs가 어떻게 동작하나요?

TaskFlow 반환 값은 자동으로 XComs로 저장돼요. 이 값들은 UI의 "XCom" 탭에서 검사할 수 있어요. 전통적인 operators의 경우 수동 xcom_pull()도 여전히 가능해요.

오류 처리와 재시도

데코레이터를 사용해 task의 재시도를 쉽게 구성할 수 있어요. 예를 들어 task 데코레이터에서 직접 최대 재시도 횟수를 설정할 수 있어요:

@task(retries=3)
def my_task(): ...

이렇게 하면 일시적인 실패가 task 실패로 이어지지 않도록 보장하는 데 도움이 돼요.

Task 파라미터화

데코레이트된 task를 여러 DAG에서 재사용하고 task_id나 retries 같은 파라미터를 재정의할 수 있어요.

start = add_task.override(task_id="start")(1, 2)

공유 모듈에서 데코레이트된 task를 가져올 수도 있어요.

다음에 살펴볼 내용

수고했어요! 이제 TaskFlow API로 첫 파이프라인을 작성했어요. 여기서 어디로 갈지 궁금한가요?

DAG에 새 task 추가하기 — 필터나 검증 단계를 넣어보세요
반환 값 수정하고 여러 출력 전달하기
.override(task_id="...")로 재시도와 재정의 탐색하기
Airflow UI를 열어 task 사이에서 데이터가 어떻게 흐르는지, task 로그와 의존성을 포함해 검사하기

참고 (See also)

다음 단계로: 간단한 데이터 파이프라인 만들기
TaskFlow API docs에서 더 배우거나, 아래의 고급 TaskFlow 패턴 계속
Core Concepts에서 Airflow 개념 읽기

고급 TaskFlow 패턴 (Advanced TaskFlow Patterns)

기본에 익숙해지면 시도해 볼 수 있는 몇 가지 강력한 기법이 있어요.

데코레이트된 Task 재사용하기

데코레이트된 task를 여러 DAG이나 Dag run에서 재사용할 수 있어요. 재사용 가능한 유틸리티나 공유 비즈니스 규칙 같은 공통 로직에 특히 유용해요. .override()를 사용해 task_id나 retries 같은 task 메타데이터를 커스터마이즈해요.

start = add_task.override(task_id="start")(1, 2)

공유 모듈에서 데코레이트된 task를 가져올 수도 있어요.

충돌하는 의존성 처리하기

때로 task는 DAG의 나머지와 다른 Python 의존성을 요구해요 — 예를 들어 특수 라이브러리나 시스템 수준 패키지 같은 것들이요. TaskFlow는 그런 의존성을 격리하기 위해 여러 실행 환경을 지원해요.

동적으로 생성된 Virtualenv

task 런타임에 임시 virtualenv를 만들어요. 실험적이거나 동적인 task에 좋지만 콜드 스타트 오버헤드가 있을 수 있어요.

/opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_python_decorator.py

@task.virtualenv(
    task_id="virtualenv_python", requirements=["colorama==0.4.0"], system_site_packages=False
)
def callable_virtualenv():
    """
    Example function that will be performed in a virtual environment.

    Importing at the module level ensures that it will not attempt to import the
    library before it is installed.
    """
    from time import sleep

    from colorama import Back, Fore, Style

    print(Fore.RED + "some red text")
    print(Back.GREEN + "and with a green background")
    print(Style.DIM + "and in dim text")
    print(Style.RESET_ALL)
    for _ in range(4):
        print(Style.DIM + "Please wait...", flush=True)
        sleep(1)
    print("Finished")

virtualenv_task = callable_virtualenv()

외부 Python 환경

사전 설치된 Python 인터프리터로 task를 실행해요 — 일관된 환경이나 공유 virtualenv에 이상적이에요.

/opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_python_decorator.py

@task.external_python(task_id="external_python", python=PATH_TO_PYTHON_BINARY)
def callable_external_python():
    """
    Example function that will be performed in a virtual environment.

    Importing at the module level ensures that it will not attempt to import the
    library before it is installed.
    """
    import sys
    from time import sleep

    print(f"Running task via {sys.executable}")
    print("Sleeping")
    for _ in range(4):
        print("Please wait...", flush=True)
        sleep(1)
    print("Finished")

external_python_task = callable_external_python()

Docker 환경

Docker 컨테이너에서 task를 실행해요. task에 필요한 모든 것을 패키징하는 데 유용하지만, 워커에 Docker가 설치되어 있어야 해요.

/opt/airflow/providers/docker/tests/system/docker/example_taskflow_api_docker_virtualenv.py

@task.docker(image="python:3.9-slim-bookworm", multiple_outputs=True)
def transform(order_data_dict: dict):
    """
    #### Transform task
    A simple Transform task which takes in the collection of order data and
    computes the total order value.
    """
    total_order_value = 0

    for value in order_data_dict.values():
        total_order_value += value

    return {"total_order_value": total_order_value}

Note

Airflow 2.2와 Docker provider가 필요해요.

KubernetesPodOperator

메인 Airflow 환경에서 완전히 격리된 Kubernetes pod 안에서 task를 실행해요. 대규모 task나 커스텀 런타임이 필요한 task에 이상적이에요.

/opt/airflow/providers/cncf/kubernetes/tests/system/cncf/kubernetes/example_kubernetes_decorator.py

@task.kubernetes(
    image="python:3.9-slim-buster",
    name="k8s_test",
    namespace="default",
    in_cluster=False,
    config_file="/path/to/.kube/config",
)
def execute_in_k8s_pod():
    import time

    print("Hello from k8s pod")
    time.sleep(2)

@task.kubernetes(image="python:3.9-slim-buster", namespace="default", in_cluster=False)
def print_pattern():
    n = 5
    for i in range(n):
        # inner loop to handle number of columns
        # values changing acc. to outer loop
        for _ in range(i + 1):
            # printing stars
            print("* ", end="")

        # ending line after each row
        print("\r")

execute_in_k8s_pod_instance = execute_in_k8s_pod()
print_pattern_instance = print_pattern()

execute_in_k8s_pod_instance >> print_pattern_instance

Note

Airflow 2.4와 Kubernetes provider가 필요해요.

Sensors 사용하기

@task.sensor를 사용해 Python 함수로 가볍고 재사용 가능한 sensors를 만들 수 있어요. 이들은 poke와 reschedule 모드를 모두 지원해요.

/opt/airflow/providers/standard/src/airflow/providers/standard/example_dags/example_sensor_decorator.py

import pendulum

from airflow.sdk import PokeReturnValue, dag, task
@dag(
    schedule=None,
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    catchup=False,
    tags=["example"],
)
def example_sensor_decorator():
    # Using a sensor operator to wait for the upstream data to be ready.
    @task.sensor(poke_interval=60, timeout=3600, mode="reschedule")
    def wait_for_upstream() -> PokeReturnValue:
        return PokeReturnValue(is_done=True, xcom_value="xcom_value")
    @task
    def dummy_operator() -> None:
        pass
    wait_for_upstream() >> dummy_operator()
tutorial_etl_dag = example_sensor_decorator()

전통적인 Task와 혼합하기

데코레이트된 task를 클래식 Operators와 결합할 수 있어요. 이는 커뮤니티 providers를 사용하거나 단계적으로 TaskFlow로 마이그레이션할 때 유용해요.

>>로 TaskFlow와 전통적인 task를 연결하거나 .output 속성으로 데이터를 전달할 수 있어요.

TaskFlow에서의 Templating

전통적인 task와 마찬가지로 데코레이트된 TaskFlow 함수는 템플릿 인자를 지원해요 — 파일에서 콘텐츠를 로드하거나 런타임 파라미터를 사용하는 것을 포함해요.

callable을 실행할 때 Airflow는 함수에서 사용할 수 있는 키워드 인자 세트를 전달해요. 이 kwarg 세트는 Jinja 템플릿에서 사용할 수 있는 것과 정확히 일치해요. 이를 위해 함수에서 수신하려는 컨텍스트 키를 키워드 인자로 추가할 수 있어요.

예를 들어 아래 코드 블록의 callable은 ti와 next_ds 컨텍스트 변수의 값을 받아요:

@task
def my_python_callable(*, ti, next_ds):
    pass

**kwargs로 전체 컨텍스트를 받기로 선택할 수도 있어요. 이렇게 하면 Airflow가 실제로 필요하지 않은 많은 것을 포함할 수 있는 전체 컨텍스트를 확장해야 하므로 약간의 성능 저하가 발생할 수 있다는 점을 유의하세요. 따라서 이전 문단에서 보여준 것처럼 명시적 인자를 사용하는 것이 더 권장돼요.

@task
def my_python_callable(**kwargs):
    ti = kwargs["ti"]
    next_ds = kwargs["next_ds"]

또한 때로는 스택 깊숙한 곳에서 컨텍스트에 접근하고 싶지만, task callable에서 컨텍스트 변수를 전달하고 싶지 않을 수 있어요. get_current_context 메서드를 통해 여전히 실행 컨텍스트에 접근할 수 있어요.

from airflow.sdk import get_current_context


def some_function_in_your_library():
    context = get_current_context()
    ti = context["ti"]

데코레이트된 함수에 전달된 인자는 자동으로 템플릿화돼요. templates_exts로 파일을 템플릿화할 수도 있어요:

@task(templates_exts=[".sql"])
def read_sql(sql): ...

조건부 실행 (Conditional Execution)

@task.run_if() 또는 @task.skip_if()를 사용해 DAG 구조를 변경하지 않고 런타임의 동적 조건에 따라 task를 실행할지 제어할 수 있어요.

@task.run_if(lambda ctx: ctx["task_instance"].task_id == "run")
@task.bash()
def echo():
    return "echo 'run'"

다음은 무엇일까요? (What's Next)

이제 TaskFlow API로 깔끔하고 유지보수 가능한 DAG을 만드는 방법을 봤으니, 여기 몇 가지 좋은 다음 단계가 있어요:

Asset-Aware Scheduling에서 asset 인식 워크플로우 살펴보기
Scheduling Options에서 스케줄링 패턴 탐구
다음 튜토리얼로: 간단한 데이터 파이프라인 만들기

TaskFlow API로 파이써닉한 DAG 작성하기