Airflow로 ML 학습 파이프라인 오케스트레이션하기

from airflow import DAG
from datetime import datetime

with DAG(
    dag_id='ml_training_pipeline',
    start_date=datetime(2026, 1, 1),
    schedule_interval='@daily',
    catchup=False,
    tags=['ml', 'training'],
) as dag:
    # Task 정의
    pass

# Push
task_instance.xcom_push(key="model_accuracy", value=0.95)

# Pull
accuracy = task_instance.xcom_pull(key="model_accuracy", task_ids="evaluate_model")

from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from kubernetes.client import models as k8s

train_model = KubernetesPodOperator(
    task_id='train_model',
    name='gpu-training-job',
    namespace='ml-training',
    image='my-registry/ml-trainer:latest',
    cmds=['python'],
    arguments=['train.py', '--epochs', '100', '--batch-size', '64'],
    container_resources=k8s.V1ResourceRequirements(
        requests={'cpu': '4', 'memory': '16Gi', 'nvidia.com/gpu': '1'},
        limits={'cpu': '8', 'memory': '32Gi', 'nvidia.com/gpu': '1'},
    ),
    get_logs=True,
    is_delete_operator_pod=True,
    do_xcom_push=True,
)

container_resources = k8s.V1ResourceRequirements(
    requests={
        'cpu': '4',
        'memory': '16Gi',
        'nvidia.com/gpu': '1'
    },
    limits={
        'cpu': '8',
        'memory': '32Gi',
        'nvidia.com/gpu': '1'
    },
)

tolerations = [
    k8s.V1Toleration(
        key='nvidia.com/gpu',
        operator='Exists',
        effect='NoSchedule',
    ),
]

node_selector = {
    'gpu-type': 'a100',
    'node-pool': 'ml-training',
}

affinity = k8s.V1Affinity(
    node_affinity=k8s.V1NodeAffinity(
        required_during_scheduling_ignored_during_execution=k8s.V1NodeSelector(
            node_selector_terms=[
                k8s.V1NodeSelectorTerm(
                    match_expressions=[
                        k8s.V1NodeSelectorRequirement(
                            key='gpu-type',
                            operator='In',
                            values=['a100', 'h100'],
                        )
                    ]
                )
            ]
        )
    )
)

volume = k8s.V1Volume(
    name='training-data',
    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(
        claim_name='ml-dataset-pvc'
    ),
)

volume_mount = k8s.V1VolumeMount(
    name='training-data',
    mount_path='/data',
    read_only=True,
)

데이터 검증 -> 전처리 -> Feature Engineering -> 학습 -> 평가 -> [조건부] 배포

from airflow import DAG
from airflow.operators.python import BranchPythonOperator
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from datetime import datetime

with DAG(
    dag_id='ml_training_pipeline',
    start_date=datetime(2026, 1, 1),
    schedule_interval='@daily',
    catchup=False,
) as dag:

    validate_data = KubernetesPodOperator(
        task_id='validate_data',
        image='ml-tools:latest',
        cmds=['python', 'validate.py'],
        # CPU만 필요
    )

    preprocess = KubernetesPodOperator(
        task_id='preprocess_data',
        image='ml-tools:latest',
        cmds=['python', 'preprocess.py'],
        # 대용량 메모리 필요
    )

    train = KubernetesPodOperator(
        task_id='train_model',
        image='ml-trainer:latest',
        cmds=['python', 'train.py'],
        # GPU 필요
        do_xcom_push=True,
    )

    evaluate = KubernetesPodOperator(
        task_id='evaluate_model',
        image='ml-trainer:latest',
        cmds=['python', 'evaluate.py'],
        do_xcom_push=True,
    )

    def check_model_quality(**context):
        metrics = context['ti'].xcom_pull(task_ids='evaluate_model')
        if metrics['accuracy'] > 0.90:
            return 'deploy_model'
        return 'notify_failure'

    branch = BranchPythonOperator(
        task_id='check_quality',
        python_callable=check_model_quality,
    )

    deploy = KubernetesPodOperator(
        task_id='deploy_model',
        image='ml-deployer:latest',
        cmds=['python', 'deploy.py'],
    )

    validate_data >> preprocess >> train >> evaluate >> branch
    branch >> [deploy]

