MLflow と SageMaker AI Model Registry の同期によるモデルガバナンス

候補モデルの実験段階から離れる瞬間に生じるギャップを埋めるため、MLflow とモデルレジストリ間の登録自動化が重要です。データサイエンティストは MLflow で数十の候補ランを追跡する一方、ガバナンス担当者は、本番環境へ投入されるモデルを検証・承認・監査するための権威あるレジストリを 1 つ必要としています。

Amazon SageMaker AI 上のマネージド MLflow では、MLflow で登録されたモデルが自動的に SageMaker AI Model Registry に同期されます。この同期機能は現在、大幅に強化されました。トレーニングメトリクスや評価結果、そして系譜情報（lineage）も引き継がれます。さらに、MLflow から駆動されるライフサイクルステージの昇格機能も追加され、データサイエンティストが実験ワークフローから離れることなく、単一の記録システムから候補モデルをガバナンスできるようになりました。

しかし、より豊富な同期機能により、モデルはレジストリに到着した時点でレビュー可能となり、ライフサイクルをスムーズに進めることができます。

MLflow アプリで Model Registry 同期を有効化すると、データサイエンティストが MLflow で登録するすべてのモデルに対応して、SageMaker AI Model Registry に「Model Package Group」と「Package version」が自動的に作成されます。同期には、トレーニングメトリクスや評価メトリクス、デプロイ可能な推論仕様、そして MLflow の実行に関連付けられた系譜情報も含まれます。

データサイエンスチームは引き続き、実験とログ化されたモデルの管理システムとして MLflow を使い続けることができます。一方、自動化機能により、組織は本番ライフサイクルを通過する登録済みモデルの管理システムとして Model Registry を利用できるようになります。

登録が自動化されることで、データサイエンティストは手動でモデルの系譜を管理する必要なく、モデルのトレーニングや実験に集中できます。一方、ガバナンス担当者は SageMaker AI Model Registry を通じて一貫性があり完全な系譜情報を得られるため、モデルライフサイクルの管理が可能になります。その結果、ガバナンス担当者はどのモデルを生産環境へ展開するかを承認・監査・制御できるようになります。

これは 2 部構成シリーズの第 1 弾です。本稿では、自動登録がどのように機能するかを紹介し、単一アカウント内でのセットアップ手順を解説します。ここではデータサイエンティストとガバナンス担当者の役割は、アカウント境界ではなく IAM のガードレールによって分離されています。Part 2 では、この仕組みを大規模で規制の厳しい組織向けの跨アカウントガバナンストポロジーへ拡張します。両記事で使用可能な作業用ノートブックは、GitHub リポジトリ で公開されています。

自動登録の仕組み

Model Registry の同期機能は、オプトイン方式の機能です。MLflow アプリケーションを作成または更新する際、モデル登録モードを AutoModelRegistrationEnabled に設定することで有効化できます（デフォルトでは無効になっています）。この同期を設定する際は、MLflow アプリの AWS Identity and Access Management (IAM) サービスロール を事前に設定しておく必要があります。このロールには、Model Package グループとバージョンの作成、タグの付与、リンケージ（系譜）情報の記録に必要な権限が含まれていることが条件です。

# Activate Model Registry sync when creating an MLflow app
aws sagemaker create-mlflow-app \
    --name my-mlflow-app \
    --artifact-store-uri s3:///mlflow \
    --role-arn arn:aws:iam:::role/my-mlflow-app-role \
    --model-registration-mode AutoModelRegistrationEnabled \
    --region

# Updating an existing MLflow app with Model Registry sync
aws sagemaker update-mlflow-app \
    --arn  \
    --model-registration-mode AutoModelRegistrationEnabled \
    --region

同期が有効になると、データサイエンティストはトレーニング実行中にモデルをログし、MLflow の 1 つの呼び出しで登録処理を完了できます。登録前に、推論仕様と評価メトリクスという 2 つのオプションアーティファクトを付与することが可能です。これらはそのまま Model Package に引き継がれます。

