MLflow と SageMaker AI Model Registry の複数アカウント連携

自動モデル登録が整った後、アカウント間でのモデルガバナンスは自然な次のステップとなります。パート1では、マネージドMLflowがAmazon SageMaker AIに登録されたモデルをSageMaker AI Model Registryに同期させる仕組みを紹介しました。ここでは、AWS Identity and Access Management (IAM) の条件キーを用いてデータサイエンティストとガバナンス担当者の役割を分離する単一アカウントのセットアップについて解説しました。

しかし大規模な組織では、開発と本番環境をアカウントレベルで完全に分離します。複数の開発用アカウントを運用し、中央集権的なガバナンス機能を担うチームが存在するのが一般的です。規制の厳しい業界ではさらに厳格な要件が課されることがあり、開発ワークロードが生産グレードのアカウントに直接書き込むことを禁止しています。

本稿では、前編で解説した構成要素を、2 つの異なるアカウント間ガバナンストポロジーに拡張して適用します。

1 つ目は「ハブ・アンド・スポーク」パターンです。AWS Resource Access Manager (AWS RAM) を活用し、開発用アカウント間で MLflow アプリケーションを共有することで、ガバナンス機能を一元化します。

2 つ目は規制の厳しい環境向けのハイブリッドパターンです。この構成では、開発用アカウントとガバナンスハブを完全に分離したまま運用します。

最後に、承認されたモデルが CI/CD (Continuous Integration and Continuous Delivery) を経由して登録からデプロイ済みエンドポイントへ移行するプロセスを示し、両トポロジーの比較を通じて最適な選択ができるよう支援します。実習用のノートブックは、GitHub リポジトリ で公開しています。

登場人物について

前編では、*データサイエンティスト*（Jupyter Notebook を使用し MLflow とのみやり取りする）と *ガバナンス担当者*（Amazon SageMaker Studio の Models UI を使用する）の 2 つの役割を紹介しました。アカウント間トポロジーを導入すると、さらに 2 つの役割が加わります。

• 管理者 (Administrator): AWS RAM の共有設定やバケットポリシー、宛先グループの設定など、一度きりの跨アカウントセットアップをコンソールまたは CLI で実行します。
• モデルオーナー (Model Owner): ハイブリッドトポロジーにおいて、各開発用アカウント内でモデルの承認を行い、その後ハブへのプロモーションを許可する役割です。

以下に紹介する各トポロジーは、これらの役割に基づくワークフローで構成されています。

事前準備

このチュートリアルを開始する前に、以下の準備を確認してください。

• パート 1 の完了: パート 1 で扱った自動モデル登録、ライフサイクルステージングの仕組み、IAM 条件キーによるゲート制御について理解している必要があります。詳しくは パート 1 をご覧ください。

