Amazon SageMaker AI を用いた MLflow へのストリーミングベンチマークおよび推薦結果の送信

生成AIモデルのベンチマークを行うチームは、本番環境への展開前に、数十種類のGPUインスタンスタイプやサービングコンテナ、並列化戦略、スペキュレーティブ・デコーディングなどの最適化手法を評価します。その結果、実務者は設定に関する判断に数週間を費やし、試したことや成功した点、その理由を手作業で整理せざるを得ない状況に陥りがちです。

この複雑さこそが、Amazon SageMaker AI における「最適化された生成AI推論の推奨機能」を導入した理由でもあります。手探りの試行錯誤から、データ駆動型のガイド付き最適化とベンチマークへとチームを移行させるためです。

本日、MLflowとの統合機能を追加しました。これにより、チームはAIベンチマークの結果や推奨事項を単一の場所にストリーミングし、すべての実験を追跡できるようになります。この統合によってデータのサイロが解消され、反復サイクルが加速するとともに、推論最適化ワークフローにおける完全な再現性が実現します。

本記事では、新しいMLflow統合機能を、Amazon SageMaker AI の「最適化された推論推奨ジョブ」や「ベンチマーク・ジョブ」と組み合わせて使用する方法を解説します。これにより、実験データを自動的に統一された追跡インターフェースへストリーミングすることが可能になります。

この統合機能により、メトリクス、パラメータ、チャートがサーバーレスの Amazon SageMaker MLflow App にリアルタイムでストリーミングされ、統一された実験追跡環境を利用できます。

ソリューションの概要

今回のリリースにより、最適化された推論推奨ジョブやベンチマークジョブを実行すると、Amazon SageMaker AI が自動的にその結果を指定した SageMaker MLflow アプリにストリーミングします。複数のジョブを同じ MLflow 実験に送信すれば、手動でデータを整形する必要なく、MLflow の実験ビューからそれらを選択して並列比較することが可能です。

以下の図は、ベンチマークおよび推奨ジョブと MLflow を連携させる設定方法を示しています：

既存の推奨ジョブやベンチマークジョブに MLflow を連携させるには、以下の手順を実行します。

• MLflow アプリの作成: AWS アカウントで Amazon SageMaker Studio を開き、MLflow に移動してCreate MLflow Appを選択してください。
• 権限の付与: 実行ロールに対して、MLflow アプリの ARN に sagemaker-mlflow:* の権限を追加します。
• MlflowConfig の指定: ベンチマークジョブまたは推奨ジョブを作成する際に、MlflowConfig を渡してください。

手動でのデータ統合を不要に: ネイティブの MLflow 連携により、複数のジョブから得られたベンチマーク結果や推奨結果は、同じ実験名の下で自動的に集約されます。これにより、別々の実行間でメトリクス、ログ、設定を手動で収集する手間がなくなります。その結果、インスタンスタイプ、モデル構成、バッチサイズ、あるいはスペキュレーティブ・ディコーディング（speculative decoding）がパフォーマンスにどう影響するかを、並列して比較検討することが可能になります。例えば、ml.g4dn.12xlarge 上で動作する qwen2-0.5b と、ml.p4d.24xlarge 上での同モデルを直接比較できます。

長時間実行されるジョブのリアルタイム監視: ベンチマークや推奨ジョブには数時間かかることも珍しくありません。完了を待つのみで進捗が見えない状況ではなく、メトリクスが MLflow UI にライブストリーミングされるのを今すぐ確認できます。各構成のテストが進むにつれてレイテンシやスループットなどのメトリクスが更新され、開始から終了までのプロセス全体を追跡可能になります。メトリクスセクションには生データが公開されており、ジョブ終了時にまとめて発表されるのではなく、時間経過とともに逐次更新されます。これらのメトリクスを用いて、推奨ジョブで想定通りのスループットが出ているか確認し、もし期待値に達していなければ早期に停止させることも可能です。

完全な監査証跡の維持: 実験の実行履歴には、ジョブのパラメータ、タイムスタンプ、チェックポイント時点でのメトリクス、生成されたアセットなど、文脈をすべて記録します。この監査証跡は数ヶ月にわたり照会可能で再現性も保たれます。どの構成変更が改善をもたらしたかを追跡したり、パフォーマンスの頭打ち（プラトー）を特定したり、過去の意思決定を完全な文脈のもとで見直したりすることが可能です。

