Amazon SageMaker AI、SDK v3 で独自モデルの Script mode を強化

2021 年に「Bring your own model with Amazon SageMaker script mode」という記事を発表しました。この投稿では、AWS が提供する管理されたフレームワークコンテナ上でスクリプトモードを活用し、独自のトレーニングや推論コードを記述する方法を紹介しています。

スクリプトモードは大きな飛躍でした。Amazon SageMaker AI 上で独自アルゴリズムを実行するために、Docker イメージの構築や維持管理が不要になったのです。

v3 SDK はゼロから再設計されたもので、Bring Your Own Model（BYOM）ワークフローを含む多くの処理をさらに簡素化しています。

新しいSDKでは、トレーニング用のフレームワーク固有の推定器クラス（SKLearn、PyTorch、XGBoost）が廃止され、代わりにトレーニングには単一の統一された ModelTrainer、デプロイには ModelBuilder が採用されています。

SDK v3 では、新しい SourceCode 設定オブジェクトを使用して、実行時にローカルのソースコードディレクトリをトレーニングジョブに同期します。コンテナイメージは、自分で構築したもの、AWS Deep Learning Container、またはサードパーティ製のものを Amazon Elastic Container Registry (Amazon ECR) から持ってきてください。SDK が実行時にコードの注入を処理します。

つまり、以下のようになります：

・イテレーションの高速化: 学習スクリプトを変更して再実行するだけで、コンテナの再構築は不要です。

・コンテナの完全な制御: システムパッケージや CUDA ライブラリをイメージ内に直接インストールできます。SDK はコンテナ内部の内容について仮定しません。

複数のフレームワークに共通の API

scikit-learn、PyTorch、Stable Diffusion、あるいは独自のカスタム C++ 推論バイナリなど、どのフレームワークでトレーニングを行う場合でも、インターフェースは統一されています。

ソリューション概要

本記事では、SageMaker Python SDK v3 のスクリプトモードがどのように機能するかを示す 2 つの完全な例をご紹介します。

• scikit-learn のランダムフォレストを学習・デプロイする – これは糖尿病データセットを用いた古典的な表形式機械学習（ML）ワークフローです。学習したモデルは、高性能なモデルサーバーである Deep Java Library (DJL) Serving を用いてリアルタイムエンドポイントにデプロイされます。

LoRA を用いた Stable Diffusion 3.5 のファインチューニング – Hugging Face Accelerate を活用したマルチ GPU 分散学習を行う生成 AI ワークフローです。

両方の例で共通して使われるのは、以下の 2 つのコアクラスです。

• ModelTrainer：v2 の Estimator ファミリーに代わるクラスです。SageMaker のトレーニングジョブを設定し、起動します。

• ModelBuilder：v2 の Model/Predictor パターンに代わるクラスです。推論ハンドラーをパッケージ化し、エンドポイントへデプロイします。

重要な概念として、SourceCode オブジェクトがあります。これは、ローカルのコードディレクトリへのパスである source_dir と、実行コマンドを指定する文字列 command のいずれかを受け取ります。

トレーニングには training を、推論には entry_script を使用します。

ジョブ起動時に、SageMaker はこのディレクトリをコンテナに同期し、コードはイメージに埋め込まれることなくコンテナ内で実行されます。

本ブログ記事のサンプルコードは、GitHub リポジトリ で確認できます。

SDK v2 から v3 への変更点

アーキテクチャの変化を以下の表にまとめました。

前提条件

本ガイドを実行するには、以下の準備が必要です。

• Amazon SageMaker AI にアクセスできる AWS アカウント。
• Amazon SageMaker AI と Amazon Simple Storage Service (Amazon S3) に対する 権限 を付与された、AWS Identity and Access Management (IAM) の実行ロール。
• SageMaker Python SDK v3 のインストール（pip install sagemaker>=3.0）。
• Amazon ECR にプッシュされたトレーニング用コンテナイメージ（本記事では、ECR へのビルドとプッシュの例を示します）。
• トレーニングデータやモデルアーティファクトを格納するための Amazon S3 バケット。
• （オプション）実験追跡用の MLflow アプリまたはトラッキングサーバー を Amazon SageMaker AI 上に設定する場合。
• （オプション）ローカルマシンではなく、Amazon SageMaker Studio の JupyterLab スペースから例のコンテナをビルド・実行する予定の場合、ドメインレベルで Docker アクセスが有効になっている必要があります。詳細は Amazon SageMaker Studio におけるローカルモードサポート を参照してください。

