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()

# 新しいモデルバージョン登録後にエイリアス設定
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")

# カナリアデプロイ：新バージョンにchallengerエイリアスを付与
client.set_registered_model_alias(model_name, "challenger", version.version)

# カナリア検証通過後に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

# サービングコード：エイリアスベースのモデルロード
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()

    # エイリアスでモデルバージョンを照会
    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}")

        # champion比でのパフォーマンス低下チェック
        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

# マルチパートアップロードチャンクサイズの調整（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