SageMaker AI MLflowを活用したモデル管理とサーバーレス推論の導入

記事タイトルとURLをコピーする

こんにちは。クロスインダストリー第1本部 松田です。

今回は、SageMaker AI の MLflow 機能を活用して、モデルの実験管理からサーバーレス推論エンドポイントへのデプロイまでを一気通貫で行う方法を検証しました。

AWS re:Invent 2025 にて、MLflow を搭載した SageMaker AI にサーバーレス機能の追加が発表されました。このブログでは、インフラストラクチャ管理不要の MLflow App の作成から SageMaker Pipelines との統合まで紹介されています。

本ブログでは、この構成をさらに掘り下げ、モデルの比較実験からサーバーレス推論エンドポイントでの推論実行までの一連の流れと、構築時にハマったポイントを紹介します。

SageMaker AI MLflow とは?

MLflow は、機械学習の実験管理やモデルレジストリとして広く利用されている OSS です。ただし、通常は MLflow サーバーを自前で構築・運用する必要があります。

SageMaker AI では、この MLflow が「MLflow App」としてマネージドに提供されています。サーバーの運用が不要になるだけでなく、SageMaker の推論基盤とシームレスに連携できる点が魅力です。さらに、サーバーレス推論と組み合わせることで、リクエストがないときにはコストが発生しない、検証用途に適した構成を実現できます。

全体の流れ

今回の検証は、以下の7つのステップで構成されています。

  1. MLflow App の作成 - SageMaker AI 上にマネージド MLflow 環境を構築
  2. モデル比較実験 - 合成データで3種の分類器を学習し、MLflow に実験結果を記録
  3. モデルレジストリへの登録 - 最良モデルを MLflow Model Registry に登録
  4. カスタムコンテナイメージのビルド - MLflow 推論用コンテナを ECR にプッシュ
  5. SageMaker モデルを作成 - ModelBuilder で SageMaker モデルを作成
  6. サーバーレス推論エンドポイントのデプロイ - boto3 で直接エンドポイントを作成
  7. 推論の実行と結果検証 - エンドポイントとローカルモデルの予測結果を比較

20260323204051

図中の番号(①〜⑥)は、後続セクションの各ステップに対応しています。

各ステップは Python スクリプトとして実装しています。

MLflow App によるモデル管理

モデル比較実験 (アーキテクチャー図の①)

まず、SageMaker AI 上に MLflow App を作成し、モデルの比較実験を行います。MLflow App の作成は AWS CLI で実行でき、S3 バケット(アーティファクトストア用)と IAM ロールを事前に準備しておけば、数分で利用可能になります。作成した MLflow App は以下にようになります。

20260323204147

今回は make_classification で生成した合成データ(20特徴量、5000サンプル)を使い、以下の3つの分類器を比較しました。

モデル 概要
DecisionTree 決定木。シンプルなベースライン
LogisticRegression ロジスティック回帰。線形モデルの代表
GradientBoosting 勾配ブースティング。アンサンブル手法

各モデルの学習・評価結果は、MLflow の start_run 内で記録しています。記録する内容は、ハイパーパラメータ、評価メトリクス(Accuracy、Precision、Recall、F1 Score、ROC AUC)、ROC曲線の画像、モデルアーティファクトです。

# ※ 以下はサンプル実装です。モデルやメトリクスは用途に応じて変更してください。
with mlflow.start_run(run_name=run_name):
    for key, value in params.items():
        mlflow.log_param(key, value)

    model.fit(X_train, y_train)
    mlflow.log_metric("f1_score", f1_score(y_test, y_pred))
    mlflow.log_artifact(roc_filename)

    # ★ log_model() でモデルをアーティファクトストアに保存する。
    #    ここで保存されたモデルが、後続のデプロイで参照される。
    mlflow.sklearn.log_model(model, "model")

MLflow UI では、3つの Run を並べて比較することで、どのモデルが最も優れているかを視覚的に確認できます。

20260323204236

実験結果は MLflow を通じて S3 バケットに保管されます。保管場所は、 MlflowClient.get_model_version() で確認することができます。

client = mlflow.tracking.MlflowClient()
mv = client.get_model_version("iris-classifier", "8")
print("Source:", mv.source)   # models:/m-70df31a022784dfba7cd3aeb2d2d3020
print("Run ID:", mv.run_id)   # b85918c706d0484dbfd57e1274ee524c

Source はモデルアーティファクトの参照先、Run ID は登録元の実験ランを示しています。