例 1: scikit-learn モデルのトレーニングとデプロイ

まずは、古典的な機械学習ワークフローから始めましょう。diabetes データセットを用いてランダムフォレスト分類器をトレーニングし、リアルタイム SageMaker エンドポイントにデプロイします。

ステップ 1: Docker コンテナの構築

トレーニング用コンテナは意図的に最小限に設計されています。実行環境とフレームワークライブラリのみを含み、トレーニングコードは一切含まれていないため、今後作成する可能性のある他の scikit-learn モデルでも再利用可能です。

scikit-learn 用の完全な Dockerfile は以下の通りです。

FROM python:3.13-slim

RUN apt-get update && apt-get install -y \
    build-essential jq git \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .

RUN pip install -r requirements.txt --no-cache-dir

このコンテナは安定しており、バージョン管理された実行環境となっています。アルゴリズム固有のコードは source_dir に配置し、SDK がランタイム時に自動的に読み込みます。

このコンテナは一度構築して Amazon ECR へプッシュすれば、Docker を再度触ることなく、トレーニングコードを何度でも修正・改善できます。

サンプルノートブックでは、2 つのシェルスクリプトを使用して Docker のビルドとプッシュコマンドを実行しています。

./build.sh --env .env.docker.sklearn
./push.sh --env .env.docker.sklearn

なお、コードサンプルを実行する環境には Docker がインストールされている必要があります。Amazon SageMaker AI 内の JupyterLab スペースで実行する場合は、ドメインレベルの設定で Docker を有効にする 必要があります。

ステップ 1a: 設定

まず、アカウントレベルの設定を自動的に検出します。以下のコードスニペットでは簡略化のためインポート文を省略していますが、完全なコードは GitHub リポジトリで確認できます。

...
boto_session = boto3.Session()
sm_session = Session(boto_session=boto_session)
AWS_REGION = boto_session.region_name
AWS_ACCOUNT_ID = boto3.client("sts").get_caller_identity()["Account"]
SAGEMAKER_EXECUTION_ROLE = sm_session.get_caller_identity_arn() # works for role or user
S3_BUCKET = sm_session.default_bucket()

以下のスニペットでは、TRAINING_IMAGE_URI に直前に構築したコンテナを指定しています。これにより、インストール済みパッケージやランタイムバージョンを完全に制御できます。

なお、自分でコンテナを構築したくない場合は、既存の管理された DJL フレームワーク用コンテナも利用可能です。プレビルドされたコンテナの詳細については、Deep Learning コンテナのイメージ一覧 をご覧ください。

# Your custom training image pushed to ECR
TRAINING_IMAGE_URI = (
    f"{AWS_ACCOUNT_ID}.dkr.ecr.{AWS_REGION}.amazonaws.com/sklearn:latest"
)

MODEL_OUTPUT_S3_PATH = f"s3://{S3_BUCKET}/random-forest/model-output"

トレーニング実行間でハイパーパラメータ、メトリクス、モデルアーティファクトを追跡したい場合は、例のトレーニングスクリプトが Amazon SageMaker AI 上の 完全に管理された MLflow に対応済みです。以下のスニペットで MLFLOW_ARN と MLFLOW_EXPERIMENT_NAME の変数を設定すると、自動的にログ記録が有効になります。これはオプションのステップであり、実験追跡をスキップする場合は値を None に設定できます。

MLFLOW_ARN = "arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_ID}:mlflow-app/{XYZ}"
MLFLOW_EXPERIMENT_NAME = "random-forest-experiment"

ステップ 2: ModelTrainer でトレーニングジョブを実行

SourceCode オブジェクトは、ローカルの source_dir と command 文字列を受け取ります。ジョブが起動すると、SageMaker は source_dir 全体をコンテナに同期し、指定された command を実行します。

