MLflow 2.x 실험 추적과 모델 레지스트리 운영 가이드

# 프로덕션 추적 서버 기동 명령
mlflow server \
  --backend-store-uri postgresql://mlflow_user:${DB_PASSWORD}@db.internal:5432/mlflow_prod \
  --default-artifact-root s3://company-mlflow-artifacts/prod/ \
  --artifacts-destination s3://company-mlflow-artifacts/prod/ \
  --host 0.0.0.0 \
  --port 5000 \
  --workers 4 \
  --gunicorn-opts "--timeout 120 --keep-alive 5"

# 클라이언트 측 설정 (올바른 방식)
import mlflow
import os

# 추적 서버 URI만 설정. 아티팩트 경로는 서버가 관리.
os.environ["MLFLOW_TRACKING_URI"] = "http://mlflow.internal:5000"

# S3 인증은 IAM Role 사용 권장 (EC2/EKS 환경)
# 로컬 개발 시에만 credential 명시
# os.environ["AWS_PROFILE"] = "mlflow-dev"

mlflow.set_experiment("/team-search/ranking-model-v3")

import mlflow

# 좋은 예: 계층적 네이밍
mlflow.set_experiment("/search-team/query-ranking/bert-finetune")
mlflow.set_experiment("/fraud-team/transaction-classifier/xgboost-baseline")
mlflow.set_experiment("/recommendation/item2vec/hyperopt-sweep")

# 나쁜 예: 플랫하고 모호한 네이밍
# mlflow.set_experiment("experiment_1")
# mlflow.set_experiment("test_model")
# mlflow.set_experiment("johns_experiment")

import mlflow
from datetime import datetime

with mlflow.start_run(run_name="bert-ranking-v3.2.1") as run:
    # 필수 태그
    mlflow.set_tag("team", "search")
    mlflow.set_tag("owner", "jane.doe@company.com")
    mlflow.set_tag("git.commit", "a1b2c3d4")
    mlflow.set_tag("data.version", "v2026.03.05")
    mlflow.set_tag("environment", "gpu-cluster-a100")
    mlflow.set_tag("purpose", "hyperparameter-sweep")

    # 파라미터 기록
    mlflow.log_params({
        "learning_rate": 2e-5,
        "batch_size": 32,
        "max_epochs": 10,
        "model_architecture": "bert-base-uncased",
        "optimizer": "AdamW",
        "warmup_steps": 500,
    })

    # 메트릭은 스텝 단위로 기록 (에포크 또는 글로벌 스텝)
    for epoch in range(10):
        train_loss = train_one_epoch(model, train_loader)
        val_ndcg = evaluate(model, val_loader)

        mlflow.log_metrics({
            "train_loss": train_loss,
            "val_ndcg@10": val_ndcg,
        }, step=epoch)

    # 최종 모델 로깅
    mlflow.pytorch.log_model(
        model,
        artifact_path="model",
        registered_model_name="search-ranking-bert",
    )

# 좋은 예 (제품/기능 중심)
fraud-detector
search-ranker
recommendation-item2vec
churn-predictor

# 나쁜 예 (알고리즘/버전 중심)
xgboost-fraud-v3
bert-search-ranking-2026
lgbm_model_final_final

from mlflow import MlflowClient

client = MlflowClient()

# 새 모델 버전 등록 후 alias 설정
model_name = "search-ranker"
version = client.create_model_version(
    name=model_name,
    source="runs:/abc123/model",
    run_id="abc123",
    description="BERT-base finetuned on 2026 Q1 query logs"
)

# 현재 champion 확인
try:
    current_champion = client.get_model_version_by_alias(model_name, "champion")
    print(f"Current champion: v{current_champion.version}")
except mlflow.exceptions.MlflowException:
    print("No champion alias set yet")

# Canary 배포: 새 버전에 challenger alias 부여
client.set_registered_model_alias(model_name, "challenger", version.version)

# Canary 검증 통과 후 champion 승격
client.set_registered_model_alias(model_name, "champion", version.version)

# 이전 champion을 archived로 이동
client.set_registered_model_alias(model_name, "previous-champion", current_champion.version)

import mlflow

# 서빙 코드: alias 기반 모델 로드
model = mlflow.pyfunc.load_model("models:/search-ranker@champion")
predictions = model.predict(input_data)

client.set_model_version_tag(model_name, version.version, "training_data", "s3://data/query-logs/2026-q1/")
client.set_model_version_tag(model_name, version.version, "training_commit", "a1b2c3d4e5f6")
client.set_model_version_tag(model_name, version.version, "validation_ndcg", "0.847")
client.set_model_version_tag(model_name, version.version, "approved_by", "jane.doe")
client.set_model_version_tag(model_name, version.version, "approval_date", "2026-03-05")

