ワークスペース間でモデルを共有する

重要

Databricks では、 Unity Catalog内のモデル を使用して、ワークスペース間でモデルを共有することをお勧めします。この記事のアプローチは非推奨です。

Databricks では、複数のワークスペース間でのモデルの共有がサポートされています。 たとえば、開発ワークスペースでモデルを開発してログに記録し、別の運用ワークスペースのモデルにアクセスして比較できます。 これは、複数のチームがモデルへのアクセスを共有する場合や、組織に開発のさまざまな段階を処理するための複数のワークスペースがある場合に便利です。 クロスワークスペース モデルの開発とデプロイの場合、Databricks では、モデル トレーニング コードを複数の環境にデプロイする コードのデプロイ アプローチをお勧めします。

マルチワークスペースの状況では、リモート モデルレジストリを使用して、Databricks ワークスペース間でモデルにアクセスできます。 たとえば、データサイエンティストは読み取り専用アクセスで本番モデルレジストリにアクセスし、開発中のモデルを現在の本番モデルと比較できます。 マルチワークスペースのセットアップ例を以下に示します。

複数のワークスペース

リモート レジストリへのアクセスはトークンによって制御されます。 アクセスが必要な各ユーザーまたはスクリプトは、リモート レジストリに個人用のアクセス トークンを作成しそのトークンをローカル ワークスペースのシークレット マネージャーにコピーします。 リモート レジストリ ワークスペースに送信される各 API リクエストには、アクセス権が含まれている必要があります。 MLflow は、モデル レジストリ操作を実行するときに使用するシークレットを指定する簡単なメカニズムを提供します。

自動化されたツール、システム、スクリプト、アプリを使用して認証する場合のセキュリティのベスト プラクティスとして、Databricks ではOAuth トークンを使用することをお勧めします。

個人のアクセス トークン認証を使用する場合、 Databricksでは、ワークスペース ユーザーではなく、サービスプリンシパルに属する個人のアクセス トークンを使用することをお勧めします。 サービスプリンシパルのトークンを作成するには、 「サービスプリンシパルのトークンの管理」を参照してください。

モデルレジストリのすべての クライアント および 流暢 な API メソッドは、リモート ワークスペースでサポートされています。

要件

ワークスペース間でモデルレジストリを使用するには、MLflow Python クライアント リリース 1.11.0 以降が必要です。

このワークフローは、MLflow クライアントのロジックから実装されます。 クライアントを実行している環境に、リモート モデルレジストリを含む Databricks ワークスペースに対してネットワーク要求を行うためのアクセス権があることを確認します。 レジストリ ワークスペースに適用される一般的な制限は IP 許可リストであり、別のワークスペースのクラスターで実行されている MLflow クライアントからの接続を禁止できます。

リモート・レジストリーの API トークンをセットアップする

  1. モデルレジストリ ワークスペースで、 アクセストークンを作成します

  2. ローカル ワークスペースで、アクセス許可とリモート ワークスペース情報を保存するためのシークレットを作成します。

    1. シークレットスコープを作成します: databricks secrets create-scope <scope>.

    2. ターゲット ワークスペースの一意の名前を選択します (ここでは<prefix>として示されています)。 次に、次の 3 つのシークレットを作成します。

      • databricks secrets put-secret <scope> <prefix>-host : モデルレジストリ ワークスペースのホスト名を入力します。 たとえば、 https://cust-success.cloud.databricks.com/.

      • databricks secrets put-secret <scope> <prefix>-token : モデルレジストリ ワークスペースからアクセストークンを入力します。

      • databricks secrets put-secret <scope> <prefix>-workspace-id : 任意のページの URL にある モデルレジストリ ワークスペースのワークスペース ID を入力します。

リモート レジストリを指定する

リモート レジストリ ワークスペース用に作成したシークレットスコープと名前プレフィックスに基づいて、次の形式のレジストリ URI を作成できます。

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

URI を使用して、最初に以下を呼び出すことで、 流暢な API メソッド のリモート レジストリを指定できます。

mlflow.set_registry_uri(registry_uri)

または、 MlflowClientをインスタンス化するときに明示的に指定することもできます。

client = MlflowClient(registry_uri=registry_uri)

次のワークフローは、両方のアプローチの例を示しています。

リモート レジストリにモデルを登録する

モデルを登録する 1 つの方法は、 mlflow.register_model API を使用することです。

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と同様の方法で、別のワークスペースの MLflow 追跡サービスを指す tracking_uri を指定することもできます。つまり、リモート ワークスペースで実行し、現在のリモート ワークスペースまたは別のリモート ワークスペースでそのモデルを登録できます。

リモートレジストリのモデルを使用する

最初にレジストリ URI を設定することで、 mlflow.<flavor>.load_model メソッドを使用してリモートレジストリにモデルバージョンを読み込んで使用できます。

mlflow.set_registry_uri(registry_uri)
model = mlflow.pyfunc.load_model(f'models:/<model-name>/Staging')
model.predict(...)

または、 models:/ URI でリモート レジストリを明示的に指定できます。

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)

サンプルノートブック: リモートモデルレジストリ

次のノートブックは、 Unity Catalogが有効になっていないワークスペースに適用されます。 現在のワークスペースから MLflow 追跡サーバーに記録済みモデルを記録し、別のワークスペースのモデルレジストリにモデルを登録する方法を示します。 Databricks では、 Unity Catalog内のモデル を使用して、ワークスペース間でモデルを共有することをお勧めします。

リモートモデルレジストリ サンプルノートブック

ノートブックを新しいタブで開く