Databricksでは、MLflowモデルレジストリを用いてモデルのステージ(None、Staging、Production、Archived)を変更することができます。これによって、モデルの現在の状態を管理することができます。
- None: なんでもない
- Staging: テスト中
- Production: 本番運用中
- Archived: お蔵入り
GUIでは承認依頼を上げることは知っていたのですが、さてPythonからはできたんだっけ？となりましたので調べてみました。MLflow APIにはありませんでしたが、Databricks REST API経由でできました。
準備
DatabricksからREST APIを呼び出すので、こちらを参考にREST APIのエンドポイントURLとトークンを取得します。
databricksURL = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiUrl().getOrElse(None)
myToken = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().getOrElse(None)
モデルステージのリクエスト
エンドポイント/api/2.0/mlflow/transition-requests/createをコールします。
import requests
header = {'Authorization': 'Bearer {}'.format(myToken)}
endpoint = '/api/2.0/mlflow/transition-requests/create'
payload = """{
"name": "wine_quality_takaakiyayoidatabrickscom",
"version": "30",
"stage": "Staging",
"comment": "このバージョンはすばらしいです！"
}
"""
resp = requests.post(
databricksURL + endpoint,
data=payload.encode("utf-8"),
headers=header
)
resp.text
{
"request": {
"creation_timestamp": 1693805509130,
"user_id": "takaaki.yayoi@databricks.com",
"activity_type": "REQUESTED_TRANSITION",
"comment": "このバージョンはすばらしいです！",
"to_stage": "Staging"
}
}
リクエストが作成されました。
モデルステージの承認
エンドポイント/api/2.0/mlflow/transition-requests/approveをコールします。
header = {'Authorization': 'Bearer {}'.format(myToken)}
endpoint = '/api/2.0/mlflow/transition-requests/approve'
payload = """{
"name": "wine_quality_takaakiyayoidatabrickscom",
"version": "30",
"stage": "Staging",
"archive_existing_versions": true,
"comment": "OKです。"
}
"""
resp = requests.post(
databricksURL + endpoint,
data=payload.encode("utf-8"),
headers=header
)
{
"activity": {
"creation_timestamp": 1693805598456,
"user_id": "takaaki.yayoi@databricks.com",
"activity_type": "APPROVED_REQUEST",
"comment": "OKです。",
"last_updated_timestamp": 1693805598456,
"from_stage": "None",
"to_stage": "Staging",
"id": ""
}
}
承認されてステージが変更されました。
これらのやりとりもすべて記録されています。
