Timetable로 Dag 스케줄링 커스터마이징하기

Timetable로 Dag 스케줄링 커스터마이징하기 (Customizing Dag Scheduling with Timetables)

Airflow의 기본 cron 표현으로는 표현하기 어려운 스케줄링 로직(예: 공휴일 제외, 각 근무일 종료 시 실행)을 커스텀 Timetable로 구현하는 방법을 설명하는 문서예요. Timetable 등록, next_dagrun_info·infer_manual_data_interval 메서드 구현, 파라미터화·직렬화, UI 표시, run_id 생성까지 코드와 함께 살펴볼게요.

출처: 문서

본문

예시로, 어떤 회사가 각 평일(weekday) 뒤에 근무 시간 동안 수집된 데이터를 처리하는 job을 실행하려 한다고 가정해 봐요. 가장 직관적인 첫 답은 schedule="0 0 * * 1-5"(월요일부터 금요일 자정)일 거예요. 하지만 이것은 금요일에 수집된 데이터가 금요일이 끝난 직후가 아니라 다음 월요일에 처리된다는 뜻이고, 그 run의 데이터 구간은 금요일 자정부터 월요일 자정까지가 돼요. 게다가 위 스케줄 문자열은 공휴일에는 처리를 건너뛸 수 없어요. 우리가 원하는 것은:

• 월요일, 화요일, 수요일, 목요일, 금요일마다 run을 스케줄한다. run의 데이터 구간은 각 날짜의 자정부터 다음 날짜의 자정까지다(예: 2021-01-01 00:00:00에서 2021-01-02 00:00:00).
• 각 run은 데이터 구간이 끝난 직후에 생성된다. 월요일을 커버하는 run은 화요일 자정에, 이런 식으로. 금요일을 커버하는 run은 토요일 자정에. 일요일과 월요일 자정에는 run이 없다.
• 정의된 공휴일에는 run을 스케줄하지 않는다.

단순화를 위해 이 예시에서는 UTC datetime만 다룰게요.

참고 (Note)

커스텀 timetable이 반환하는 모든 datetime 값은 반드시 "aware"여야 해요. 즉 시간대 정보를 포함해야 해요. 게다가 pendulum의 datetime과 timezone 타입을 사용해야 해요.

Timetable 등록 (Timetable Registration)

Timetable은 Timetable의 서브클래스여야 하고, plugin의 일부로 등록되어야 해요. 다음은 새 timetable을 구현하기 위한 골격(skeleton)이에요:

from airflow.plugins_manager import AirflowPlugin
from airflow.timetables.base import Timetable

class AfterWorkdayTimetable(Timetable):
    pass

class WorkdayTimetablePlugin(AirflowPlugin):
    name = "workday_timetable_plugin"
    timetables = [AfterWorkdayTimetable]

다음으로 AfterWorkdayTimetable에 코드를 넣기 시작할게요. 구현이 끝나면 Dag 파일에서 timetable을 사용할 수 있어야 해요:

import pendulum

from airflow.sdk import DAG
from airflow.example_dags.plugins.workday import AfterWorkdayTimetable

with DAG(
    dag_id="example_after_workday_timetable_dag",
    start_date=pendulum.datetime(2021, 3, 10, tz="UTC"),
    schedule=AfterWorkdayTimetable(),
    tags=["example", "timetable"],
):
    ...

스케줄링 로직 정의하기 (Define Scheduling Logic)

Airflow의 스케줄러가 Dag를 만나면 두 메서드 중 하나를 호출해 Dag의 다음 run을 언제 스케줄할지 알아내요.

• next_dagrun_info: 스케줄러가 timetable의 정규 스케줄(우리 예시의 "매 근무일마다 하나, 끝에서 실행" 부분)을 아는 데 사용해요.
• infer_manual_data_interval: Dag run이 수동으로 트리거될 때(예: 웹 UI에서), 스케줄러는 이 메서드로 스케줄 외 run의 데이터 구간을 역추론하는 방법을 알아요.

둘 중 더 쉬운 infer_manual_data_interval부터 시작할게요:

airflow/example_dags/plugins/workday.py [source]

def infer_manual_data_interval(self, run_after: DateTime) -> DataInterval:
    start = DateTime.combine((run_after - timedelta(days=1)).date(), Time.min).replace(tzinfo=UTC)
    # Skip backwards over weekends and holidays to find last run
    start = self.get_next_workday(start, incr=-1)
    return DataInterval(start=start, end=(start + timedelta(days=1)))