これにより、コードとコンテナイメージの依存関係が解消されます。スクリプトを変更しても、コンテナを再構築する必要なく再度起動できます。

source_code = SourceCode(
    source_dir="./train/random_forest",
    command=(
        "python random_forest.py"
        " --n_jobs 4 --max_depth 10 --n_estimators 120"
        f" --mlflow_arn {MLFLOW_ARN}"
        f" --mlflow_experiment_name {MLFLOW_EXPERIMENT_NAME}"
    ),
)

compute = Compute(
    instance_type="ml.m5.2xlarge",
    instance_count=1,
    volume_size_in_gb=30,
    keep_alive_period_in_seconds=3600, # Warm pool for faster re-runs
)

stopping_condition = StoppingCondition(max_runtime_in_seconds=3600)
output_config = OutputDataConfig(s3_output_path=MODEL_OUTPUT_S3_PATH)

model_trainer = ModelTrainer(
    training_image=TRAINING_IMAGE_URI,
    source_code=source_code,
    compute=compute,
    output_data_config=output_config,
    stopping_condition=stopping_condition,
    role=SAGEMAKER_EXECUTION_ROLE,
    base_job_name="random-forest-training",
    sagemaker_session=Session(),
)

model_trainer.train(wait=True)

いくつか注意点があります：

source_dir には、ユーティリティモジュールや設定ファイル、シェルスクリプトなどのファイルを格納できます。これらはコンテナに同期されます。

command は、コンテナ内で実行されるシェルコマンドです。Python スクリプトや Bash スクリプトを呼び出すほか、コンテナがサポートする他の処理も可能です。

keep_alive_period_in_seconds を設定すると、SageMaker のウォームプール機能が有効になります。インスタンスは 1 時間温存状態を保つため、反復的な再実行でも数秒で起動し、従来の数分を要するプロセスから大幅に短縮されます。

OutputDataConfig は、トレーニングジョブが完了した際に SageMaker が学習結果をアップロードする S3 の宛先を設定します。スクリプトで /opt/ml/model（環境変数 SM_MODEL_DIR）に保存されたすべてのデータはここに格納されます。

変数はこのパスの下で model.tar.gz としてパッケージ化され、これがステップ 3 でデプロイするモデルアーティファクトとなります。

詳細については、フォルダ構造についてはSageMaker 学習および推論ツールキットの使用を、SageMaker の環境変数については以下のドキュメントをご覧ください。

セットアップについては、こちら をご覧ください。

ステップ 3: ModelBuilder を使用してリアルタイムエンドポイントへデプロイする

トレーニングが完了したら、モデルの成果物を SageMaker のリアルタイムエンドポイントにデプロイします。ModelBuilder は推論ハンドラをパッケージ化し、モデルの成果物と組み合わせてエンドポイントを数行で作成します。

トレーニングジョブから取得したメタデータを使って最終的なモデル成果物の S3 パスを確認し、それを ModelBuilder オブジェクトに渡すことができます。サービス提供には、カスタムコンテナを構築するのではなく、事前に用意された AWS Deep Learning コンテナを使用しています。必要に応じて独自のコンテナを利用することも可能です。

その他の事前用意されたコンテナについては、利用可能な Deep Learning コンテナイメージ をご覧ください。

# Locate the trained model artifact
sm_session = Session()
sm_client = sm_session.boto_session.client("sagemaker")
training_job_desc = sm_client.describe_training_job(
    TrainingJobName=training_job_name
)
model_artifact_s3_uri = training_job_desc["ModelArtifacts"]["S3ModelArtifacts"]

# Define inference source code
inference_source_code = SourceCode(
    source_dir="./deploy/random_forest",
    entry_script="inference.py",
)

# Build and deploy
model_builder = ModelBuilder(
    image_uri=INFERENCE_IMAGE_URI,
    model_server=ModelServer.DJL_SERVING,
    source_code=inference_source_code,
    s3_model_data_url=model_artifact_s3_uri,
    env_vars={"OPTION_ENTRYPOINT": "code/inference.py"},
    sagemaker_session=Session(),
    role_arn=SAGEMAKER_EXECUTION_ROLE,
)