モデルの実体は S3 上の s3://<artifact-bucket>/<experiment_id>/<run_id>/artifacts/model/ に格納されており、Model Registry はこの S3 パスへの参照を管理しているだけで、登録時にデータのコピーは発生しません。

モデルレジストリへの登録 (アーキテクチャー図の②)

比較実験の結果から、F1 Score が最も高い Run を検索し、MLflow Model Registry に登録します。

# ※ 検索条件やモデル名はサンプルです。
runs = mlflow.search_runs(
    experiment_names=["model-comparison-experiment"],
    order_by=["metrics.f1_score DESC"],
    max_results=1,
)

# ★ register_model() を呼ぶだけで、モデル名の新規作成とバージョン採番が自動で行われる。
#    登録後は models:/<モデル名>/<バージョン> の URI でデプロイ時に参照できる。
mlflow.register_model(model_uri=f"runs:/{best_run_id}/model", name="iris-classifier")

mlflow.register_model() を呼ぶだけで、モデル名が存在しない場合は自動的に新規作成され、バージョン番号も自動採番されます。 登録されたモデルは models:/<モデル名>/<バージョン> という URI で参照でき、後続のデプロイスクリプトからこの URI を使ってモデルを取得します。

20260323204330

サーバーレス推論エンドポイントの構築

方法1(失敗): ModelBuilder が自動選択するプリビルドの推論コンテナを利用

MLflow モデルを SageMaker 推論エンドポイントにデプロイするには、推論コンテナが必要です。最初に検討したのは、SageMaker が提供する sklearn プリビルトコンテナを使う方法でした。ModelBuilderimage_uri を指定しなければ、SDK が自動的にプリビルトコンテナを選択してくれます。

ModelBuilder を使用して Amazon SageMaker AI でモデルを作成する https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/how-it-works-modelbuilder-creation.html

しかし、私のケースではプリビルトコンテナは Python 3.9 ベースであり、ローカル環境(Python 3.14)との間で依存ライブラリのバージョン互換性の問題が発生しました。 どうやら最も新しいバージョンであっても Python 3.10 となるようで、互換性の問題は避けられない状況でした。

Amazon SageMaker AI で Scikit-learn を使用するためのリソース https://docs.aws.amazon.com/ja_jp/sagemaker/latest/dg/sklearn.html

この問題に対して、まず ModelBuilderdependencies パラメータで依存ライブラリのバージョンを明示的に上書きする方法を試みました。

model_builder = ModelBuilder(
    mode=Mode.SAGEMAKER_ENDPOINT,
    schema_builder=SchemaBuilder(sample_input=SAMPLE_INPUT, sample_output=SAMPLE_OUTPUT),
    role_arn=IAM_ROLE_ARN,
    model_metadata={
        "MLFLOW_MODEL_PATH": mlflow_model_path,
        "MLFLOW_TRACKING_ARN": MLFLOW_APP_ARN,
    },
    # ★ auto=True(デフォルト)の場合、ModelBuilder は build() 時にローカル Python 環境の
    #    パッケージをスキャンして requirements.txt を自動生成する(ベストエフォート)。
    #    auto=False でこの自動収集を無効化し、custom で互換バージョンを明示指定する。
    #    優先順位は custom > requirements > auto で、同一パッケージは custom が最優先される。
    dependencies={
        "auto": False,
        "custom": [
            "mlflow>=2.17,<3.2",       # Python 3.9 互換の最新
            "jsonschema>=4.18,<4.26",  # Python 3.9 互換
        ],
    },
)

https://aws.amazon.com/jp/blogs/machine-learning/package-and-deploy-classical-ml-and-llms-easily-with-amazon-sagemaker-part-1-pysdk-improvements/

上記 AWS ブログにあるように、ModelBuilderdependencies パラメータは autorequirementscustom の3つのキーを持ちます。ここでは auto=False で自動収集を無効化しましたが、この方法では根本的な問題は解決しませんでした。

調査したところ ModelBuilder はローカルの Python バージョンを LOCAL_PYTHON 環境変数に記録しますが、コンテナは Python 3.9 の sklearn プリビルドイメージを選択しました。そして依存関係インストールする際に、Python 3.10+ が必要なパッケージ(mlflow, jsonschema)がインストールできずに失敗します。

そこで今回は、MLflow が提供する mlflow sagemaker build-and-push-container コマンドでカスタムコンテナをビルドする方式を採用しました。 ローカル環境と同じ Python バージョン・依存ライブラリをコンテナに含めることで、学習環境と推論環境の一致が保証されます。依存ライブラリのバージョン管理も不要になります。

