[DAY 117] 최종 프로젝트_ Airflow + MLflow 연동

[천재교육] 프로젝트 기반 빅데이터 서비스 개발자 양성 과정 9기
학습일 : 2025.01.06

📕 프로젝트 작업 내역

• Airflow + MLflow 연동 필요성
• 연동 방법
• 코드 예시

 

📗 수행 결과

1. Airflow + MLflow 연동 필요성

• 각각의 역할
  • Airflow : 모델 학습 과정을 자동으로 스케줄링
  • MLflow : 모델 실험의 파라미터·성능 지표·결과를 기록하고 버전 관리
• 이 두 가지를 연동하면 얻게 되는 이점
  • 모델 학습을 주기적으로 자동 실행하면서
  • 각 실험의 결과를 MLflow에 자동으로 기록하여 추적 가능
  • 성능 비교, 베스트 모델 관리, 재학습 여부 판단이 쉬워짐

 

2. 연동 방법

[Airflow DAG 실행]
      ↓
[PythonOperator 내에서 MLflow Tracking API 호출]
      ↓
[실험명 설정 → 파라미터 로깅 → 메트릭 기록 → 모델 저장]
      ↓
[MLflow UI에서 모든 실험 결과 조회 가능]

 

3. 코드 예시

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
import mlflow
import mlflow.sklearn

def run_training():
    mlflow.set_tracking_uri("http://localhost:5000")  # 또는 S3, EC2 내부 주소 등
    mlflow.set_experiment("ocr_model_experiment")

    with mlflow.start_run(run_name="run_20250105") as run:
        # 파라미터 기록
        mlflow.log_param("learning_rate", 0.001)
        mlflow.log_param("batch_size", 64)

        # 모델 학습 코드 삽입
        accuracy = 0.92  # 예시

        # 메트릭 기록
        mlflow.log_metric("val_accuracy", accuracy)

        # 모델 저장
        mlflow.sklearn.log_model(sk_model="model.pkl", artifact_path="model")

        print("MLflow 연동 완료")

with DAG(
    dag_id="ml_training_with_mlflow",
    start_date=datetime(2025, 1, 1),
    schedule_interval="@monthly",
    catchup=False,
) as dag:
    train = PythonOperator(
        task_id="train_and_log_model",
        python_callable=run_training
    )

 

 

📙 내일 일정

• Airflow 수행 결과