• AWS アカウントの用意: 開発用（スポーク）アカウントとガバナンス用（ハブ）アカウント、計 2 つの AWS アカウントが必要です。それぞれのアカウントに対して、mlops-spoke や mlops-hub のように名前を付けた AWS CLI プロファイルを設定してください。セットアップの詳細については、関連するリポジトリの「Account and profile setup (REPO_LINK#account-and-profile-setup)」セクションをご参照ください。

両方のアカウントにデプロイされる AWS CloudFormation スack に対応しています。テンプレート（cfn/sagemaker-studio-mlflow.yaml (REPO_LINK)）は、SageMaker AI Studio のドメイン、ユーザープロファイル、実行ロール、そして AutoModelRegistrationEnabled を有効にしたマネージド MLflow アプリをセットアップします。アカウントごとにスタックを 1 つずつデプロイしてください。正確なコマンドについては、リポジトリの README にある「環境のプロビジョニング」 (REPO_LINK#provision-the-environment-cloudformation) を参照してください。

CloudFormation のデプロイ、スタック出力の確認、環境変数の設定を含む完全なプロビジョニングの手順については、対応するリポジトリ (REPO_LINK) をご覧ください。

トポロジー 1: ハブ・アンド・スポーク型中央ガバナンス

複数のチームを抱える大規模組織では、開発アカウントを複数用意し、ガバナンスを一元管理するアカウントを設ける構成が一般的です。このトポロジーにより、各開発アカウント内でワークロードの分離を保ちつつ、MLflow アプリと中央モデルレジストリをホストする一元化されたアカウントを通じて共有サービスを提供できます。

一元化されたハブアカウントは、AWS RAM を用いて MLflow アプリを 1 つ以上のスプーク（開発）アカウントと共有します。このパターンは、SageMaker Model Registry と AWS RAM の共有によるモデルガバナンスの一元化で説明されたアプローチを拡張したものです。

以下の図にワークフローを示します：

• ハブが MLflow アプリ向けの AWS RAM リソース共有を作成し、スプーク側が招待を受け入れます。外部プリンシパルを使用することで、アカウントが同じ組織内にない場合でも共有機能を利用できます。

• スプーク内のデータサイエンティストが、共有された MLflow アプリに対してモデルを登録します。自動登録により、register コールと同時にハブアカウントに Model Package Group とバージョンが同期して作成されます。

• ハブ側でグループにリソースポリシーを付与し、AllowDeploy マネージド権限を使ってスプークへ共有します。これにより、スプーク側はモデルの記述やデプロイが可能になります。

• ガバナンス担当者が同期されたメトリクスと系譜を検証し、ハブ内で中央集権的にモデルを承認します。

承認イベントが発生すると CI/CD パイプラインが起動し、機械学習（ML）エンジニアが承認されたモデルをスポークアカウント内の Amazon SageMaker エンドポイントにデプロイします。

図 1: 集中型ガバナンスを持つハブ・アンド・スポークトポロジー

管理者：クロスアカウントアクセスの設定（初回のみ）

管理者は、いずれの役割も作業を開始する前に、ハブと共有環境の準備を完了させておく必要があります。前提条件として、モデルアーティファクトはハブの MLflow アーティファクトストア、つまりハブアカウント内の S3 バケットに存在している必要があります。クロスアカウントによる S3 アクセス権限は、両方のアカウント側で設定しておくことが必須です。

• リソース側: ハブのアーティファクトバケットに対して、スプークアカウントが s3:GetObject、s3:PutObject、s3:ListBucket、および s3:GetBucketLocation を実行できるようにするバケットポリシーを設定します。

これがないと、登録処理が行われる前にデータサイエンティストの log_model コマンドが S3 アクセスエラーで失敗してしまいます。

アイデンティティ側では、スプークの実行ロールがデプロイ時にアーティファクトを読み取るために、ハブバケットに対して s3:GetObject 権限が必要です。AmazonSageMakerFullAccess ポリシーは、名前に sagemaker を含むバケットをカバーしています。もしスプーク側のロールの権限制御がより厳格な場合は、ハブバケットを明示的に追加してください。

前提条件が整った上で、管理者は次の手順を実行します：

ハブの MLflow アプリに対して、スプーク側が受け入れるための AWS RAM リソース共有を作成します。

初回の登録後、同期されたモデルパッケージグループにリソースポリシーを付与し、Managed Permission の AllowDeploy を使用してスプーク（支店）へ RAM 共有を行います。

ハブの AWS RAM コンソールでは、両方の共有が「アクティブ」状態として表示されます。

図 2：AWS RAM コンソールのハブで「アクティブ」として表示される MLflow アプリのリソース共有

図 3：AWS RAM コンソールのハブで「アクティブ」として表示されるモデルパッケージグループのリソース共有

データサイエンティスト：ノートブックからスプーク（子）側で登録する

このワークフローはパート 1 と同様ですが、1 つだけ異なる点があります。それは、ノートブックが MLflow の追跡 URI をハブアプリの Amazon Resource Name (ARN) に設定していることです。自動登録により、モデルパッケージグループとバージョンが、登録呼び出しと同時にハブ内に作成されます。

候補は、スプーク側のランから取得したメトリクス、評価カード、および系譜情報とともに、ハブレジストリに承認待ちとして表示されます（以下の図参照）：

図 4：メトリクスと系譜情報を備えた、承認待ちの同期済み候補がハブレジストリに表示される様子

計画する際に注意すべき点が 2 つあります。まず、自動登録ではグループ名に短いハッシュ接尾辞が付加されます。例えば my-model は my-model- のようになります。このグループは list_model_package_groups(NameContains=...) コマンドで検索し、名前ではなく完全な ARN を使用して他アカウントから参照してください。

次に、系譜情報はハブに記録されます。同期処理がハブ側で実行されるためです。ガバナンス担当者はハブの Studio で完全な系譜グラフを確認できますが、AWS RAM の共有機能によって系譜エンティティがスプーク側に自動で戻ってくることはありません。他アカウント間での系譜情報の共有は自動的に実現されないのです。

ガバナンス担当者：ハブ側の Studio UI 上で中央集権的に承認する

ガバナンス担当者は、ハブの「Studio Models」ビューで作業を行います。これはすべてのスプークから候補が一元管理される制御ポイントです。同期されたメトリクスとラインージを確認し、本番環境へのプロモーションを実行して承認ステータスを Approved に設定します。

以下の図は、SageMaker AI Model Registry に登録されたバージョンの現在のライフサイクルステータスを示しています。

Figure 5: Lifecycle status of the registered version in the SageMaker AI Model Registry

承認後の処理については、「承認からデプロイまで」のセクションで解説します。

Topology 2: Hub-and-spoke hybrid governance

規制の厳しい組織の中には、ハブを本番グレードのアカウントとして扱い、データサイエンティストが直接・間接を問わず書き込みを行わないようにするケースがあります。こうした顧客には、各開発アカウントを独立した環境に保つハイブリッドトポロジーをお勧めします。

以下の図はワークフローの流れを示しています。

• データサイエンティストが、開発アカウント固有の MLflow アプリに対してモデルを登録します。
• 自動同期により、そのモデルが開発アカウント内の Model Registry（ローカル）に登録されます。ハブ側には一切触れません。
• 開発アカウントのモデルオーナーが、開発レジストリ上でモデルを検証・承認します。これによりコピーワークフローがトリガーされます。
• ワークフローは、承認されたモデルをクロスアカウントでハブへ転送します。具体的には、モデルアーティファクトをハブ所有のバケットに複製し、推論仕様の記述を書き換えた上で、ハブの共有宛先グループに登録パッケージとして登録します。

ハブのガバナンス担当者は、コピーされたモデルを再検証し、中央のガバナンス記録として承認します。

ML エンジニアは、ローカルで承認されたモデルを CI/CD パイプラインを通じてスポークのエンドポイントにデプロイします。この際、アカウント間でのアーティファクト転送は不要です。

図 6: ハイブリッドガバナンスを持つハブ・アンド・スポークトポロジー

管理者：ハブの宛先グループを公開する（1 回限り）

ハブ管理者は、宛先のモデルパッケージグループを作成し、開発アカウントがその中に CreateModelPackage を実行できるリソースポリシーを付与します。また、RAM 経由で AllowRegister という管理権限付きで共有設定を行います。

さらに管理者は、コピーワークフローによってモデルアーティファクトが書き込まれる「ハブ所有のアーティファクトバケット」を作成（または指定）します。

バケットポリシーにより、開発アカウントに対して s3:PutObject および s3:ListBucket の権限が付与されています。これにより、コピー処理は開発アカウントのロールの下で実行可能です。なお、本記事で使用している AWS 管理フレームワークのイメージはリージョンごとに公開されているため、コンテナイメージの複製は不要です。

開発アカウントのプライベートリポジトリにコンテナイメージが格納されている場合は、そのイメージをハブへ複製するか、Amazon Elastic Container Registry (Amazon ECR) のリポジトリポリシーを追加してください。

データサイエンティスト：開発アカウント内でノートブックから登録する

開発アカウント内だけで完結する、パート 1 と同じ開発者体験です。データサイエンティストが開発環境の MLflow アプリでモデルを登録すると、その登録されたモデルは自動的に開発用の Model Registry に同期されます。ハブ側には一切触れません。

開発アカウントのモデル所有者：ローカルで承認する

モデル所有者は開発レジストリ内の候補モデルを確認し、承認します。このローカルでの承認が、ハブへのプロモーションをトリガーするきっかけとなります（以下の図参照）。

Figure 7: The development account model owner approving a model version locally

承認によるコピーと、ハブ側でのガバナンス担当者の再検証

承認が行われると、ワークフローが起動します。これは、本番環境における Model Package の状態変更を検知する Amazon EventBridge ルールによって実行されます。このプロセスでは、承認されたパッケージをハブへコピーします。

コピー処理では、モデルアーティファクトをハブ所有のバケットに複製し、推論仕様に書き換えてハブ側のコピーを指すように設定します。さらに、ソース側へのポインタを含む形でハブ側にパッケージを再作成します。その結果、ハブ側のパッケージは完全な自己完結型となり、開発アカウントへのランタイム依存関係は一切不要となります。

コピー処理では、3 つのアクションが実行されます。まず開発用バケットからモデルアーティファクトを HUB 所有のバケットへ複製し、推論仕様を書き換えて HUB 内のコピーを指すようにします。最後に、CustomerMetadataProperties にソースパッケージの ARN とアカウント情報を記録しながら、HUB の宛先グループにパッケージを登録します。これにより、HUB 側のパッケージは開発用アカウントへのランタイム依存を持たない完全な自己完結型となります。実際の実装についてはリポジトリをご参照ください。

ガバナンス担当者は、HUB の Studio で複製されたパッケージを確認し、ソースへ遡るプロベナンスメタデータに基づいて再検証を行った上で、独立して承認を行います。ガバナンス担当者が確認できるメタデータの例を以下の図に示します。

Figure 8: The copied package in the hub with provenance metadata pointing to the source

承認からデプロイへ

どちらのトポロジーにおいても、レジストリにおける承認がゲート機能を持ちます。承認ステータスが変更されると Amazon EventBridge にイベントが発行され、これを消費して CI/CD デプロイパイプラインをトリガーします。具体的には、ModelApprovalStatus: Approved を条件にマッチする Amazon EventBridge ルールが起動し、承認された Model Package からエンドポイントの作成または更新を行います。

パイプラインの実行場所とエンドポイントの配置場所は、組織的な判断に委ねられます。一般的な構成では、パイプラインを共有サービスアカウントで実行し、専用デプロイメントアカウントへ展開します。よりシンプルな設定では、開発（スプーク）アカウントの両方で動作させることも可能ですが、これは付随するサンプルで示されている通りです。

いずれのアカウントからデプロイを行う場合でも、モデルアーティファクトへの読み取りアクセス権限が必要です。中央集約型トポロジーの場合、これはハブ側のアーティファクトバケット（直前の管理者必須要件を参照）へのアクセスを意味します。一方、ハイブリッド型トポロジーでは、ハブのコピーがすでに自己完結型の状態となっています。

別のデプロイメントアカウントを導入する際の考慮点として、モデルの系譜情報は同期が記録されたアカウント内に残存し続けます。ガバナンスレビューは、ハブ（または開発）アカウントを軸に計画する必要があります。