Skip to content
Published on

MLflow完全ガイド:実験追跡からModel Registry、プロダクションデプロイまで

Authors
  • Name
    Twitter

MLflowとは?

MLflowはMLライフサイクルを管理するオープンソースプラットフォームです。4つの主要コンポーネントで構成されています:

  • MLflow Tracking: 実験のパラメータ、メトリクス、アーティファクトを記録
  • MLflow Projects: 再現可能なMLコードのパッケージング
  • MLflow Models: 様々なフレームワークのモデルを統一形式でパッケージング
  • MLflow Model Registry: モデルのバージョン管理とデプロイワークフロー

インストールとサーバー設定

基本インストール

# pipインストール
pip install mlflow

# 追加フレームワークサポート
pip install mlflow[extras]  # sklearn, tensorflow, pytorchなど

# サーバー起動(ローカル)
mlflow server --host 0.0.0.0 --port 5000

# PostgreSQL + S3バックエンドでプロダクションサーバー
mlflow server \
  --backend-store-uri postgresql://mlflow:password@localhost:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/ \
  --host 0.0.0.0 --port 5000

Docker Composeでデプロイ

# docker-compose.yml
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.18.0
    ports:
      - '5000:5000'
    environment:
      - MLFLOW_BACKEND_STORE_URI=postgresql://mlflow:password@postgres:5432/mlflow
      - MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts/
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:password@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0 --port 5000
    depends_on:
      - postgres

  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: password
      POSTGRES_DB: mlflow
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

実験追跡(Tracking)

基本的な使い方

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score

# トラッキングサーバーの設定
mlflow.set_tracking_uri("http://localhost:5000")

# 実験の作成/設定
mlflow.set_experiment("iris-classification")

# データの準備
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 実験の実行
with mlflow.start_run(run_name="rf-baseline"):
    # パラメータの記録
    params = {
        "n_estimators": 100,
        "max_depth": 5,
        "min_samples_split": 2,
        "random_state": 42
    }
    mlflow.log_params(params)

    # モデル学習
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 予測とメトリクス
    y_pred = model.predict(X_test)
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1_macro": f1_score(y_test, y_pred, average="macro"),
        "precision_macro": precision_score(y_test, y_pred, average="macro")
    }
    mlflow.log_metrics(metrics)

    # タグ
    mlflow.set_tag("model_type", "random_forest")
    mlflow.set_tag("dataset", "iris")

    # モデルの保存
    mlflow.sklearn.log_model(
        model,
        artifact_path="model",
        registered_model_name="iris-classifier"
    )

    # カスタムアーティファクト(グラフ、レポートなど)
    import matplotlib.pyplot as plt
    from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

    cm = confusion_matrix(y_test, y_pred)
    fig, ax = plt.subplots()
    ConfusionMatrixDisplay(cm).plot(ax=ax)
    fig.savefig("confusion_matrix.png")
    mlflow.log_artifact("confusion_matrix.png")

    print(f"Run ID: {mlflow.active_run().info.run_id}")
    print(f"Metrics: {metrics}")

ハイパーパラメータチューニングの追跡

import optuna
import mlflow

def objective(trial):
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 50, 500),
        "max_depth": trial.suggest_int("max_depth", 2, 20),
        "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 5),
    }

    with mlflow.start_run(nested=True, run_name=f"trial-{trial.number}"):
        mlflow.log_params(params)

        model = RandomForestClassifier(**params, random_state=42)
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)

        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric("accuracy", accuracy)

        return accuracy

# Optunaスタディの実行
with mlflow.start_run(run_name="hyperparameter-tuning"):
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=50)

    # 最適結果の記録
    mlflow.log_params(study.best_params)
    mlflow.log_metric("best_accuracy", study.best_value)
    mlflow.set_tag("best_trial", study.best_trial.number)

PyTorchモデルの追跡

import torch
import torch.nn as nn
import mlflow.pytorch

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