より良いコラボレーションとガバナンスの実現：共有された MLflow の実験は、最適化活動における唯一の信頼できる情報源となります。チームメンバーは、何を試したか、何が成功し、何が失敗したかをレビューできるため、重複作業を減らし、シフトやチーム間の引き継ぎをより確実なものにできます。

技術的な解説：SageMaker AI のベンチマークと推奨事項を MLflow で追跡する

この例では、顧客が Qwen/Qwen2-0.5B-Instruct を Amazon SageMaker AI のリアルタイムエンドポイントにデプロイしています。これは ml.g6.12xlarge インスタンス上で OpenAI 互換のサービングスタックを使用しています。チームは以下の 2 つの質問への回答を求めています。

1. 既存の Qwen エンドポイントは、軽量なチャットボットのワークロードに対してどのように動作するか？
2. より長いプロンプトとレスポンスを扱う大規模なワークロードには、どのような設定を使用すべきか？

評価結果の再現性を確保するため、ベンチマークジョブと推奨ジョブの両方で、メトリクス、パラメータ、チャートが同じ SageMaker MLflow 実験に書き込まれます。ベンチマークジョブは既存の SageMaker エンドポイントを評価し、推奨ジョブは Amazon Simple Storage Service (Amazon S3) に保存された Qwen モデルアーティファクトに対するデプロイオプションを評価して、ランク付けされた設定を返します。

ここでは、Qwen/Qwen2-0.5B-Instruct を ml.g6.12xlarge インスタンス上で使用し、レイテンシの最適化を実現するための MLflow 設定方法を、ウォークスルーノートブックを通じてご紹介します。

エンドツーエンドのウォークスルーには、エンドポイントの準備状況、モデルサイズ、ワークロード設定、インスタンスの利用可能性、推薦ジョブの検索空間によって45〜120 分ほどかかります。

前提条件

この ノートブック では、ベンチマークと推薦ジョブの実験追跡のために MLflow アプリをセットアップする方法を解説しています。このノートブックを実行して始めるには、以下の手順を完了する必要があります。

• Amazon SageMaker AI 上で開発環境として Amazon SageMaker Studio を設定します。
• SageMaker Unified Studio（SageMaker Unified Studio → MLflow → Create MLflow App）から SageMaker MLflow アプリ を作成します。
• ベンチマークジョブ用に、OpenAI 互換モデルを実行しているSageMaker エンドポイント（すでに InService 状態）が必要です。以下の図に示すように、Amazon SageMaker JumpStart Studio UI から Qwen 2 0.5B モデルをデプロイして始められます。
• 推薦ジョブ用の Amazon S3 に保存されたモデルアーティファクトを用意します。
• AmazonSageMakerFullAccess と MLflow およびエンドポイントへの呼び出しアクセス権を持つインラインポリシーを含む、AWS Identity and Access Management (IAM) の実行ロールが必要です。IAM ロールの作成手順については、同梱のノートブックを参照してください。

重要な考慮事項

この統合機能は SageMaker MLflow Apps に対応していますが、自己ホスト型の MLflow 追跡サーバーへ結果をストリーミングするものではありません。

MLflow のネストされた実行（nested run）サポートを利用するには、tooling.version を 0.8.0 以降に設定してください。

推薦ジョブは内部で独自のエンドポイントをプロビジョニングするため、検索空間を制限する特別な理由がない限り ComputeSpec.InstanceTypes を指定する必要はありません。実行ロールには、SageMaker、MLflow、Amazon S3、およびエンドポイント呼び出しに必要な権限を含める必要があります。また、S3 出力バケットはジョブと同じリージョンに配置する必要があります。

AIRecommendationJob を含む追加の例も掲載された完全なエンドツーエンドのノートブックについては、ステップバイステップチュートリアルをご覧ください。

ステップ 1: 環境の設定

まず、環境変数を設定し、SageMaker クライアントを初期化します。

import json
import time
import uuid

import boto3
import botocore.config

# Update these values for your environment
REGION = "us-west-2"
ACCOUNT_ID = ""
MLFLOW_APP_ARN = "arn:aws:sagemaker:us-west-2::mlflow-app/"
ENDPOINT_NAME = "my-model-endpoint"

sm = boto3.client(
    "sagemaker",
    region_name=REGION,
    config=botocore.config.Config(
        retries={"mode": "standard", "max_attempts": 15}
    ),
)

