Airflow Issue #62945 — forward reference warning 버그 분석
분석 계기
Apache Airflow 공식 저장소의 good first issue 레이블을 훑다가 이 이슈를 발견했다. 재현 코드도 명확하고 범위가 좁아서 분석해봤다.
https://github.com/apache/airflow/issues/62945
문제 상황
Airflow에는 TaskFlow API라는 기능이 있는데, @task 데코레이터를 붙이면 일반 Python 함수가 Airflow Task로 동작한다.
Airflow는 함수의 반환 타입을 기반으로 multiple_outputs 여부를 자동으로 추론한다.
multiple_outputs는 반환값이 딕셔너리처럼 여러 출력으로 나뉘는지 여부를 나타내는 옵션이다.
Airflow 2.x에서는 TYPE_CHECKING 블록 안에 있는 타입을 반환 타입 힌트로 사용해도 조용히 넘어갔는데, 3.x에서 아래 warning이 새로 생겼다.
1
2
3
Cannot infer multiple_outputs for TaskFlow function 'dataframe_task'
with forward type references that are not imported.
(Error was name 'DataFrame' is not defined) category=UserWarning
재현 코드는 이렇다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
from typing import TYPE_CHECKING
from airflow.sdk import dag, task
if TYPE_CHECKING:
from pandas import DataFrame
@task(task_id="df_task")
def dataframe_task() -> "DataFrame":
import pandas as pd
df = pd.DataFrame({"a": [1, 2, 3]})
return df
@dag()
def simple_dag():
_ = dataframe_task()
simple_dag()
-> "DataFrame" 처럼 따옴표로 감싸는 건 올바른 Python 패턴인데, Airflow가 불필요한 warning을 발생시키는 것이다.
배경 지식
typing 모듈
Python 타입 힌트에 필요한 도구들을 모아둔 표준 라이브러리다. TYPE_CHECKING, List, Dict, Optional 같은 것들이 여기에 있다.
TYPE_CHECKING
런타임에는 무시되고 타입 검사할 때만 사용되는 조건부 import 패턴이다. TYPE_CHECKING 변수는 런타임에 항상 False이고, VS Code가 코드를 분석할 때만 True가 된다. DAG 파일 읽을 때마다 무거운 라이브러리를 로드하지 않아도 된다는 장점이 있다.
1
2
3
4
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import pandas as pd # 실제 실행에는 무시, 타입 검사할 때만
forward reference
TYPE_CHECKING 블록 안에 있는 타입은 런타임에 존재하지 않는다. 그래서 타입 힌트로 쓸 때는 따옴표로 감싸야 한다 (-> "DataFrame") 따옴표로 감싸면 Python이 런타임에 실제 타입으로 평가하지 않고 문자열로만 본다. 런타임에 존재하는 타입은 따옴표가 불필요하다. (-> str)
원인 분석
에러 메시지는 코드에 문자열로 정의되어 있기 때문에 grep으로 검색하면 에러를 직접 재현하지 않아도
해당 메시지를 출력하는 코드 위치를 찾을 수 있다.
1
grep -r "Cannot infer" ~/projects/airflow-dev/airflow/ --include="*.py" -l
1
2
task-sdk/src/airflow/sdk/bases/decorator.py ← 버그 있는 실제 코드
providers/standard/tests/unit/standard/decorators/test_python.py ← 테스트 코드
decorator.py의 _infer_multiple_outputs 메서드를 보면 이렇다.
1
2
3
4
5
6
7
8
9
10
11
try:
def fake(): ...
fake.__annotations__ = {"return": self.function.__annotations__["return"]}
return_type = typing_extensions.get_type_hints(fake, self.function.__globals__).get("return", Any)
except NameError as e:
warnings.warn(
f"Cannot infer multiple_outputs for TaskFlow function {self.function.__name__!r} with forward"
f" type references that are not imported. (Error was {e})",
stacklevel=4,
)
return False
Airflow는 @task 함수의 반환 타입을 보고 multiple_outputs 여부를 자동으로 추론하는데, 이때 get_type_hints()가 "DataFrame" 문자열을 런타임에 실제 타입으로 평가하려 한다. DataFrame은 TYPE_CHECKING 블록 안에 있어서 런타임에 import가 안 된 상태라 NameError가 발생하고, 그 시점에 경고를 출력한다. forward reference는 Python에서 일반적으로 사용하는 패턴인데, Airflow가 이 타입을 평가하려다 실패하면서 불필요한 경고가 발생한다. 테스트 파일에도 이미 @pytest.mark.xfail이 붙어있었는데, 이건 이 테스트가 현재 실패 상태라는 표시다.
수정 방향 비교
decorator.py
기존 코드는 NameError가 발생하면 warning을 출력하고 return False한다. 내가 생각한 방법은 get_type_hints() 호출 전에 반환 타입이 문자열인지 먼저 확인하는 것이었다. forward reference면 평가 시도 자체를 건너뛰는 방식이다.
1
2
3
4
5
6
7
8
9
return_annotation = self.function.__annotations__["return"]
if isinstance(return_annotation, str): # forward reference면 skip
return False
try:
return_type = typing_extensions.get_type_hints(...)
except NameError as e:
warnings.warn(...)
return False
PR에서 제안된 방법은 NameError 발생 자체는 놔두고, catch 시점에 warning 없이 조용히 return False하는 방식이다. import warnings도 더 이상 안 쓰니까 같이 제거했다.
1
2
3
4
except NameError:
# Forward references using TYPE_CHECKING-only imports are valid Python patterns.
# We cannot infer multiple_outputs when the type is not available at runtime.
return False
두 방법의 결과는 동일하지만 차이가 있다. isinstance(return_annotation, str) 체크는 문자열로 감싼 forward reference만 걸러낼 수 있는데, NameError가 발생하는 케이스는 이것 말고도 있다. 타입 힌트에 쓴 이름이 오타거나 import를 빠뜨린 경우처럼 forward reference가 아닌데도 NameError가 날 수 있고, 내 방법은 그런 케이스를 걸러내지 못하고 get_type_hints() 호출까지 진행하게 된다. PR 방식이 NameError 케이스를 더 넓게 처리할 수 있는 구조라고 생각이 들었다.
test_python.py
기존 테스트는 warning이 발생하는지 검증하고 있었고, xfail이 붙어있었다.
1
2
3
4
5
@pytest.mark.xfail(reason="TODO AIP72: ...")
def test_infer_multiple_outputs_forward_annotation(self):
...
with pytest.warns(UserWarning, match="Cannot infer multiple_outputs.*t3") as recwarn:
...
PR에서는 xfail을 제거하고 warning 검증 코드를 전부 지웠다. 버그를 수정했으니 warning이 발생하지 않는 게 정상 동작이고, 그걸 검증하는 것으로 바꾼 것이다.
1
2
3
4
def test_infer_multiple_outputs_forward_annotation(self):
...
# No warning should be raised for TYPE_CHECKING-only forward references
assert t3(5, 5).operator.multiple_outputs is False
정리
이번 이슈는 코드 수정 자체는 단순했지만 Airflow 내부 동작과 Python 타입 시스템이 같이 얽혀 있는 문제였다. 오픈소스 코드를 읽다 보니 에러 메시지를 grep으로 찾아 들어가는 방식이 생각보다 유용하다는 것도 다시 느꼈다.