MLflow: 4. Model Registry を使った実験管理

はじめに
- 環境情報
前回までのMLflowシリーズ
MLflowとは
Model Registryとは
Model Registryでできること
Model Registryのワークフロー
UIでの基本的な操作
APIでの基本的な操作
まとめ
参考

はじめに

今回はMLflowシリーズの最後の機能であるMLflow Model Registryについてまとめる。

環境情報

Databricks RunTime: 10.2 ML (includes Apache Spark 3.2.0, Scala 2.12)

前回までのMLflowシリーズ

MLflow Tracking

ktksq.hatenablog.com

MLflow Projects

ktksq.hatenablog.com

MLflow Models

ktksq.hatenablog.com

MLflowとは

前回の記事でまとめたので、そちらをご参照ください。 ktksq.hatenablog.com

Model Registryとは

MLflowモデルの全ライフサイクルをメンバー間で共同管理するための一元的なモデルストア、APIのセット、および UI のこと。モデルレジストリはチームでMLモデルを共有し、実験からオンラインテスト、本番まで共同で作業し、承認とガバナンスのワークフローと統合し、MLのデプロイメントとそのパフォーマンスを監視するコラボレーションハブとして機能する。

Model Registryでできること

モデルレジストリはセントラル・リポジトリーとしてMLflowモデルを登録できる。登録されたモデルは、一意の名前、バージョン、ステージ、およびコメントやタグなどのメタデータを持つ。モデルレジストリでは、以下の情報を管理することができる。

モデルのバージョン管理

同じモデル名のモデルが更新された場合、自動的にモデルのバージョンを更新し管理する。

モデルのステージの遷移管理

各モデルのバージョンに「ステージング」や「プロダクション」のようなステージを割り当てることで、モデルのライフサイクルを管理する。新規登録イベントや変更をアクティビティとして記録し、ユーザー、変更、コメントなどの追加メタデータを自動的にログに記録できる。ステージ遷移、変更の要求、レビュー、承認の一連のフローを記録することでコントロールとガバナンスを向上させる。

モデルのコード管理

モデルがどの MLflowエクスペリメントとランから作成されたか記録する。

モデルの概要管理

アルゴリズムの説明や採用したデータセット、方法論など、チームにとって有用な関連情報をコメントやタグで記録する。

モデルのデプロイ管理

特定のモデルバージョンを要求した本番ジョブはどれかなど、どのモデルをデプロイしたか管理する。

モデルサービング

モデルレジストリに登録したMLflowモデルをRESTエンドポイントとして公開する。

Model Registryのワークフロー

1.モデルの登録

モデルに対応するモデルフレーバーのmlflow.<model_flavor>.log_model()を使用して、MLflowモデルを作成する。記録されたMLflowモデルはモデルレジストリに登録することができる。一度モデルが記録されると、UIやAPIを通してモデルレジストリ内のモデルを追加、変更、更新、移行、または削除が可能になる。

2.モデルのバージョン登録

登録された各モデルは、1つまたは複数のバージョンを持つことができる。新しいモデルがモデルレジストリに追加される場合は、バージョン1として追加される。同じモデル名で新しいモデルが登録されると、バージョン番号が自動的に増加する。

3.モデルステージ遷移

各バージョンのモデルは、任意の時点で1つのステージを割り当てることができる。ステージの種類にはステージング、プロダクション、アーカイブがあり、各ステージ間の移動の際に変更のリクエスト、レビュー、承認のフローを経ることでモデルの管理とガバナンスを可能にしている。

4.モデルサービング

登録したMLflowモデルをRESTエンドポイントとして公開できる。

UIでの基本的な操作

1.モデルを登録する

ランの詳細ページから、アーティファクションとして記録された MLflowモデルを選択する。その後「モデルを登録」を押下する。

2.モデルのバージョンを登録する

モデルを登録すると自動でバージョンが登録されるため、特に操作は必要ない。

3.モデルのステージを登録する

バージョン詳細ページで、モデルのバージョンの詳細と現在のステージを確認することができる。右上のステージのドロップダウンをクリックすると、モデルバージョンを他の有効なステージに移行させることが可能である。

4.モデルサービング

APIでの基本的な操作

モデルを登録する

APIでモデルを登録する方法は3種類ある。

