MLflow와 Airflow를 활용한 모델 관리 및 자동화 파이프라인 구축

본 프로젝트에서는 MLflow를 활용한 모델 버전 관리 및 Airflow를 이용한 자동화 파이프라인을 구축하여, 데이터 학습 및 모델 관리의 효율성을 극대화하였습니다. 이를 통해 학습된 모델의 실험 관리, 저장, 버전 관리뿐만 아니라, 일일 주기로 자동화된 모델 학습 및 배포가 가능합니다.

 

시스템 구성

MLflow를 활용한 모델 버전 관리

MLflow는 머신러닝 실험을 추적하고 모델을 효율적으로 관리하기 위한 도구입니다. 본 프로젝트에서는 MLflow를 활용하여 다음과 같은 기능을 구현하였습니다.

• 자동 실험 로깅: mlflow.sklearn.autolog()을 사용하여 모델 학습 과정에서 하이퍼파라미터 및 성능 지표를 자동으로 기록합니다.
• 실험 관리 및 저장: initialize_experiment 함수를 활용하여 모델 실험을 자동으로 생성하고 관리합니다.
• 모델 저장 및 불러오기: 학습이 완료된 모델은 save_model_to_pickle 메서드를 통해 로컬 디렉토리에 저장되며, 추후 필요할 때 불러올 수 있습니다.

 

Airflow 기반의 자동화 파이프라인 구축

 

Airflow는 워크플로우 자동화를 지원하는 오픈소스 툴로, 본 프로젝트에서는 모델 학습 및 배포를 자동화하는 데 활용되었습니다. 주요 기능은 다음과 같습니다.

 

• DAG(Directed Acyclic Graph) 구성: DAG('model_training_selection', schedule_interval=timedelta(days=1))을 통해 모델 학습 및 선택 작업을 하루 단위로 자동 실행합니다.
• PythonOperator를 이용한 태스크 관리: 각 모델 학습 및 저장 과정이 독립적인 PythonOperator로 구성되어 있으며, 이를 통해 CSU, STS, FTS 등 다양한 시스템 헬스 모델이 자동으로 학습됩니다.
• 모델 학습 및 저장 자동화: train_and_save_model 함수를 활용하여 모델을 학습한 후 저장하며, 모든 과정이 Airflow DAG에 의해 자동으로 실행됩니다.

 

실행 과정

1. 데이터 전처리 및 학습 데이터 생성
   • preprocess_sensor_data 메서드를 통해 학습에 필요한 데이터를 전처리합니다.
   • generate_train_data 메서드를 활용하여 학습 데이터셋을 구성합니다.
2. 모델 학습 및 성능 평가
   • train_xgboost_regression_model을 통해 XGBoost 회귀 모델을 학습합니다.
   • 학습 과정에서 MLflow를 활용하여 실험 정보를 자동으로 기록합니다.
3. 모델 저장 및 버전 관리
   • 학습이 완료된 모델은 save_model_to_pickle을 통해 저장됩니다.
   • MLflow에서 실험 정보와 함께 모델이 관리되므로, 향후 모델 재학습 및 배포가 용이합니다.
4. Airflow를 이용한 자동 실행
   • DAG가 주기적으로 실행되며, 각 모델 학습 및 저장 프로세스를 자동으로 수행합니다.
   • PythonOperator를 통해 run_csu_model, run_sts_model 등의 개별 학습 태스크를 관리합니다.

 

결론

 

MLflow와 Airflow를 활용하여 모델 실험 및 버전 관리뿐만 아니라, 전 과정을 자동화하는 파이프라인을 구축하였습니다. 이를 통해 모델 관리 및 배포의 효율성이 증가하며, 데이터 변화에 따라 지속적으로 업데이트되는 최적의 모델을 운영할 수 있습니다. 향후 추가적인 모델 최적화 및 파이프라인 확장을 통해 더욱 안정적이고 효과적인 AI 시스템을 구축할 예정입니다.