from airflow.utils.task_group import TaskGroup

with TaskGroup('data_preparation') as data_prep:
    validate_data >> preprocess >> feature_engineering

with TaskGroup('model_training') as training:
    train >> evaluate

data_prep >> training >> branch

# train.py (KubernetesPodOperator 내부에서 실행)
import json

metrics = {
    'accuracy': 0.9534,
    'f1_score': 0.9421,
    'loss': 0.0312,
    'model_path': 's3://ml-models/experiment-42/model.pt',
    'run_id': 'exp-42-20260301',
}

# /airflow/xcom/return.json에 JSON 작성
with open('/airflow/xcom/return.json', 'w') as f:
    json.dump(metrics, f)

value = task_instance.xcom_pull(task_ids='train_model')

@task(multiple_outputs=True)
def evaluate_model(**context):
    return {
        'accuracy': 0.95,
        'f1_score': 0.94,
        'model_path': 's3://models/latest.pt',
    }

[core]
xcom_backend = airflow.providers.common.io.xcom.backend.XComObjectStorageBackend

from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator

hyperparams = [
    {'lr': '0.001', 'batch_size': '32', 'optimizer': 'adam'},
    {'lr': '0.01', 'batch_size': '64', 'optimizer': 'adam'},
    {'lr': '0.001', 'batch_size': '32', 'optimizer': 'sgd'},
    {'lr': '0.0001', 'batch_size': '128', 'optimizer': 'adamw'},
]

train_tasks = KubernetesPodOperator.partial(
    task_id='hyperparam_training',
    image='ml-trainer:latest',
    namespace='ml-training',
    get_logs=True,
    is_delete_operator_pod=True,
    do_xcom_push=True,
    container_resources=k8s.V1ResourceRequirements(
        requests={'nvidia.com/gpu': '1'},
        limits={'nvidia.com/gpu': '1'},
    ),
).expand(
    arguments=[
        ['train.py', '--lr', hp['lr'], '--batch-size', hp['batch_size'], '--optimizer', hp['optimizer']]
        for hp in hyperparams
    ],
)

@task
def generate_hyperparams():
    """Grid Search 또는 Random Search로 하이퍼파라미터 생성"""
    import itertools

    learning_rates = [0.001, 0.01, 0.0001]
    batch_sizes = [32, 64, 128]
    optimizers = ['adam', 'sgd']

    combinations = list(itertools.product(learning_rates, batch_sizes, optimizers))
    return [
        {'lr': str(lr), 'batch_size': str(bs), 'optimizer': opt}
        for lr, bs, opt in combinations
    ]

hp_list = generate_hyperparams()

@task
def train(lr: float, batch_size: int):
    # 학습 로직
    pass

train.expand(lr=[0.001, 0.01], batch_size=[32, 64])
# 2 x 2 = 4개의 Task Instance 생성

@task
def select_best_model(results):
    """모든 하이퍼파라미터 조합의 결과 중 최선을 선택"""
    best = max(results, key=lambda x: x['accuracy'])
    return best

training_results = train_task.expand(params=hyperparams)
best = select_best_model(training_results)

from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor

wait_for_data = S3KeySensor(
    task_id='wait_for_training_data',
    bucket_name='ml-datasets',
    bucket_key='daily/{{ ds }}/training_data.parquet',
    aws_conn_id='aws_default',
    poke_interval=300,  # 5분마다 확인
    timeout=3600 * 6,   # 최대 6시간 대기
    mode='reschedule',  # Worker slot 반환
    deferrable=True,    # Triggerer를 사용한 비동기 대기
)

from airflow.sensors.external_task import ExternalTaskSensor

wait_for_preprocessing = ExternalTaskSensor(
    task_id='wait_for_preprocessing',
    external_dag_id='data_preprocessing_pipeline',
    external_task_id='final_validation',
    allowed_states=['success'],
    execution_delta=timedelta(hours=0),
    timeout=3600,
    poke_interval=60,
    mode='reschedule',
)

wait_for_data >> validate_data >> preprocess >> train >> evaluate >> deploy

# train.py (컨테이너 내부에서 실행)
import mlflow
import mlflow.pytorch

mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment("daily-training")

