DatabricksにおけるMLflow 3.0のサポート - 従来の機械学習編

こちらのマニュアルに沿ってサンプルノートブックを動かします。

https://docs.databricks.com/aws/ja/mlflow/mlflow-3-install

:::note info
**注意**
執筆時点では[ベータ版](https://docs.databricks.com/aws/ja/release-notes/release-types)です。
:::

# MLflow 3.0

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

https://qiita.com/taka_yayoi/items/8ebc58c04f7889438eb1

# 有効化

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

![Screenshot 2025-04-09 at 21.15.18.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/54771c50-54cf-40c5-9ada-215335bb440a.png)

# 従来のMLのワークフロー

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

https://docs.databricks.com/aws/ja/mlflow/mlflow3-ml-workflow

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

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

```py
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, ...>, ...]
```

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

![Screenshot 2025-04-09 at 21.16.17.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/e1109526-322f-4ff5-8853-e76703e3e4f7.png)

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

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

```py
# テストデータセット評価ジョブを表すランを開始
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
  )
```

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

![Screenshot 2025-04-09 at 21.17.21.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/cd40ac7b-c425-4c7b-85f1-ec5e92da3def.png)

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

```py
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、パラメータ、およびメトリクスも確認できます。

```py
# カタログに対する `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のモデルバージョンページでモデルバージョンとすべての集中化されたパフォーマンスデータを確認できます。

![Screenshot 2025-04-09 at 21.18.05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/823def41-14cc-4b65-81c7-b5157fd3c661.png)

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

```py
# モデルバージョンを取得
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'>
```

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

https://qiita.com/taka_yayoi/items/651ad6078a9c338cd482

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)