方法2(成功): カスタムの推論コンテナを ECR にプッシュし、ModelBuilder からイメージ URI を指定 (アーキテクチャー図の③)

ECR へのプッシュは以下ステップで行います。

20260323204432

まず、MLflow CLI でベースイメージをローカルにビルドし、それをカスタム Dockerfile で拡張して ECR にプッシュします。

# Step 1: MLflow ベースイメージをローカルにビルド(mlflow-pyfunc というイメージ名で作成される)
mlflow sagemaker build-and-push-container --build --no-push

# Step 2: カスタム Dockerfile でビルド
docker build -t <your-image-name> -f docker/Dockerfile .

# Step 3-4: ECR にプッシュ
aws ecr get-login-password --region <region> | docker login --username AWS --password-stdin <account>.dkr.ecr.<region>.amazonaws.com
docker tag <your-image-name>:latest <account>.dkr.ecr.<region>.amazonaws.com/<your-repo>:latest
docker push <account>.dkr.ecr.<region>.amazonaws.com/<your-repo>:latest

Step 2 で利用した Dockerfile は以下のようになっています。

# mlflow sagemaker build-and-push-container --build がローカルに作成するイメージ名
FROM mlflow-pyfunc

# ダミーの /opt/activate を作成
RUN echo '#!/bin/bash' > /opt/activate && chmod +x /opt/activate

この /opt/activate の追加は重要なポイントです。 詳細は後述で解説します。

SageMaker Model の作成 ― ModelBuilder を活用 (アーキテクチャー図の④)

boto3 の create_model() で直接 SageMaker Model を作成することも可能ですが、ModelBuilder の利用を推奨します。ModelBuilder は MLflow のモデル管理と SageMaker の推論基盤をシームレスに繋ぐ役割を果たしており、具体的には以下のようなメリットがあります。

  • MLflow モデル URI の自動解決: models:/<モデル名>/<バージョン> という MLflow の URI を model_metadata に渡すだけで、SDK がモデルアーティファクトの S3 パスを自動解決してくれる。boto3 で直接 create_model() を呼ぶ場合は、S3 パスへの変換処理を自前で実装する必要がある
  • 推論コンテナの環境変数の自動設定: MLFLOW_TRACKING_ARNMLFLOW_MODEL_PATH など、MLflow 推論コンテナが必要とする環境変数を model_metadata から適切に設定してくれる。手動で設定する場合、環境変数名の誤りや設定漏れによるデプロイ失敗が起きやすい
  • SchemaBuilder による入出力スキーマの定義: サンプルの入出力データを渡すだけで、推論リクエスト・レスポンスのスキーマを自動生成してくれる。エンドポイントの入出力仕様がコード上で明示的に管理される
# ※ image_uri や model_metadata の値はサンプルです。自身の環境に合わせて変更してください。
model_builder = ModelBuilder(
    mode=Mode.SAGEMAKER_ENDPOINT,
    schema_builder=SchemaBuilder(sample_input=SAMPLE_INPUT, sample_output=SAMPLE_OUTPUT),
    role_arn=IAM_ROLE_ARN,
    image_uri=image_uri,  # ECR のカスタムコンテナ
    # ★ model_metadata に MLflow の URI と ARN を渡すだけで、
    #    SDK が MLflow Registry からモデルアーティファクトの S3 パスを自動解決し、
    #    推論コンテナに必要な環境変数も自動設定してくれる。
    model_metadata={
        "MLFLOW_MODEL_PATH": "models:/<モデル名>/<バージョン>",
        "MLFLOW_TRACKING_ARN": MLFLOW_APP_ARN,
    },
    # ★ カスタムコンテナでは仮想環境の再構築をスキップする必要がある。
    #    詳細は「ハマりポイント」を参照。
    env_vars={"MLFLOW_DISABLE_ENV_CREATION": "true"},
)
model_builder.build()

build() による SageMaker Model の作成までを ModelBuilder に任せることで、MLflow と SageMaker の接続部分を SDK に委譲できます。これにより、デプロイスクリプトの実装をシンプルに保つことができています。

補足: MLFLOW_DISABLE_ENV_CREATION の内部動作と /opt/activate

SageMaker が起動する MLflow 推論コンテナは、起動時に以下の処理を行います。

  1. ENTRYPOINT が _serve_pyfunc() を呼び出す
  2. _install_pyfunc_deps() が pyenv で仮想環境を構築し、/opt/activate を生成する
  3. source /opt/activate で仮想環境を有効化し、推論サーバーを起動する