with mlflow.start_run() as run:
    # 하이퍼파라미터 로깅
    mlflow.log_params({
        'learning_rate': 0.001,
        'batch_size': 64,
        'epochs': 100,
    })

    # 모델 학습
    model = train_model(...)

    # 메트릭 로깅
    mlflow.log_metrics({
        'accuracy': 0.95,
        'f1_score': 0.94,
    })

    # 모델 등록
    mlflow.pytorch.log_model(model, "model")

    # run_id를 XCom으로 전달
    import json
    with open('/airflow/xcom/return.json', 'w') as f:
        json.dump({'run_id': run.info.run_id}, f)

@task
def register_model(run_id: str, accuracy: float):
    import mlflow
    from mlflow.tracking import MlflowClient

    client = MlflowClient("http://mlflow-server:5000")

    # 모델 등록
    model_uri = f"runs:/{run_id}/model"
    mv = mlflow.register_model(model_uri, "production-model")

    # 정확도가 기준을 넘으면 Production 스테이지로 전환
    if accuracy > 0.93:
        client.transition_model_version_stage(
            name="production-model",
            version=mv.version,
            stage="Production",
        )

@task
def compare_with_baseline():
    from mlflow.tracking import MlflowClient

    client = MlflowClient("http://mlflow-server:5000")

    # Production 모델의 메트릭 조회
    prod_versions = client.get_latest_versions("production-model", stages=["Production"])
    if prod_versions:
        prod_run = client.get_run(prod_versions[0].run_id)
        baseline_accuracy = float(prod_run.data.metrics['accuracy'])
        return {'baseline_accuracy': baseline_accuracy}
    return {'baseline_accuracy': 0.0}

from airflow.sdk import dag, task
from datetime import datetime

@dag(
    schedule='@daily',
    start_date=datetime(2026, 1, 1),
    catchup=False,
)
def ml_training_taskflow():

    @task
    def extract_data():
        """데이터 추출"""
        return {'data_path': 's3://bucket/data/2026-03-01.parquet', 'num_rows': 100000}

    @task(multiple_outputs=True)
    def preprocess(data_info: dict):
        """데이터 전처리"""
        return {
            'train_path': f"{data_info['data_path']}/train.parquet",
            'test_path': f"{data_info['data_path']}/test.parquet",
            'feature_count': 128,
        }

    @task
    def train_model(train_path: str, feature_count: int):
        """모델 학습"""
        # 실제로는 KubernetesPodOperator로 GPU 학습 실행
        return {
            'model_path': 's3://models/latest.pt',
            'accuracy': 0.95,
        }

    @task
    def evaluate_and_deploy(model_info: dict):
        """모델 평가 및 조건부 배포"""
        if model_info['accuracy'] > 0.90:
            return f"Deployed model from {model_info['model_path']}"
        return "Model quality below threshold, skipping deployment"

    # 자동 의존성 연결
    data = extract_data()
    processed = preprocess(data)
    model = train_model(
        train_path=processed['train_path'],
        feature_count=processed['feature_count'],
    )
    evaluate_and_deploy(model)

ml_training_taskflow()

@task
def log_execution_info(task_instance=None, dag_run=None):
    print(f"Run ID: {task_instance.run_id}")
    print(f"DAG Run: {dag_run.dag_id}")

"""
ML Training Pipeline DAG
- S3 데이터 도착 감지
- 데이터 전처리 (KubernetesPodOperator)
- Dynamic Task Mapping을 활용한 하이퍼파라미터 튜닝
- 최적 모델 선택 및 MLflow 등록
- 조건부 배포
"""

from __future__ import annotations

from datetime import datetime, timedelta
from airflow import DAG
from airflow.sdk import task
from airflow.operators.python import BranchPythonOperator
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
from airflow.utils.task_group import TaskGroup
from kubernetes.client import models as k8s

# -- 공통 설정 --
GPU_RESOURCES = k8s.V1ResourceRequirements(
    requests={'cpu': '4', 'memory': '16Gi', 'nvidia.com/gpu': '1'},
    limits={'cpu': '8', 'memory': '32Gi', 'nvidia.com/gpu': '1'},
)

GPU_TOLERATIONS = [
    k8s.V1Toleration(
        key='nvidia.com/gpu',
        operator='Exists',
        effect='NoSchedule',
    ),
]