import mlflow, sagemaker_mlflow

# Log the trained model during the run.
model_info = mlflow.sklearn.log_model(model, name="sklearn-model")

# Optionally log evaluation metrics
sagemaker_mlflow.evaluate(model_info, data=dataset, model_type="regressor")

# Optionally log an inference specification:
# enables direct deployment from the registry
sagemaker_mlflow.log_inference_specification(
    model_info.model_id, inference_specification=inference_spec
)

# Register the logged model:
# creates the Model Package Group and version automatically
mv = mlflow.register_model(f"models:/{model_info.model_id}", "my-model")

登録後、同期機能は Amazon SageMaker Studio の Model Registry へ 4 つのカテゴリのメタデータを転送し、ガバナンス担当者が参照可能にします。

• 実行メタデータ — モデルのパラメータ、トレーニングで得られたメトリクス、トレーニングデータの場所、およびモデルアーティファクトのパス。

• 評価メトリクス — Model Package のバージョンにモデルカードとして付与され、Studio の「Evaluate」タブに表示されます。これにより、パフォーマンス測定値をモデルと併せてレビューすることが可能になります。

推論仕様は、コンテナイメージ、モデルデータの場所、およびサポートされるインスタンスタイプを定義し、レジストリから直接デプロイできるようにします。パッケージを実際にデプロイ可能にするためには、モデルアーティファクトに推論ハンドラー（トレーニングや登録時にモデルとパッケージ化された inference.py など）を含める必要があります。推論仕様のコンテナ定義は、SAGEMAKER_PROGRAM 環境変数を通じてこれを参照し、関連するノートブックではトレーニング中にハンドラーをパッケージ化します。

ラインージスは、MLflow の実験、モデルバージョン、コンテナイメージ、およびモデルパッケージグループ間の関係を記録します。

ガバナンス担当者は、Model Registry のステージング構成 を活用してモデルのライフサイクルを管理できます。組織は、事前に用意されたステージング構成テンプレートを利用するか、ユースケースに合わせて独自の構成を定義し、IAM の条件キーと組み合わせることで、適切な権限を持つユーザーのみがライフサイクル全体を通じてモデルのプロモーションを行えるように制御することが可能です。

MLflow アリases に sagemakerlifecycle-{stage}-{status} という命名規則を適用することで、ライフサイクルステージとステータスを設定できます。ここで、stage には staging を指定します。

または、production でステータスが pending または active の場合（使用時）

事前に用意されたテンプレートを使用すると、対応する SageMaker AI Model Package のライフサイクルが自動的に更新されます。

client = mlflow.MlflowClient()
# Move a candidate to staging
client.set_registered_model_alias(
  "my-model",
  "sagemakerlifecycle-staging-pending",
  version
)

モデルライフサイクルの更新を適切なロールに限定するため、2 つのガバナンス制御を適用できます。

1 つ目は IAM 条件キーです。これにより、ライフサイクルの移行を制限できます。sagemaker:ModelLifeCycle/stage および sagemaker:ModelLifeCycle/stageStatus を適切な値で定義することで、誰がモデルのステージを更新できるか、またどのターゲットステージとステータスへ移行可能かを制御できます。

2 つ目はリソースタグ条件です。承認後に Model Package Group をロックし、MLflow 経由でのさらなる更新が流れ込まないようにします。

ライフサイクルの変更は Amazon EventBridge にイベントとして発行され、監査証跡として記録されます。これらのイベントを消費することで、承認プロセスをダウンストリームのサードパーティ製ガバナンスツールに連携させることが可能です。

sagemaker_mlflow に関する注記

本記事で解説するワークフローは、sagemaker-mlflow プラグインを基盤としています。バージョン 0.5.0 では、特筆すべき機能が 2 つ追加されました。1 つ目は推論イメージのログ記録機能です。これにより、モデルを登録する際に、そのモデルに対応する推論コンテナのイメージも同時に記録できるようになります。

