[Airflow - 5단계: PythonOperator + MLflow Tracking 연동]

이 글에서 다루는 것

PythonOperator로 MLflow 실험을 자동화하여, 하이퍼파라미터·메트릭·모델을 자동 기록하는 Airflow + MLflow 연동 패턴을 다룹니다.

선수지식

• Airflow 4단계: BashOperator로 외부 학습 스크립트 실행 — 외부 스크립트 분리 패턴

이 단계에서 해결하려는 문제

모델을 학습할 때마다 하이퍼파라미터, 정확도, 모델 아티팩트를 수동으로 기록하면 실험 재현이 불가능합니다. MLflow Tracking은 이 모든 정보를 자동 기록하고, Airflow DAG에서 호출하면 학습 파이프라인 자체를 자동화할 수 있습니다.

실습 코드: GitHub (Airflow_and_MLflow)

🧭 실습 전체 흐름 요약

[1단계] MLflow 실험 스크립트 작성
[2단계] Airflow DAG 구성
[3단계] DAG 실행 및 파라미터/메트릭 확인
[4단계] 모델 저장 및 로깅 상태 점검

📁 디렉토리 구조

airflow/
├── dags/
│   ├── train_with_mlflow.py        ← DAG 파일
├── ml_code/
│   ├── train_mlflow.py             ← MLflow 연동 학습 스크립트
└── mlruns/                         ← MLflow 로깅 결과 저장 폴더 (자동 생성)

🧪 1단계: MLflow 학습 스크립트 (핵심 부분)

# airflow/ml_code/train_mlflow.py

def run_experiment():
    mlflow.set_tracking_uri("file:/opt/airflow/mlruns")
    mlflow.set_experiment("airflow_mlflow_example")

    with mlflow.start_run():
        data = load_iris()
        X, y = data.data, data.target

        model = RandomForestClassifier(n_estimators=50, max_depth=3)
        model.fit(X, y)
        preds = model.predict(X)

        # 학습 데이터로 평가 — 실제 운영에서는 train_test_split 필수
        acc = accuracy_score(y, preds)

        mlflow.log_param("n_estimators", 50)
        mlflow.log_param("max_depth", 3)
        mlflow.log_metric("accuracy", acc)
        mlflow.sklearn.log_model(model, "model")

전체 코드: GitHub (train_mlflow.py)

🧪 2단계: Airflow DAG 작성

# airflow/dags/train_with_mlflow.py

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
import sys
sys.path.append("/opt/airflow/ml_code")

from train_mlflow import run_experiment

with DAG(
    dag_id='mlflow_tracking_dag',
    start_date=datetime(2023, 1, 1),
    schedule_interval=None,
    catchup=False,
) as dag:

    run_mlflow = PythonOperator(
        task_id='run_mlflow_training',
        python_callable=run_experiment,
    )

✅ 실행 절차

1. train_mlflow.py 작성
2. train_with_mlflow.py DAG 등록
3. Airflow UI에서 DAG 실행
4. Task 로그에서 파라미터/메트릭/모델 기록 확인

❓ 모듈 에러가 발생하면?

1. docker-compose.yaml에서 볼륨 확인:

   volumes:
     - ./ml_code:/opt/airflow/ml_code
   

2. requirements.txt에 MLflow 추가:

   mlflow
   

3. Dockerfile 작성 후 재빌드:

   docker-compose down
   docker-compose build
   docker-compose up -d
   

🔍 확인 포인트

설계 판단 (Why This Way?)

MLflow 함수를 직접 호출해야 하므로 PythonOperator가 적합하며, 로컬 file URI 대신 프로덕션에서는 PostgreSQL+S3 백엔드의 MLflow 서버를 사용해 팀 단위 실험 추적과 아티팩트 관리를 수행합니다.

다음에 읽을 글

→ MLOps 플랫폼 구축 1단계: 인프라 설계 및 환경 준비 — Level 2: Helm 기반 MLOps 플랫폼 구축 시작