DatabricksにおけるMLflowモデルレジストリWebhook

[MLflow Model Registry Webhooks on Databricks \| Databricks on AWS](https://docs.databricks.com/applications/mlflow/model-registry-webhooks.html) [2022/2/1時点]の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

> **プレビュー**
本機能は[パブリックプレビュー](https://docs.databricks.com/release-notes/release-types.html)です。

Webhookを用いることで、モデルレジストリのイベントに基づいて、アクションを自動で起動できるようになります。既存のCI/CDツールやワークフローと、皆様の機械学習パイプラインを統合し、自動化することができます。例えば、新規モデルバージョンが作成された際にCIビルドを実行したり、モデルのプロダクションへの移行がリクエストされた際にSlackで通知を送ることが可能となります。

Webhookは[Databricks REST API](https://docs.databricks.com/dev-tools/api/latest/mlflow.html)、[PyPI](https://pypi.org/project/databricks-registry-webhooks/)からインストールできるPythonクライアント`databricks-registry-webhooks`から利用することができます。

# Webhookのイベント

以下のイベントをWebhookに指定することができます。複数指定も可能です。

- **MODEL_VERSION_CREATED:** 新規モデルバージョンが作成された。
- **MODEL_VERSION_TRANSITIONED_STAGE:** モデルバージョンのステージが変更された。
- **TRANSITION_REQUEST_CREATED:** ユーザーがモデルのステージ変更をリクエストした。
- **COMMENT_CREATED:** 登録モデルにユーザーがコメントをした。
- **REGISTERED_MODEL_CREATED:** 新規登録モデルが作成された。モデル名を作成リクエストで指定しない場合に作成されるレジストリに対するWebhookでのみ、このイベントタイプを指定できます。
- **MODEL_VERSION_TAG_SET:** モデルバージョンにユーザーがタグを設定した。
- **MODEL_VERSION_TRANSITIONED_TO_STAGING:** モデルバージョンがステージングに遷移した。
- **MODEL_VERSION_TRANSITIONED_TO_PRODUCTION:** モデルバージョンがプロダクションに遷移した。
- **MODEL_VERSION_TRANSITIONED_TO_ARCHIVED:** モデルバージョンがアーカイブされた。
- **TRANSITION_REQUEST_TO_STAGING_CREATED:** ユーザーがモデルバージョンのステージングへの遷移をリクエストした。
- **TRANSITION_REQUEST_TO_PRODUCTION_CREATED:** ユーザーがモデルバージョンのプロダクションへの遷移をリクエストした。
- **TRANSITION_REQUEST_TO_ARCHIVED_CREATED:** ユーザーがモデルバージョンのアーカイブをリクエストした。

# Webhookのタイプ

トリガーするターゲットに応じて2つのタイプのWebhookが存在します。

- **HTTPエンドポイントを用いたWebhook:** HTTPエンドポイントにトリガーを送信します。
- **ジョブトリガーを用いたWebhook:** Databricksワークスペースのジョブを起動します。ジョブがあるワークスペースでIPアクセスリストが有効化されている場合、モデルレジストリのワークスペースのIPアドレスを許可する必要があります。詳細は[ジョブによるレジストリWebhookに対するIPアドレスの許可](#セキュリティ)を参照ください。

また、スコープによって2つのタイプのWebhookが存在します。

- **モデル固有のWebhook:** 特定の登録モデルに適用されるWebhookです。
    - アクセスコントロール: モデル固有のWebhookを作成、編集、削除、テストするために登録モデルに対する*Can Manage*権限が必要となります。
- **レジストリ規模のWebhook:** 登録モデルの作成を含め、ワークスペースの全ての登録モデルのイベントで起動するWebhookです。レジストリ規模のWebhookを作成するには、作成時に`model_name`フィールドを削除します。
    - アクセスコントロール: レジストリ規模のWebhookを作成、編集、削除、テストするためにワークスペース管理者権限が必要となります。

# Webhookのペイロード

それぞれのイベントトリガーには、Webhookエンドポイントに対するリクエストのペイロードに含まれる最小限のフィールドがあります。

- アーティファクトのパスのようなセンシティブな情報は除外されます。適切なACLを持つユーザー、プリンシパルが、このような情報に対するモデルレジストリに対するクエリーを行うためにクライアントやREST APIを使用することができます。
- Slackとの連携で`text`フィールドを活用します。Slackメッセージを送信するためには、Slack WebhookエンドポイントをWebhook URLとして指定します。
- ペイロードは暗号化されません。DatabricksがWebhookのソースであることを、どのように検証するのかについては[セキュリティ](#セキュリティ)を参照ください。
- ジョブのレジストリWebhookのペイロードは、ターゲットワークスペースのjobs/run-nowエンドポイントに対して、複数のフィールドを`event_message `という単一のフィールドにラップします。このJSONオブジェクトの値はDatabricksノートブックでアンパックすることができます。

```py:Python
import json

event_message = dbutils.widgets.get("event_message")
event_message_dict = json.loads(event_message)
```

**サンプルペイロード**

- イベント: `MODEL_VERSION_TRANSITIONED_STAGE`

```:レスポンス
POST
/your/endpoint/for/event/model-versions/stage-transition
--data {
  "event": "MODEL_VERSION_TRANSITIONED_STAGE",
  "webhook_id": "c5596721253c4b429368cf6f4341b88a",
  "event_timestamp": 1589859029343,
  "model_name": "Airline_Delay_SparkML",
  "version": "8",
  "to_stage": "Production",
  "from_stage": "None",
  "text": "Registered model 'someModel' version 8 transitioned from None to Production."
}
```

- イベント: `MODEL_VERSION_TAG_SET`

```:レスポンス
POST
/your/endpoint/for/event/model-versions/tag-set
--data {
  "event": "MODEL_VERSION_TAG_SET",
  "webhook_id": "8d7fc634e624474f9bbfde960fdf354c",
  "event_timestamp": 1589859029343,
  "model_name": "Airline_Delay_SparkML",
  "version": "8",
  "tags": [{"key":"key1","value":"value1"},{"key":"key2","value":"value2"}],
  "text": "example@yourdomain.com set version tag(s) 'key1' => 'value1', 'key2' => 'value2' for registered model 'someModel' version 8."
}
```

- イベント: `COMMENT_CREATED`

```:レスポンス
POST
/your/endpoint/for/event/comments/create
--data {
  "event": "COMMENT_CREATED",
  "webhook_id": "8d7fc634e624474f9bbfde960fdf354c",
  "event_timestamp": 1589859029343,
  "model_name": "Airline_Delay_SparkML",
  "version": "8",
  "comment": "Raw text content of the comment",
  "text": "A user commented on registered model 'someModel' version 8."
}
```

# セキュリティ

セキュリティに関しては、Databricksはペイロードから生成されるヘッダーに`X-Databricks-Signature`を埋め込み、[HMAC with SHA\-256 algorithm](https://en.wikipedia.org/wiki/HMAC)を用いてWebhookに紐づけられたシークレットキーを共有します。

加えて、Webhookの`HttpUrlSpec`に標準的な認証ヘッダーを指定することで、認証ヘッダーを搬出するリクエストに追加することができます。

- クライアントによる検証: 共有シークレットが設定されている場合、クライアントは共有シークレットを用いてペイロードが使用しているHMAC-encodingし、ヘッダーのX-Databricks-Signatureと比較することで、HTTPリクエストのソースを検証する必要があります。特に、`enable_ssl_verification`が`false`に設定されておりSSL証明書の検証が無効化されている場合には、この検証を行うことを強く推奨します。

```py:Python
import hmac
import hashlib
import json

secret = shared_secret.encode('utf-8')
signature_key = 'X-Databricks-Signature'

def validate_signature(request):
  if not request.headers.has_key(signature_key):
    raise Exception('No X-Signature. Webhook not be trusted.')

  x_sig = request.headers.get(signature_key)
  body = request.body.encode('utf-8')
  h = hmac.new(secret, body, hashlib.sha256)
  computed_sig = h.hexdigest()

  if not hmac.compare_digest(computed_sig, x_sig.encode()):
    raise Exception('X-Signature mismatch. Webhook not be trusted.')
```

- HTTPレジストリWebhookのヘッダーの認証: 認証ヘッダーが設定されている場合、クライアントは認証ヘッダーのbearerトークン/認証クレディンシャルを検証することでHTTPリクエストのソースを検証すべきです。
- ジョブのレジストリWebhookのIPアドレス許可: IPアクセスリストが有効化されている別のワークスペースでジョブを起動するためにWebhookを使うには、リクエストを受け付けるためにWebhoookが存在するリージョンNATを[許可する](https://docs.databricks.com/security/network/ip-access-list.html)必要があります。許可するIPを特定するには、Databricksアカウントチームにコンタクトください。Webhookとジョブが同じワークスペースにある場合には、IPアドレスを許可する必要はありません。

# 監査ログ

ワークスペースで監査ログが有効化されている場合、以下のイベントが監査ログに記録されます。

- Webhookの作成
- Webhookの更新
- Webhookの一覧
- Webhookの削除
- Webhookのテスト
- Webhookのトリガー
    - HTTPリクエストによるWebhookにおいては、Webhookに指定されたURLに送信されたHTTPリクエストとURL、`enable_ssl_verification`の値が記録されます。
    - ジョブトリガーによるWebhookにおいては、`job_id`と`workspace_url`が記録されます。

# サンプル

このセクションでは、HTTPレジストリWebhookのサンプルワークフロー、ジョブによるレジストリWebhookのサンプルワークフロー、サンプルノートブックを説明します。

## HTTPによるレジストリWebhookサンプルワークフロー

### 1. Webhookの作成

HTTPSエンドポイントがWebhookのイベントリクエストを受け取る準備ができているのであれば、Webhooks Databricks REST APIを用いてWebhookを作成することができます。例えば、Slackのチャンネルにメッセージを投稿するためにWebhookのURLにSlackを指定することができます。

```bash:Bash
$ curl -X POST -H "Authorization: Bearer <access_token>" -d \
'{"model_name": "<model-name>",
  "events": ["MODEL_VERSION_CREATED"],
  "description": "Slack notifications",
  "status": "TEST_MODE",
  "http_url_spec": {
    "url": "https://hooks.slack.com/services/...",
    "secret": "anyRandomString"
    "authorization": "Bearer AbcdEfg1294"}}' https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/create
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient, HttpUrlSpec

http_url_spec = HttpUrlSpec(
  url="https://hooks.slack.com/services/...",
  secret="secret_string",
  authorization="Bearer AbcdEfg1294"
)
http_webhook = RegistryWebhooksClient().create_webhook(
  model_name="<model-name>",
  events=["MODEL_VERSION_CREATED"],
  http_url_spec=http_url_spec,
  description="Slack notifications",
  status="TEST_MODE"
)
```

```:レスポンス
{"webhook": {
   "id":"1234567890",
   "creation_timestamp":1571440826026,
   "last_updated_timestamp":1582768296651,
   "status":"TEST_MODE",
   "events":["MODEL_VERSION_CREATED"],
   "http_url_spec": {
     "url": "https://hooks.slack.com/services/...",
     "enable_ssl_verification": True
}}}
```

### 2. Webhookのテスト

上述のWebhookは`TEST_MODE`で作成されたので、指定されたURLにリクエストを送信するためにダミーのイベントを起動することができます。しかし、実際のイベントではWebhookは起動されません。テストのエンドポイントは、指定されたURLで受信したステータスコードと本文を返却します。

```bash:Bash
$ curl -X POST -H "Authorization: Bearer <access_token>" -d \
'{"id": "1234567890"}' \
https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/test
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient

http_webhook = RegistryWebhooksClient().test_webhook(
  id="1234567890"
)
```

```:レスポンス
{
 "status":200,
 "body":"OK"
}
```

### 3. Webhookをアクティブに更新

実際のイベントに対してWebhookが動作するようにするには、他のプロパティを更新することもできる更新コールを用いてステータスを`ACTIVE`に設定する必要があります。

```bash:Bash
$ curl -X PATCH -H "Authorization: Bearer <access_token>" -d \
'{"id": "1234567890", "status": "ACTIVE"}' \
https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/update
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient

http_webhook = RegistryWebhooksClient().update_webhook(
  id="1234567890",
  status="ACTIVE"
)
```

```:レスポンス
{"webhook": {
   "id":"1234567890",
   "creation_timestamp":1571440826026,
   "last_updated_timestamp":1582768296651,
   "status": "ACTIVE",
   "events":["MODEL_VERSION_CREATED"],
   "http_url_spec": {
     "url": "https://hooks.slack.com/services/...",
     "enable_ssl_verification": True
}}}
```

### 4. Webhookの削除

Webhookを無効化するには、(上と同じコマンドを用いて)ステータスを`DISABLED`に更新するか削除します。

```bash:Bash
$ curl -X DELETE -H "Authorization: Bearer <access_token>" -d \
'{"id": "1234567890"}' \
https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/delete
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient

http_webhook = RegistryWebhooksClient().delete_webhook(
  id="1234567890"
)
```

```:レスポンス
{}
```

## ジョブによるレジストリWebhookサンプルワークフロー

ジョブレジストリWebhookを管理するワークフローは、HTTPレジストリWebhookと類似しています。唯一の違いは`http_url_spec`フィールドが`job_spec`フィールドに変わるという点です。

Webhookを用いて、同じワークスペースのジョブあるいは、異なるワークスペースのジョブを起動することができます。`workspace_url`パラメーターでワークスペースを指定することができます。`workspace_url`が指定されない場合、デフォルトの挙動では同じワークスペースのジョブを起動します。

### 要件

- 既存の[ジョブ](https://qiita.com/taka_yayoi/items/b3275a1983c51a8bbe1a)
- [パーソナルアクセストークン](https://docs.databricks.com/dev-tools/api/latest/authentication.html)。APIから返却されるWebhookオブジェクトにはトークンは含まれないことに注意してください。

### ジョブレジストリWebhookの作成

```bash:Bash
$ curl -X POST -H "Authorization: Bearer <access_token>" -d \ '{"model_name": "<model-name>",
  "events": ["TRANSITION_REQUEST_CREATED"],
  "description": "Job webhook trigger",
  "status": "TEST_MODE",
  "job_spec": {
    "job_id": "1",
    "workspace_url": "https://my-databricks-workspace.com",
    "access_token": "dapi12345..."}}'
https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/create
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient, JobSpec

job_spec = JobSpec(
  job_id="1",
  workspace_url="https://my-databricks-workspace.com",
  access_token="dapi12345..."
)
job_webhook = RegistryWebhooksClient().create_webhook(
  model_name="<model-name>",
  events=["TRANSITION_REQUEST_CREATED"],
  job_spec=job_spec,
  description="Job webhook trigger",
  status="TEST_MODE"
)
```

```:レスポンス
{"webhook": {
   "id":"1234567891",
   "creation_timestamp":1591440826026,
   "last_updated_timestamp":1591440826026,
   "status":"TEST_MODE",
   "events":["TRANSITION_REQUEST_CREATED"],
   "job_spec": {
     "job_id": "1",
     "workspace_url": "https://my-databricks-workspace.com"
}}}
```

### レジストリWebhookを一覧するサンプル

```bash:Bash
$ curl -X GET -H "Authorization: Bearer <access_token>" -d \ '{"model_name": "<model-name>"}'
https://<databricks-instance>/api/2.0/mlflow/registry-webhooks/list
```

```py:Python
from databricks_registry_webhooks import RegistryWebhooksClient

webhooks_list = RegistryWebhooksClient().list_webhooks(model_name="<model-name>")
```

```:レスポンス
{"webhooks": [{
   "id":"1234567890",
   "creation_timestamp":1571440826026,
   "last_updated_timestamp":1582768296651,
   "status": "ACTIVE",
   "events":["MODEL_VERSION_CREATED"],
   "http_url_spec": {
     "url": "https://hooks.slack.com/services/...",
     "enable_ssl_verification": True
}},
{
   "id":"1234567891",
   "creation_timestamp":1591440826026,
   "last_updated_timestamp":1591440826026,
   "status":"TEST_MODE",
   "events":["TRANSITION_REQUEST_CREATED"],
   "job_spec": {
     "job_id": "1",
     "workspace_url": "https://my-databricks-workspace.com"
}}]}
```

## サンプルノートブック

### MLflowモデルレジストリWebhook REST APIサンプルノートブック

https://docs.databricks.com/_static/notebooks/mlflow/mlflow-model-registry-webhooks-rest-api-example.html

### MLflowモデルレジストリWebhook Pythonクライアントサンプルノートブック

https://docs.databricks.com/_static/notebooks/mlflow/mlflow-model-registry-webhooks-python-client-example.html

### Databricks 無料トライアル

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