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

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

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

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

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

必要条件

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

リモート レジストリの API トークンを設定する

1. モデル レジストリ ワークスペースで、アクセス トークンを作成します。
2. ローカル ワークスペースで、アクセス トークンとリモート ワークスペースの情報を格納するシークレットを作成します。
   1. シークレットのスコープ databricks secrets create-scope <scope> を作成します。
   2. <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)

その他のモデルの登録方法の例については、このページの最後にあるノートブックを参照してください。

また、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 のモデルを使用して、ワークスペース間でモデルを共有することが推奨されています。

リモート モデル レジストリの例のノートブック

ノートブックを入手