GPU_NODE_SELECTOR = {'gpu-type': 'a100', 'node-pool': 'ml-training'}

DATA_VOLUME = k8s.V1Volume(
    name='shared-data',
    persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(
        claim_name='ml-shared-pvc'
    ),
)

DATA_VOLUME_MOUNT = k8s.V1VolumeMount(
    name='shared-data',
    mount_path='/shared',
)

default_args = {
    'owner': 'ml-team',
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
    'execution_timeout': timedelta(hours=4),
}


with DAG(
    dag_id='ml_full_training_pipeline',
    default_args=default_args,
    start_date=datetime(2026, 1, 1),
    schedule_interval='@daily',
    catchup=False,
    tags=['ml', 'training', 'production'],
    max_active_runs=1,
    doc_md="""
    ## ML Full Training Pipeline
    매일 새로운 데이터로 모델을 재학습하고,
    품질 기준을 충족하면 자동 배포하는 파이프라인.
    """,
) as dag:

    # ===== Stage 1: 데이터 도착 감지 =====
    wait_for_data = S3KeySensor(
        task_id='wait_for_training_data',
        bucket_name='ml-datasets',
        bucket_key='daily/{{ ds }}/features.parquet',
        aws_conn_id='aws_default',
        poke_interval=300,
        timeout=3600 * 6,
        mode='reschedule',
        deferrable=True,
    )

    # ===== Stage 2: 데이터 전처리 =====
    with TaskGroup('data_preparation') as data_prep:

        validate_data = KubernetesPodOperator(
            task_id='validate_data',
            name='data-validation',
            namespace='ml-training',
            image='ml-tools:latest',
            cmds=['python', 'validate.py'],
            arguments=['--date', '{{ ds }}', '--source', 's3://ml-datasets/daily/{{ ds }}/'],
            env_vars=[
                k8s.V1EnvVar(name='DATA_DATE', value='{{ ds }}'),
            ],
            container_resources=k8s.V1ResourceRequirements(
                requests={'cpu': '2', 'memory': '8Gi'},
                limits={'cpu': '4', 'memory': '16Gi'},
            ),
            get_logs=True,
            is_delete_operator_pod=True,
            do_xcom_push=True,
        )

        preprocess = KubernetesPodOperator(
            task_id='preprocess_data',
            name='data-preprocessing',
            namespace='ml-training',
            image='ml-tools:latest',
            cmds=['python', 'preprocess.py'],
            arguments=[
                '--input', 's3://ml-datasets/daily/{{ ds }}/',
                '--output', 's3://ml-processed/{{ ds }}/',
            ],
            container_resources=k8s.V1ResourceRequirements(
                requests={'cpu': '4', 'memory': '32Gi'},
                limits={'cpu': '8', 'memory': '64Gi'},
            ),
            volumes=[DATA_VOLUME],
            volume_mounts=[DATA_VOLUME_MOUNT],
            get_logs=True,
            is_delete_operator_pod=True,
            do_xcom_push=True,
        )

        validate_data >> preprocess

    # ===== Stage 3: 하이퍼파라미터 생성 및 병렬 학습 =====
    @task
    def generate_hyperparams():
        """학습할 하이퍼파라미터 조합 생성"""
        import itertools
        learning_rates = [0.001, 0.0001]
        batch_sizes = [32, 64]
        optimizers = ['adam', 'adamw']

        combos = list(itertools.product(learning_rates, batch_sizes, optimizers))
        return [
            [
                'train.py',
                '--lr', str(lr),
                '--batch-size', str(bs),
                '--optimizer', opt,
                '--data-path', 's3://ml-processed/{{ ds }}/',
                '--experiment-name', 'daily-training-{{ ds }}',
            ]
            for lr, bs, opt in combos
        ]

    hp_args = generate_hyperparams()

    with TaskGroup('hyperparameter_tuning') as hp_tuning:

        train_tasks = KubernetesPodOperator.partial(
            task_id='train_with_hp',
            name='gpu-hp-training',
            namespace='ml-training',
            image='ml-trainer:latest',
            cmds=['python'],
            container_resources=GPU_RESOURCES,
            tolerations=GPU_TOLERATIONS,
            node_selector=GPU_NODE_SELECTOR,
            volumes=[DATA_VOLUME],
            volume_mounts=[DATA_VOLUME_MOUNT],
            env_vars=[
                k8s.V1EnvVar(name='MLFLOW_TRACKING_URI', value='http://mlflow:5000'),
                k8s.V1EnvVar(name='CUDA_VISIBLE_DEVICES', value='0'),
            ],
            get_logs=True,
            is_delete_operator_pod=True,
            do_xcom_push=True,
            startup_timeout_seconds=600,
        ).expand(arguments=hp_args)

    # ===== Stage 4: 최적 모델 선택 =====
    @task
    def select_best_model(training_results):
        """모든 학습 결과 중 최적 모델 선택"""
        valid_results = [r for r in training_results if r is not None]
        if not valid_results:
            raise ValueError("No valid training results found")

        best = max(valid_results, key=lambda x: x.get('accuracy', 0))
        print(f"Best model: accuracy={best['accuracy']}, run_id={best['run_id']}")
        return best

    best_model = select_best_model(train_tasks)

    # ===== Stage 5: 모델 평가 =====
    @task(multiple_outputs=True)
    def evaluate_model(model_info: dict):
        """최적 모델에 대한 상세 평가"""
        import mlflow
        from mlflow.tracking import MlflowClient

        client = MlflowClient("http://mlflow:5000")

        # 현재 Production 모델과 비교
        try:
            prod_versions = client.get_latest_versions(
                "production-model", stages=["Production"]
            )
            if prod_versions:
                prod_run = client.get_run(prod_versions[0].run_id)
                baseline_accuracy = float(prod_run.data.metrics.get('accuracy', 0))
            else:
                baseline_accuracy = 0.0
        except Exception:
            baseline_accuracy = 0.0

        new_accuracy = model_info['accuracy']
        improvement = new_accuracy - baseline_accuracy

        return {
            'should_deploy': improvement > 0.005 and new_accuracy > 0.90,
            'new_accuracy': new_accuracy,
            'baseline_accuracy': baseline_accuracy,
            'improvement': improvement,
            'run_id': model_info['run_id'],
            'model_path': model_info.get('model_path', ''),
        }

    eval_result = evaluate_model(best_model)

    # ===== Stage 6: 조건부 배포 =====
    def decide_deployment(**context):
        should_deploy = context['ti'].xcom_pull(
            task_ids='evaluate_model', key='should_deploy'
        )
        if should_deploy:
            return 'deploy_model'
        return 'skip_deployment'

    deployment_branch = BranchPythonOperator(
        task_id='deployment_decision',
        python_callable=decide_deployment,
    )

    deploy = KubernetesPodOperator(
        task_id='deploy_model',
        name='model-deployment',
        namespace='ml-serving',
        image='ml-deployer:latest',
        cmds=['python', 'deploy.py'],
        arguments=[
            '--run-id', "{{ ti.xcom_pull(task_ids='evaluate_model', key='run_id') }}",
            '--model-name', 'production-model',
        ],
        container_resources=k8s.V1ResourceRequirements(
            requests={'cpu': '2', 'memory': '4Gi'},
            limits={'cpu': '4', 'memory': '8Gi'},
        ),
        get_logs=True,
        is_delete_operator_pod=True,
    )

    @task(task_id='skip_deployment')
    def skip_deployment():
        print("Model quality does not meet deployment criteria. Skipping.")

    skip = skip_deployment()

    # ===== Stage 7: 알림 =====
    @task(trigger_rule='none_failed_min_one_success')
    def send_notification(**context):
        """학습 결과 알림 전송"""
        eval_data = context['ti'].xcom_pull(task_ids='evaluate_model', key='return_value')
        print(f"Pipeline completed. Accuracy: {eval_data.get('new_accuracy')}")
        print(f"Improvement: {eval_data.get('improvement')}")
        # Slack / Email 알림 로직

    notification = send_notification()

    # ===== DAG 의존성 정의 =====
    wait_for_data >> data_prep >> hp_args >> hp_tuning
    hp_tuning >> best_model >> eval_result >> deployment_branch
    deployment_branch >> [deploy, skip] >> notification