AI 플랫폼 스택 설계: Kubeflow, MLflow, KServe 통합 운영

"""
Kubeflow Pipelines v2 component로 학습을 수행하면서
MLflow에 실험 결과를 기록하는 예시.
"""
from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model, Metrics

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=[
        "mlflow==2.17.2",
        "scikit-learn==1.5.2",
        "pandas==2.2.3",
        "boto3==1.35.0",
    ],
)
def train_model(
    training_data: Input[Dataset],
    model_output: Output[Model],
    metrics_output: Output[Metrics],
    mlflow_tracking_uri: str,
    mlflow_experiment_name: str,
    n_estimators: int = 200,
    max_depth: int = 10,
    learning_rate: float = 0.1,
):
    import mlflow
    import pandas as pd
    from sklearn.ensemble import GradientBoostingClassifier
    from sklearn.model_selection import cross_val_score
    import json
    import os

    # MLflow 연결
    mlflow.set_tracking_uri(mlflow_tracking_uri)
    mlflow.set_experiment(mlflow_experiment_name)

    # 데이터 로드
    df = pd.read_parquet(training_data.path)
    X = df.drop(columns=["label"])
    y = df["label"]

    with mlflow.start_run() as run:
        # 하이퍼파라미터 기록
        mlflow.log_params({
            "n_estimators": n_estimators,
            "max_depth": max_depth,
            "learning_rate": learning_rate,
            "feature_count": X.shape[1],
            "training_rows": X.shape[0],
        })

        # 학습
        clf = GradientBoostingClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            learning_rate=learning_rate,
            random_state=42,
        )
        cv_scores = cross_val_score(clf, X, y, cv=5, scoring="f1_weighted")
        clf.fit(X, y)

        # 메트릭 기록
        mlflow.log_metrics({
            "cv_f1_mean": float(cv_scores.mean()),
            "cv_f1_std": float(cv_scores.std()),
        })

        # 모델 저장 (MLflow Model Registry 형식)
        mlflow.sklearn.log_model(
            clf,
            artifact_path="model",
            registered_model_name="recommendation-classifier",
        )

        # Kubeflow Metrics에도 기록 (UI 표시용)
        metrics_output.log_metric("cv_f1_mean", float(cv_scores.mean()))
        metrics_output.log_metric("cv_f1_std", float(cv_scores.std()))
        metrics_output.log_metric("mlflow_run_id", run.info.run_id)

        # 모델 경로를 output으로 전달 (다음 component에서 사용)
        model_output.uri = f"runs:/{run.info.run_id}/model"
        model_output.metadata["mlflow_run_id"] = run.info.run_id

from kfp import dsl, compiler

@dsl.pipeline(
    name="recommendation-training-pipeline",
    description="데이터 검증 -> 학습 -> 모델 등록 -> 서빙 배포",
)
def training_pipeline(
    data_path: str = "gs://ml-data/recommendation/2026-03-04/",
    mlflow_tracking_uri: str = "http://mlflow.ml-platform.svc:5000",
    mlflow_experiment: str = "recommendation-v3",
    serving_namespace: str = "model-serving",
):
    # Step 1: 데이터 검증
    validate_task = validate_data(data_path=data_path)

    # Step 2: 전처리
    preprocess_task = preprocess_data(
        raw_data=validate_task.outputs["validated_data"],
    )

    # Step 3: 학습 + MLflow 추적
    train_task = train_model(
        training_data=preprocess_task.outputs["processed_data"],
        mlflow_tracking_uri=mlflow_tracking_uri,
        mlflow_experiment_name=mlflow_experiment,
        n_estimators=300,
        max_depth=8,
        learning_rate=0.05,
    )
    train_task.set_cpu_limit("4").set_memory_limit("16Gi")
    train_task.set_accelerator_type("nvidia.com/gpu").set_accelerator_limit(1)

    # Step 4: 모델 품질 게이트
    gate_task = quality_gate(
        metrics=train_task.outputs["metrics_output"],
        f1_threshold=0.80,
    )

    # Step 5: KServe 배포 (품질 게이트 통과 시)
    with dsl.Condition(gate_task.outputs["passed"] == "true"):
        deploy_task = deploy_to_kserve(
            model=train_task.outputs["model_output"],
            serving_namespace=serving_namespace,
        )

