Databricksワークスペース間における機械学習モデルの共有

[Share models across workspaces — Databricks Documentation](https://docs.databricks.com/applications/machine-learning/manage-model-lifecycle/multiple-workspaces.html) [2021/11/9時点]の翻訳です。

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

Databricksは複数のワークスペース間でのモデル共有をサポートしています。例えば、あなた自身のワークスペースでモデルを構築し、ロギングしたのち、集中管理されているモデルレジストリにモデルを登録することができます。これは、あなたの組織において、異なる開発ステージに対応するために複数のワークスペースが存在しているケースで、複数のチームがモデルにアクセスする際に有益です。

このようなケースにおいては、Databricksは、それぞれのユーザーがアクセスするためのアカウントが設定された、中央モデルレジストリ専用のワークスペースを作成することをお勧めします。これには、以下のマルチワークスペース環境の図に示すように、モデルをロギングし、登録を行うデータサイエンティスト、モデルの管理、デプロイを行うプロダクションユーザーが含まれます。
![](https://docs.databricks.com/_images/multiworkspace1.png)

中央モデルレジストリへのアクセスはトークンによって制御されます。アクセスが必要なそれぞれのユーザーあるいはスクリプトは、中央モデルレジストリワークスペースにおいて[パーソナルアクセストークンを作成](https://docs.databricks.com/dev-tools/api/latest/authentication.html#token-management)し、彼らが属するワークスペースの[シークレットマネージャーにコピー](https://docs.databricks.com/dev-tools/cli/secrets-cli.html)します。中央モデルレジストリへのAPIリクエストはアクセストークンを伴う必要があります。MLflowはモデルレジストリに対する操作を行う際に、シークレットを指定するメカニズムを提供しています。

モデルレジストリに対する全ての[クライアントAPI](https://www.mlflow.org/docs/latest/python_api/mlflow.tracking.html)、[フルーエントAPI](https://www.mlflow.org/docs/latest/python_api/mlflow.html)はリモートワークスペースをサポートしています。

本記事では以下を説明します：

- [要件](#要件)
- [リモートレジストリに対するAPIトークンの設定](#リモートレジストリに対するapiトークンの設定)
- [リモートレジストリの指定](#リモートレジストリの指定)
- [リモートレジストリへのモデル登録](#リモートレジストリへのモデル登録)
- [リモートレジストリのモデルの利用](#リモートレジストリのモデルの利用)
- [リモートレジストリのモデルの管理](#リモートレジストリのモデルの管理)
- [ノートブック](#ノートブック)

# 要件
ワークスペースにまたがるモデル登録を行うためには、1.11.0リリース以降のMLflow Pythonクライアントが必要となります。

> **注意**
このワークフローはMLflowクライアントのロジックを用いて実装されています。クライアントを実行する環境が、集中管理されたモデルレジストリを有するDatabricksワークスペースへのネットワークリクエストが可能であることを確認してください。レジストリワークスペースに課せられる一般的な制限はIPアクセスリストであり、別のワークスペースのクラスターで稼働するMLflowクライアントからの接続を拒否することができます。

# リモートレジストリに対するAPIトークンの設定

1. 中央モデルレジストリのワークスペースで[アクセストークンを作成](https://docs.databricks.com/dev-tools/api/latest/authentication.html#token-management)します。
2. ローカルのワークスペースで、アクセストークン及びリモートワークスペースの情報を格納するための[シークレットを作成](https://docs.databricks.com/dev-tools/cli/secrets-cli.html)します。
    1. シークレットのスコープを定義します: `databricks secrets create-scope --scope <scope>`
    2. ターゲットとなるワークスペースに対して一意となる名前をつけます。ここでは、`<prefix>`とします。そして、以下のコマンドを実行して3つのシークレットを作成します。
        - `databricks secrets put --scope <scope> --key <prefix>-host`: 実行後、中央モデルレジストリのホスト名を入力します。(例)`https://cust-success.cloud.databricks.com/`
        - `databricks secrets put --scope <scope> --key <prefix>-token`: 中央モデルレジストリのワークスペースで作成したトークンを入力します。
        - `databricks secrets put --scope <scope> --key <prefix>-workspace-id`: 中央モデルレジストリのワークスペースIDを指定します。ワークスペースIDは中央モデルレジストリのワークスペースにアクセスした際に表示される[URLから特定](https://docs.databricks.com/workspace/workspace-details.html#workspace-instance-names-urls-and-ids)できます。

> **注意** [ワークスペースごとにシークレットのスコープ数の上限](https://docs.databricks.com/security/secrets/secret-scopes.html)があるため、他のユーザーと[シークレットスコープを共有](https://docs.databricks.com/security/access-control/secret-acl.html)するケースがあるかもしれません。

# リモートレジストリの指定

中央レジストリワークスペースに対するシークレットスコープ及び名称プレフィクスに基づき以下の形式でレジストリのURIを作成できます：

```python
registry_uri = f'databricks://<scope>:<prefix>'
```

[フルーエントAPI](https://www.mlflow.org/docs/latest/python_api/mlflow.html)呼び出しの最初にURIを指定することでリモートレジストリを指定できます：

```python
mlflow.set_registry_uri(registry_uri)
```

あるいは、`MLflowClient`作成時に明示的に指定することも可能です:

```python
client = MlflowClient(registry_uri=registry_uri)
```

以降のワークフローでは両方のアプローチで説明します。

# リモートレジストリへのモデル登録

モデルを登録する一つ目の方法は、`mlflow.register_model `APIです：

```python
mlflow.set_registry_uri(registry_uri)
mlflow.register_model(model_uri=f'runs:/<run_id>/<artifact_path>', name=model_name)
```

別の登録方法に関しては、記事の最後にあるノートブックで確認できます。

> **注意** リモートワークスペースへのモデル登録の際、リモートワークスペースのDBFSに一時的にモデルアーティファクトのコピーが作成されます。モデルバージョンが**READY**になった後は、このコピーを削除して構いません。一時ファイルは`/dbfs/databricks/mlflow/tmp-external-source/<run_id>`フォルダーにあります。

また、`registry_uri`で指定したのと同じように、`tracking_uri`を指定することで、リモートワークスペースのMLflowトラッキングサービスを指定することも可能です。

# リモートレジストリのモデルの利用

モデルレジストリのURIを指定した後で、`mlflow.<flavor>.load_model`を呼び出すことで、リモートレジストリのモデルバージョンをロードし、利用することができます：

```python
mlflow.set_registry_uri(registry_uri)
model = mlflow.pyfunc.load_model(f'models:/<model_name>/Staging')
model.predict(...)
```

あるいは、URIに`models:/`を用いて、明示的にモデルを指定することもできます：

```python
model = mlflow.pyfunc.load_model(f'models://<scope>:<prefix>@databricks/<model_name>/Staging')
model.predict(...)
```

この他にもモデルファイルにアクセスするためのヘルパーメソッドがサポートされています。例えば：

```python
client.get_latest_versions(model_name)
client.get_model_version_download_uri(model_name, version)
```

# リモートレジストリのモデルの管理

要求されるアクセス権を持っている限り、リモートレジストリに対する任意のアクションが可能です。例えば、モデルに対して**Can Manage**権限を持っているのであれば、`MlflowClient`メソッドを利用して、モデルのバージョンのステータスを変更したり削除したりすることができます：

```python
client = MlflowClient(tracking_uri=None, registry_uri=registry_uri)
client.transition_model_version_stage(model_name, version, 'Archived')
client.delete_registered_model(model_name)
```

# ノートブック

https://github.com/taka-yayoi/public_repo/tree/main/model_sharing

### Databricks 無料トライアル

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