AI 플랫폼 모델 레지스트리와 A/B 배포 파이프라인 설계 2026

"""
MLflow Model Registry: 모델 등록, 버전 관리, Alias 설정
MLflow 2.19+ 기준
"""
import mlflow
from mlflow.tracking import MlflowClient

# MLflow Tracking Server 연결
mlflow.set_tracking_uri("http://mlflow.internal.example.com:5000")
client = MlflowClient()

# 1. 학습 실행과 모델 로깅
with mlflow.start_run(run_name="fraud-detection-v3") as run:
    # 학습 메트릭 기록
    mlflow.log_params({
        "model_type": "xgboost",
        "n_estimators": 500,
        "max_depth": 8,
        "learning_rate": 0.05,
        "training_data_version": "2026-03-01",
        "feature_count": 47,
    })
    mlflow.log_metrics({
        "auc_roc": 0.9847,
        "precision_at_95_recall": 0.912,
        "f1_score": 0.943,
        "latency_p99_ms": 12.3,
    })

    # 모델 artifact 로깅 + 레지스트리 등록을 한 번에 수행
    model_info = mlflow.xgboost.log_model(
        xgb_model=trained_model,
        artifact_path="model",
        registered_model_name="fraud-detection",
        input_example=sample_input,
        signature=mlflow.models.infer_signature(X_test, y_pred),
    )
    print(f"Model URI: {model_info.model_uri}")
    # 출력: models:/fraud-detection/3

# 2. 등록된 모델 버전에 Alias 할당
# 새 모델을 challenger로 지정
client.set_registered_model_alias(
    name="fraud-detection",
    alias="challenger",
    version=3,
)

# 3. 검증 완료 후 champion으로 승격
def promote_to_champion(model_name: str, version: int):
    """모델을 champion으로 승격하고 이전 champion에 대한 정보를 기록한다."""
    # 현재 champion 확인
    try:
        current_champion = client.get_model_version_by_alias(
            name=model_name, alias="champion"
        )
        # 이전 champion에 archived 태그 추가
        client.set_model_version_tag(
            name=model_name,
            version=current_champion.version,
            key="previously_champion",
            value="true",
        )
        print(f"Previous champion: v{current_champion.version}")
    except mlflow.exceptions.MlflowException:
        print("No existing champion found")

    # 새 champion 지정
    client.set_registered_model_alias(
        name=model_name, alias="champion", version=version
    )
    # challenger alias 제거
    client.delete_registered_model_alias(
        name=model_name, alias="challenger"
    )
    print(f"Promoted v{version} to champion")

promote_to_champion("fraud-detection", version=3)

"""
Alias 기반으로 모델을 로드하여 서빙에 활용하는 패턴.
배포 파이프라인에서 champion 모델 URI를 자동으로 resolve한다.
"""
import mlflow

# champion 모델 로드 (Alias 기반)
champion_model = mlflow.pyfunc.load_model("models:/fraud-detection@champion")

# 특정 버전 로드 (디버깅/비교용)
v2_model = mlflow.pyfunc.load_model("models:/fraud-detection/2")

# 모델 URI를 KServe storageUri로 변환하는 유틸리티
def resolve_model_storage_uri(model_name: str, alias: str) -> str:
    """MLflow Alias에서 실제 artifact storage URI를 추출한다."""
    client = MlflowClient()
    mv = client.get_model_version_by_alias(name=model_name, alias=alias)
    run = client.get_run(mv.run_id)
    artifact_uri = run.info.artifact_uri
    # 예: s3://ml-artifacts/12/abc123/artifacts/model
    return f"{artifact_uri}/model"

storage_uri = resolve_model_storage_uri("fraud-detection", "champion")
print(f"KServe storageUri: {storage_uri}")
# 출력: s3://ml-artifacts/12/abc123/artifacts/model

"""
모델 등록 시 필수 메타데이터를 검증하는 게이트키퍼.
CI/CD 파이프라인의 등록 단계에 삽입하여 불완전한 모델이
레지스트리에 들어가는 것을 방지한다.
"""
from dataclasses import dataclass
from typing import Optional