이 메서드는 인자 run_after(Dag가 외부에서 트리거된 시점을 나타내는 pendulum.DateTime 객체)를 받아요. 우리 timetable은 각 완전한 근무일에 데이터 구간을 만들므로, 여기서 추론된 데이터 구간은 보통 run_after 하루 전의 자정에 시작해야 해요. 하지만 run_after가 일요일이나 월요일(즉 전날이 토요일이나 일요일)이라면 이전 금요일로 더 뒤로 밀려야 해요. 구간의 시작을 알면 끝은 단순히 그로부터 하루 전체 뒤예요. 그런 다음 이 구간을 설명하는 DataInterval 객체를 만들어요.

다음은 next_dagrun_info의 구현이에요:

airflow/example_dags/plugins/workday.py [source]

def next_dagrun_info(
    self,
    *,
    last_automated_data_interval: DataInterval | None,
    restriction: TimeRestriction,
) -> DagRunInfo | None:
    if last_automated_data_interval is not None:  # There was a previous run on the regular schedule.
        last_start = last_automated_data_interval.start
        next_start = DateTime.combine((last_start + timedelta(days=1)).date(), Time.min)
    # Otherwise this is the first ever run on the regular schedule...
    elif (earliest := restriction.earliest) is None:
        return None  # No start_date. Don't schedule.
    elif not restriction.catchup:
        # If the DAG has catchup=False, today is the earliest to consider.
        next_start = max(earliest, DateTime.combine(Date.today(), Time.min, tzinfo=UTC))
    elif earliest.time() != Time.min:
        # If earliest does not fall on midnight, skip to the next day.
        next_start = DateTime.combine(earliest.date() + timedelta(days=1), Time.min)
    else:
        next_start = earliest
    # Skip weekends and holidays
    next_start = self.get_next_workday(next_start.replace(tzinfo=UTC))

    if restriction.latest is not None and next_start > restriction.latest:
        return None  # Over the DAG's scheduled end; don't schedule.
    return DagRunInfo.interval(start=next_start, end=(next_start + timedelta(days=1)))

이 메서드는 두 인자를 받아요. last_automated_data_interval은 이 Dag의 이전 비-수동 트리거 run의 데이터 구간을 나타내는 DataInterval 인스턴스이거나, Dag가 처음으로 스케줄되는 경우에는 None이에요. restriction은 Dag와 그 태스크가 스케줄을 어떻게 지정하는지를 담고 있으며 세 가지 속성을 포함해요:

• earliest: Dag가 스케줄될 수 있는 가장 이른 시간. Dag와 그 태스크의 모든 start_date 인자에서 계산된 pendulum.DateTime이거나, start_date 인자가 전혀 없으면 None.
• latest: earliest와 유사하게 end_date 인자에서 계산된, Dag가 스케줄될 수 있는 가장 늦은 시간.
• catchup: Dag의 catchup 인자를 반영하는 불리언. 기본값은 False.

참고 (Note)

earliest와 latest 둘 다 Dag run의 logical date(데이터 구간의 시작)에 적용되지, run이 스케줄될 시점(보통 데이터 구간의 끝 이후)에는 적용되지 않아요.

참고 (Note)

last_automated_data_interval은 Dag가 Dag processor에 처음 포착될 때만 None이에요: 첫 run은 어떤 Dag run도 존재하기 전에 파싱 시간에 계산되고, 그 결과는 Dag에 저장돼요. 스케줄링 중에는 next_dagrun_info가 항상 이전 run의 데이터 구간이 채워진 상태로 호출되므로, Dag가 처음 unpause될 때 스케줄러 로그에는 None 경우가 나타나지 않아요.

이전에 스케줄된 run이 있었다면, 후속 날짜를 반복해 토요일·일요일·미국 공휴일이 아닌 날을 찾아 다음 비공휴일 평일을 스케줄해야 해요. 반면 이전 스케줄된 run이 없었다면, restriction.earliest 이후의 다음 비공휴일 근무일 자정을 고른다. restriction.catchup도 고려해야 해요—False라면 start_date 값이 과거여도 현재 시간 이전에는 스케줄할 수 없어요. 마지막으로 계산된 데이터 구간이 restriction.latest보다 늦다면 그것을 존중하고 None을 반환해 run을 스케줄하지 않아야 해요.

run을 스케줄하기로 결정했다면 DagRunInfo로 그것을 설명해야 해요. 이 타입에는 두 인자와 속성이 있어요:

• data_interval: 다음 run의 데이터 구간을 설명하는 DataInterval 인스턴스.
• run_after: Dag run이 언제 스케줄될 수 있는지 스케줄러에 알려주는 pendulum.DateTime 인스턴스.

DagRunInfo는 이렇게 만들 수 있어요:

info = DagRunInfo(
    data_interval=DataInterval(start=start, end=end),
    run_after=run_after,
)

보통 데이터 구간이 끝나자마자 run을 스케줄하고 싶으므로, 위의 end와 run_after는 일반적으로 같아요. DagRunInfo는 이를 위한 단축키를 제공해요:

