Databricksにおけるモデル管理機能について

本記事は、Japan Digital Design Advent Calendar 2023 の10日目の記事になります。

adventar.org

---

三菱UFJフィナンシャル・グループ（以下MUFG）の戦略子会社であるJapan Digital Design（以下JDD）のTechnology ＆ Development Division所属の町田です。普段はシステムアーキテクチャ策定やデータ基盤の開発・運用に従事しております。

今回の記事では、AIモデル開発プラットフォームDatabricksにおけるモデル管理機能について記載いたします。

はじめに

Databricksのモデル管理機能として、以下の２つの機能があります。

• Workspace Model Registry

• Unity Catalog

Workspace Model Registryは従来から用いられてきた方法であるのに対して、Unity Catalogによるモデル管理は比較的新しい機能です。Unity Catalogモデル管理機能の登場により、当初、混乱した部分もあったので、これらの機能差異や使い方を比較・説明しました。

なお、これら２つのモデル管理機能は並行利用することも可能です。またUnity Catalogモデル管理機能は、Unity Catalogを有効化したWorkspaceのみ利用可能な機能です。DatabricksのデータガバナンスソリューションUnity Catalogに関しては、以下の解説記事をご参照ください。

DatabricksのUnity Catalogとは？ #Databricks - Qiita

Workspace Model RegistryとUnity Catalogのモデル管理を、以下の観点で整理・比較してみました。

Unity Catalogモデル管理では、テーブルデータ同様にモデルのガバナンスを効かせられる点は大きなメリットと感じました。一方で従来のWorkspace Model Registryの移行要求・承認機能がUnity Catalogで提供されていない点は、今後のアップデートに期待したいところです。

モデル管理機能の比較

事前準備①（モデル作成）

Scikit-LearnのIrisデータセットを利用して、LinearSVCモデルを作成します。

# import library
import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.svm import LinearSVC

# load dataset
iris_data = load_iris()

# split data
train_df, test_df, train_target, test_target = train_test_split(iris_data.data, iris_data.target, test_size=0.33)

# model
model = LinearSVC()
model.fit(train_df, train_target)

上記コードを実行すると、Logged 1 run to an experiment in MLflow. と出力される。「1 run」のリンクをクリックすると、以下のようにMLflow Experimentの画面が表示されるので、左上のRun IDを控えておきます。

事前準備②（Unity Catalogカタログ・スキーマ・権限付与）

Unity Catalogでは、<catalog-name>.<schema-name>.<object-name>という３階層の名前空間でオブジェクトを一意に特定します。今回はカタログ名production、スキーマ名group_a、モデルオブジェクト名irisとして進めていきます。

Unity Catalogモデル管理の準備として、カタログ・スキーマ作成と権限付与を行います。

%sql
CREATE CATALOG IF NOT EXISTS production;
USE CATALOG production;
CREATE SCHEMA IF NOT EXISTS group_a;
USE SCHEMA group_a;

GRANT USE CATALOG ON CATALOG production TO `account users`;
GRANT USE SCHEMA ON SCHEMA group_a TO `account users`;
GRANT CREATE MODEL ON SCHEMA group_a TO `account users`;%

これらの作業は、以下のとおりCatalog Explorerで実施することも可能です。

モデルの登録

Workspace Model RegistryとUnity Catalogどちらでモデル管理を行うのか明確化するため、モデル管理作業前にレジストリURIを指定する必要があります。下表のとおり、mlflow.set_registry_uri()でレジストリURI設定をします（mlflow.get_registry_uri()で確認可能）。以降は、適切にレジストリURIが設定されている前提で話を進めます。

以下のとおり、MLflowのregister_model関数にRun IDを指定して、モデルを登録します。

• Workspace Model Registry

import mlflow
model_name = "iris"
mlflow.set_registry_uri("databricks")
mlflow.register_model("runs:/3633a1156f774b4bb5c4030b93f8d961/model", model_name)

• Unity Catalog

import mlflow
catalog = "production"
schema = "group_a"
model_name = "iris"
mlflow.set_registry_uri("databricks-uc")
mlflow.register_model("runs:/3633a1156f774b4bb5c4030b93f8d961/model", f"{catalog}.{schema}.{model_name}")