しかし、mlflow sagemaker build-and-push-container でビルドしたカスタムコンテナには pyenv がインストールされていないため、ステップ2で仮想環境の構築に失敗します。

この問題に対して、MLFLOW_DISABLE_ENV_CREATION=true を設定することでステップ2をスキップし、コンテナにプリインストールされた依存ライブラリをそのまま使用できます。サーバーレス推論のような環境では起動時の依存関係の再構築は不要であり、コールドスタートの短縮にもつながります。

ただし、ステップ2をスキップしても、ステップ3の source /opt/activate はスキップされません。/opt/activate が生成されないまま source が実行されるため、"No such file or directory" エラーでコンテナが起動に失敗します。 これを回避するため、ダミーの /opt/activate をあらかじめコンテナに配置しています。

サーバーレス推論エンドポイント の作成 (アーキテクチャー図の⑤)

現状 ModelBuilder.deploy() は MLflow モデルに対してサーバーレス推論をサポートしていません(sagemaker SDK v3 の制限)。そこで、build() で SageMaker Model を作成した後、boto3 で直接 EndpointConfig と Endpoint を作成しています。

# ※ メモリサイズや同時実行数はサンプル値です。ユースケースに応じて調整してください。
sm_client.create_endpoint_config(
    EndpointConfigName=endpoint_config_name,
    ProductionVariants=[{
        "VariantName": "AllTraffic",
        "ModelName": sagemaker_model_name,
        # ★ ServerlessConfig を指定することでサーバーレス推論エンドポイントになる。
        #    ModelBuilder.deploy() ではこの指定ができないため、boto3 で直接作成する。
        "ServerlessConfig": {
            "MemorySizeInMB": 2048,
            "MaxConcurrency": 5,
        },
    }],
)
sm_client.create_endpoint(
    EndpointName=ENDPOINT_NAME,
    EndpointConfigName=endpoint_config_name,
)

デプロイには数分かかり、エンドポイントのステータスが InService になるまでポーリングで待機します。

無事に、エンドポイントがサーバーレスタイプで起動しました。

20260323204534

推論の実行と結果検証 (アーキテクチャー図の⑥)

デプロイしたエンドポイントに推論リクエストを送信します。MLflow カスタムコンテナは MLflow Serving Format(dataframe_split 形式)を受け付けます。これは MLflow の推論サーバー仕様 で定義されている入力形式で、pandas DataFrame を JSON にシリアライズする標準的な方法です。

# ※ カラム名やデータ値はサンプルです。モデルの入力スキーマに合わせてください。
response = runtime_client.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="application/json",
    # ★ MLflow カスタムコンテナは dataframe_split 形式を受け付ける。
    #    columns でカラム名、data で入力データを指定する。
    Body=json.dumps({
        "dataframe_split": {
            "columns": ["f0", "f1", ..., "f19"],
            "data": [[0.5, -1.2, 0.8, ...]],
        }
    }),
)

検証では、学習時と同じパラメータで再生成したテストデータ(1000件)を使い、以下の2つの観点で結果を確認しました。

  • 正解ラベルとの比較: エンドポイントの予測結果と正解ラベルの評価メトリクスを算出
  • ローカルモデルとの比較: MLflow からローカルにロードしたモデルの予測結果と、エンドポイントの予測結果が完全に一致することを確認

ローカルモデルとエンドポイントの予測結果が一致したことで、モデルが正しくデプロイされていることを確認できました。

まとめ

今回は、SageMaker AI の MLflow App を活用して、モデルの比較実験からサーバーレス推論エンドポイントへのデプロイまでの一連の流れを検証しました。

  • MLflow App はマネージドで提供されるため、MLflow サーバーの運用が不要
  • MLflow の実験トラッキングとモデルレジストリにより、モデルのバージョン管理が容易
  • サーバーレス推論はリクエストがないときにコストが発生しないため、検証用途に適している
  • プリビルトコンテナの Python バージョン互換性の問題は、カスタムコンテナ方式で学習環境と推論環境を一致させることで解決できる
  • MLFLOW_DISABLE_ENV_CREATION=true 設定時の /opt/activate 問題や SDK のサーバーレス推論非対応など、MLflow コンテナの内部動作を理解した上での回避策が必要だった

サーバーレス推論はコールドスタート(数十秒〜数分の遅延)がある点には留意が必要ですが、検証・開発フェーズやリクエスト頻度が低いユースケースではコスト効率の良い選択肢になると思います。 今後は、リアルタイム推論エンドポイントとのコスト比較や、モデルの自動再学習パイプラインとの連携についても検証していきたいと思います。

参考