Apache Airflow

# dags/ml_training_pipeline.py
# ML 모델 학습 파이프라인 DAG

from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.providers.amazon.aws.operators.s3 import S3CreateObjectOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.utils.trigger_rule import TriggerRule

# TaskFlow API를 사용한 현대적인 DAG 작성
@dag(
    dag_id='ml_training_pipeline',
    schedule_interval='@weekly',
    start_date=datetime(2024, 1, 1),
    catchup=False,
    default_args={
        'owner': 'ml-team',
        'retries': 2,
        'retry_delay': timedelta(minutes=10),
    },
    tags=['ml', 'training'],
)
def ml_training_pipeline():

    @task
    def prepare_training_data(execution_date: str) -> dict:
        """학습 데이터 준비"""
        import pandas as pd
        from sklearn.model_selection import train_test_split

        # 데이터 로드 (예시)
        s3_hook = S3Hook(aws_conn_id='aws_default')
        data_key = f'raw/features/{execution_date}/data.parquet'

        # 데이터 전처리
        # df = pd.read_parquet(...)
        # X_train, X_test, y_train, y_test = train_test_split(...)

        return {
            'train_path': f's3://ml-bucket/processed/{execution_date}/train.parquet',
            'test_path': f's3://ml-bucket/processed/{execution_date}/test.parquet',
            'feature_count': 150,
            'train_samples': 100000,
        }

    @task
    def train_model(data_info: dict) -> dict:
        """모델 학습"""
        import mlflow
        from sklearn.ensemble import RandomForestClassifier

        mlflow.set_tracking_uri('http://mlflow:5000')
        mlflow.set_experiment('weekly-training')

        with mlflow.start_run():
            # 모델 학습
            model = RandomForestClassifier(n_estimators=100)
            # model.fit(X_train, y_train)

            # 메트릭 기록
            mlflow.log_params({
                'n_estimators': 100,
                'feature_count': data_info['feature_count'],
            })
            mlflow.log_metrics({
                'accuracy': 0.95,
                'f1_score': 0.93,
            })

            # 모델 저장
            model_uri = mlflow.sklearn.log_model(model, 'model').model_uri

        return {
            'model_uri': model_uri,
            'accuracy': 0.95,
            'f1_score': 0.93,
        }

    @task
    def evaluate_model(model_info: dict) -> dict:
        """모델 평가 및 검증"""
        # A/B 테스트용 기준 모델과 비교
        baseline_accuracy = 0.90

        if model_info['accuracy'] > baseline_accuracy:
            return {
                'approved': True,
                'model_uri': model_info['model_uri'],
                'improvement': model_info['accuracy'] - baseline_accuracy,
            }
        else:
            raise ValueError(f"Model accuracy {model_info['accuracy']} below baseline {baseline_accuracy}")

    @task.branch
    def check_deployment_decision(evaluation: dict) -> str:
        """배포 여부 결정"""
        if evaluation['approved'] and evaluation['improvement'] > 0.02:
            return 'deploy_to_production'
        elif evaluation['approved']:
            return 'deploy_to_staging'
        else:
            return 'notify_failure'

    @task
    def deploy_to_production(evaluation: dict):
        """프로덕션 배포"""
        import mlflow.sagemaker

        print(f"Deploying model {evaluation['model_uri']} to production")
        # mlflow.sagemaker.deploy(...)

    @task
    def deploy_to_staging(evaluation: dict):
        """스테이징 배포"""
        print(f"Deploying model {evaluation['model_uri']} to staging")

    @task(trigger_rule=TriggerRule.ONE_FAILED)
    def notify_failure():
        """실패 알림"""
        from airflow.providers.slack.hooks.slack_webhook import SlackWebhookHook

        hook = SlackWebhookHook(slack_webhook_conn_id='slack_ml_alerts')
        hook.send(text="ML Training Pipeline Failed!")

    @task(trigger_rule=TriggerRule.ONE_SUCCESS)
    def notify_success(**context):
        """성공 알림"""
        from airflow.providers.slack.hooks.slack_webhook import SlackWebhookHook

        hook = SlackWebhookHook(slack_webhook_conn_id='slack_ml_alerts')
        hook.send(text=f"ML Training Pipeline Completed! Run: {context['run_id']}")

    # DAG 흐름 정의
    data_info = prepare_training_data(execution_date='{{ ds }}')
    model_info = train_model(data_info)
    evaluation = evaluate_model(model_info)
    decision = check_deployment_decision(evaluation)

    # 조건부 분기
    prod_deploy = deploy_to_production(evaluation)
    staging_deploy = deploy_to_staging(evaluation)
    failure_notify = notify_failure()
    success_notify = notify_success()

    decision >> [prod_deploy, staging_deploy, failure_notify]
    [prod_deploy, staging_deploy] >> success_notify

# DAG 인스턴스 생성
dag = ml_training_pipeline()