model_builder.build(model_name="random-forest-endpoint", mode=Mode.SAGEMAKER_ENDPOINT)

build() ステップは、インフラを起動することなく実行可能なモデルを組み立てます。ModelBuilder は推論ハンドラとモデルアーティファクトを受け取り、選択したモデルサーバー（ここでは DJL Serving）の規約に従ってこれらをパッケージ化します。その後、Amazon S3 上の推論イメージと再パッケージ化したアーティファクトを指す SageMaker モデルが登録されます。

ModelBuilder はここで紹介している以上の機能も提供します。例えば、コンテナの自動選択や依存関係の自動キャプチャ、生フレームワークモデルからのシリアライズコード生成などです。詳細は ModelBuilder を使用して Amazon SageMaker AI でモデルを作成する をご覧ください。

モデルが完成したら、deploy() を呼び出してリアルタイムエンドポイントを作成します。これにより Endpoint インターフェースが返されます。

predictor = model_builder.deploy(
    endpoint_name="random-forest-endpoint",
    initial_instance_count=1,
)

推論においても同様に SourceCode パターンを使用します。これはハンドラを含むローカルディレクトリを指し、entry_script を指定するものです。SDK はハンドラをモデルアーカイブに再パッケージ化するため、DJL Serving がランタイム時にこれを検出できるようになります。

上記のコードスニペットに関するいくつかの補足です：

inference.py スクリプトは、DJL Python モードのドキュメント に従って、1 つの handle(inputs) 関数を実装しています。これは SageMaker がリクエストごとに呼び出すものです。

推論ワーカーが最初に起動すると、一度限りのモデル読み込み処理を完了させるために、空のリクエストがハンドラーに送信されます。

この概念は、モデルを一度読み込んでから、その後のリクエストを予測にマッピングするものです。デフォルトでは、モデルは /opt/ml/model に配置され、これは SM_MODEL_DIR 環境変数に対応しています。

初期のモデル読み込み後、各 incoming リクエストに対して推論スクリプトが Content-Type を判定し、ペイロードをシリアライズ解除して、予測結果を JSON オブジェクトとして返します。

コールドスタート時にモデルを一度読み込み、リクエスト間で再利用してください。毎回再読み込みすると、すべての呼び出しでレイテンシが増加するためです。また、リクエストの Content-Type を検証し、予期しない入力に対しては明確かつ即座にエラーを返すエンドポイントとして動作させることも重要です。

model_server 引数は、ModelBuilder がモデルをパッケージ化してエンドポイント内で実行するサービングランタイムを指定します。このモデルサーバーは、モデルを読み込み、SageMaker が期待するエンドポイントを公開し、各リクエストをハンドラーに配信するプロセスです。そのため、この設定は inference.py の記述方法と対応しています。

ここでは、汎用的な Python 推論や大規模モデルの提供に適した、柔軟性が高く高性能なサーバーである ModelServer.DJL_SERVING を選択します。これが、前述の DJL の handle(inputs) コントラクトに従ってハンドラーが実装されている理由でもあります。

ModelServer が提供する他のモデル提供の選択肢については、ModelServer API リファレンス をご覧ください。

• mode パラメータは、モデルが実行される場所を制御します。ここでは Mode.SAGEMAKER_ENDPOINT を使用し、完全に管理されたリアルタイムエンドポイントへデプロイします。

ModelBuilder は、テストやデバッグのために、ローカルの Docker コンテナで実行する Mode.LOCAL_CONTAINER や、現在の Python プロセス内で直接実行する Mode.IN_PROCESS もサポートしています。

ハンドラーをローカルで反復して開発する。

この場合、リアルタイムエンドポイントへのデプロイを行います。ワークロードに応じて、単一のモデルを独立したエンドポイントでホストすることもできますし、推論コンポーネントを用いて複数のモデルを1 つのエンドポイントにまとめて配置することも可能です。これにより、各モデルのリソース割り当てとスケーリングを個別に行うことができます。詳細は リアルタイム推論 をご覧ください。