MLflow 3.0

こちらにもまとめていますが、エクスペリメントページにおけるモデルタブの追加、Logged Modelエンティティの追加、モデルレジストリUIの改善などが含まれています。特に大きな変更点は、情報アーキテクチャが見直されたことだと思います。これまでは、MLflowラン(トレーニングの実行単位)にモデルやアーティファクト、メトリクス、パラメータなどが紐づけられていました。生成AIのように一つのモデルに繰り返し評価が行われる場合や、ディープラーニングのように複数のモデルチェックポイントができる場合には、この情報アーキテクチャでは困難さが生じていたため、今回の変更でモデルに複数のメトリクス、アーティファクトが紐づけられるようになっています。

有効化

ワークスペースのプレビューメニューにアクセスし、MLflow 3を有効化します。

従来のMLのワークフロー

こちらのサンプルを動かします。

このノートブックは最初にモデルのトレーニングジョブを実行し、MLflow Runとして追跡され、トレーニングされたモデルを生成します。モデルはMLflow Logged Modelとして追跡されます。

%pip install mlflow  --upgrade --pre
dbutils.library.restartPython()

import pandas as pd
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

import mlflow
import mlflow.sklearn
from mlflow.entities import Dataset

# メトリクスを計算するためのヘルパー関数
def compute_metrics(actual, predicted):
    rmse = mean_squared_error(actual, predicted)
    mae = mean_absolute_error(actual, predicted)
    r2 = r2_score(actual, predicted)
    return rmse, mae, r2


# Irisデータセットをロードし、DataFrameを準備
iris = load_iris()
iris_df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
iris_df['quality'] = (iris.target == 2).astype(int)  # 簡単のためにバイナリターゲットを作成

# トレーニングとテストデータセットに分割
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)

# トレーニングジョブを表すランを開始
with mlflow.start_run() as training_run:
    # トレーニングデータセットをMLflowでロード。トレーニングメトリクスをこのデータセットにリンク
    train_dataset: Dataset = mlflow.data.from_pandas(train_df, name="train")
    train_x = train_dataset.df.drop(["quality"], axis=1)
    train_y = train_dataset.df[["quality"]]

    # トレーニングデータセットにモデルをフィット
    lr = ElasticNet(alpha=0.5, l1_ratio=0.5, random_state=42)
    lr.fit(train_x, train_y)

    # モデルをログし、ElasticNetのパラメータ（alpha, l1_ratio）を指定
    # 新機能として、LoggedModelエンティティはその名前とパラメータにリンク
    model_info = mlflow.sklearn.log_model(
        sk_model=lr,
        name="elasticnet",
        params={
            "alpha": 0.5,
            "l1_ratio": 0.5,
        },
        input_example = train_x
    )

    # LoggedModelとそのプロパティを確認
    logged_model = mlflow.get_logged_model(model_info.model_id)
    print(logged_model.model_id, logged_model.params)
    # m-fa4e1bca8cb64971bce2322a8fd427d3, {'alpha': '0.5', 'l1_ratio': '0.5'}

    # トレーニングデータセットでモデルを評価し、メトリクスをログ
    # これらのメトリクスはLoggedModelエンティティにリンク
    predictions = lr.predict(train_x)
    (rmse, mae, r2) = compute_metrics(train_y, predictions)
    mlflow.log_metrics(
        metrics={
            "rmse": rmse,
            "r2": r2,
            "mae": mae,
        },
        model_id=logged_model.model_id,
        dataset=train_dataset
    )

    # メトリクスを持つLoggedModelを確認
    logged_model = mlflow.get_logged_model(model_info.model_id)
    print(logged_model.model_id, logged_model.metrics)
    # m-fa4e1bca8cb64971bce2322a8fd427d3, [<Metric: dataset_name='train', key='rmse', model_id='m-fa4e1bca8cb64971bce2322a8fd427d3, value=0.7538635773139717, ...>, ...]

これでランが作成されますが、新たなエンティティとしてモデルが記録されます。エクスペリメントページのモデルタブでも確認できます。

しばらくして、最新の本番データに基づいた新しい評価データセットを取得したら、新しいモデル評価ジョブを実行して、この新しいデータセットに対するモデルのパフォーマンスを測定できます。このジョブは新しいMLflow Runとして追跡されます。

この例では、2つのMLflow Run（training_run と evaluation_run）と1つのMLflow LoggedModel（elasticnet）が生成されます。生成された LoggedModel から、すべてのパラメータとメタデータ、およびトレーニングと評価のランからリンクされたすべてのメトリクスを確認できます。

# テストデータセット評価ジョブを表すランを開始
with mlflow.start_run() as evaluation_run:
  # テストデータセットをMLflowでロード。テストメトリクスをこのデータセットにリンク
  test_dataset: mlflow.entities.Dataset = mlflow.data.from_pandas(test_df, name="test")
  test_x = test_dataset.df.drop(["quality"], axis=1)
  test_y = test_dataset.df[["quality"]]

  # モデルをロード
  model = mlflow.sklearn.load_model(f"models:/{logged_model.model_id}")

  # トレーニングデータセットでモデルを評価し、メトリクスをログ。モデルにリンク
  predictions = model.predict(test_x)
  (rmse, mae, r2) = compute_metrics(test_y, predictions)
  mlflow.log_metrics(
    metrics={
      "rmse": rmse,
      "r2": r2,
      "mae": mae,
    },
    dataset=test_dataset,
    model_id=logged_model.model_id
  )