この画像メタデータは、モデルとともにモデルレジストリへ転送され、ここで説明されている「レジストリから直接デプロイする」パスが強化されます。また、セッションの注入（use_session() および set_session() を介して）により、カスタムの boto3 セッションを指定できるため、プラグインはコンテキスト固有の認証情報を用いてリクエストに署名できます。

本稿で扱う単一アカウント構成は必須ではありませんが、これは後続の Part 2 で解説する跨アカウントガバナンストポロジーを構築するための基礎となります。

2 つのペルソナ

この仕組みを始めるには、異なるツールを使う 2 つの役割が必要です。データサイエンティストは Jupyter Notebook で作業し、MLflow を使ってモデルの学習、ログ記録、登録、ステージングを行います。一方、ガバナンス担当者は SageMaker Studio の UI、特に「Models」ビューで作業し、同期されたメトリクスや系譜情報を確認したり、本番環境へのプロモーションを承認したりします。プラットフォーム管理者が、これら 2 つの役割を分けるための初期セットアップを一度だけ行います。以降では、それぞれの役割について順を追って解説していきます。

事前準備

開始前に、以下の準備を整えてください：

SageMaker AI リソースの作成と管理に必要な権限を持つ AWS アカウントを用意してください。

Studio アクセスが設定された Amazon SageMaker AI ドメインが必要です。

ドメイン内に作成した Managed MLflow アプリケーションを使用します。セットアップ手順については、MLflow アプリの作成 をご覧ください。

MLflow アプリで Model Registry の同期機能を有効化してください。

実験環境に sagemaker-mlflow プラグイン（バージョン 0.5.0 以降）がインストールされている必要があります。インストールコマンドは以下の通りです：

pip install sagemaker-mlflow>=0.5.0

モデルの登録、ライフサイクルステージの更新、Model Registry への書き込み権限を持つ IAM ロールが必要です。本記事で説明するガバナンス制御を行うには、IAM 条件キーやリソースタグポリシーを付与する権限も別途必要となります。

単一アカウントでのセットアップ

単一アカウント構成では、データサイエンティストとガバナンス担当者が同じ AWS アカウント内で作業を行います。これは、まだアカウントの分離が必要でない小規模チームや初期段階のプロジェクトに適しています。この構成では、開発が行われる同一アカウント内で、ガバナンス担当者がモデルの検証と承認を担当します。

以下の図は、モデル登録からガバナンスによるプロモーションを経てデプロイに至るまでの、単一 AWS アカウント内でのエンドツーエンドのワークフローを示しています。

図 1 – 単一アカウント構成

上記の図は、以下のワークフローを示しています： （原文の技術表記: pip install sagemaker-mlflow>=0.5.0）

モデルのトレーニングが完了したら、データサイエンティストは MLflow アプリにモデルをログ記録し、登録します。この MLflow アプリでは「Model Registry sync」機能が有効になっています。

自動登録により、モデルパッケージグループとバージョンが SageMaker AI Model Registry に追加されます。これにはメトリクス、評価結果、推論仕様、そして系譜情報（lineage）も含まれます。

IAM の条件キーによるゲートが設定されているため、データサイエンティストはモデルを「staging」ステータスに設定できます。ただし、そのロールでは本番環境へのプロモーションは拒否されます。

代わりに、本番環境へのプロモーションを許可されたガバナンスオフィサーのロールが、モデルの承認と本番環境へのプロモーションを行うことができます。

機械学習（ML）エンジニアは、承認されたモデルを SageMaker AI エンドポイントにデプロイします。通常は、継続的インテグレーションおよび継続的デリバリー（CI/CD）パイプラインを通じて行われます。

プラットフォーム管理者：ガードレールの設定