# The S3 output bucket must be in the same Region as the job
S3_OUTPUT_BUCKET = f"mlflow-sagemaker-{REGION}-{ACCOUNT_ID}"

ステップ 2: 実行ロールの取得

Amazon SageMaker AI は、ベンチマークジョブおよび推薦ジョブを実行するために実行ロールを使用します。このロールには、SageMaker へのアクセス権限、エンドポイントの呼び出し権限、ジョブ出力を Amazon S3 に書き込む権限、および SageMaker MLflow App への追跡情報の公開権限が必要です。IAM ロールの作成については、ノートブック を参照してください。

ロールには以下の権限を含める必要があります：

ROLE_NAME = "tutorial-exec-role"

iam = boto3.client("iam")
role = iam.get_role(RoleName=ROLE_NAME)["Role"]
ROLE_ARN = role["Arn"]

print(f"Role ARN: {ROLE_ARN}")

SageMaker がロールを仮定できるように、実行ロールの信頼ポリシーを設定してください。環境によっては、ロールの作成方法や引き渡し方に応じて、アカウントのルートプリンシパルを含める必要がある場合もあります。

ステップ 3: MLflow アプリの確認と実験名の定義

次に、SageMaker の MLflow アプリが利用可能か確認します。その後、ベンチマークジョブとレコメンデーションジョブで共有する実験名と、それぞれ固有のラン名を定義してください。

両方のジョブで同じ MlflowExperimentName を使用することで、結果を一元管理して比較しやすくなります。

# Verify that the MLflow App is available
app = sm.describe_mlflow_app(Arn=MLFLOW_APP_ARN)
MLFLOW_APP_URL = app["Url"]

assert app["Status"] in ("Created", "Updated")

print(f"MLflow URL: {MLFLOW_APP_URL}")

# Define a shared experiment and unique run names
MLFLOW_EXPERIMENT = "gemma-experiments"
SESSION_UID = uuid.uuid4().hex[:6]

BENCH_RUN_NAME = f"benchmark-{SESSION_UID}"
REC_RUN_NAME = f"recommendation-{SESSION_UID}"

ステップ 4: ベンチマークジョブの提出

SageMaker AI のベンチマークジョブは、既存のエンドポイントをワークロード仕様に従って評価するものです。今回の例では、NVIDIA AIPerf を使用し、OpenAI と互換性のある API スキーマを採用しています。

OutputConfig 内の MlflowConfig が重要な設定項目です。これにより、SageMaker はベンチマークのメトリクス、パラメータ、およびチャートを SageMaker MLflow App に直接ストリーミングできるようになります。

# Tokenizer MUST match the deployed model.
# tooling.version=0.8.0 is the minimum for MLflow nested run support.
TOKENIZER = "Qwen/Qwen2-0.5B-Instruct"
bench_workload = {
    "benchmark": {"type": "aiperf"},
    "parameters": {
        "model": TOKENIZER,
        "concurrency": 1,
        "request_count": 3,
        "streaming": True,
        "prompt_input_tokens_mean": 32,
        "prompt_input_tokens_stddev": 8,
        "output_tokens_mean": 16,
        "output_tokens_stddev": 4,
        "random_seed": 42,
    },
    "tooling": {
        "api_standard": "openai",
        "version": "0.8.0",
    },
}

BENCH_CONFIG = f"mlflow-bench-cfg-{SESSION_UID}"
BENCH_JOB = f"mlflow-bench-job-{SESSION_UID}"

sm.create_ai_workload_config(
    AIWorkloadConfigName=BENCH_CONFIG,
    AIWorkloadConfigs={
        "WorkloadSpec": {
            "Inline": json.dumps(bench_workload)
        }
    },
)

sm.create_ai_benchmark_job(
    AIBenchmarkJobName=BENCH_JOB,
    AIWorkloadConfigIdentifier=BENCH_CONFIG,
    RoleArn=ROLE_ARN,
    BenchmarkTarget={
        "Endpoint": {
            "Identifier": ENDPOINT_NAME
        }
    },
    OutputConfig={
        "S3OutputLocation": f"s3://{S3_OUTPUT_BUCKET}/mlflow-tutorial/{BENCH_JOB}/",
        "MlflowConfig": {
            "MlflowResourceArn": MLFLOW_APP_ARN,
            "MlflowExperimentName": MLFLOW_EXPERIMENT,
            "MlflowRunName": BENCH_RUN_NAME,
        },
    },
)

print(f"Benchmark job submitted: {BENCH_JOB}")