# .github/workflows/train-and-register.yml
name: Train and Register Model

on:
  push:
    paths:
      - 'models/search-ranker/**'
    branches: [main]
  workflow_dispatch:
    inputs:
      experiment_name:
        description: 'MLflow experiment name'
        required: true
        default: '/search-team/query-ranking/scheduled-retrain'

env:
  MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_TRACKING_URI }}
  AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
  AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}

jobs:
  train:
    runs-on: [self-hosted, gpu]
    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        run: |
          pip install -r requirements.txt
          pip install mlflow[extras]

      - name: Train model
        run: |
          python models/search-ranker/train.py \
            --experiment-name "${{ github.event.inputs.experiment_name || '/search-team/query-ranking/ci-train' }}" \
            --run-name "ci-${{ github.sha }}" \
            --register-model search-ranker

      - name: Validate model
        run: |
          python models/search-ranker/validate.py \
            --model-uri "models:/search-ranker@challenger" \
            --threshold-ndcg 0.82

      - name: Promote to champion
        if: success()
        run: |
          python scripts/promote_model.py \
            --model-name search-ranker \
            --from-alias challenger \
            --to-alias champion

# scripts/validate_model.py
import mlflow
import sys
from mlflow import MlflowClient

def validate_model(model_name: str, alias: str, threshold: float) -> bool:
    """모델 성능이 임계값을 충족하는지 검증한다."""
    client = MlflowClient()

    # alias로 모델 버전 조회
    model_version = client.get_model_version_by_alias(model_name, alias)
    run = client.get_run(model_version.run_id)

    # 검증 메트릭 확인
    val_ndcg = run.data.metrics.get("val_ndcg@10")
    if val_ndcg is None:
        print(f"ERROR: val_ndcg@10 metric not found in run {model_version.run_id}")
        return False

    # 현재 champion과 비교
    try:
        champion = client.get_model_version_by_alias(model_name, "champion")
        champion_run = client.get_run(champion.run_id)
        champion_ndcg = champion_run.data.metrics.get("val_ndcg@10", 0)
        print(f"Champion v{champion.version} NDCG: {champion_ndcg:.4f}")
        print(f"Challenger v{model_version.version} NDCG: {val_ndcg:.4f}")

        # 챔피언 대비 성능 저하 체크
        if val_ndcg < champion_ndcg * 0.98:  # 2% 이상 하락 시 실패
            print("FAIL: Challenger performs worse than champion by more than 2%")
            return False
    except Exception:
        print("No existing champion found. Proceeding with threshold check only.")

    # 절대 임계값 체크
    if val_ndcg < threshold:
        print(f"FAIL: NDCG {val_ndcg:.4f} below threshold {threshold:.4f}")
        return False

    print(f"PASS: NDCG {val_ndcg:.4f} meets threshold {threshold:.4f}")
    return True

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-name", required=True)
    parser.add_argument("--alias", default="challenger")
    parser.add_argument("--threshold", type=float, default=0.80)
    args = parser.parse_args()

    if not validate_model(args.model_name, args.alias, args.threshold):
        sys.exit(1)

# 인증 활성화된 서버 기동
mlflow server \
  --backend-store-uri postgresql://mlflow:${DB_PASSWORD}@db:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/ \
  --app-name basic-auth \
  --host 0.0.0.0 \
  --port 5000

# 팀별 실험 접두어로 논리적 격리
TEAM_PREFIX = {
    "search": "/search-team",
    "fraud": "/fraud-team",
    "recommendation": "/rec-team",
}

def get_experiment_name(team: str, project: str, experiment: str) -> str:
    """팀 접두어가 포함된 실험 이름을 생성한다."""
    prefix = TEAM_PREFIX.get(team)
    if prefix is None:
        raise ValueError(f"Unknown team: {team}. Allowed: {list(TEAM_PREFIX.keys())}")
    return f"{prefix}/{project}/{experiment}"

# 사용 예
experiment = get_experiment_name("search", "query-ranking", "bert-v4-sweep")
mlflow.set_experiment(experiment)  # "/search-team/query-ranking/bert-v4-sweep"

from mlflow import MlflowClient
from datetime import datetime, timedelta