REQUIRED_TAGS = [
    "training_data_version",
    "feature_schema_hash",
    "evaluation_dataset",
    "code_commit_sha",
    "owner_team",
]

REQUIRED_METRICS = [
    "auc_roc",
    "latency_p99_ms",
]

@dataclass
class RegistrationValidation:
    is_valid: bool
    missing_tags: list[str]
    missing_metrics: list[str]
    message: str

def validate_before_registration(
    run_id: str,
    min_auc: float = 0.95,
    max_latency_ms: float = 50.0,
) -> RegistrationValidation:
    """모델 등록 전 필수 조건을 검증한다."""
    client = MlflowClient()
    run = client.get_run(run_id)

    # 필수 태그 확인
    existing_tags = set(run.data.tags.keys())
    missing_tags = [t for t in REQUIRED_TAGS if t not in existing_tags]

    # 필수 메트릭 확인
    existing_metrics = set(run.data.metrics.keys())
    missing_metrics = [m for m in REQUIRED_METRICS if m not in existing_metrics]

    # 성능 기준선 확인
    threshold_failures = []
    if "auc_roc" in existing_metrics:
        auc = run.data.metrics["auc_roc"]
        if auc < min_auc:
            threshold_failures.append(
                f"auc_roc={auc:.4f} < min={min_auc}"
            )
    if "latency_p99_ms" in existing_metrics:
        latency = run.data.metrics["latency_p99_ms"]
        if latency > max_latency_ms:
            threshold_failures.append(
                f"latency_p99_ms={latency:.1f} > max={max_latency_ms}"
            )

    is_valid = (
        len(missing_tags) == 0
        and len(missing_metrics) == 0
        and len(threshold_failures) == 0
    )

    messages = []
    if missing_tags:
        messages.append(f"Missing tags: {missing_tags}")
    if missing_metrics:
        messages.append(f"Missing metrics: {missing_metrics}")
    if threshold_failures:
        messages.append(f"Threshold failures: {threshold_failures}")

    return RegistrationValidation(
        is_valid=is_valid,
        missing_tags=missing_tags,
        missing_metrics=missing_metrics,
        message=" | ".join(messages) if messages else "All checks passed",
    )

1. 모델 등록 (MLflow Registry)
   └─ 필수 메타데이터 검증 통과
   └─ @challenger alias 할당

2. 사전 검증 (Pre-deployment Validation)
   └─ 오프라인 평가 데이터셋으로 정확도 확인
   └─ 추론 레이턴시 벤치마크 (p50, p95, p99)
   └─ 입출력 스키마 호환성 검증

3. 카나리 배포 시작
   └─ KServe InferenceService 업데이트
   └─ canaryTrafficPercent: 5 설정
   └─ 모니터링 대시보드 활성화

4. 점진적 트래픽 증가
   └─ 5% -> 10% -> 25% -> 50% -> 100%
   └─ 각 단계에서 메트릭 자동 비교
   └─ 이상 감지 시 자동 롤백

5. 승격 또는 롤백
   └─ 성공: @champion alias 전환, 이전 버전 @archived
   └─ 실패: canaryTrafficPercent: 0, 원인 분석 태그 기록

# kserve-canary-deploy.yaml
# KServe InferenceService: 카나리 배포 설정
# 새 모델에 10%의 트래픽을 할당하고 점진적으로 증가시킨다.
apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: fraud-detection
  namespace: ml-serving
  annotations:
    # 최소 스케일 설정 (콜드 스타트 방지)
    autoscaling.knative.dev/min-scale: '2'
    # 배포 메타데이터
    mlflow.model.version: '3'
    mlflow.model.alias: 'challenger'
    deploy.pipeline.run-id: 'deploy-20260304-001'
spec:
  predictor:
    # 카나리 트래픽 비율 (0-100)
    canaryTrafficPercent: 10
    # 새 모델 버전 (카나리)
    model:
      modelFormat:
        name: xgboost
      runtime: kserve-xgbserver
      storageUri: 's3://ml-artifacts/fraud-detection/v3/model'
      resources:
        requests:
          cpu: '2'
          memory: '4Gi'
        limits:
          cpu: '4'
          memory: '8Gi'
    # Autoscaler 설정
    scaleTarget: 10
    scaleMetric: concurrency
  # 트래픽 라우팅 설정 (선택)
  transformer:
    containers:
      - name: feature-enricher
        image: registry.example.com/ml/feature-enricher:v2.1
        resources:
          requests:
            cpu: '1'
            memory: '2Gi'