info = DagRunInfo.interval(start=start, end=end)
assert info.data_interval.end == info.run_after  # Always True.

참고로, 우리의 플러그인과 Dag 파일을 전체로 보면 다음과 같아요:

airflow/example_dags/plugins/workday.py [source]

from pendulum import UTC, Date, DateTime, Time

from airflow.plugins_manager import AirflowPlugin
from airflow.timetables.base import DagRunInfo, DataInterval, Timetable

if TYPE_CHECKING:
    from airflow.timetables.base import TimeRestriction

log = logging.getLogger(__name__)

class AfterWorkdayTimetable(Timetable):
    _NOT_LOADED = object()
    _holiday_calendar = _NOT_LOADED

    @classmethod
    def _get_holiday_calendar(cls):
        if cls._holiday_calendar is cls._NOT_LOADED:
            try:
                from pandas.tseries.holiday import USFederalHolidayCalendar

                cls._holiday_calendar = USFederalHolidayCalendar()
            except ImportError:
                log.warning("Could not import pandas. Holidays will not be considered.")
                cls._holiday_calendar = None

        return cls._holiday_calendar

    def get_next_workday(self, d: DateTime, incr=1) -> DateTime:
        holiday_calendar = self._get_holiday_calendar()
        next_start = d
        while True:
            if next_start.weekday() not in (5, 6):  # not on weekend
                if holiday_calendar is None:
                    holidays = set()
                else:
                    holidays = holiday_calendar.holidays(start=next_start, end=next_start).to_pydatetime()
                if next_start not in holidays:
                    break
            next_start = next_start.add(days=incr)
        return next_start
    def infer_manual_data_interval(self, run_after: DateTime) -> DataInterval:
        start = DateTime.combine((run_after - timedelta(days=1)).date(), Time.min).replace(tzinfo=UTC)
        # Skip backwards over weekends and holidays to find last run
        start = self.get_next_workday(start, incr=-1)
        return DataInterval(start=start, end=(start + timedelta(days=1)))
    def next_dagrun_info(
        self,
        *,
        last_automated_data_interval: DataInterval | None,
        restriction: TimeRestriction,
    ) -> DagRunInfo | None:
        if last_automated_data_interval is not None:  # There was a previous run on the regular schedule.
            last_start = last_automated_data_interval.start
            next_start = DateTime.combine((last_start + timedelta(days=1)).date(), Time.min)
        # Otherwise this is the first ever run on the regular schedule...
        elif (earliest := restriction.earliest) is None:
            return None  # No start_date. Don't schedule.
        elif not restriction.catchup:
            # If the DAG has catchup=False, today is the earliest to consider.
            next_start = max(earliest, DateTime.combine(Date.today(), Time.min, tzinfo=UTC))
        elif earliest.time() != Time.min:
            # If earliest does not fall on midnight, skip to the next day.
            next_start = DateTime.combine(earliest.date() + timedelta(days=1), Time.min)
        else:
            next_start = earliest
        # Skip weekends and holidays
        next_start = self.get_next_workday(next_start.replace(tzinfo=UTC))

        if restriction.latest is not None and next_start > restriction.latest:
            return None  # Over the DAG's scheduled end; don't schedule.
        return DagRunInfo.interval(start=next_start, end=(next_start + timedelta(days=1)))

class WorkdayTimetablePlugin(AirflowPlugin):
    name = "workday_timetable_plugin"
    timetables = [AfterWorkdayTimetable]

import pendulum

from airflow.sdk import DAG
from airflow.example_dags.plugins.workday import AfterWorkdayTimetable
from airflow.providers.standard.operators.empty import EmptyOperator

with DAG(
    dag_id="example_workday_timetable",
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    schedule=AfterWorkdayTimetable(),
    tags=["example", "timetable"],
):
    EmptyOperator(task_id="run_this")

파라미터화된 Timetables (Parameterized Timetables)

때로는 timetable에 런타임 인자를 전달해야 해요. AfterWorkdayTimetable 예시를 이어서, 서로 다른 시간대에서 실행되는 Dags가 있고 자정 대신 다음 날 오전 8시에 일부 Dags를 스케줄하고 싶을 수 있어요. 목적마다 별도 timetable을 만들기보다는 이렇게 하고 싶을 거예요:

class SometimeAfterWorkdayTimetable(Timetable):
    def __init__(self, schedule_at: Time) -> None:
        self._schedule_at = schedule_at

    def next_dagrun_info(self, last_automated_dagrun, restriction):
        ...
        end = start + timedelta(days=1)
        return DagRunInfo(
            data_interval=DataInterval(start=start, end=end),
            run_after=DateTime.combine(end.date(), self._schedule_at).replace(tzinfo=UTC),
        )