compiler.Compiler().compile(
    pipeline_func=training_pipeline,
    package_path="recommendation_pipeline.yaml",
)

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["kubernetes==31.0.0", "mlflow==2.17.2"],
)
def deploy_to_kserve(
    model: Input[Model],
    serving_namespace: str,
    canary_traffic_percent: int = 20,
):
    """
    MLflow에 등록된 모델을 KServe InferenceService로 배포한다.
    canary 방식으로 신규 버전에 트래픽 일부를 할당한다.
    """
    from kubernetes import client, config
    import json
    import mlflow

    config.load_incluster_config()
    api = client.CustomObjectsApi()

    mlflow_run_id = model.metadata.get("mlflow_run_id")
    model_uri = model.uri  # runs:/<run_id>/model

    # MLflow에서 GCS 경로 추출
    mlflow_client = mlflow.tracking.MlflowClient()
    run = mlflow_client.get_run(mlflow_run_id)
    artifact_uri = run.info.artifact_uri  # gs://mlflow-artifacts/<exp>/<run>/artifacts

    storage_uri = f"{artifact_uri}/model"

    inference_service = {
        "apiVersion": "serving.kserve.io/v1beta1",
        "kind": "InferenceService",
        "metadata": {
            "name": "recommendation-classifier",
            "namespace": serving_namespace,
            "labels": {
                "mlflow-run-id": mlflow_run_id,
                "pipeline": "recommendation-training",
            },
            "annotations": {
                "serving.kserve.io/deploymentMode": "Serverless",
                "serving.kserve.io/autoscalerClass": "hpa",
                "serving.kserve.io/metrics": "cpu",
                "serving.kserve.io/targetUtilizationPercentage": "70",
            },
        },
        "spec": {
            "predictor": {
                "canaryTrafficPercent": canary_traffic_percent,
                "minReplicas": 2,
                "maxReplicas": 10,
                "model": {
                    "modelFormat": {"name": "mlflow"},
                    "storageUri": storage_uri,
                    "resources": {
                        "requests": {"cpu": "1", "memory": "2Gi"},
                        "limits": {"cpu": "2", "memory": "4Gi"},
                    },
                },
            },
        },
    }

    try:
        api.patch_namespaced_custom_object(
            group="serving.kserve.io",
            version="v1beta1",
            namespace=serving_namespace,
            plural="inferenceservices",
            name="recommendation-classifier",
            body=inference_service,
        )
        print(f"Updated InferenceService with canary {canary_traffic_percent}%")
    except client.exceptions.ApiException as e:
        if e.status == 404:
            api.create_namespaced_custom_object(
                group="serving.kserve.io",
                version="v1beta1",
                namespace=serving_namespace,
                plural="inferenceservices",
                body=inference_service,
            )
            print("Created new InferenceService")
        else:
            raise

"""
Canary 배포 후 일정 시간 동안 메트릭을 관찰하여
자동으로 승격하거나 롤백하는 판정 스크립트.
CronJob 또는 Argo Workflow step으로 실행한다.
"""
import requests
import time
from dataclasses import dataclass
from typing import Optional

@dataclass
class CanaryConfig:
    service_name: str
    namespace: str
    prometheus_url: str
    observation_minutes: int = 15
    check_interval_seconds: int = 60
    error_rate_threshold: float = 0.02    # 2%
    p99_latency_threshold_ms: float = 500  # 500ms
    min_request_count: int = 100           # 최소 관찰 요청 수

def query_prometheus(config: CanaryConfig, query: str) -> Optional[float]:
    """Prometheus에서 메트릭 값을 조회한다."""
    resp = requests.get(
        f"{config.prometheus_url}/api/v1/query",
        params={"query": query},
        timeout=10,
    )
    results = resp.json().get("data", {}).get("result", [])
    if results:
        return float(results[0]["value"][1])
    return None

def evaluate_canary(config: CanaryConfig) -> dict:
    """Canary 버전의 error rate와 latency를 평가한다."""
    error_rate_query = (
        f'sum(rate(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}",response_code=~"5.."}}[5m])) / '
        f'sum(rate(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}"}}[5m]))'
    )
    p99_query = (
        f'histogram_quantile(0.99, sum(rate(kserve_request_duration_seconds_bucket'
        f'{{service_name="{config.service_name}",namespace="{config.namespace}"}}'
        f'[5m])) by (le)) * 1000'
    )
    request_count_query = (
        f'sum(increase(kserve_request_total{{service_name="{config.service_name}",'
        f'namespace="{config.namespace}"}}[{config.observation_minutes}m]))'
    )

    error_rate = query_prometheus(config, error_rate_query) or 0.0
    p99_latency = query_prometheus(config, p99_query) or 0.0
    request_count = query_prometheus(config, request_count_query) or 0

    return {
        "error_rate": error_rate,
        "p99_latency_ms": p99_latency,
        "request_count": request_count,
        "error_rate_ok": error_rate < config.error_rate_threshold,
        "latency_ok": p99_latency < config.p99_latency_threshold_ms,
        "sufficient_traffic": request_count >= config.min_request_count,
    }