・mlflow.<model_flavor>.log_model()を使用する方法

from random import random, randint
from sklearn.ensemble import RandomForestRegressor

import mlflow
import mlflow.sklearn

with mlflow.start_run(run_name="YOUR_RUN_NAME") as run:
    params = {"n_estimators": 5, "random_state": 42}
    sk_learn_rfr = RandomForestRegressor(**params)

    # MLflowのAPIを使用してパラメータとメトリクスをログに記録する
    mlflow.log_params(params)
    mlflow.log_param("param_1", randint(0, 100))
    mlflow.log_metrics({"metric_1": random(), "metric_2": random() + 1})

    # sklearn モデルをログに記録し、バージョン1として登録する
    mlflow.sklearn.log_model(
        sk_model=sk_learn_rfr,
        artifact_path="sklearn-model",
        registered_model_name="sk-learn-random-forest-reg-model"
    )

・ mlflow.register_model()を使用する方法

runs:URIの引数の一部としてrun_idを使用してモデルを登録する。

result = mlflow.register_model(
    "runs:/d16076a3ec534311817565e6527539c0/sklearn-model",
    "sk-learn-random-forest-reg "
)

・create_registered_model() を使って新しいモデルを作成する方法

# バージョンに関連付けされていない空の登録済みモデルを作成する
from mlflow.tracking import MlflowClient

client = MlflowClient()
client.create_registered_model("sk-learn-random-forest-reg-model")

# モデルの新しいバージョンを作成する
client = MlflowClient()
result = client.create_model_version(
    name="sk-learn-random-forest-reg-model",
    source="mlruns/0/d16076a3ec534311817565e6527539c0/artifacts/sklearn-model",
    run_id="d16076a3ec534311817565e6527539c0"
)

モデルレジストリからMLflowモデルを取得する

MLflow のモデルを登録したら、mlflow.<model_flavor>.load_model()を使って、任意のモデルを取得することができる。

・モデル URI の一部としてそのバージョン番号を指定して、特定のモデルのバージョンを取得する

# 特定のモデルのバージョンを取得する
import mlflow.pyfunc

model_name = "sk-learn-random-forest-reg-model"
model_version = 1

model = mlflow.pyfunc.load_model(
    model_uri=f "models:/{model_name}/{model_version}"
)
model.predict(data)

・モデルURIの一部としてモデルステージを指定して、指定ステージの最新バージョンのモデルを取得する

# ステージごとにモデルのバージョンを取得する
import mlflow.pyfunc

model_name = "sk-learn-random-forest-reg-model"
stage = 'staging'

model = mlflow.pyfunc.load_model(
    model_uri=f "models:/{model_name}/{stage}"
)
model.predict(data)

モデルレジストリからMLflowモデルをデプロイする

MLflowモデルを登録したら、モデルをホスト上のサービスとしてデプロイすることができる。

#!/usr/bin/env sh

# Model Registryが存在するトラッキングURLの環境変数を設定する
export MLFLOW_TRACKING_URI=http://localhost:5000

# モデルレジストリから本番モデルを配信する
mlflow models serve -m "models:/sk-learn-random-forest-reg-model/Production"

MLflowモデル情報の追加と更新

・モデルのバージョン情報を更新する

client = MlflowClient()
client.update_model_version(
    name="sk-learn-random-forest-reg-model",
    version=1,
    description="このモデルのバージョンは、100本の決定木を含むscikit-learnランダムフォレストです"
)

・MLflowモデル名を変更する

特定のバージョンのモデルの説明を追加または更新するのと同様に、rename_registered_model()を使用して既存の登録済みモデルの名前を変更することができる。

client = MlflowClient()
client.rename_registered_model(
    name="sk-learn-random-forest-reg-model",
    new_name="sk-learn-random-forest-reg-model-100"
)

・MLflowモデルのステージを移行する

モデルのライフサイクルの中で、モデルはステージング、プロダクション、アーカイブとステージが変化していく。ライフサイクルに合わせて、登録されたモデルは任意のステージに移行させることができる。

client = MlflowClient()
client.transition_model_version_stage(
    name="sk-learn-random-forest-reg-model",
    version=3,
    stage="Production"
)
# <stage>に指定できる値: Staging|Archived|Production|None