커스텀 schedule_at 값에 대해 AfterWorkdayTimetable의 첫 run 로직을 적용한다면, 후보 시간을 self._schedule_at과 비교해요. 이전 예시의 자정 특정 체크는 run이 00:00에 스케줄될 때만 올바르지만, 예를 들어 earliest 시간이 06:00이면 같은 날 08:00 run은 여전히 허용되어야 하고, earliest 시간이 09:00이면 다음 근무일로 이동해야 해요.

그러나 timetable은 Dag의 일부이므로, __init__에서 제공하는 context로 어떻게 직렬화할지 Airflow에 알려줘야 해요. 이는 timetable 클래스에 두 개의 추가 메서드를 구현해 이루어져요:

class SometimeAfterWorkdayTimetable(Timetable):
    ...

    def serialize(self) -> dict[str, Any]:
        return {"schedule_at": self._schedule_at.isoformat()}

    @classmethod
    def deserialize(cls, value: dict[str, Any]) -> Timetable:
        return cls(Time.fromisoformat(value["schedule_at"]))

Dag가 직렬화될 때 serialize가 호출되어 JSON 직렬화 가능한 값을 얻어요. 그 값은 스케줄러가 직렬화된 Dag에 접근해 timetable을 재구성할 때 deserialize로 전달돼요.

UI에서의 Timetable 표시 (Timetable Display in UI)

기본적으로 커스텀 timetable은 UI에서 클래스 이름으로 표시돼요(예: "dags" 테이블의 Schedule 열). summary 속성을 오버라이드해 이를 커스터마이징할 수 있어요. 이는 특히 파라미터화된 timetable이 __init__에서 제공된 인자를 포함시키는 데 유용해요. 예를 들어 SometimeAfterWorkdayTimetable 클래스는 이렇게 할 수 있어요:

@property
def summary(self) -> str:
    return f"after each workday, at {self._schedule_at}"

따라서 이런 Dag에 대해:

with DAG(
    schedule=SometimeAfterWorkdayTimetable(Time(8)),  # 8am.
    ...,
):
    ...

Schedule 열은 after each workday, at 08:00:00라고 말할 거예요.

더 보기 (See also)

Module airflow.timetables.base

공개 인터페이스는 서브클래스가 무엇을 구현해야 하는지를 설명하도록 상세히 문서화되어 있어요.

UI에서의 Timetable 설명 표시 (Timetable Description Display in UI)

description 속성을 오버라이드해 Timetable 구현에 대한 설명을 제공할 수도 있어요. 이는 구현에 대한 포괄적인 설명을 UI에서 제공하는 데 특히 유용해요. SometimeAfterWorkdayTimetable 클래스는 예를 들어 이렇게 할 수 있어요:

description = "Schedule: after each workday"

설명을 파생하고 싶다면 이 로직을 __init__ 안에 감쌀 수도 있어요:

def __init__(self) -> None:
    self.description = "Schedule: after each workday, at f{self._schedule_at}"

이는 summary 속성과 다른 포괄적인 설명을 제공하고 싶을 때 특히 유용해요.

따라서 이런 Dag에 대해:

with DAG(
    schedule=SometimeAfterWorkdayTimetable(Time(8)),  # 8am.
    ...,
):
    ...

i 아이콘은 Schedule: after each workday, at 08:00:00을 보여줄 거예요.

더 보기 (See also)

Module airflow.timetables.interval

UI에서 포괄적인 cron 설명을 제공하는 CronDataIntervalTimetable 설명 구현을 확인하세요.

생성된 run_id 변경하기 (Changing generated run_id)

2.4 버전에 추가됨 (Added in version 2.4).

Airflow 2.4부터 Timetable은 DagRun의 run_id 생성을 담당해요.

예를 들어 Run ID가 run이 시작된 시점의 "사람 친화적인" 날짜(즉 데이터 구간의 끝. 현재 사용되는 시작이 아니라)를 보여주게 하려면 커스텀 timetable에 이와 같은 메서드를 추가할 수 있어요:

def generate_run_id(
    self,
    *,
    run_type: DagRunType,
    logical_date: DateTime,
    data_interval: DataInterval | None,
    **extra,
) -> str:
    if run_type == DagRunType.SCHEDULED and data_interval:
        return data_interval.end.format("YYYY-MM-DD dddd")
    return super().generate_run_id(
        run_type=run_type, logical_date=logical_date, data_interval=data_interval, **extra
    )

RunID는 250자로 제한되고 Dag 안에서 고유해야 한다는 것을 기억하세요.

더 알아보기 (Learn more)

• Operator 사용하기
• 커스텀 View 플러그인
• Timetable API 참조
• Airflow 최신 버전 문서 홈