def run_canary_judgment(config: CanaryConfig) -> str:
    """
    observation_minutes 동안 메트릭을 관찰하고 승격/롤백을 결정한다.
    Returns: "promote" or "rollback"
    """
    checks_passed = 0
    total_checks = config.observation_minutes * 60 // config.check_interval_seconds

    for i in range(total_checks):
        result = evaluate_canary(config)
        print(f"Check {i+1}/{total_checks}: {result}")

        if not result["sufficient_traffic"]:
            print("Insufficient traffic, waiting...")
            time.sleep(config.check_interval_seconds)
            continue

        if result["error_rate_ok"] and result["latency_ok"]:
            checks_passed += 1
        else:
            # 즉시 롤백 조건: error rate가 임계치의 3배 초과
            if result["error_rate"] > config.error_rate_threshold * 3:
                print(f"Immediate rollback: error_rate={result['error_rate']}")
                return "rollback"

        time.sleep(config.check_interval_seconds)

    pass_ratio = checks_passed / max(total_checks, 1)
    decision = "promote" if pass_ratio >= 0.9 else "rollback"
    print(f"Decision: {decision} (pass_ratio={pass_ratio:.2f})")
    return decision

증상: Pipeline이 성공(green)으로 표시되지만 MLflow UI에 해당 Run이 없음
에러 로그 (Pipeline pod):
  requests.exceptions.ConnectionError: HTTPConnectionPool(host='mlflow.ml-platform.svc', port=5000):
  Max retries exceeded with url: /api/2.0/mlflow/runs/create

원인: Kubeflow Pipeline pod의 ServiceAccount에 MLflow service에 대한
      NetworkPolicy 접근이 허용되지 않음

해결:
  1. NetworkPolicy에서 pipeline runner namespace -> mlflow namespace 허용
  2. MLflow tracking URI가 cluster-internal DNS인지 확인
  3. Pipeline component에 retry 로직 추가

증상: InferenceService가 "FailedCreate" 상태로 멈춤
에러 로그 (kserve-container):
  mlflow.exceptions.MlflowException: Could not find an "MLmodel" configuration file
  at "gs://mlflow-artifacts/3/abc123def/artifacts/model"

원인: MLflow log_model의 artifact_path와 KServe storageUri의 경로 불일치
      (artifact_path="model" 이지만 storageUri에 "/model" 이 빠짐)

해결:
  storageUri를 artifact_uri + "/model" 형태로 조합해야 한다.
  잘못된 예: gs://mlflow-artifacts/3/abc123def/artifacts
  올바른 예: gs://mlflow-artifacts/3/abc123def/artifacts/model

증상: canaryTrafficPercent을 0으로 설정해도 이전 버전으로 안 돌아감
에러 로그 (kserve-controller):
  RevisionFailed: Revision "recommendation-classifier-predictor-prev"
  has no ready pods

원인: KServe가 이전 Revision의 pod를 scale-to-zero한 상태에서
      cold start 시간이 readiness probe timeout을 초과

해결:
  1. 서빙 서비스에 minReplicas >= 1 설정 (중요 서비스는 scale-to-zero 비활성화)
  2. readinessProbe timeout을 모델 로딩 시간 기준으로 늘림
  3. 롤백 시에는 canary 비율 조정이 아닌 storageUri를 이전 버전으로 교체

# 롤백 시 적용할 KServe 매니페스트
apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: recommendation-classifier
  namespace: model-serving
spec:
  predictor:
    canaryTrafficPercent: 0 # canary 트래픽 제거
    minReplicas: 2 # scale-to-zero 방지
    model:
      modelFormat:
        name: mlflow
      # 이전 안정 버전의 model URI로 교체
      storageUri: gs://mlflow-artifacts/3/previous-stable-run/artifacts/model
      resources:
        requests:
          cpu: '1'
          memory: '2Gi'
        limits:
          cpu: '2'
          memory: '4Gi'
    containerSpec:
      readinessProbe:
        initialDelaySeconds: 30
        timeoutSeconds: 10
        periodSeconds: 5