GUI画面でモデルを確認すると、以下のように表示されます。

• Workspace Model Registry：Modelsから参照可能

• Unity Catalog：Catalogから参照可能

モデルの移行

２つのモデル管理方法で、稼働中・テスト中などのモデルの状態を識別する手法が異なります。Workspace Model Registryでは「ステージ」、Unity Catalogでは「エイリアス」を用いてモデルの状態を管理します。

• Workspace Model Registry

Workspace Model RegistyではModels画面から以下のようにステージの移行（リクエスト）を行います。特定モデルバージョンにProduction/Stagingステージいずれかを付与できますが、Archivedステージは利用しなくなった複数モデルバージョンに付与することが可能です（Production/Stagingステージが外れると自動的にArchivedステージへ移行）。

APIでもステージの移行は可能です。ただし、MLflowClientにはステージ移行リクエストのAPIが提供されていないため、ステージ移行リクエストはREST APIで対応します（REST APIを利用する場合はPersonal Access Tokenの生成が必要）。

# Transition to Production stage
from mlflow import MlflowClient
MlflowClient().transition_model_version_stage(model_name, version=3, stage="Production")

# Make a transition request
import requests
DOMAIN = "<workspace-host-name>"
TOKEN = "<personal-access-token>"
response = requests.post(
  'https://%s/api/2.0/mlflow/transition-requests/create' % (DOMAIN),
  headers={'Authorization': 'Bearer %s' % TOKEN},
  json={
      "name": "iris",
      "version": "2",
      "stage": "Staging",
      "comment": "Great!"
  }
)

• Unity Catalog

Unity Catalogでは任意のエイリアス名を設定し、モデルの状態移行を行うことが可能です。特定モデルバージョンに複数のエイリアスを付与することも可能ですので、Workspace Model Registryより柔軟なモデル状態管理が可能になります。なお、Workspace Model Registryとは異なり、Archivedステージ相当のエイリアスや、エイリアス移行を要求・承認する機能はありません。

from mlflow import MlflowClient
MlflowClient().set_registered_model_alias(f"{catalog}.{schema}.{model_name}", "prod", 3)

モデルのロード

モデルのロードは、以下のようにモデル名を指定してロードします。

• Workspace Model Registry

stage = "Production"
model = mlflow.pyfunc.load_model(f"models:/{model_name}/{stage}")
version = "2"
model = mlflow.pyfunc.load_model(f"models:/{model_name}/{version}")

• Unity Catalog

stage = "prod"
model = mlflow.pyfunc.load_model(f"models:/{catalog}.{schema}.{model_name}@{stage}")
version = "2"
model = mlflow.pyfunc.load_model(f"models:/{catalog}.{schema}.{model_name}/{version}")

モデルのアクセス制御

• Workspace Model Registry

Workspace Model Registryに登録した各モデルに対して、No Permissions, Can Read, Can Edit, Can Manage Staging Versions, Can Manage Production Versions, Can Manageという6段階の権限レベルを設定することが可能です（権限レベルの詳細は以下リンクご参照）。

Share MLflow Models | Databricks on AWS

import requests
DOMAIN = "<workspace-host>"
TOKEN = "<personal-access-token>"
response = requests.put(
  'https://%s/api/2.0/permissions/registered-models/<registered_model_id>' % (DOMAIN),
  headers={'Authorization': 'Bearer %s' % TOKEN},
  json={
      "access_control_list": [{
          "group_name": "<group-name>",
          "permission_level": "CAN_MANAGE"
          }]
      }
)

• Unity Catalog

Unity Catalogでは、テーブル同様の管理スキームで、モデルのアクセス権限を管理することが可能です。カタログに対するCREATE MODEL権限（モデル作成権限）、スキーマに対するCREATE MODEL権限（モデル作成権限）、モデルに対するAPPLY TAG, EXECUTE権限（モデルのタグ編集・実行権限）のレベルでアクセス制御することが可能です。以下のようにCatalog ExplorerまたはGRANT/REVOKEステートメントで権限設定を行います。なお、エイリアス編集権限はGRANTすることは出来ず、モデル所有者に付与されます。

