← 학습 카테고리

Learn

Airflow

151개 모듈 · 현재 18번째

Airflow 모듈 18/151 airflow-learn-18

TaskFlow API — @task로 Dag 작성하기

Apache Airflow Official Documentation (in-repo snapshot) — Apache Software Foundation core-concepts/taskflow.rst

이 모듈을 다 읽으면

  • TaskFlow가 XCom과 다운스트림 의존성을 자동으로 처리하는 방식을 설명할 수 있다
  • TaskFlow 함수 인자로 임의 객체를 전달하기 위한 직렬화 요구사항(dataclass/attr.define, 커스텀 serialize/deserialize)을 설명할 수 있다
  • TaskFlow가 PythonOperator에 비해 어떤 문제(보일러플레이트)를 줄여주는지 설명할 수 있다

TaskFlow API(Airflow 2.0에서 추가)는 ``@task`` 데코레이터를 통해 순수 Python 코드로 Dag를 깔끔하게 작성할 수 있게 해준다. XCom을 이용한 입출력 전달과 의존성 계산을 자동화하며, 컨텍스트 변수 접근, 로깅, 임의 객체 전달을 위한 직렬화 규칙을 함께 다룬다.

XComArg와 자동 의존성 계산

Operator보다 순수 Python 코드로 대부분의 Dag를 작성한다면, TaskFlow API가 불필요한 보일러플레이트 없이 깔끔한 Dag를 작성하게 해준다. TaskFlow는 XCom을 이용해 태스크 간 입출력을 자동으로 옮겨주며, 의존성도 자동으로 계산한다 — Dag 파일 안에서 TaskFlow 함수를 호출하면, 그 함수가 즉시 실행되는 것이 아니라 결과의 XCom을 나타내는 객체(``XComArg``)를 돌려받게 되고, 이를 다운스트림 태스크나 오퍼레이터의 입력으로 그대로 쓸 수 있다.

예를 들어 ``get_ip``라는 ``@task``가 IP를 반환하고, ``compose_email``이라는 ``@task(multiple_outputs=True)``가 그 IP를 받아 제목·본문 딕셔너리를 반환하며, 마지막으로 전통적인 ``EmailOperator``가 ``compose_email``의 반환값을 파라미터로 쓴다고 하자. 이 경우 ``compose_email``은 ``get_ip``의 반환값을 자동으로 전달받으면서 동시에 ``get_ip``의 다운스트림으로 자동 선언되고, ``send_email_notification``도 ``compose_email``의 반환값을 쓰는 순간 자동으로 그 다운스트림이 된다. 즉 TaskFlow 함수든 전통적 오퍼레이터든, 다른 태스크의 반환값을 파라미터로 사용하는 것만으로 XCom 전달과 의존성 선언이 동시에 이루어진다.

일반 값이나 변수를 TaskFlow 함수 호출에 그대로 전달할 수도 있다(물론 Dag가 실제로 실행되기 전까지는 코드가 실행되지 않고, 그 값은 Dag가 실행될 때까지 태스크 파라미터로 유지된다).

핵심 포인트

  • TaskFlow 함수를 Dag 파일에서 호출하면 즉시 실행되는 대신 XComArg(결과의 XCom을 나타내는 객체)를 반환하며, 이를 다운스트림에 넘기면 XCom 전달과 의존성 선언이 자동으로 함께 이루어진다
  • TaskFlow 함수와 전통적 Operator는 서로 자유롭게 조합할 수 있으며, 어느 쪽이든 다른 태스크의 반환값을 파라미터로 쓰면 자동으로 다운스트림이 된다

컨텍스트, 로깅, 임의 객체 전달

Airflow 컨텍스트 변수는 태스크 함수의 키워드 인자로 추가하면 접근할 수 있고, 또는 함수 시그니처에 ``**kwargs``를 추가하면 모든 컨텍스트 변수를 ``kwargs`` 딕셔너리 안에서 접근할 수 있다. 태스크 함수 안에서 로깅을 하려면 Python의 표준 logging 시스템을 그대로 쓰면 되며(``logger = logging.getLogger("airflow.task")``), 이렇게 만든 모든 로그 라인은 태스크 로그에 기록된다.

TaskFlow는 XCom으로 태스크 간에 변수를 전달하므로, 인자로 쓰이는 변수는 직렬화 가능해야 한다. Airflow는 기본적으로 int, str 같은 내장 타입은 물론, ``@dataclass``나 ``@attr.define``으로 데코레이트된 객체도 지원한다(Airflow 2.5.0에서 추가). ``Asset``(``@attr.define``으로 정의됨)을 입력 인자로 쓰면 자동으로 inlet으로 등록되고, 태스크의 반환값이 ``Asset``이나 ``list[Asset]``이면 자동으로 outlet으로 등록된다는 부가 이점도 있다.

커스텀 객체를 전달하고 싶다면 보통 ``@dataclass``나 ``@attr.define``으로 클래스를 데코레이트하면 Airflow가 알아서 처리하지만, 직렬화를 직접 제어하고 싶다면 클래스에 ``serialize()`` 메서드와 정적 메서드 ``deserialize(data: dict, version: int)``를 추가하면 된다. 직렬화되는 객체는 버전을 관리하는 것이 좋은 관행이며, 이를 위해 클래스에 ``__version__: ClassVar[int] = <x>``를 추가한다. Airflow는 클래스가 하위 호환된다고 가정하므로(버전 2가 버전 1도 역직렬화할 수 있어야 함), 커스텀 역직렬화 로직이 필요하다면 ``deserialize(data: dict, version: int)``에서 버전별로 분기하면 된다.

핵심 포인트

  • 태스크 함수는 키워드 인자 또는 **kwargs로 Airflow 컨텍스트 변수에 접근할 수 있고, 표준 logging.getLogger("airflow.task")로 태스크 로그에 기록할 수 있다
  • TaskFlow 인자로 임의 객체를 전달하려면 직렬화 가능해야 하며, @dataclass/@attr.define 데코레이터나 커스텀 serialize()/deserialize() 메서드로 이를 지원할 수 있다
  • Asset을 인자나 반환값으로 쓰면 자동으로 inlet/outlet으로 등록되며, 커스텀 직렬화 객체는 __version__: ClassVar[int]로 버전을 관리하는 것이 권장된다