MLflow完全ガイド｜Databricksの実験管理・モデル管理基盤

MLflowは、機械学習のライフサイクル全体（実験管理・モデル管理・デプロイ）を統合管理するオープンソースプラットフォームです。 Databricksが開発を主導しており、Databricks上ではManaged版としてシームレスに利用できます。 ML Associate/Professional試験の中核トピックであり、Data Engineer系でも基本概念が問われます。

MLflowの4コンポーネント

MLflow Tracking

Trackingは「誰が、いつ、どのパラメータで、どのデータで訓練し、結果はどうだったか」を記録する仕組みです。 1回の訓練をRunと呼び、複数のRunをExperimentにまとめます。

import mlflow

# Experimentの指定（なければ自動作成）
mlflow.set_experiment("/Users/taro/churn_prediction")

# Runの開始
with mlflow.start_run(run_name="xgboost_v1"):
    # パラメータの記録
    mlflow.log_param("max_depth", 5)
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("n_estimators", 100)

    # モデル訓練（例）
    model = train_xgboost(X_train, y_train, max_depth=5, lr=0.01)

    # メトリクスの記録
    accuracy = evaluate(model, X_test, y_test)
    mlflow.log_metric("accuracy", accuracy)
    mlflow.log_metric("f1_score", f1)

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

    # アーティファクトの保存（図表、設定ファイルなど）
    mlflow.log_artifact("confusion_matrix.png")

Databricks上では mlflow.autolog() を呼ぶだけで、多くのフレームワーク（scikit-learn、XGBoost、 LightGBM、PyTorch、TensorFlow等）のパラメータ・メトリクス・モデルが自動記録されます。

# Autologging（推奨）
mlflow.autolog()

# これだけで訓練時のパラメータ・メトリクス・モデルが自動記録される
model = RandomForestClassifier(n_estimators=100, max_depth=5)
model.fit(X_train, y_train)

Model Registry

Model Registryは、訓練済みモデルのバージョン管理とライフサイクル管理を行う仕組みです。 「どのモデルバージョンがどのステージにあるか」を一元管理し、本番デプロイの承認フローを構築できます。

# RunからModel Registryに登録
mlflow.register_model(
    model_uri="runs:/abc123def456/model",
    name="churn_prediction_model"
)

# バージョンの遷移（Unity Catalog Model Registryではエイリアスを使用）
# 旧Model Registry: Staging → Production
# Unity Catalog: エイリアス "champion" / "challenger" を設定
from mlflow import MlflowClient
client = MlflowClient()
client.set_registered_model_alias(
    name="churn_prediction_model",
    alias="champion",
    version=3
)

モデルのデプロイ

Model Registryに登録したモデルは、Databricksの Model Serving でリアルタイム推論エンドポイントとしてデプロイできます。 バッチ推論の場合は、mlflow.pyfunc.spark_udf() でSpark UDFとして利用します。

# バッチ推論: Spark UDFとしてモデルを適用
import mlflow
model_uri = "models:/churn_prediction_model@champion"
predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri)

predictions = (spark.table("silver.customers")
  .withColumn("churn_probability", predict_udf("age", "tenure", "monthly_charges"))
)

Experiment管理の実務パターン

• Experiment名はプロジェクト名 + モデル目的で命名する（例: /Projects/churn/xgboost）
• Run名には手法とバージョンを入れる（例: xgboost_v3_tuned）
• データのバージョン（Delta TableのVersion番号）もパラメータとして記録する
• Compare機能で複数Runのメトリクスを並べて比較する
• 最良のRunをModel Registryに登録し、エイリアスでchampionを指定する

試験で問われるポイント

• Experiment / Run / Model の階層関係
• mlflow.log_param() / log_metric() / log_model() / log_artifact() の使い分け
• mlflow.autolog() の挙動（対応フレームワーク、何が自動記録されるか）
• Model Registryのステージ遷移（旧: Staging→Production、新: エイリアス）
• Unity Catalog Model Registryと旧Workspace Model Registryの違い
• spark_udf() でのバッチ推論パターン

問題で確認

よくある質問