管理者は MLflow アプリで Model Registry sync を有効化し（前節を参照）、MLflow アプリのサービスロールに登録と系譜情報の権限を付与します。さらに、ライフサイクルのガードレール（IAM ポリシー）をデータサイエンティストのロールにアタッチします。このガードレールは、sagemaker:ModelLifeCycle/stage 条件キーを使用して、本番環境へのプロモーションを拒否します。

{
  "Version": "2012-10-17",
  "Statement": [{
    "Sid": "DenyProductionPromotion",
    "Effect": "Deny",
    "Action": "sagemaker:UpdateModelPackage",
    "Resource": "*",
    "Condition": {
      "StringEquals": {"sagemaker:ModelLifeCycle/stage": "production"}
    }
  }]
}

ガバナンスオフィサーのロールにはこの拒否設定が含まれていないため、本番環境へのプロモーションが可能なのはその役割を持つ人物だけです。これは一度きりのセットアップ作業です。完了後、2 つの役割はそれぞれ独立して業務を遂行します。

データサイエンティスト：ノートブックからの実験と登録

データサイエンティストは、Jupyter ノートブック（SageMaker Studio またはローカル環境）で作業し、MLflow とのみ対話します。トレーニング完了後、評価指標と推論仕様をログ記録し、モデルを登録してライフサイクルエイリアスを使ってステージングへ移行します。

mv = mlflow.register_model(f"runs:/{run_id}/sklearn-model", "my-model")
client = mlflow.MlflowClient()
client.set_registered_model_alias("my-model", "sagemakerlifecycle-staging-pending", mv.version)

自動登録により、Model Registry 内に Model Package Group と Model Package バージョンが作成されます。エイリアスの設定でステージングまたは保留状態に移行しますが、データサイエンティストが生産用エイリアスを指定しようとすると、ガードレールによって移行は拒否されます。

以下に、登録されたモデル実行の MLflow UI における Artifacts タブのスクリーンショットを示します。

図 2 – 推論仕様を含む MLflow アーティファクト

Artifacts パネルには、モデルのアーティファクトが格納されたディレクトリが表示されます。ここには inference.py（推論ハンドラ）や input_example.json も含まれています。

モデルディレクトリの下には、2 つの追加アーティファクトが強調表示されています。sagemaker_inference_specification.json と serving_input_example.json です。右側のペインには、推論仕様の内容が表示されます。

コンテナイメージのURI、Amazon S3 上のモデルアーティファクトを指す ModelDataUrl、およびサポートされるリアルタイム推論用インスタンスタイプを定義します。

また、環境変数も設定されます。SAGEMAKER_PROGRAM には inference.py が指定され、SAGEMAKER_SUBMIT_DIRECTORY にはモデルコードの場所が指し示されます。この仕様が、レジストリから直接モデルパッケージをデプロイ可能にする理由です。

ガバナンス担当者：Studio UI でレビューと承認を行う

ガバナンス担当者は、SageMaker Studio のModelsビューで作業を行います。同期された Model Package バージョンには、トレーニングメトリクス、評価モデルカード、および系譜グラフ（lineage graph）が既に含まれているため、レビューは一つの場所で完結します。

以下のスクリーンショットでは、トレーニングと評価の過程で計算されたモデルメトリクスと、その系譜情報を示しています。

図 3 – Model Package バージョンに同期されたトレーニングメトリクスを表示する SageMaker Studio

図 4 – SageMaker Studio の「Evaluate」タブに表示される評価メトリクス

図 5 – モデルバージョンの系譜（リンネージ）

メトリクスと系譜の確認が完了したら、ガバナンス担当者はモデルを本番環境へプロモートし、承認ステータスを Approved に設定します。承認ステータスはライフサイクルステージとは別の属性であるため、プロモーションの決定とデプロイの準備は別々の判断事項として扱われます。

以下の 2 つのスクリーンショットは、担当者の操作前後におけるモデルバージョンの状態を示しています。

図 6 – ガバナンス担当者が本番環境へプロモートする前の、ステージング中の Model Package バージョン