Databricks: MLflow Model Registryを別のWorkspaceに移す (legacy)

注意：今はUnity Catalogがあるため、もっと簡単手軽に共有移動が可能です。

目的

Model Registry に登録されているモデルを別の Workspace に移します。
モチベーションとしては、以下のようなケースを想定しています：

• 非常に優秀なモデルがあるが、再現性がなく、学習コードも紛失してしまっているため再生成が困難
• 学習コストや時間が膨大で、同じモデルを再度作るのが現実的ではない

このような場合に、既存モデルをそのまま別 Workspace に移す方法を紹介します。

方法は何種類かあります。

そもそもMLflow自身が mlflow.register_model の引数に取れるのは、run であったり、modelのpathもいけます。
正直 pathの方は only dbfs locations are currently supported.というエラーが回避できなかったので、迂回策をとっています。

結論

runを一回生成して、そのrunに対してModelを生成する。また、その間に MLflow registryを切り替える。

まずDatabrikcsは異なるワークスペース間でモデルを共有できます。(legacy)
https://docs.databricks.com/aws/en/machine-learning/manage-model-lifecycle/multiple-workspaces
コードは２行です。

registry_uri = f'databricks://<scope>:<prefix>'
mlflow.set_registry_uri(registry_uri)

コード自体はシンプルですが、これを実行する前に共有するモデルが置かれているWorkspaceに紐づく三つのSecretを命名規則に従って用意する必要があります。

• <scope> <prefix>-workspace-id
• <scope> <prefix>-token
• <scope> <prefix>-host
  これでモデルは共有可能です。
  そしてこれを利用して、今のWorkspaceに対してデータをコピーします。

# develop環境にMLflowを接続する
mlflow.set_registry_uri("databricks://mlflow-env:develop")

# modelを変数modelに代入する
logged_model = f"models:/{params.model_name}/Production"
model = mlflow.pytorch.load_model(logged_model)

# 現環境にMLflowを接続する
mlflow.set_registry_uri("databricks")

# experimentに対して runを作ってその上でそのrunに対して log_modelで modelを登録
with mlflow.start_run(run_name=model_name, experiment_id="<experiment_id>") as run:
    # model typeやmodelの格納先のdirectoryは自由に
    mlflow.pytorch.log_model(model, "model")

# 上記で走らせたrunを利用して modelを再度格納
model_details = mlflow.register_model(f"runs:/{run.info.run_id}/model", model_name)

# productionから取得したので、ついでにProductionに昇格
client.transition_model_version_stage(
    name=model_name, version=model_details.version, stage="Production"
)

まとめ

• モデルの再生成が難しい場合、既存モデルを別 Workspace にコピーする方法が有効
• run を使ってモデルを log し、再度 Registry に登録することで移動可能
• Registry URI の切り替えがポイント

いいねいただけると励みになります。お読みいただきありがとうございました。