Airflow와 MLflow를 활용한 Mlops 구축기

현재 회사에서 강의/컨텐츠 추천 시스템 관련하여 연구, 개발을 진행하고 있습니다.

데이터 전처리와 모델 학습, 인퍼런스를 위한 코딩은 모두 끝냈지만 새로운 아이템이 빠르게 자주 생성되는 강의/컨텐츠 특성 상 전처리부터 배포까지 Batch로 전체 파이프라인을 돌릴 필요가 있어서 MLops를 구축하게 되었습니다.

쿠버네티스를 활용하여 전체 인프라를 구축하고 싶은 욕심이 있었으나, 기간과 리소스가 충분하지 않았기에 쿠버네티스는 오버스펙이라고 생각하여 Airflow와 MLflow를 활용하기로 결정하였습니다.

먼저 전체적인 Workflow 제어를 위해 Airflow를 활용하였고, MLflow는 모델의 모니터링과 버전 관리 역할로 활용하였습니다. 전체 구조를 간략하게 그림으로 표현하면 아래와 같습니다. 각 모듈은 Airflow DAG로 작성되어 순차적으로 실행되게 되고, Model Training 부분에서는 모델의 성능 결과와 모델 파일을 MLflow를 통해 추적할 수 있도록 하였습니다.

Airflow와 MLflow 각각에 대해 정리한 글도 추후에 작성할 예정입니다.

Airflow에는 다양한 모듈이 존재하는데 이러한 모듈을 한 번에 컨테이너로 올릴 수 있는 Docker-compose 파일을 제공하고 있어서 매우 수월하게 진행할 수 있었습니다. 저희는 GPU를 활용하기 때문에 airflow-worker 부분에서 gpu 사용을 아래 부분을 추가했습니다.

deploy:
  resources:
    reservations:
      devices:
        - driver: nvidia
          count: 1
          capabilities: [gpu]

추가적으로 MLflow 서버와 DAG 등의 작업을 위한 Jupyter 서버도 같이 올라갈 수 있도록 yaml파일을 수정하였습니다.

mlflow:
  image: mlflow:latest
  command: mlflow server --backend-store-uri sqlite:///mlflow.db --default-artifact-root (pwd)/artifacts --host 0.0.0.0
  ports:
    - 25000:5000
  healthcheck:
    test: ["CMD", "curl", "--fail", "http://localhost:5000/"]
    interval: 10s
    timeout: 10s
    retries: 5
  restart: always

jupyter:
  image: mlflow:latest
  command :
    - /bin/bash
    - -ic
    - |
      conda run -n mlflow jupyter lab --allow-root --ip=0.0.0.0 --port=8888
  deploy:
    resources:
      reservations:
        devices:
          - driver: nvidia
            count: 0
            capabilities: [ gpu ]

위에 있는 mlflow 이미지는 anaconda를 베이스 이미지로 해서 새로운 conda env를 추가하고 bash에서 바로 적용될 수 있도록 수정한 이미지 입니다. 빌드를 위한 Dockerfile은 아래와 같고 필요한 conda environment를 env.yaml로 작성하시면 바로 사용하실 수 있습니다. 이때 env.yaml에서 name을 mlflow로 맞춰주시거나 밑에 RUN echo "conda activate mlflow" 부분을 env.yaml에서 name으로 설정해주신 걸로 변경해주시면 됩니다.

FROM continuumio/anaconda3

ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
ENV PATH /opt/conda/bin:$PATH

USER root
# Install Dependencies of Miniconda
RUN apt-get update --fix-missing && \
    apt-get install -y wget bzip2 curl git && \
    apt-get clean && \
    apt install -y python3 && \
    apt install -y python3-pip

COPY env.yaml /environment.yml
RUN conda env create -f /environment.yml

RUN echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
    echo "conda activate mlflow" >> ~/.bashrc

ENV PATH /opt/conda/envs/env/bin:$PATH

RUN pip install -U pip setuptools wheel \
    && pip install pytz \
    && pip install pyOpenSSL \
    && pip install ndg-httpsclient \
    && pip install pyasn1 \
    && pip install 'redis>=2.10.5,<3' \
    && pip install mlflow

EXPOSE 5000