[Airflow] Airflow Timeout 설정

 

 이번 글에서는 Airflow에서 Timeout을 설정하는 방법에 대해 간단하게 살펴보고자 합니다. 지난 MLflow Timeout 오류 관련 글을 작성하면서, Airflow에서도 비슷한 경험을 겪었던 것이 생각나 작성하게 되었습니다. 우선 Airflow에 대해 간략하게 설명하면, airbnb에서 만든 워크플로우 관리 툴로 대표적인 스케쥴링 툴 중의 하나입니다. 저 역시도 실무에서 AI 모델 학습과 추론 과정에서 스케쥴링 툴로 사용하고 있습니다.

 

 Airflow의 기본적인 Dag 구조는 다음과 같습니다.

 

with DAG (dag_id='test', schedule=None) as dag:

    t1 = BashOperator(
        task_id='hello',
        bash_command='echo "hello"'
    )
    
    t2 = BashOperator(
        task_id='bye',
        bash_command='echo "Bye"'
    )
    
    t1 >> t2

 

 위의 코드처럼 DAG를 작성하면, t1과 t2 task가 순차적으로 문제없이 실행될 것입니다. 문제는 task에서 실행할 파일이 오래 걸릴 경우입니다. 오래 걸리게 되면 다음과 같은 오류가 발생합니다.

 

airflow.exceptions.AirflowException: SSH command timed out

 

 이를 해결하기 위해서는 BashOperator 내 인자 중 cmd_timeout을 늘려주면 됩니다. default는 10으로 되어있는데, 이는 10초를 의미합니다. 따라서 실행하고자 하는 task에 따라 적절하게 초 단위로 설정해주면 됩니다.

 

 설정한 예시는 다음과 같습니다. 전처리 과정은 10분(600초), 학습 과정은 60분(3600초)으로 설정하였습니다. 

 

with DAG (dag_id='test', schedule=None) as dag:

    t1 = BashOperator(
        task_id='preprocess',
        bash_command='python3 preprocess.py',
        cmd_timeout=600
    )
    
    t2 = BashOperator(
        task_id='train',
        bash_command='python3 train.py',
        cmd_timeout=3600
    )
    
    t1 >> t2