def cleanup_old_runs(experiment_name: str, days_old: int = 90, dry_run: bool = True):
    """지정 기간이 지난 실패/취소된 run의 아티팩트를 정리한다."""
    client = MlflowClient()
    experiment = client.get_experiment_by_name(experiment_name)

    if experiment is None:
        print(f"Experiment '{experiment_name}' not found")
        return

    cutoff_ts = int((datetime.now() - timedelta(days=days_old)).timestamp() * 1000)

    runs = client.search_runs(
        experiment_ids=[experiment.experiment_id],
        filter_string=f"attributes.end_time < {cutoff_ts} AND attributes.status != 'RUNNING'",
        order_by=["attributes.end_time ASC"],
        max_results=500,
    )

    deleted_count = 0
    for run in runs:
        # 태그로 보존 여부 확인
        if run.data.tags.get("keep", "false").lower() == "true":
            continue

        # 레지스트리에 등록된 모델이 있는 run은 건너뜀
        if run.data.tags.get("mlflow.registeredModelName"):
            continue

        if dry_run:
            print(f"[DRY RUN] Would delete run {run.info.run_id} "
                  f"(ended: {datetime.fromtimestamp(run.info.end_time / 1000)})")
        else:
            client.delete_run(run.info.run_id)
            deleted_count += 1

    print(f"{'Would delete' if dry_run else 'Deleted'} {deleted_count} runs "
          f"out of {len(runs)} found")

# 사용: 먼저 dry_run으로 확인 후 실제 삭제
cleanup_old_runs("/search-team/query-ranking/bert-finetune", days_old=60, dry_run=True)

{
  "Rules": [
    {
      "ID": "MoveOldArtifactsToIA",
      "Status": "Enabled",
      "Filter": {
        "Prefix": "prod/"
      },
      "Transitions": [
        {
          "Days": 90,
          "StorageClass": "STANDARD_IA"
        },
        {
          "Days": 365,
          "StorageClass": "GLACIER"
        }
      ]
    },
    {
      "ID": "DeleteTempArtifacts",
      "Status": "Enabled",
      "Filter": {
        "Prefix": "tmp/"
      },
      "Expiration": {
        "Days": 7
      }
    }
  ]
}

# 서버 기동 시 연결 풀 파라미터 조정
mlflow server \
  --backend-store-uri "postgresql://mlflow:pass@db:5432/mlflow?pool_size=20&max_overflow=40" \
  --default-artifact-root s3://artifacts/ \
  --workers 8

# 업로드 타임아웃 연장
export MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT=600

# multipart upload 청크 크기 조정 (S3)
export MLFLOW_S3_UPLOAD_EXTRA_ARGS='{"ServerSideEncryption": "aws:kms"}'

from mlflow import MlflowClient

client = MlflowClient()

# 고착된 run 강제 종료
stuck_runs = client.search_runs(
    experiment_ids=["1"],
    filter_string="attributes.status = 'RUNNING'",
)

for run in stuck_runs:
    end_time = run.info.end_time
    # end_time이 없고 시작 시간이 24시간 이상 전인 경우
    if end_time is None or end_time == 0:
        start_time = run.info.start_time
        if (datetime.now().timestamp() * 1000 - start_time) > 86400000:  # 24h
            client.set_terminated(run.info.run_id, status="FAILED")
            print(f"Force-terminated stuck run: {run.info.run_id}")

import redis
import time

def safe_promote_model(model_name: str, version: str, alias: str, redis_url: str):
    """분산 잠금을 사용한 안전한 모델 승격."""
    r = redis.from_url(redis_url)
    lock_key = f"mlflow:promote:{model_name}:{alias}"

    # 30초 TTL의 분산 잠금 획득
    lock = r.lock(lock_key, timeout=30)
    if lock.acquire(blocking=True, blocking_timeout=10):
        try:
            client = MlflowClient()
            client.set_registered_model_alias(model_name, alias, version)
            print(f"Successfully promoted {model_name} v{version} to @{alias}")
        finally:
            lock.release()
    else:
        raise RuntimeError(f"Failed to acquire lock for {model_name}@{alias}")

from mlflow import MlflowClient

def rollback_model(model_name: str):
    """champion 모델을 previous-champion으로 롤백한다."""
    client = MlflowClient()

    try:
        previous = client.get_model_version_by_alias(model_name, "previous-champion")
    except Exception:
        print("ERROR: No previous-champion alias found. Manual intervention required.")
        return False

    current = client.get_model_version_by_alias(model_name, "champion")

    # 롤백 실행
    client.set_registered_model_alias(model_name, "champion", previous.version)
    client.set_registered_model_alias(model_name, "rolled-back", current.version)

    # 롤백 사유 태깅
    client.set_model_version_tag(
        model_name, current.version, "rollback_reason", "performance_degradation"
    )
    client.set_model_version_tag(
        model_name, current.version, "rolled_back_at", datetime.now().isoformat()
    )

    print(f"Rolled back {model_name}: v{current.version} -> v{previous.version}")
    return True

# 아티팩트 가비지 컬렉션
mlflow gc \
  --backend-store-uri postgresql://mlflow:pass@db:5432/mlflow \
  --older-than 30d