ステップ 5: レコメンデーションジョブの実行

SageMaker AI のレコメンデーションジョブは、デプロイ構成を評価し、ワークロードとパフォーマンス目標に基づいて推奨オプションを提示します。ベンチマークジョブとは異なり、レコメンデーションジョブでは評価プロセス中にエンドポイントが内部で自動的にプロビジョニングされます。

以下の例では、ベンチマークジョブと同じ MLflow 実験名を使用しています。これにより、両方の実行結果は SageMaker MLflow アプリ上でまとめて表示されます。

rec_workload = {
    "benchmark": {"type": "aiperf"},
    "parameters": {
        "prompt_input_tokens_mean": 1600,
        "prompt_input_tokens_stddev": 200,
        "output_tokens_mean": 600,
        "output_tokens_stddev": 100,
        "request_count": 100,
    },
    "tooling": {
        "api_standard": "openai",
        "version": "0.8.0",
    },
}

REC_CONFIG = f"mlflow-rec-cfg-{SESSION_UID}"
REC_JOB = f"mlflow-rec-job-{SESSION_UID}"

sm.create_ai_workload_config(
    AIWorkloadConfigName=REC_CONFIG,
    AIWorkloadConfigs={
        "WorkloadSpec": {
            "Inline": json.dumps(rec_workload)
        }
    },
)

sm.create_ai_recommendation_job(
    AIRecommendationJobName=REC_JOB,
    ModelSource={
        "S3": {
            "S3Uri": MODEL_S3_URI
        }
    },
    OutputConfig={
        "S3OutputLocation": f"s3://{S3_OUTPUT_BUCKET}/mlflow-tutorial/{REC_JOB}/",
        "MlflowConfig": {
            "MlflowResourceArn": MLFLOW_APP_ARN,
            "MlflowExperimentName": MLFLOW_EXPERIMENT,
            "MlflowRunName": REC_RUN_NAME,
        },
    },
    RoleArn=ROLE_ARN,
    AIWorkloadConfigIdentifier=REC_CONFIG,
    PerformanceTarget={
        "Constraints": [
            {
                "Metric": "ttft-ms"
            }
        ]
    },
    OptimizeModel=False,
)

print(f"Recommendation job submitted: {REC_JOB}")

ステップ 6：MLflow で結果を比較する

ジョブが完了したら、SageMaker MLflow App を開き、設定した実験（experiment）へ移動してください。以下のコードは、ノートブック内で MLflow アプリの URL を取得するためのものです。

print(f"Open: {MLFLOW_APP_URL}")

実験ビューでは、親ランは子ランとの関係性を示すための空のプレースホルダーとして表示されます。ネストされたメトリクスを確認するには、各親ランを展開する必要があります。

benchmark-        -> Single run with AIPerf benchmark metrics
recommendation-   -> Nested run tree with ranked recommendation configurations

子ランにはモデルのメトリクスが格納されています。推薦ジョブのメトリクスを確認するには、MLflow Runs ページへ行き、推薦ジョブの親ランを展開して、並列実行（concurrency）の行が表示されるまでネストされた子ランを順次展開してください。ベンチマークのメトリクスは、以下の図に示す通り、各並列実行ランの「Metrics」カラムに表示されます。

推薦結果のアートファクトを検証するには、Concurrency=4 のランを開き、Artifacts タブを選択してください。アートファクトのファイルツリーで exports を展開し、outputs.json を選択します。プレビューパネルには、アートファクトのファイルパスと JSON 出力が表示され、セッション ID、リクエストのタイミング、出力トークン数、リクエストレイテンシ、レスポンステキストといったリクエストレベルの詳細フィールドも含まれています（以下の図参照）。

Runs ページで、比較したいベンチマーク実行（Concurrency=32、Concurrency=16、Concurrency=8、Concurrency=4 など）を選択し、「Compare」をクリックします。すると、以下に示すように、実行時間、スループット、出力トークン数、リクエストレイテンシの各メトリクス間の差異を確認できます。

Metrics セクションでは、「Show diff only」トグルをオンにすると、変更されたメトリクスのみを表示できます。以下に示すように、この表では選択したベンチマークまたは推奨ジョブの結果間で、ベンチマークの実行時間やエンドツーエンドの出力トークンスループットなどのメトリクスが比較されます。

リソースのクリーンアップ

継続的な課金を防ぐため、以下のリソースを削除してください。