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": ""
    }
}

承認されてステージが変更されました。

これらのやりとりもすべて記録されています。

Databricksクイックスタートガイド

Databricksクイックスタートガイド

Databricks無料トライアル

Databricks無料トライアル