[Airflow] Docker 환경에서 Airflow와 Wandb 같이 활용하기

이번 포스팅에서는 Airflow를 docker 컨테이너에서 사용할 경우 WandB를 설정하는 방법을 정리하였다.

 

 

 

Airflow

 

왜 Airflow를 WandB와 같이 사용하려는가?

 

Airflow는 일반적으로 데이터 파이프라인 구축하여 스케줄링 하고 모니터링 하기 위해 주로 사용된다. 실시간으로 로그가 쌓이는 작업에서 데이터 파이프라인을 관리하기 위해서 보통 많이 쓰이는 것으로 알고 있다. 이런 과정을 위해서 쓴다면 로그들은 ETL를 거쳐 데이터베이스에 저장되고 ELT 과정을 통해 의미있는 데이터들을 얻어낸 다음 머신러닝 학습을 한 후 얻게 된 정보를 바탕으로 서비스를 제공하는 방식으로 운영이 될 것이다.

 

만약, 데이터베이스에 저장될 수 있는 데이터가 아닌 이미지 데이터를 위한 머신러닝 모델을 활용하는 task가 이 과정에서 포함된다면 모델의 학습과정이나 모델의 성능을 시각화하여 확인하는 과정이 필요할 것이다. 

 

이를 위한 여러가지 visualization tool들이 존재하지만 이번 글에서는 WandB를 선택하였고 그 이유와 필요한 세팅과정을 정리해보았다. 아래는 한 사이트에서 Airflow와 WandB를 비교하여 기능에 따라 점수화한 자료이다.

 

https://www.saashub.com/compare-airflow-vs-weights-and-biases

 

 

시각화 툴 중에서 WandB의 장점?

 

WandB에 대한 장점은 구글링을 통해서 충분히 찾아볼 수 있기 때문에 개인적으로 생각한 장점을 적어보면 다음과 같다.

 

https://wandb.ai/site/experiment-tracking

 

WandB는 특정 포트를 브라우저에 입력해서 확인하는 것이 아니라 WandB 페이지에 계정을 통해서 확인할 수 있다. 따라서, 장소에 불문하고 WandB 계정을 통해 결과를 확인할 수 있는 장점이 존재한다. (단적인 예로 모델 학습을 시켜놓은 뒤 식사를 하러가서 밥을 먹으며 학습이 잘되고 있는지 휴대폰으로도 확인할 수 있다!)

 

또한, 프로젝트 단위의 실험 관리가 편하다. 아래의 왼쪽의 사진에서 여러 개의 프로젝트들을 테이블 형태로 확인할 수 있는 것이 보이고 오른쪽의 사진에서 그 프로젝트에 대한 내용을 정리할 것을 확인할 수 있다.

 

WandB 사용예시

 

이를 통해  Airflow의 DAG를 WandB의 프로젝트에 대응하여 활용할 수도 있고, DAG를 구성하는 Task들을 프로젝트에 대응하여 활용할 수 있다.

 

일반적인 local 환경에 직접 Airflow를 설치하여 구성한다면 WandB를 일반적인 방법으로 잘 활용하면 된다. 하지만, Docker 환경에서 Airflow를 활용할 경우 두 개의 user(airflow, root)로 관리하다보니 wandb에 관한 로그가 남게 되는 디렉토리를 airflow의 디렉토리 안에 지정을 해주어야 에러가 발생하지 않는다. 

 

 

Docker환경에서 Airflow 사용 시 Wandb를 활용하기 위한 설정

 

이제 본론으로 들어가서 docker 환경에서 Airflow와 WandB 설정하는 방법을 다음과 같다. 먼저, 아래의 명령어로 docker-compose.yaml 파일을 다운로드한 후

 

curl -LfO 'https://airflow.apache.org/docs/apache-airflow/2.5.1/docker-compose.yaml'

 

파일을 내용을 확인해보면 

 

docker-compose.yaml 파일 내용 일부분

 

42번째 줄부터 74번째 줄까지는  Airflow와 관련된 모든 컨테이너에 대한 공통적인 설정을 해주는 부분이다. 이 부분의 environment 부분에 맨 마지막줄 (63번째 줄) 이후에 아래의 설정들을 넣어준다.

 

WANDB_DATA_DIR: '/opt/airflow/logs/wandb'
WANDB_CONFIG_DIR: '/opt/airflow/logs/wandb'
WANDB_CACHE_DIR: '/opt/airflow/logs/wandb'
WANDB_DIR: '/opt/airflow/logs/wandb'

 

위의 설정들은 WandB를 사용하면서 생겨나는 파일들의 경로들을 airflow 컨테이너에서 airflow 유저에게 권한이 있는 /opt/airflow 안으로 하여 에러가 발생하지 않도록 해준다. 추가적으로 맨 처음 컨테이너에 접속했을 때 로그인을 해주지 않아도 되도록 WandB의 api를 미리 설정해줄 수 있다.

 

 

