동적 태스크 매핑
동적 태스크 매핑 (Dynamic Task Mapping)
DAG을 짤 때 "데이터가 몇 개 올지는 모르지만, 올 때마다 하나씩 처리하고 싶어" 하는 경우가 참 많아요. 파일이 10개 올지 100개 올지는 그때그때 달라지는데, DAG 작성자가 미리 태스크를 정해 둘 수는 없잖아요. 이럴 때 쓰는 기능이 바로 동적 태스크 매핑이에요. 실행 시점에 상위 태스크의 결과를 보고 태스크 개수를 알아서 늘려 주니까, 데이터가 얼마나 오는지 몰라도 되는 거죠.
본문
동적 태스크 매핑은 DAG 파일에서 for 루프로 태스크를 만드는 것과 비슷해 보이지만 결정적인 차이가 있어요. for 루프는 DAG 파싱 시점에 목록을 직접 알아야 하지만, 매핑은 스케줄러가 상위 태스크의 출력을 기준으로 태스크 인스턴스의 개수를 실행 시점에 정해요. 매핑된 태스크가 실행되기 직전에 입력 하나당 태스크 복사본 하나씩, 그러니까 입력 개수만큼 태스크를 만들어 주는 구조예요. 매핑된 태스크들의 결과를 모아서 처리하는 리듀스 역할도 가능하죠.
단순 매핑
가장 간단한 형태는 DAG 파일에 정의된 리스트를 대상으로 expand()를 쓰는 거예요. 태스크를 직접 호출하는 대신 expand()로 감싸면 돼요.
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Example DAG demonstrating the usage of dynamic task mapping."""
from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, task, task_group
with DAG(
dag_id="example_dynamic_task_mapping", schedule=None, start_date=datetime(2022, 3, 4), tags=["example"]
):
@task
def add_one(x: int):
return x + 1
@task
def sum_it(values):
total = sum(values)
print(f"Total was {total}")
added_values = add_one.expand(x=[1, 2, 3])
sum_it(added_values)
with DAG(
dag_id="example_task_mapping_second_order",
schedule=None,
catchup=False,
start_date=datetime(2022, 3, 4),
tags=["example"],
):
@task
def get_nums():
return [1, 2, 3]
@task
def times_2(num):
return num * 2
@task
def add_10(num):
return num + 10
_get_nums = get_nums()
_times_2 = times_2.expand(num=_get_nums)
add_10.expand(num=_times_2)
with DAG(
dag_id="example_task_group_mapping",
schedule=None,
catchup=False,
start_date=datetime(2022, 3, 4),
tags=["example"],
):
@task_group
def op(num):
@task
def add_1(num):
return num + 1
@task
def mul_2(num):
return num * 2
return mul_2(add_1(num))
op.expand(num=[1, 2, 3])
이 DAG을 실행하면 태스크 로그에 Total was 9가 찍혀요. 그리드 뷰(gird view)에서도 매핑된 태스크가 세부 정보 패널에 잘 드러나요.
여기서 add_one의 매핑 결과 3개가 모여 sum_it의 values로 들어가는 걸 볼 수 있는데요, 주의할 점이 하나 있어요. values는 보통 리스트가 아니라 lazy proxy예요. 매핑된 태스크 인스턴스가 몇 개나 될지 미리 알 수 없기 때문에, 값이 필요해질 때 하나씩 가져오는 지연 시퀀스인 거죠. 그래서 print(values)를 하면 LazySelectSequence([15 items])처럼 보여요.
LazySelectSequence([15 items])
이 객체에는 일반 시퀀스 문법(values[0])이나 for 루프를 그대로 쓸 수 있어요. list(values)로 진짜 리스트를 얻을 수도 있는데, 그러면 참조된 모든 상위 매핑 태스크의 값을 한 번에 eager하게 불러오므로 매핑 수가 크다면 성능 영향을 염두에 둬야 해요. 이 lazy proxy를 XCom으로 넘길 때도 Airflow가 자동으로 값을 강제 변환하긴 하지만 경고를 뱉어요.
@task
def forward_values(values):
return values # This is a lazy proxy!
이 경우 이런 경고가 나와요.
Coercing mapped lazy proxy return value from task forward_values to list, which may degrade
performance. Review resource requirements for this operation, and call list() explicitly to suppress this message. See Dynamic Task Mapping documentation for more information about lazy proxy objects.
경고를 없애고 싶다면 태스크 안에서 명시적으로 list(values)를 호출하면 돼요.
@task
def forward_values(values):
return list(values)
꼭 리듀스 태스크가 있어야 하는 건 아니에요. 아래처럼 매핑된 태스크만 두고 그 후속 태스크가 없어도 매핑 태스크는 정상 실행돼요.
태스크가 생성하는 매핑 (Task-generated Mapping)
지금까지 예시는 DAG 파일에 리스트를 적어 두고 매핑했는데, 동적 태스크 매핑의 진짜 힘은 태스크가 그 리스트를 만들어 내는 데 있어요. API 호출이든 DB 조회든, 결과를 리스트나 딕셔너리로 만들어 XCom 백엔드에 저장할 수만 있으면 뭐든 반복 대상으로 쓸 수 있어요.
@task
def make_list():
# This can also be from an API call, checking a database, -- almost anything you like, as long as the
# resulting list/dictionary can be stored in the current XCom backend.
return [1, 2, {"a": "b"}, "str"]
@task
def consumer(arg):
print(arg)
with DAG(dag_id="dynamic-map", start_date=datetime(2022, 4, 2)) as dag:
consumer.expand(arg=make_list())
make_list는 일반 태스크로 먼저 실행되고 [1, 2, {"a":"b"}, "str"]를 돌려줘요. 그러면 consumer가 이 값 4개 각각에 대해 한 번씩, 총 네 번 호출돼요.
한 가지 제약이 있어요. 태스크가 생성하는 매핑에서는 TriggerRule.ALWAYS를 쓸 수 없어요. 확장할 파라미터가 태스크 실행 시점에 정의되지 않기 때문인데, 이 제약은 DAG 파싱 시점에 적용되어 위반하면 에러가 나요.
반복 매핑 (Repeated mapping)
한 매핑 태스크의 결과를 또 다른 매핑 태스크의 입력으로 쓸 수도 있어요.
with DAG(dag_id="repeated_mapping", start_date=datetime(2022, 3, 4)) as dag:
@task
def add_one(x: int):
return x + 1
first = add_one.expand(x=[1, 2, 3])
second = add_one.expand(x=first)
이 경우 second의 결과는 [3, 4, 5]가 돼요.
여러 파라미터 매핑
expand에 파라미터를 여러 개 넘기면 각 조합에 대해 태스크를 호출하는 **교차곱(cross product)**이 만들어져요.
@task
def add(x: int, y: int):
return x + y
added_values = add.expand(x=[2, 4, 8], y=[5, 10])
# This results in the add function being called with
# add(x=2, y=5)
# add(x=2, y=10)
# add(x=4, y=5)
# add(x=4, y=10)
# add(x=8, y=5)
# add(x=8, y=10)
add는 총 6번 호출돼요. 다만 확장 순서는 보장되지 않아요. 만약 x에만 값을 매핑하고 y는 고정하고 싶다면 partial()을 쓰면 돼요.
@task
def add(x: int, y: int):
return x + y
added_values = add.partial(y=10).expand(x=[1, 2, 3])
# This results in add function being expanded to
# add(x=1, y=10)
# add(x=2, y=10)
# add(x=3, y=10)
네임드 매핑 (Named mapping)
기본적으로 매핑된 태스크에는 정수 인덱스가 붙어요. map_index_template으로 Jinja 템플릿을 주면, 태스크 입력에 기반한 이름으로 인덱스를 대체할 수 있어요. 확장이 .expand(<property>=...) 형태라면 보통 map_index_template="{{task.<property>}}"처럼 쓰죠. 이 템플릿은 각 확장 태스크가 실행된 뒤 태스크 컨텍스트를 이용해 렌더링돼요.
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
# The two expanded task instances will be named "2024-01-01" and "2024-01-02".
SQLExecuteQueryOperator.partial(
...,
sql="SELECT * FROM data WHERE date = %(date)s",
map_index_template="""{{ task.parameters['date'] }}""",
).expand(
parameters=[{"date": "2024-01-01"}, {"date": "2024-01-02"}],
)
위 예시에서 확장된 태스크 인스턴스는 각각 2024-01-01, 2024-01-02로 이름 붙어서 UI에 인덱스 0, 1 대신 표시돼요.
템플릿은 메인 실행 블록 뒤에 렌더링되므로, 렌더링 컨텍스트에 값을 동적으로 주입하는 것도 가능해요. Jinja 문법으로 원하는 이름 표현이 어려울 때, 특히 TaskFlow 함수 안에서 유용하죠.
from airflow.sdk import get_current_context
@task(map_index_template="{{ my_variable }}")
def my_task(my_value: str):
context = get_current_context()
context["my_variable"] = my_value * 3
... # Normal execution...
# The task instances will be named "aaa" and "bbb".
my_task.expand(my_value=["a", "b"])
TaskFlow 아닌 오퍼레이터와의 매핑
클래식 스타일 오퍼레이터에서도 partial과 expand를 쓸 수 있어요. 단, task_id·queue·pool처럼 BaseOperator에 속한 대부분의 인자는 매핑할 수 없고 반드시 partial()에 넘겨야 해요.
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Example DAG demonstrating the usage of dynamic task mapping with non-TaskFlow operators."""
from __future__ import annotations
from datetime import datetime
from airflow.sdk import DAG, BaseOperator
class AddOneOperator(BaseOperator):
"""A custom operator that adds one to the input."""
def __init__(self, value, **kwargs):
super().__init__(**kwargs)
self.value = value
def execute(self, context):
return self.value + 1
class SumItOperator(BaseOperator):
"""A custom operator that sums the input."""
template_fields = ("values",)
def __init__(self, values, **kwargs):
super().__init__(**kwargs)
self.values = values
def execute(self, context):
total = sum(self.values)
print(f"Total was {total}")
return total
with DAG(
dag_id="example_dynamic_task_mapping_with_no_taskflow_operators",
schedule=None,
start_date=datetime(2022, 3, 4),
catchup=False,
tags=["example"],
):
# map the task to a list of values
add_one_task = AddOneOperator.partial(task_id="add_one").expand(value=[1, 2, 3])
# aggregate (reduce) the mapped tasks results
sum_it_task = SumItOperator(task_id="sum_it", values=add_one_task.output)
클래식 오퍼레이터의 결과를 매핑하려면, 오퍼레이터 자체가 아니라 output을 명시적으로 참조해야 해요.
# Create a list of data inputs.
extract = ExtractOperator(task_id="extract")
# Expand the operator to transform each input.
transform = TransformOperator.partial(task_id="transform").expand(input=extract.output)
# Collect the transformed inputs, expand the operator to load each one of them to the target.
load = LoadOperator.partial(task_id="load").expand(input=transform.output)
expand_kwargs로 여러 인자 지정
상위 태스크가 하위 오퍼레이터에 여러 인자를 한꺼번에 지정해야 한다면 expand_kwargs를 써요. 매핑 대상이 되는 매핑(딕셔너리) 시퀀스를 받아요.
BashOperator.partial(task_id="bash").expand_kwargs(
[
{"bash_command": "echo $ENV1", "env": {"ENV1": "1"}},
{"bash_command": "printf $ENV2", "env": {"ENV2": "2"}},
],
)
이 코드는 런타임에 각각 1과 2를 출력하는 태스크 인스턴스 두 개를 만들어요. op_kwargs처럼 대부분의 오퍼레이터 인자와 섞어 쓸 수도 있어요.
def print_args(x, y):
print(x)
print(y)
return x + y
PythonOperator.partial(task_id="task-1", python_callable=print_args).expand_kwargs(
[
{"op_kwargs": {"x": 1, "y": 2}, "show_return_value_in_logs": True},
{"op_kwargs": {"x": 3, "y": 4}, "show_return_value_in_logs": False},
]
)
expand와 마찬가지로 expand_kwargs도 딕셔너리 리스트를 돌려주는 XCom, 또는 딕셔너리를 돌려주는 XCom의 리스트를 대상으로 매핑할 수 있어요. 예를 들어 파일 확장자에 따라 서로 다른 버킷으로 복사하는 "브랜칭" 로직을 매핑 태스크로 표현할 수 있죠.
list_filenames = S3ListOperator(...) # Same as the above example.
@task
def create_copy_kwargs(filename):
if filename.rsplit(".", 1)[-1] not in ("json", "yml"):
dest_bucket_name = "my_text_bucket"
else:
dest_bucket_name = "my_other_bucket"
return {
"source_bucket_key": filename,
"dest_bucket_key": filename,
"dest_bucket_name": dest_bucket_name,
}
copy_kwargs = create_copy_kwargs.expand(filename=list_filenames.output)
# Copy files to another bucket, based on the file's extension.
copy_filenames = S3CopyObjectOperator.partial(
task_id="copy_files", source_bucket_name=list_filenames.bucket
).expand_kwargs(copy_kwargs)
태스크 그룹 매핑
@task_group으로 장식된 함수에도 expand나 expand_kwargs를 호출해 태스크 그룹을 통째로 매핑할 수 있어요.
@task_group
def file_transforms(filename):
return convert_to_yaml(filename)
file_transforms.expand(filename=["data1.json", "data2.json"])
여기서 convert_to_yaml 태스크는 런타임에 두 개의 인스턴스로 확장되고, 각각 data1.json과 data2.json을 입력으로 받아요.
태스크 그룹 함수는 주의할 점이 있어요. 태스크 그룹에는 연결된 워커가 없어서, 그룹 함수 안에서 전달받은 인자의 실제 값을 해석할 수 없어요. 값은 그 참조가 태스크로 넘어가 실행될 때 비로소 결정돼요. 그래서 그룹 안에서 값에 대한 분기 로직을 돌리면 예상과 다르게 동작해요.
@task
def my_task(value):
print(value)
@task_group
def my_task_group(value):
if not value: # DOES NOT work as you'd expect!
task_a = EmptyOperator(...)
else:
task_a = PythonOperator(...)
task_a << my_task(value)
my_task_group.expand(value=[0, 1, 2])
my_task_group이 실행될 때 value는 여전히 참조일 뿐이라 if not value 분기가 의도대로 동작하지 않아요. 값을 실제로 이용해 어떤 로직을 돌리고 싶다면 반드시 태스크 안에서 해야 해요. 조건 분기에는 @task.branch(또는 BranchPythonOperator), 루프에는 태스크 매핑 메서드를 쓰는 게 정석이에요.
참고로 매핑된 태스크 그룹 안에 태스크 매핑을 중첩하는 것은 아직 허용되지 않아요. UI 복잡도가 크게 늘고 일반적인 용도에 필요하지 않다고 판단해 의도적으로 빠뜨린 기능이에요.
깊이 우선 실행 (Depth-first execution)
매핑된 태스크 그룹이 여러 태스크를 담고 있으면, 그룹 안의 모든 태스크가 같은 입력에 대해 "함께" 확장돼요.
@task_group
def file_transforms(filename):
converted = convert_to_yaml(filename)
return replace_defaults(converted)
file_transforms.expand(filename=["data1.json", "data2.json"])
그룹 file_transforms가 두 개로 확장되면 convert_to_yaml과 replace_defaults도 각각 두 인스턴스가 돼요. 두 태스크를 따로 확장해도 비슷한 효과는 얻을 수 있어요.
converted = convert_to_yaml.expand(filename=["data1.json", "data2.json"])
replace_defaults.expand(filename=converted)
차이는 의존성의 범위예요. 태스크 그룹을 쓰면 안의 각 태스크가 자기 "관련 입력"에만 의존해요. 위 예시에서 replace_defaults는 같은 확장 그룹의 convert_to_yaml에만 의존하지, 다른 그룹의 같은 태스크 인스턴스에는 관여하지 않아요. 이 전략을 깊이 우선 실행이라고 부르는데, 태스크 분리가 논리적이고 세밀한 의존성 규칙과 정확한 리소스 할당이 가능해져요. 첫 번째 replace_defaults는 convert_to_yaml("data2.json")이 끝나기 전에도 돌 수 있는 거죠.
매핑된 태스크에서 항목 필터링
매핑된 태스크가 None을 돌려주면, 그 요소는 하위 태스크로 전달되지 않아요. 특정 확장자의 파일만 복사하고 싶은 상황을 떠올려 보면 create_copy_kwargs를 이렇게 바꿀 수 있어요.
@task
def create_copy_kwargs(filename):
# Skip files not ending with these suffixes.
if filename.rsplit(".", 1)[-1] not in ("json", "yml"):
return None
return {
"source_bucket_key": filename,
"dest_bucket_key": filename,
"dest_bucket_name": "my_other_bucket",
}
매핑 제한 걸기
매핑 태스크를 제한하는 방법은 두 가지예요. 하나는 확장으로 만들어지는 인스턴스의 수를 제한하는 것이고, 다른 하나는 동시에 실행되는 복사본 수를 제한하는 거예요.
- 매핑 태스크 수 제한:
[core]의max_map_length설정이expand로 만들 수 있는 최대 태스크 수예요. 기본값은1024. 상위 태스크가 이보다 긴 리스트를 돌려주면 그 태스크가 실패로 처리돼요. - 병렬 복사본 제한: 큰 매핑 태스크가 모든 runner 슬롯을 독점하지 않도록
max_active_tis_per_dag로 동시 실행 수를 제한할 수 있어요. 단, 이 값은 특정 DAG Run이 아니라 모든 활성 DAG Run에 걸쳐 그 태스크의 모든 복사본에 적용된다는 점을 기억하세요.
@task(max_active_tis_per_dag=16)
def add_one(x: int):
return x + 1
BashOperator.partial(task_id="my_task", max_active_tis_per_dag=16).expand(bash_command=commands)
길이가 0인 매핑 자동 스킵
입력이 비어 있으면(길이 0) 새 태스크는 만들어지지 않고 매핑 태스크는 SKIPPED로 표시돼요. DAG이 런타임에 할 일을 찾는데 가끔은 할 일이 없을 때 유용하죠. 예를 들어 스캔-복구 DAG이 복구할 게 없으면 빈 리스트를 돌려주고, 매핑 태스크는 스킵되며, 하위 요약 태스크는 스킵된 상위 태스크를 허용하는 트리거 규칙을 쓰면 성공적인 no-op으로 처리할 수 있어요.
from airflow.sdk import TriggerRule, task
@task
def find_work_items():
# Return an empty list when no files, records, or partitions need repair.
return []
@task
def repair(item): ...
@task(trigger_rule=TriggerRule.NONE_FAILED)
def summarize(repaired_items):
if not repaired_items:
print("No work found; nothing to repair.")
return
print(f"Repaired {len(repaired_items)} item(s).")
repaired_items = repair.expand(item=find_work_items())
summarize(repaired_items)
템플릿 필드와 매핑 인자의 관계
오퍼레이터의 모든 인자는 매핑될 수 있어요. 템플릿 파라미터를 받지 않는 인자라도 매핑 대상이 되는데, 중요한 예외가 하나 있어요. 어떤 필드가 템플릿 필드로 표시되어 있으면서 동시에 매핑되면, 그 필드는 템플릿화되지 않아요. 아래 예시는 날짜 스탬프가 아니라 {{ds}}를 그대로 출력해요.
@task
def make_list():
return ["{{ ds }}"]
@task
def printer(val):
print(val)
printer.expand(val=make_list())
값을 보간하고 싶다면 task.render_template를 직접 호출하거나 탬플릿을 이용하면 돼요.
@task
def make_list(ds=None):
return [ds]
@task
def make_list(**context):
return [context["task"].render_template("{{ ds }}", context)]