with mlflow.start_run(run_name="pytorch-model"):
    model = SimpleNet(4, 32, 3)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    mlflow.log_params({
        "hidden_dim": 32,
        "learning_rate": 0.001,
        "optimizer": "Adam",
        "epochs": 100
    })

    for epoch in range(100):
        # 学習ロジック...
        loss = criterion(model(X_tensor), y_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # エポックごとのメトリクス記録
        mlflow.log_metric("train_loss", loss.item(), step=epoch)

    # PyTorchモデルの保存
    mlflow.pytorch.log_model(model, "model")

Model Registry

モデルの登録とバージョン管理

from mlflow import MlflowClient

client = MlflowClient()

# モデル登録(log_modelでregistered_model_nameを使用すると自動登録)
# または手動登録:
result = client.create_registered_model(
    name="iris-classifier",
    description="Iris花分類モデル"
)

# 特定の実行のモデルをバージョンとして登録
model_version = client.create_model_version(
    name="iris-classifier",
    source=f"runs:/{run_id}/model",
    run_id=run_id,
    description="RandomForest baseline v1"
)

print(f"Model Version: {model_version.version}")

Aliasを活用したデプロイ管理

# MLflow 2.xではAliasを使用(Stageは非推奨)
client = MlflowClient()

# プロダクションaliasの設定
client.set_registered_model_alias(
    name="iris-classifier",
    alias="champion",
    version=3
)

# チャレンジャーモデルの設定
client.set_registered_model_alias(
    name="iris-classifier",
    alias="challenger",
    version=5
)

# Aliasでモデルをロード
champion_model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")
challenger_model = mlflow.pyfunc.load_model("models:/iris-classifier@challenger")

# A/Bテスト
champion_pred = champion_model.predict(X_test)
challenger_pred = challenger_model.predict(X_test)

print(f"Champion accuracy: {accuracy_score(y_test, champion_pred)}")
print(f"Challenger accuracy: {accuracy_score(y_test, challenger_pred)}")

モデルタグの活用

# モデルバージョンにタグを追加
client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="validation_status",
    value="approved"
)

client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="approved_by",
    value="data-science-lead"
)

# タグでモデルを検索
from mlflow import search_model_versions

approved_versions = search_model_versions(
    "name='iris-classifier' AND tag.validation_status='approved'"
)

モデルサービング

MLflow内蔵サービング

# ローカルREST APIサービング
mlflow models serve \
  -m "models:/iris-classifier@champion" \
  --port 8080 \
  --no-conda

# テストリクエスト
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

FastAPIカスタムサービング

from fastapi import FastAPI
import mlflow.pyfunc
import numpy as np

app = FastAPI()

# モデルのロード(サーバー起動時に1回)
model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")

@app.post("/predict")
async def predict(features: list[list[float]]):
    predictions = model.predict(np.array(features))
    return {
        "predictions": predictions.tolist(),
        "model_version": "champion"
    }

@app.get("/health")
async def health():
    return {"status": "healthy", "model": "iris-classifier@champion"}

実験の比較と分析

MLflow UIでの比較

# 実験検索(CLI)
mlflow runs list --experiment-id 1

# メトリクスベースの検索
mlflow runs list \
  --experiment-id 1 \
  --filter "metrics.accuracy > 0.95" \
  --order-by "metrics.accuracy DESC"

Python APIでの分析

import mlflow
import pandas as pd

# 実験の全実行を照会
runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10
)

# DataFrameで分析
print(runs[["run_id", "params.n_estimators", "params.max_depth", "metrics.accuracy"]])

# 最適な実行を見つける
best_run = runs.iloc[0]
print(f"Best run: {best_run.run_id}, Accuracy: {best_run['metrics.accuracy']}")

プロダクションチェックリスト

□ バックエンドストアをPostgreSQL/MySQLに設定
□ アーティファクトストアをS3/GCS/MinIOに設定
□ 認証/認可の設定(OIDC、Basic Auth)
□ 自動実験記録(autolog)の設定
□ Model Registryのalias規則を策定
□ CI/CDでモデル検証を自動化
□ モデルサービングのヘルスチェックを設定
□ 実験整理ポリシーの定義(古い実行のアーカイブ)

確認クイズ(6問)

Q1. MLflowの4つの主要コンポーネントは?

Tracking、Projects、Models、Model Registry

Q2. mlflow.log_paramsとmlflow.log_metricsの違いは?

log_paramsは学習ハイパーパラメータ(文字列)を記録し、log_metricsは性能指標(数値)を記録します。メトリクスはstepパラメータでエポックごとの追跡が可能です。

Q3. MLflow 2.xでモデルデプロイ管理に使用する概念は?

Alias(例:@champion、@challenger)。Stageは非推奨になりました。

Q4. nested=Trueパラメータはいつ使用しますか?

ハイパーパラメータチューニングのように、親の実行の中で複数の子実行を記録する際に使用します。

Q5. アーティファクトストアにS3を使用する理由は?

モデルファイルやグラフなどの大容量アーティファクトをスケーラブルなオブジェクトストレージに保存し、チーム間の共有とバージョン管理を容易にするためです。

Q6. mlflow.autolog()の長所と短所は?

長所:コード変更なしで自動的にパラメータ/メトリクス/モデルを記録。短所:不要な情報が多く記録される場合があり、カスタムメトリクスは別途記録が必要です。