WANDB_API_KEY: 'your wandb api'

 

 

이 역시 위의 설정들과 같은 위치에 입력해준다. WandB의 api는 https://wandb.ai/ 에서 로그인 후 User settings의 Danger Zone에서 확인이 가능하다.

 

추가적으로 airflow dag 예제들을 안보이도록 61번째 줄 AIRFLOW__CORE__LOAD_EXAMPLES를 true에서 false로 변경해준다.

 

AIRFLOW__CORE__LOAD_EXAMPLES: 'false' # 'ture'

 

 

간단한 테스트

 

설정이 잘 되었는지 확인을 위해 간단한 DAG를 생성 후 확인할 수 있다. 간단한 테스트를 위해 docker-compose.yaml 파일이 있는 위치에 여러 디렉토리들을 생성해주고 권한을 airflow 유저에게 할당한다.

 

mkdir -p ./dags ./logs ./plugins ./config
echo -e "AIRFLOW_UID=$(id -u)" > .env

 

이제 아래의 명령어로 docker-compose 파일로 airflow와 관련된 컨테이너들을 띄워준다.

 

docker-compose up airflow-init
docker-compose up -d

 

이후, (docker-compose파일이 존재하는 디렉토리명)-airflow-scheduler-1 의 컨테니어 접속한다. 필자는 docker-compose파일이 존재하는 디렉토리명이 airflow이기에 다음과 같은 명령어로 접속하였다. 이후, wandb 및 필요한 라이브러리를 설치해준다. 

 

docker exec -it airflow-airflow-scheduler-1 bash
pip install opencv-python wandb

 

생성해준 dags 폴더 안에 테스트할 파일 test_dag.py을 만들고 아래의 테스트 코드를 입력해준다.

 

from airflow                          import DAG
from airflow.operators.python         import PythonOperator

from datetime                         import datetime
import numpy                          as np

import random
import wandb

def print_hello():
    print('Hello')

def train_wandb():
    epochs=10
    lr=0.01

    run = wandb.init(
        # Set the project where this run will be logged
        project="my-awesome-project",
        # Track hyperparameters and run metadata
        config={
            "learning_rate": lr,
            "epochs": epochs,
        })

    offset = random.random() / 5

    # simulating a training run
    for epoch in range(2, epochs):
        acc = 1 - 2 ** -epoch - random.random() / epoch - offset
        loss = 2 ** -epoch + random.random() / epoch + offset
        wandb.log({"accuracy": acc, "loss": loss})



test_wandb = DAG(
    dag_id = f'test_wandb',
    start_date = datetime(2022,6,11),
    catchup = False,
    tags = ['hello', 'wandb'],
    schedule = '@once',
)

print_hello = PythonOperator(
    task_id         = 'print_hello',
    python_callable = print_hello,
    dag             = test_wandb
)

train_wandb = PythonOperator(
    task_id         = 'train_wandb',
    python_callable = train_wandb,
    dag             = test_wandb
)

print_hello >> train_wandb

 

이제 인터넷 브라우저에 127.0.0.1:8080을 입력하고 Airflow에 로그인을 해준다. username과 password는 모두 airflow이다.

 

airflow

 

성공적으로 DAG 생성되었다면 DAG를 실행시킨 후 결과를 확인한다.

 

 

만약, DAG 실행시켰는데 실패하였다면

현재의 컨테이너의 상태를 도커 이미지에 commit해준 후 Airflow를 재가동해주면 된다. 이유는 잘 모르겠으나 경험적으로 알아낸 방법이다.

 

도커 airflow-scheduler 컨테이너의 상태를 apache/airflow:wandb라는 도커 이미지 형태로 저장

docker commit airflow-airflow-scheduler1 apache/airflow:wandb

 

docker-compose.yaml 파일에서 49번째 줄 image 부분을 아래와 같이 수정

image:  ${AIRFLOW_IMAGE_NAME:-apache/airflow:wandb} # ${AIRFLOW_IMAGE_NAME:-apache/airflow:2.5.1}

 

그 후 다시 docker-compose를 통해 Airflow 가동

docker-compose up airflow-init
docker-compose up -d

 

 

이제 결과를 확인하면 아래와 같이 DAG가 성공적으로 수행되었음을 확인할 수 있다.

 

airflow dags

 

또한, 본인의 WandB에 접속해보면 my-awesome-project라는 project가 생성되었으며 프로젝트를 클릭해보면 다음과 같이 차트가 입력되어 있음을 확인할 수 있다.

 

my-awesome-project

 

이와 같은 방식으로 Comtuper Vision과 관련된 머신러닝 학습의 파이프라인도 Airflow를 통하여 자동화하고 WandB를 통하여 확인할 수 있을 것이다.