%sql
GRANT APPLY TAG ON MODEL production.group_a.iris TO `account users`;
GRANT EXECUTE ON MODEL production.group_a.iris TO `account users`;

CI/CD

• Workspace Model Registry

Workspace Model Registryでは、モデルのWebhook機能を利用することが可能です。モデルの移行要求の手動承認などのイベントをトリガーに、ジョブ実行やHTTPエンドポイント連携することが可能です。以下はモデルのStaging/Production移行をトリガーにジョブを稼働させる設定例です。

MLflow モデルレジストリ Webhooks on Databricks | Databricks on AWS

from databricks_registry_webhooks import RegistryWebhooksClient, JobSpec

job_spec = JobSpec(
  job_id="<job-id>",
  workspace_url="<workspace-url>",
  access_token="<personal-access-token>"
)
job_webhook = RegistryWebhooksClient().create_webhook(
  model_name="iris",
  events=["MODEL_VERSION_TRANSITIONED_TO_STAGING", "MODEL_VERSION_TRANSITIONED_TO_PRODUCTION"],
  job_spec=job_spec,
  description="Job webhook trigger",
  status="TEST_MODE"
)

• Unity Catalog

Unity Catalogでは、モデルのWebhook機能を提供していませんが、ジョブのWebhook機能を利用して、トレーニングジョブの正常終了などをトリガーにHTTPエンドポイントへ連携することは可能です。

また、Unity Catalogモデル管理では、モデル移行要求・手動承認機能はないため、CI/CDツール等へ連携して機械学習パイプラインに手動承認処理を追加することができます。

機械学習ワークフローを Unity Catalog のターゲットモデルにアップグレードする | Databricks on AWS

モデル共用

• Workspace Model Registry

Workspace間でモデルを共用するには、まず接続先となるWorkspaceでPersonal Access Tokenを発行し、モデルのアクセス権限を付与しておきます。接続元Workspaceでは、以下のとおり接続先のアクセス情報をsecretsに格納し、リモートレジストリURI指定でモデルへアクセスします。

ワークスペース間でモデルを共有する | Databricks on AWS

!databricks secrets create-scope --scope <scope>
!databricks secrets put --scope <scope> --key <prefix>-host
!databricks secrets put --scope <scope> --key <prefix>-token
!databricks secrets put --scope <scope> --key <prefix>-workspace-id

import mlflow.pyfunc
model = mlflow.pyfunc.load_model(f'models://{scope}:{key}@databricks/{model_name}/Production')

• Unity Catalog

Unity Catalogでは、テーブルデータ同様にWorkspace横断でモデルへアクセスすることが可能です。ただし、モデルが所属するカタログをWorkspaceに割り当てておく必要はあります。カタログの割り当てと権限付与を行っておけば、どのWorkspaceからも簡単かつ同一コードで、モデル利用できるのは便利です。

カタログの作成と管理 | Databricks on AWS

さいごに

今回の記事では、Databricksのモデル管理機能について比較・整理しました。

Workspace Model RegistryとUnity Catalogそれぞれで細かい機能差がありますので、足りない部分をどのように補っていくか検討していく参考になれば幸いです。Databricksでは、Unity Catalogモデル管理機能の利用を推奨しており、今後のアップデートにも注目です。

Databricksは機能リリースが多いので、今後もさまざまなアップデートを追いかけていきたいと思います。

参考リンク

Unity Catalogでモデルのライフサイクルを管理する | Databricks on AWS

ワークスペースモデルレジストリを使用したモデルのライフサイクルの管理 | Databricks on AWS

最後までご覧いただきありがとうございました。

---

Japan Digital Design株式会社では、一緒に働いてくださる仲間を募集中です。カジュアル面談も実施しておりますので下記リンク先からお気軽にお問合せください。

採用情報｜Japan Digital Design 株式会社

この記事に関するお問い合わせはこちらにお願いします。

お問い合わせ｜Japan Digital Design 株式会社