Unity Catalog でモデルがターゲットにされるように ML ワークフローをアップグレードする

  • [アーティクル]

この記事では、ML ワークフローを移行して Unity Catalog のターゲット モデルにアップグレードする方法について説明します。

要件

開始する前に、「要件」に記載されている要件を満たしていることを確認します。 特に、モデル トレーニング、デプロイ、推論のワークフローの実行に使用するユーザーまたはプリンシパルが、Unity Catalog の登録済みモデルに対して必要な権限を持っていることを確認します。

  • トレーニング: 登録済みモデルの所有権 (新しいモデル バージョンを作成するために必要です)、および外側のカタログとスキーマに対する USE CATALOG 権限と USE SCHEMA 権限。
  • デプロイ: 登録済みモデルの所有権 (モデルにエイリアスを設定するために必要です)、および外側のカタログとスキーマに対する USE CATALOG 権限と USE SCHEMA 権限。
  • 推論: 登録済みモデルに対する EXECUTE 権限 (モデル バージョンで推論を読み取り、実行するために必要です)、および外側のカタログとスキーマに対する USE CATALOG 'USE SCHEMA 権限。

並列トレーニング、デプロイ、ワークフローの作成

モデル トレーニングと推論ワークフローを Unity Catalog にアップグレードするために、Databricks では、Unity Catalog でモデルを活用する並列トレーニング、デプロイ、推論パイプラインを作成する段階的なアプローチをお勧めしています。 Unity Catalog を使用した結果に満足している場合は、ダウンストリーム コンシューマーを切り替えてバッチ推論の出力を読み取るか、エンドポイントを提供する際に Unity Catalog のモデルにルーティングされるトラフィックを増やすことができます。

モデル トレーニングのワークフロー

モデル トレーニングのワークフローを複製します。 次に、以下を確認します。

  1. ワークフロー クラスターに Unity Catalog へのアクセス権があり、「要件」で説明されている要件が満たされている。
  2. ワークフローを実行しているプリンシパルに Unity Catalog の登録済みモデルに対する必要なアクセス許可がある。

次に、複製したワークフロー内のモデル トレーニング コードを変更します。 ワークフローによって実行されるノートブックを複製するか、複製したワークフローに新しい Git ブランチを作成してターゲットにする必要がある場合があります。 これらの手順に従って、必要なバージョンの MLflow をインストールし、トレーニング コードで Unity Catalog をターゲットとするようにクライアントを構成し、モデルが Unity Catalog に登録されるようにモデル トレーニング コードを更新します。

モデル デプロイのワークフロー

モデル デプロイのワークフローを複製します。「モデル トレーニングのワークフロー」と同様の手順に従って、Unity Catalog へのアクセスを有効になるようにそのコンピューティング構成を更新します。

複製したワークフローを所有するプリンシパルに必要なアクセス許可があることを確認します。 デプロイのワークフローにモデル検証ロジックがある場合は、UC からモデル バージョンを読み込むようにそれを更新します。 別名を使用して、運用モデルのロールアウトを管理します。

モデル推論のワークフロー

バッチ推論のワークフロー

「モデル トレーニングのワークフロー」と同様の手順に従って、バッチ推論のワークフローを複製し、Unity Catalog へのアクセスを有効になるようにそのコンピューティング構成を更新します。 複製したバッチ推論ジョブを実行しているプリンシパルに、推論用のモデルを読み込むために必要なアクセス許可があることを確認します。

モデル提供のワークフロー

Mosaic AI Model Serving を使用している場合、既存のエンドポイントを複製する必要はありません。 代わりに、トラフィック分割機能を利用して、トラフィックのごく一部を Unity Catalog のモデルにルーティングできます。

まず、モデル提供のエンドポイントを所有するプリンシパルに、推論用のモデルを読み込むために必要なアクセス許可があることを確認します。 次に、トラフィックのごく一部が Unity Catalog のモデル バージョンに割り当てられるように、複製したモデル デプロイのワークフローを更新します。

環境間でモデルを昇格する

Databricks では、ML パイプラインをコードとしてデプロイすることをお勧めしています。 こうすると、運用環境で自動化されたトレーニング ワークフローを通じてあらゆる運用モデルを作成できるため、環境間でモデルを昇格する必要がなくなります。

ただし、場合によっては、環境間でのモデルの再トレーニングのコストがきわめて高くなることがあります。 代わりに、Unity Catalog に登録されているモデル間でモデル バージョンをコピーして、環境間で昇格させることができます。

以下のコード例を実行するには、次の権限が必要です。

  • staging および prod カタログに対する USE CATALOG。
  • staging.ml_team および prod.ml_team スキーマに対する USE SCHEMA。
  • staging.ml_team.fraud_detection に対する EXECUTE。

さらに、登録済みモデル prod.ml_team.fraud_detection の所有者である必要があります。

次のコード スニペットでは、MLflow バージョン 2.8.0 以上で利用できる copy_model_version MLflow Client API を使用しています。

import mlflow
mlflow.set_registry_uri("databricks-uc")

client = mlflow.tracking.MlflowClient()
src_model_name = "staging.ml_team.fraud_detection"
src_model_version = "1"
src_model_uri = f"models:/{src_model_name}/{src_model_version}"
dst_model_name = "prod.ml_team.fraud_detection"
copied_model_version = client.copy_model_version(src_model_uri, dst_model_name)

モデル バージョンが運用環境に昇格された後、必要なデプロイ前の検証を実行できます。 次に、別名を使用して、デプロイ用のモデル バージョンをマークできます。

client = mlflow.tracking.MlflowClient()
client.set_registered_model_alias(name="prod.ml_team.fraud_detection", alias="Champion", version=copied_model_version.version)

上記の例では、登録済みモデル staging.ml_team.fraud_detection から読み取り、登録済みモデル prod.ml_team.fraud_detection に書き込むことのできるユーザーのみが、ステージング モデルを運用環境に昇格させることができます。 また、同一のユーザーが別名を使用して、運用環境内にデプロイされるモデル バージョンを管理することもできます。 モデルの昇格とデプロイを管理するために、他のルールやポリシーを構成する必要はありません。

このフローをカスタマイズして、セットアップに一致する複数の環境 (dev、qa、prod など) 間でモデル バージョンを昇格させることができます。 アクセスの制御は、各環境で構成されているとおりに適用されます。

モデル デプロイの手動承認にジョブ Webhook を使用する

Databricks では、モデル デプロイ プロセス中に適切な検査とテストを使用して、可能であればモデル デプロイを自動化することをお勧めしています。 ただし、運用モデルをデプロイするために手動承認を実行する必要がある場合は、モデル トレーニング ジョブが正常に完了した後で、ジョブ Webhook を使用して外部 CI/CD システムに呼び出し、モデルをデプロイするために手動承認を要求することができます。 手動承認が行われた後、CI/CD システムで (たとえば、“Champion” という別名を設定するなどして) モデル バージョンをデプロイして、トラフィックを提供することができます。