ワークスペース間でモデルを共有する
重要
Databricks では、Unity Catalog のモデルを使用して、ワークスペース間でモデルを共有することが推奨されています。 この記事へのアプローチは非推奨となっています。
Azure Databricks では、複数のワークスペース間でのモデルの共有がサポートされています。 たとえば、開発ワークスペースでモデルを開発してログに記録してからそれにアクセスし、別の運用ワークスペース内のモデルと比較できます。 これは、複数のチームがモデルへのアクセスを共有するときや、組織に開発の異なる段階を処理する複数のワークスペースがあるときに便利です。 ワークスペース間のモデルの開発とデプロイの場合、Databricks では、モデル トレーニング コードが複数の環境にデプロイされるデプロイ コード アプローチが推奨されます。
複数ワークスペースの状況では、リモート モデル レジストリを使用して Azure Databricks ワークスペース全体でモデルにアクセスできます。 たとえば、データ サイエンティストは、読み取り専用アクセス権で運用モデル レジストリにアクセスして、開発中のモデルと現在の運用モデルを比較できます。 マルチワークスペースのセットアップの例を下に示します。
リモート レジストリへのアクセスは、トークンによって制御されます。 アクセスを必要とする各ユーザーまたはスクリプトは、リモート レジストリに個人用アクセス トークンを作成し、そのトークンをローカル ワークスペースのシークレット マネージャーにコピーします。 リモート レジストリ ワークスペースに送信される各 API 要求には、アクセス トークンが含まれている必要があります。MLflow には、モデル レジストリの操作を実行するときに使用するシークレットを指定するための単純なメカニズムが用意されています。
注意
セキュリティのベスト プラクティスとして、自動化ツール、システム、スクリプト、アプリを使用して認証する場合、Databricks では、ワークスペース ユーザーではなくサービス プリンシパルに属する個人用アクセス トークンを使用することを推奨しています。 サービス プリンシパルのトークンを作成するには、「サービス プリンシパルのトークンを管理する」をご覧ください。
リモート ワークスペースでは、モデル レジストリのすべてのクライアントおよび fluent API メソッドがサポートされています。
必要条件
ワークスペース間でモデル レジストリを使用するには、MLflow Python クライアント リリース 1.11.0 以降が必要です。
注意
このワークフローは、MLflow クライアントのロジックから実装されます。 クライアントを実行している環境に、リモート モデル レジストリが含まれる Azure Databricks ワークスペースに対してネットワーク要求を行うためのアクセス許可があることを確認します。 レジストリ ワークスペースに対する一般的な制限は IP 許可一覧であり、別のワークスペースのクラスターで実行されている MLflow クライアントからの接続を禁止することができます。
リモート レジストリの API トークンを設定する
- モデル レジストリ ワークスペースで、アクセス トークンを作成します。
- ローカル ワークスペースで、アクセス トークンとリモート ワークスペースの情報を格納するシークレットを作成します。
- シークレットのスコープ
databricks secrets create-scope <scope>
を作成します。 <prefix>
に示すように、ターゲット ワークスペースの一意の名前を選択します。 次に、次の 3 つのシークレットを作成します。databricks secrets put-secret <scope> <prefix>-host
: モデル レジストリ ワークスペースのホスト名を入力します。 たとえば、https://westus.azuredatabricks.net/
またはhttps://adb-5555555555555555.19.azuredatabricks.net/
です。databricks secrets put-secret <scope> <prefix>-token
: モデル レジストリ ワークスペースのアクセス トークンを入力します。databricks secrets put-secret <scope> <prefix>-workspace-id
: モデル レジストリ ワークスペースのワークスペース ID を入力します。これは、任意のページのURL 内にあります。
- シークレットのスコープ
注意
ワークスペースごとのシークレット スコープの数に制限があるため、シークレット スコープを他のユーザーと共有することができます。
リモート レジストリを指定する
リモート レジストリ ワークスペース用に作成したシークレットのスコープと名前のプレフィックスに基づいて、次のような形式のレジストリの URI を作成できます。
registry_uri = f'databricks://<scope>:<prefix>'
最初に以下を呼び出すことで、URI を使用して fluent 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)
リモート レジストリでモデルを管理する
必要なアクセス許可がある限り、リモート レジストリ内のモデルに対して任意のアクションを実行できます。 たとえば、モデルに対して管理可能アクセス許可がある場合は、モデル バージョン ステージを切り替えたり、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 追跡サーバーにモデルをロギングし、モデルを別のワークスペースの Model Registry に登録する方法を示します。 Databricks では、Unity Catalog のモデルを使用して、ワークスペース間でモデルを共有することが推奨されています。