"""
KServe 카나리 트래픽을 단계적으로 증가시키는 자동화 스크립트.
각 단계에서 핵심 메트릭을 확인하고, 기준을 미달하면 즉시 롤백한다.
"""
import time
import subprocess
import json
from dataclasses import dataclass

@dataclass
class CanaryStage:
    traffic_percent: int
    observation_minutes: int
    max_error_rate: float
    max_latency_p99_ms: float

CANARY_STAGES = [
    CanaryStage(traffic_percent=5,   observation_minutes=10, max_error_rate=0.01, max_latency_p99_ms=50),
    CanaryStage(traffic_percent=10,  observation_minutes=15, max_error_rate=0.01, max_latency_p99_ms=50),
    CanaryStage(traffic_percent=25,  observation_minutes=20, max_error_rate=0.005, max_latency_p99_ms=45),
    CanaryStage(traffic_percent=50,  observation_minutes=30, max_error_rate=0.005, max_latency_p99_ms=45),
    CanaryStage(traffic_percent=100, observation_minutes=0,  max_error_rate=0.005, max_latency_p99_ms=45),
]

def update_canary_traffic(
    isvc_name: str,
    namespace: str,
    traffic_percent: int,
) -> bool:
    """KServe InferenceService의 canaryTrafficPercent를 업데이트한다."""
    patch = {
        "spec": {
            "predictor": {
                "canaryTrafficPercent": traffic_percent
            }
        }
    }
    cmd = [
        "kubectl", "patch", "inferenceservice", isvc_name,
        "-n", namespace,
        "--type=merge",
        f"-p={json.dumps(patch)}",
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Failed to patch: {result.stderr}")
        return False
    print(f"Updated canaryTrafficPercent to {traffic_percent}%")
    return True

def check_canary_health(
    isvc_name: str,
    namespace: str,
    stage: CanaryStage,
    prometheus_url: str = "http://prometheus.monitoring:9090",
) -> bool:
    """Prometheus에서 카나리 메트릭을 확인한다."""
    import requests

    # 에러율 확인
    error_query = (
        f'sum(rate(revision_request_count'
        f'{{service_name="{isvc_name}",response_code!="200",'
        f'namespace="{namespace}"}}[5m])) / '
        f'sum(rate(revision_request_count'
        f'{{service_name="{isvc_name}",'
        f'namespace="{namespace}"}}[5m]))'
    )
    resp = requests.get(
        f"{prometheus_url}/api/v1/query",
        params={"query": error_query},
    )
    data = resp.json()
    if data["data"]["result"]:
        error_rate = float(data["data"]["result"][0]["value"][1])
        if error_rate > stage.max_error_rate:
            print(f"Error rate {error_rate:.4f} exceeds {stage.max_error_rate}")
            return False

    # p99 레이턴시 확인
    latency_query = (
        f'histogram_quantile(0.99, '
        f'sum(rate(revision_request_latencies_bucket'
        f'{{service_name="{isvc_name}",'
        f'namespace="{namespace}"}}[5m])) by (le))'
    )
    resp = requests.get(
        f"{prometheus_url}/api/v1/query",
        params={"query": latency_query},
    )
    data = resp.json()
    if data["data"]["result"]:
        latency_ms = float(data["data"]["result"][0]["value"][1])
        if latency_ms > stage.max_latency_p99_ms:
            print(f"p99 latency {latency_ms:.1f}ms exceeds {stage.max_latency_p99_ms}ms")
            return False

    return True

def run_canary_rollout(isvc_name: str, namespace: str):
    """카나리 롤아웃을 단계별로 실행한다."""
    for i, stage in enumerate(CANARY_STAGES):
        print(f"\n--- Stage {i+1}/{len(CANARY_STAGES)}: "
              f"{stage.traffic_percent}% traffic ---")

        if not update_canary_traffic(isvc_name, namespace, stage.traffic_percent):
            print("ABORT: Failed to update traffic")
            rollback(isvc_name, namespace)
            return False

        if stage.observation_minutes > 0:
            print(f"Observing for {stage.observation_minutes} minutes...")
            # 관찰 기간 동안 주기적으로 메트릭 확인
            checks = stage.observation_minutes // 2  # 2분마다 확인
            for check in range(checks):
                time.sleep(120)  # 2분 대기
                if not check_canary_health(isvc_name, namespace, stage):
                    print(f"ROLLBACK: Health check failed at stage {i+1}")
                    rollback(isvc_name, namespace)
                    return False
                print(f"  Check {check+1}/{checks}: OK")

    print("\nCanary rollout completed successfully!")
    return True

def rollback(isvc_name: str, namespace: str):
    """카나리 트래픽을 0%로 설정하여 즉시 롤백한다."""
    update_canary_traffic(isvc_name, namespace, 0)
    print("Rolled back to previous stable version")

# istio-ab-routing.yaml
# 헤더 기반 A/B 트래픽 라우팅
# x-model-variant 헤더에 따라 다른 모델 버전으로 라우팅한다.
apiVersion: networking.istio.io/v1beta1
kind: VirtualService
metadata:
  name: fraud-detection-ab
  namespace: ml-serving
spec:
  hosts:
    - fraud-detection.ml-serving.svc.cluster.local
  http:
    # 규칙 1: 명시적 variant B 요청 -> 새 모델
    - match:
        - headers:
            x-model-variant:
              exact: 'B'
      route:
        - destination:
            host: fraud-detection-predictor.ml-serving.svc.cluster.local
            port:
              number: 80
          headers:
            request:
              set:
                x-model-version: 'v3'
          weight: 100
    # 규칙 2: 특정 사용자 세그먼트 -> 새 모델 (쿠키 기반)
    - match:
        - headers:
            cookie:
              regex: '.*ab_group=treatment.*'
      route:
        - destination:
            host: fraud-detection-predictor.ml-serving.svc.cluster.local
            port:
              number: 80
          weight: 100
    # 규칙 3: 기본 트래픽 -> 가중치 기반 분할
    - route:
        - destination:
            host: fraud-detection-predictor.ml-serving.svc.cluster.local
            subset: stable
          weight: 90
        - destination:
            host: fraud-detection-predictor.ml-serving.svc.cluster.local
            subset: canary
          weight: 10
---
# DestinationRule로 서브셋 정의
apiVersion: networking.istio.io/v1beta1
kind: DestinationRule
metadata:
  name: fraud-detection-subsets
  namespace: ml-serving
spec:
  host: fraud-detection-predictor.ml-serving.svc.cluster.local
  subsets:
    - name: stable
      labels:
        model-version: 'v2'
    - name: canary
      labels:
        model-version: 'v3'

"""
A/B 배포 모니터링 및 자동 롤백 컨트롤러.
Prometheus에서 메트릭을 수집하고, 기준 위반 시 자동 롤백을 트리거한다.
"""
import time
import requests
import logging
from dataclasses import dataclass, field
from enum import Enum

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("ab-monitor")

class RollbackReason(Enum):
    ERROR_RATE = "error_rate_exceeded"
    LATENCY = "latency_exceeded"
    PREDICTION_DRIFT = "prediction_distribution_drift"
    BUSINESS_KPI = "business_kpi_degradation"

@dataclass
class MonitorConfig:
    prometheus_url: str = "http://prometheus.monitoring:9090"
    isvc_name: str = "fraud-detection"
    namespace: str = "ml-serving"
    check_interval_sec: int = 60
    # 임계값 설정
    max_error_rate: float = 0.01
    max_latency_p99_ms: float = 50.0
    max_psi: float = 0.2  # PSI > 0.2이면 유의미한 분포 변화
    # 연속 실패 횟수 기준
    consecutive_failures_to_rollback: int = 3
    # 알림 설정
    slack_webhook: str = ""
    pagerduty_key: str = ""

@dataclass
class MonitorState:
    consecutive_failures: int = 0
    total_checks: int = 0
    failure_reasons: list = field(default_factory=list)

def query_prometheus(config: MonitorConfig, query: str) -> float | None:
    """Prometheus 쿼리를 실행하고 단일 값을 반환한다."""
    try:
        resp = requests.get(
            f"{config.prometheus_url}/api/v1/query",
            params={"query": query},
            timeout=10,
        )
        data = resp.json()
        if data["status"] == "success" and data["data"]["result"]:
            return float(data["data"]["result"][0]["value"][1])
    except Exception as e:
        logger.error(f"Prometheus query failed: {e}")
    return None

def check_model_health(config: MonitorConfig) -> tuple[bool, list[RollbackReason]]:
    """모델의 건강 상태를 종합적으로 확인한다."""
    reasons = []

    # 1. 에러율 확인
    error_rate = query_prometheus(config, (
        f'sum(rate(revision_request_count'
        f'{{service_name="{config.isvc_name}",'
        f'response_code=~"5..",namespace="{config.namespace}"}}[5m]))'
        f' / sum(rate(revision_request_count'
        f'{{service_name="{config.isvc_name}",'
        f'namespace="{config.namespace}"}}[5m]))'
    ))
    if error_rate is not None and error_rate > config.max_error_rate:
        logger.warning(f"Error rate: {error_rate:.4f} > {config.max_error_rate}")
        reasons.append(RollbackReason.ERROR_RATE)

    # 2. p99 레이턴시 확인
    latency = query_prometheus(config, (
        f'histogram_quantile(0.99, sum(rate('
        f'revision_request_latencies_bucket'
        f'{{service_name="{config.isvc_name}",'
        f'namespace="{config.namespace}"}}[5m])) by (le))'
    ))
    if latency is not None and latency > config.max_latency_p99_ms:
        logger.warning(f"p99 latency: {latency:.1f}ms > {config.max_latency_p99_ms}ms")
        reasons.append(RollbackReason.LATENCY)

    # 3. 예측 분포 변화 (PSI) 확인
    psi = query_prometheus(config, (
        f'ml_prediction_psi'
        f'{{model="{config.isvc_name}",namespace="{config.namespace}"}}'
    ))
    if psi is not None and psi > config.max_psi:
        logger.warning(f"PSI: {psi:.4f} > {config.max_psi}")
        reasons.append(RollbackReason.PREDICTION_DRIFT)

    is_healthy = len(reasons) == 0
    return is_healthy, reasons

def send_alert(config: MonitorConfig, message: str):
    """Slack과 PagerDuty로 알림을 전송한다."""
    if config.slack_webhook:
        requests.post(config.slack_webhook, json={
            "text": f":rotating_light: [AB Deploy Monitor] {message}",
            "channel": "#ml-alerts",
        })
    logger.info(f"Alert: {message}")

def trigger_rollback(config: MonitorConfig, reasons: list[RollbackReason]):
    """자동 롤백을 실행한다."""
    import subprocess, json
    reason_str = ", ".join([r.value for r in reasons])
    logger.critical(f"Triggering rollback. Reasons: {reason_str}")

    # KServe canaryTrafficPercent를 0으로 패치
    patch = json.dumps({
        "spec": {"predictor": {"canaryTrafficPercent": 0}}
    })
    subprocess.run([
        "kubectl", "patch", "inferenceservice", config.isvc_name,
        "-n", config.namespace,
        "--type=merge", f"-p={patch}",
    ], check=True)

    send_alert(config, f"Auto-rollback executed for {config.isvc_name}. Reasons: {reason_str}")

def run_monitor(config: MonitorConfig):
    """모니터링 루프를 실행한다."""
    state = MonitorState()
    logger.info(f"Starting AB deploy monitor for {config.isvc_name}")

    while True:
        is_healthy, reasons = check_model_health(config)
        state.total_checks += 1

        if is_healthy:
            state.consecutive_failures = 0
            state.failure_reasons.clear()
            if state.total_checks % 10 == 0:
                logger.info(f"Check #{state.total_checks}: Healthy")
        else:
            state.consecutive_failures += 1
            state.failure_reasons.extend(reasons)
            logger.warning(
                f"Check #{state.total_checks}: Unhealthy "
                f"({state.consecutive_failures}/{config.consecutive_failures_to_rollback})"
            )

            if state.consecutive_failures >= config.consecutive_failures_to_rollback:
                trigger_rollback(config, reasons)
                break

        time.sleep(config.check_interval_sec)

# 실행 예시
if __name__ == "__main__":
    config = MonitorConfig(
        isvc_name="fraud-detection",
        namespace="ml-serving",
        slack_webhook="https://hooks.slack.com/services/xxx/yyy/zzz",
    )
    run_monitor(config)

"""
모델 승격 파이프라인: 카나리 배포 성공 후 champion 전환 자동화.
MLflow Alias 업데이트 + KServe 카나리 정리 + 알림을 처리한다.
"""
import mlflow
from mlflow.tracking import MlflowClient
import subprocess
import json
from datetime import datetime

def promote_model_pipeline(
    model_name: str,
    new_version: int,
    isvc_name: str,
    namespace: str,
    ab_test_results: dict | None = None,
):
    """
    모델 승격 전체 파이프라인.

    Args:
        model_name: MLflow 등록 모델 이름
        new_version: 승격할 모델 버전 번호
        isvc_name: KServe InferenceService 이름
        namespace: 쿠버네티스 네임스페이스
        ab_test_results: A/B 테스트 결과 메트릭 딕셔너리
    """
    client = MlflowClient()
    timestamp = datetime.utcnow().isoformat()

    # Step 1: 현재 champion 정보 백업
    try:
        old_champion = client.get_model_version_by_alias(
            name=model_name, alias="champion"
        )
        old_version = int(old_champion.version)
        print(f"Current champion: v{old_version}")

        # 이전 champion에 메타데이터 기록
        client.set_model_version_tag(
            name=model_name, version=old_version,
            key="demoted_at", value=timestamp,
        )
        client.set_model_version_tag(
            name=model_name, version=old_version,
            key="demoted_by", value=f"v{new_version}",
        )
        # 이전 champion을 archived로 이동
        client.set_registered_model_alias(
            name=model_name, alias="archived", version=old_version,
        )
    except mlflow.exceptions.MlflowException:
        old_version = None
        print("No existing champion found")

    # Step 2: 새 모델을 champion으로 승격
    client.set_registered_model_alias(
        name=model_name, alias="champion", version=new_version,
    )
    # challenger alias 제거
    try:
        client.delete_registered_model_alias(
            name=model_name, alias="challenger",
        )
    except mlflow.exceptions.MlflowException:
        pass

    # 승격 메타데이터 기록
    client.set_model_version_tag(
        name=model_name, version=new_version,
        key="promoted_to_champion_at", value=timestamp,
    )
    if ab_test_results:
        for k, v in ab_test_results.items():
            client.set_model_version_tag(
                name=model_name, version=new_version,
                key=f"ab_result_{k}", value=str(v),
            )

    # Step 3: KServe에서 카나리 설정 정리
    # canaryTrafficPercent를 제거하여 새 모델이 100% 트래픽을 받도록 한다
    patch = json.dumps({
        "spec": {
            "predictor": {
                "canaryTrafficPercent": 100
            }
        }
    })
    subprocess.run([
        "kubectl", "patch", "inferenceservice", isvc_name,
        "-n", namespace,
        "--type=merge", f"-p={patch}",
    ], check=True)

    print(f"Promotion complete: {model_name} v{new_version} is now champion")

    return {
        "model_name": model_name,
        "new_champion_version": new_version,
        "old_champion_version": old_version,
        "promoted_at": timestamp,
    }

# 실행 예시
result = promote_model_pipeline(
    model_name="fraud-detection",
    new_version=3,
    isvc_name="fraud-detection",
    namespace="ml-serving",
    ab_test_results={
        "auc_improvement": 0.012,
        "latency_p99_delta_ms": -2.1,
        "error_rate_delta": -0.003,
        "test_duration_hours": 24,
        "traffic_percentage": 50,
    },
)