Share models across workspaces — Databricks Documentation [2021/11/9時点]の翻訳です。
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
Databricksは複数のワークスペース間でのモデル共有をサポートしています。例えば、あなた自身のワークスペースでモデルを構築し、ロギングしたのち、集中管理されているモデルレジストリにモデルを登録することができます。これは、あなたの組織において、異なる開発ステージに対応するために複数のワークスペースが存在しているケースで、複数のチームがモデルにアクセスする際に有益です。
このようなケースにおいては、Databricksは、それぞれのユーザーがアクセスするためのアカウントが設定された、中央モデルレジストリ専用のワークスペースを作成することをお勧めします。これには、以下のマルチワークスペース環境の図に示すように、モデルをロギングし、登録を行うデータサイエンティスト、モデルの管理、デプロイを行うプロダクションユーザーが含まれます。
中央モデルレジストリへのアクセスはトークンによって制御されます。アクセスが必要なそれぞれのユーザーあるいはスクリプトは、中央モデルレジストリワークスペースにおいてパーソナルアクセストークンを作成し、彼らが属するワークスペースのシークレットマネージャーにコピーします。中央モデルレジストリへのAPIリクエストはアクセストークンを伴う必要があります。MLflowはモデルレジストリに対する操作を行う際に、シークレットを指定するメカニズムを提供しています。
モデルレジストリに対する全てのクライアントAPI、フルーエントAPIはリモートワークスペースをサポートしています。
本記事では以下を説明します：
要件
ワークスペースにまたがるモデル登録を行うためには、1.11.0リリース以降のMLflow Pythonクライアントが必要となります。
注意
このワークフローはMLflowクライアントのロジックを用いて実装されています。クライアントを実行する環境が、集中管理されたモデルレジストリを有するDatabricksワークスペースへのネットワークリクエストが可能であることを確認してください。レジストリワークスペースに課せられる一般的な制限はIPアクセスリストであり、別のワークスペースのクラスターで稼働するMLflowクライアントからの接続を拒否することができます。
リモートレジストリに対するAPIトークンの設定
- 中央モデルレジストリのワークスペースでアクセストークンを作成します。
- ローカルのワークスペースで、アクセストークン及びリモートワークスペースの情報を格納するためのシークレットを作成します。
- シークレットのスコープを定義します:
databricks secrets create-scope --scope <scope> - ターゲットとなるワークスペースに対して一意となる名前をつけます。ここでは、
<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から特定できます。
-
- シークレットのスコープを定義します:
注意 ワークスペースごとにシークレットのスコープ数の上限があるため、他のユーザーとシークレットスコープを共有するケースがあるかもしれません。
リモートレジストリの指定
中央レジストリワークスペースに対するシークレットスコープ及び名称プレフィクスに基づき以下の形式でレジストリのURIを作成できます：
registry_uri = f'databricks://<scope>:<prefix>'
フルーエントAPI呼び出しの最初にURIを指定することでリモートレジストリを指定できます：
mlflow.set_registry_uri(registry_uri)
あるいは、MLflowClient作成時に明示的に指定することも可能です:
client = MlflowClient(registry_uri=registry_uri)
以降のワークフローでは両方のアプローチで説明します。
リモートレジストリへのモデル登録
モデルを登録する一つ目の方法は、mlflow.register_modelAPIです：
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を呼び出すことで、リモートレジストリのモデルバージョンをロードし、利用することができます：
mlflow.set_registry_uri(registry_uri)
model = mlflow.pyfunc.load_model(f'models:/<model_name>/Staging')
model.predict(...)
あるいは、URIにmodels:/を用いて、明示的にモデルを指定することもできます：
model = mlflow.pyfunc.load_model(f'models://<scope>:<prefix>@databricks/<model_name>/Staging')
model.predict(...)
この他にもモデルファイルにアクセスするためのヘルパーメソッドがサポートされています。例えば：
client.get_latest_versions(model_name)
client.get_model_version_download_uri(model_name, version)
リモートレジストリのモデルの管理
要求されるアクセス権を持っている限り、リモートレジストリに対する任意のアクションが可能です。例えば、モデルに対してCan Manage権限を持っているのであれば、MlflowClientメソッドを利用して、モデルのバージョンのステータスを変更したり削除したりすることができます：
client = MlflowClient(tracking_uri=None, registry_uri=registry_uri)
client.transition_model_version_stage(model_name, version, 'Archived')
client.delete_registered_model(model_name)