改めてモデルを確認すると、最初のトレーニングにおけるランとメトリクス、評価のランとメトリクスをモデルから確認できるようになっています。

get_logged_modelでも情報を確認できます。

print(mlflow.get_logged_model(logged_model.model_id).to_dictionary())

{'artifact_location': 'dbfs:/databricks/mlflow-tracking/3126308776674605/logged_models/m-c86bf92aab5146fda4b404b0fedae336/artifacts', 'creation_timestamp': 1744200949778, 'experiment_id': '3126308776674605', 'last_updated_timestamp': 1744200955164, 'metrics': [<Metric: dataset_digest='d95ce242', dataset_name='train', key='rmse', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.13264009364619966>, <Metric: dataset_digest='d95ce242', dataset_name='train', key='r2', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.395372792495956>, <Metric: dataset_digest='d95ce242', dataset_name='train', key='mae', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.3209158276037644>, <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='rmse', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.13442354610166166>, <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='r2', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.4211426244426051>, <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='mae', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.3256178886187552>], 'model_id': 'm-c86bf92aab5146fda4b404b0fedae336', 'model_type': '', 'name': 'elasticnet', 'params': {'alpha': '0.5', 'l1_ratio': '0.5'}, 'source_run_id': 'b41bd03fb85c4648a50b8f13ac4249a7', 'status': 2, 'status_message': '', 'tags': {'mlflow.databricks.cluster.id': '0409-121238-ckjb5rwi-v2n', 'mlflow.databricks.notebook.commandID': '1744200760564_8172306878656848519_7ccb8f652f2640609f3486a1a2b0f332', 'mlflow.databricks.notebookID': '3126308776674605', 'mlflow.databricks.notebookPath': '/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/mlflow3-ml-example', 'mlflow.databricks.webappURL': 'https://xxxxx.cloud.databricks.com', 'mlflow.databricks.workspaceID': '5099015744649857', 'mlflow.databricks.workspaceURL': 'https://xxxxx.cloud.databricks.com', 'mlflow.source.name': '/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/mlflow3-ml-example', 'mlflow.source.type': 'NOTEBOOK', 'mlflow.user': 'spark-f6ea5246-5ebd-427d-9450-f0'}}

次に、モデルをUCに登録します。UCモデルバージョンページでモデルID、パラメータ、およびメトリクスも確認できます。

# カタログに対する `USE CATALOG` 権限、およびスキーマに対する `USE SCHEMA` 権限が必要です。
# 必要に応じて、ここでカタログ名とスキーマ名を変更してください。

CATALOG = "takaakiyayoi_catalog"
SCHEMA = "mlflow"
MODEL = "ml_model"
MODEL_NAME = f"{CATALOG}.{SCHEMA}.{MODEL}"

uc_model_version = mlflow.register_model(model_info.model_uri, name=MODEL_NAME)

これで、Unity Catalogのモデルバージョンページでモデルバージョンとすべての集中化されたパフォーマンスデータを確認できます。

次のセルに示すように、APIを使用して同じ情報を取得することもできます。

# モデルバージョンを取得
from mlflow import MlflowClient
client = MlflowClient()
model_version = client.get_model_version(name=MODEL_NAME, version=uc_model_version.version)
print(model_version)

<ModelVersion: aliases=[], creation_timestamp=1744201072223, current_stage=None, deployment_job_state=<ModelVersionDeploymentJobState: current_task_name='', job_id='', job_state='DEPLOYMENT_JOB_CONNECTION_STATE_UNSPECIFIED', run_id='', run_state='DEPLOYMENT_JOB_RUN_STATE_UNSPECIFIED'>, description='', last_updated_timestamp=1744201073758, metrics=[<Metric: dataset_digest='d95ce242', dataset_name='train', key='rmse', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.13264009364619966>,
 <Metric: dataset_digest='d95ce242', dataset_name='train', key='r2', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.395372792495956>,
 <Metric: dataset_digest='d95ce242', dataset_name='train', key='mae', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='b41bd03fb85c4648a50b8f13ac4249a7', step=0, timestamp=1744200955317, value=0.3209158276037644>,
 <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='rmse', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.13442354610166166>,
 <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='r2', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.4211426244426051>,
 <Metric: dataset_digest='cc28e2f0', dataset_name='test', key='mae', model_id='m-c86bf92aab5146fda4b404b0fedae336', run_id='c526255db4f34c169ad46b8529b47a4b', step=0, timestamp=1744201004520, value=0.3256178886187552>], model_id='m-c86bf92aab5146fda4b404b0fedae336', name='takaakiyayoi_catalog.mlflow.ml_model', params=[<LoggedModelParameter: key='alpha', value='0.5'>,
 <LoggedModelParameter: key='l1_ratio', value='0.5'>], run_id='b41bd03fb85c4648a50b8f13ac4249a7', run_link=None, source='models:/m-c86bf92aab5146fda4b404b0fedae336', status='READY', status_message='', tags={}, user_id='takaaki.yayoi@databricks.com', version='1'>

こちらで、デプロイメントジョブと連携しています。

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル