MLOps/Airflow

[Airflow] @task 데코레이터 사용하기

monkeykim 2024. 10. 31. 01:27

Airflow에서는 파이썬 함수를 간단히 데코레이터를 통해 task로 변환할 수 있는 @task 데코레이터를 제공합니다. 이를 통해 PythonOperator를 직접 정의하는 것보다 효율적이고 간결하게 DAG 내의 task를 작성할 수 있습니다.


데코레이터(Decorator)란?

데코레이터는 함수를 '감싸서' 기능을 추가하는 파이썬 기능입니다. 함수를 인자로 전달하거나 함수 내부에 함수를 정의하는 것이 가능하기 때문에, 함수를 쉽게 확장할 수 있습니다. @데코레이터명으로 사용하며, Airflow에서는 이를 활용해 task를 정의합니다.

 

Airflow에서의 @task 데코레이터 활용하기

Airflow의 @task 데코레이터를 사용하면, 파이썬 함수 정의만으로 손쉽게 task를 생성할 수 있습니다. 이는 Airflow 공식 문서에서도 권장하는 방식으로, 코드 가독성을 높이고 관리가 간편해지는 장점이 있습니다.

 

@task를 활용한 DAG 작성 예시

from airflow import DAG
from airflow.decorators import task
import pendulum

with DAG(
    dag_id="dag_python_decorator",
    schedule="30 6 * * *",
    start_date=pendulum.datetime(2024, 3, 1, tz="Asia/Seoul"),
    catchup=False
) as dag:
    @task(task_id="python_task_1")
    def print_context(arg):
        print(arg)

    python_task_1 = print_context("############ task decorator 실행 ############")

코드 설명

  • DAG 설정: dag_python_decorator라는 DAG은 매일 오전 6시 30분에 실행됩니다.
  • @task 데코레이터: print_context 함수 위에 @task(task_id="python_task_1")를 선언하여 함수가 Airflow의 task로서 동작하도록 합니다.
  • Task 실행: print_context 함수에 인자를 전달해 task가 실행될 때 "############ task decorator 실행 ############" 메시지가 출력됩니다.

정리

@task 데코레이터를 사용하면 PythonOperator를 직접 정의할 필요 없이 간단하게 Airflow task를 생성할 수 있어서 DAG 파일이 훨씬 깔끔해집니다.