Amazon SageMaker におけるカスタムコンテナ実装パターン詳説〜推論編〜

みなさんこんにちは！AWS Japan のソリューションアーキテクトの辻です。このブログでは、Amazon SageMaker のカスタムコンテナの実装方法の選択肢が多すぎて、自分たちのユースケースにどれが向いているのか分からない！、という意見にお答えして、カスタムコンテナの実装方法のパターンを整理・解説しています。
SageMaker ではカスタムコンテナが利用可能な場面はいくつかあるのですが、このブログではホスティングによる推論 (SageMaker Endpoint) で利用する場合について説明します。学習編もこちらで公開しているので、是非一読ください。

はじめに

このブログは、既に SageMaker の概要を知っている方を対象としています。SageMaker Endpoint という SageMaker の推論環境に関する単語を知っており、SageMaker Python SDK や Boto3 での推論の実行方法がイメージできていると躓かずに読み進められると思います。もしこれらの領域を復習または新規に学習されたい方は AWS 公式 YouTube チャンネルにある SageMaker の動画シリーズが非常に良くまとまっているためオススメです。SageMaker 以外では Docker の概要や Dockerfile の記述方法についても基礎的な知識があるとなお良いです。

Amazon SageMaker は機械学習のあらゆるワークロードを支援するフルマネージドサービスです。SageMaker の数多ある機能の中で汎用性の高いものとして、推論用のモデルをホスティング環境をマネージドで提供する機能があります。この環境は、コンテナでの処理が前提であり、ユーザーはデプロイ時にコンテナイメージを指定します。ほとんどのユースケースでは AWS が提供しているコンテナイメージをそのまま利用できますが、ビルドが必要なライブラリをインストールしたいなど、どうしてもカスタマイズが必要になってくる場合があります。その際にカスタムコンテナを定義することになりますが、このカスタムコンテナの実装方法は複数あり、どれを選択すれば良いかが判断が難しいという相談をよく受けていました。そのような相談にお答えする形で、このブログでは推論に絞った上で、カスタムコンテナの実装パターンとその詳細について説明します。

SageMaker の推論コンテナを構成する 3 つのレイヤー

本題のカスタムコンテナの実装パターンの説明に入る前に一旦俯瞰して、どのようなアーキテクチャを前提としているのか説明します。通常、ML 推論コンテナは、アプリケーションコンテナとは分離したコンポーネントとしてデプロイします。ML 推論コンテナには、入出力データの前後処理と推論処理などの ML のロジックが含まれます。一方でここで言うアプリケーションコンテナには、ユーザー認証認可、他のサービスへの連携、データベースへのアクセスなどのビジネスロジックが含まれます。この 2 つのロジックは、デプロイのライフサイクル（変更が必要となるタイミング）が異なるため、開発境界の明確化やデプロイ起因の障害の極小化の観点で、コンポーネントを分離したアーキテクチャを採用する方が良いといえます。このブログでもこのような分離アーキテクチャを前提とします。ML ロジックは SageMaker を利用し、ビジネスロジックはまた別のサービスを利用します。具体例として下図のような形です。この図の後段部分の SageMaker Endpoint のコンテナを、如何にして実装するかがこのブログでの本題です。

一般的に推論、とくに HTTP リクエストを捌くようなホスティングによる推論は、学習と比べて複雑な処理を行う必要があります。当然自前ですべて実装することも可能ですが、開発運用難易度は低くはありません。幸いにも世の中には多くの開発者の協力によって、そのような処理を実現するためのライブラリが OSS として提供されています。ライブラリごとに処理全体を構築できるものや、一部分だけの機能を提供しているものなど様々です。SageMaker Endpoint へのデプロイを前提とした場合、これらのライブラリは大雑把に次の 3 つのレイヤーのいずれか、もしくは複数の機能を実装していることが多いかと思います。

推論ランタイム
モデルサービング
SageMaker 連携

1 つ目の推論ランタイムは、1 回の推論を実行するために必要なレイヤーです。例としては PyTorch の torch パッケージです。ローカルの開発環境で推論を試すだけならば、このレイヤーだけでも問題はありません。しかし本番環境などで秒間数十、数百、数千リクエストを捌くことを考えると、リクエストのたびにプロセスを起動し、ライブラリとモデルを読み込み、推論処理し、結果を返し、プロセスを終了するのはリソースの無駄です。そこで登場するのが 2 つ目のモデルサービングのレイヤーです。このレイヤーはプロセスプーリングや HTTP リクエストの終端、リクエストのルーティングなど行います。SageMaker 以外で推論環境を構築する場合は、1 つ目のレイヤーと 2 つ目のレイヤーだけでコンテナは構成されます。しかし SageMaker Endpoint で構築する場合は、加えて SageMaker と連携するレイヤーが必要です。これが 3 つ目です。SageMaker のコントロールプレーンから送信されるヘルスチェックへの応答や、入出力データの serialize/deserialize、コンテナ起動時・終了時の処理などを担います。

例として、PyTorch のモデルを SageMaker Endpoint にデプロイする場合を見てみます。コンテナイメージは AWS が提供しているものをそのまま利用します。PyTorch であれば、このコンテナイメージは TorchServe をベースにしたものが利用されます。 1 つ目のレイヤーは PyTorch のランタイムです（TorchServe のドキュメント参照）。2 つ目のレイヤーは TorchServe です。そして 3 つ目のレイヤーは SageMaker PyTorch Inference Toolkit が利用されます。

3 つのレイヤーとベースイメージの関係を表したのが下の図です。

カスタムコンテナの実装する際、自分たちがどのレイヤーをカスタマイズする必要があるのかを意識することが重要です。

カスタムコンテナの実装パターン

実装パターンを列挙する前に、まずはカスタムコンテナの利用目的（= どこまでカスタマイズ、または持ち込みたいのか）を整理してみます。基本的には次の 3 つのいずれかになるはずです。

AWS 提供のコンテナイメージに追加のライブラリを追加したい
独自の推論ランタイムを持ち込みたい（カスタマイズしたい）
独自のモデルサービングのライブラリを持ち込みたい（カスタマイズしたい）

各利用目的ごとに実装方法を整理しパターン化していき、本セクションの最後に全パターンを列挙します。

追加のライブラリをインストールしたい場合

まずは追加のライブラリ等をインストールしたい場合です。例えば、ベースのコンテナイメージは AWS が提供しているもので問題ないが、推論処理のために追加のライブラリなどが必要な場合に選択します。この場合、シンプルに AWS が提供しているコンテナイメージをベースイメージとして、レイヤーを追加すれば OK です。これがパターン 1 です。AWS が提供しているコンテナイメージをベースイメージとするため、実装もデプロイも比較的に簡単に実現できます。簡単である一方で pip install でインストールできるライブラリのほとんどは、デプロイ時に requirements.txt ファイルのあるディレクトリを指定して動的にインストールできるため、その場合はカスタムコンテナイメージではなく、デプロイ時指定を利用することをオススメします。フレームワークごとに細かい指定方法が異なってくるため、利用したいフレームワークごとにドキュメントを確認ください。例えば PyTorch の場合のドキュメントはこちらです。

推論ランタイムだけをカスタマイズしたい場合

続いて推論ランタイムだけをカスタマイズしたい場合です。例えば、特殊なランタイムだけで動作するようなモデルファイルだけあり、このモデルを SageMaker 上にデプロイしたい場合に選択します。この場合、モデルサービングのライブラリ（レイヤー 2）として MMS (Multi Model Server) が採用可能な否かで、さらに場合分けされます（MMS の詳細は後ほど説明します）。MMS が採用できる場合、SageMaker 連携のライブラリ（レイヤー 3）として SageMaker Inference Toolkit を利用します。この MMS と SageMaker Inference Toolkit の組み合わせがパターン 2 です。先程の 3 つのレイヤーのレイヤー 1 だけをカスタマイズする形です。一方で、MMS が採用できない場合、例えば、ランタイムが推奨しているモデルサービングのライブラリがあり、それを採用する場合などは、後述する「モデルサービングをカスタマイズしたい場合」と実装方法は同じになります。

モデルサービングをカスタマイズしたい場合

最後にモデルサービングをカスタマイズしたい場合です。例えば、オンプレミスや Amazon Elastic Container Service (ECS) など他のプラットフォームで既に動作しているコンテナイメージを、大きく変更せずに SageMaker に持ち込みたい場合に選択します。この場合、実装方法はモデルサービングのライブラリ依存となります。モデルサービングのライブラリを持ち込む場合、追加で SageMaker 連携機能の実装が必要です。学習編ブログを読まれた方ならば、SageMaker Training Job で SageMaker Training Toolkit をインストールして対応したように、SageMaker Endpoint でも SageMaker Inference Toolkit をインストールして対応できるのでは、と思われるかもしれませんが、残念ながら同じようには対応できません。これは SageMaker 連携機能が一般的にモデルサービングのライブラリに強く依存しているためです。例えば SageMaker Inference Toolkit は MMS との組み合わせでなければ基本的には動作しません。SageMaker Training Toolkit が汎用的なライブラリであったのと対称的といえます。具体的な実装方法は、モデルサービングのライブラリ依存ですが、強いてパターン化するならば次の 2 つに分けることができます。

モデルサービングのライブラリと同じライブラリのコンテナイメージを AWS が提供している
モデルサービングのライブラリと同じライブラリのコンテナイメージを AWS が提供していない

1 について。例えば、持ち込みたいモデルサービングのライブラリが TorchServe や TensorFlow Serving など SageMaker が対応しているもので、かつ既に他のプラットフォームでコンテナをホスティングしており、そのコンテナイメージをそのまま持ち込みたい場合が 1 です。この場合、SageMaker 対応のコンテナイメージの Dockerfile を参考しつつ実装します。これがパターン 3 です。パターン 3 では、そもそもを実装前に本当にコンテナイメージを「そのまま」持ち込む必要があるか検討すべきです。例えば TorchServe のコンテナイメージを持ち込むことを検討している場合、AWS が提供している PyTorch の推論コンテナイメージを利用しても問題ないことが多いはずです。なので複数プラットフォームで同じコンテナイメージ（ベースイメージ）をどうしても利用せざるを得ない状況で初めて選択する実装パターンとなります。
続いて 2 について。この場合は参考になる Dockerfile もないためイチからの実装が必要になります。SageMaker の入出力の仕様を確認しつつ自前で実装します。これがパターン 4 です。

パターン 3 と 4 どちらも、前述した他のパターンと比べて実装工数は多くなりがちです。実装の中心は SageMaker 連携機能となるので、果たして本当に SageMaker Endpoint を利用することがアーキテクチャとして適切なのかは検討することをオススメします。例えば Amazon ECS や AWS App Runner でも持ち込んだコンテナイメージをデプロイすることが可能です。

4 つの実装パターン

上で説明した実装パターンをまとめます。SageMaker Endpoint のカスタムコンテナの実装方法は、大きく 4 つのパターンがあります。

AWS が提供しているコンテナイメージを拡張する方法
独自のランタイムのコンテナイメージに MMS と SageMaker Inference Toolkit をインストールする方法
独自のモデルサービングのコンテナイメージに、SageMaker 連携機能を実装する方法（SageMaker 対応の Dockerfile を参考に実装）
独自のモデルサービングのコンテナイメージに、SageMaker 連携機能を実装する方法（スクラッチで実装）

それぞれのパターンのコンテナイメージで、どこまでのレイヤーが AWS から提供され、どこからユーザーが定義するのかを表したのが下記の図です。

次のセクションでそれぞれのパターンについて詳しく実装方法を説明していきます。

パターン 1: AWS が提供しているコンテナイメージを拡張する方法

パターン 1 で重要となるのは、ベースとなるコンテナイメージを特定することです。AWS が提供しているコンテナイメージの URI は学習同様こちらとこちらのドキュメントから辿ることができますが、少々手間がかかるため、オススメは一度目的のフレームワークのバージョンなどを指定して Model を作成し、URI を serving_image_uri() で取得してしまう方法です。この Model は URI 取得用としてのみ利用する形です。

from sagemaker.pytorch import PyTorchModel
model = PyTorchModel(framework_version='1.12', py_version='py38', ...)
model.serving_image_uri(...)

URI を取得できれば、それを Dockerfile の FROM に記述しビルドをプッシュを行います。ビルドとプッシュを行う際 CLI から aws ecr get-login-password | docker login ... を実行し認証しますが、ビルドとプッシュで異なる認証情報が必要なことに注意していください。ビルド時は SageMaker がオーナーの Amazon Elastic Container Registry (ECR) のリポジトリに対して認証が必要で、プッシュ時は自身がオーナーの ECR のリポジトリに対して認証が必要なためです。プッシュ後は URI を Model クラスのコンストラクタの image_uri に指定し、その後 deploy() を実行すればデプロイ完了です。

model = PyTorchModel(
   image_uri='<CUSTOM_IMAGE_URI_HERE>',
   model_data='s3://DOC-EXAMPLE-BUCKET/model.tar.gz',
   entry_point='entry_point.py',
   role=sagemaker.get_execution_role(),
)
predictor = model.deploy(initial_instance_count=1, instance_type='ml.m5.xlarge')
predictor.predict(...)

ベースイメージとして AWS が提供しているものを利用しているので、SageMaker との連携機能はカスタムでない場合と同じように利用可能です。

パターン 1 のサンプルコードとして Amazon SageMaker Examples の bring your own pipe-mode algorithm が参考になります。

パターン 2: 独自のランタイムのコンテナイメージに MMS と SageMaker Inference Toolkit をインストールする方法

パターン 2 は AWS が中心となって開発している MMS (Multi Model Server) と SageMaker Inference Toolkit を利用します。このパターンを実装するにあたり、これら 2 つのツールの機能の概要の理解が必要です。まずはそれぞれのツールの概要を説明します。

MMS は OSS のモデルサービングのライブラリです。SageMaker に限らずオンプレミスや ECS 上など様々なプラットフォームで動作します。MMS は特定のフレームワークに依存しないライブラリです。なので任意のランタイムでモデルサービングの機能を実現することができます。MMS と同じようにフレームワークに依存しないライブラリとして他にも BentoML などがあります。MMS は AWS が中心となって開発しているツールですが、MMS 自体には SageMaker との連携機能は備わっていません。MMS が SageMaker と連携するためには、SageMaker Inference Toolkit が必要になります。

SageMaker Inference Toolkit は MMS 専用の SageMaker 連携用ライブラリで SageMaker Endpoint のコンテナ内で利用されます。SageMaker との様々な連携機能が実装されており、MMS を利用している限りは、スクラッチでこれら機能を実装することは工数が大幅に増えるため非推奨です。具体的に以下の機能は、独自に処理を記述しない場合、SageMaker Inference Toolkit がなければ正しく動作しません。

ローカルの推論コードを entry_point として指定できる機能
ヘルスチェックへの応答や推論エラーなどをサービスのコントロールプレーンに通知する機能
などなど

SageMaker が提供しているいくつかのフレームワークのコンテナイメージ、例えば MXNet, Hugging Face, Scikit-learn は MMS と SageMaker Inference Toolkit の組み合わせを利用しています。

パターン 2 で重要となるのは、MMS と SageMaker Inference Toolkit に対応したコンテナイメージと推論コードを作成・実装することです。具体的に、SageMaker Endpoint をデプロイするための次の 4 つの手順それぞれについて実装方法の概要を説明します。

モデルの作成
コンテナイメージの作成（重要）
推論コードの実装（重要）
デプロイと推論呼び出し

手順 1

まずモデルの作成です。この手順は MMS と SageMaker Inference Toolkit とは独立しています。今回はなるべく話をシンプルにするために、2 x 2 行列をモデルとみなし、NumPy で実装します。

import numpy as np
import tarfile

w = np.identity(2, dtype=np.float32) * 10
np.save('model.npy', w)
with tarfile.open('model.tar.gz', 'w:gz') as f:
    f.add('model.npy')

上のモデル (w) は 2 次元ベクトルの入力を 10 倍して返します。上のコードを実行するとモデルファイルを圧縮した model.tar.gz が得られます。これを適当な S3 のパスへアップロードしておきます（ここでは s3://DOC-EXAMPLE-BUCKET/model.tar.gz とします）。

手順 2

次はコンテナイメージの作成です。この手順は MMS と SageMaker Inference Toolkit に対応したコンテナイメージの作成および設定が必要です。作成の基本的な手順は SageMaker Inference Toolkit の GitHub の README に書かれている次の 4 ステップです（日本語に翻訳し、少し補足を加えました）。

モデルの読み込み、一回の推論時のモデルの入出力変換を実装し、そのファイルをコピーする (Handler の実装とコピー)
コンテナ起動時、リクエスト受取時の挙動を実装し、そのファイルをコピーする (HandlerService の実装とコピー)
コンテナ起動時のエントリーポイントを実装し、そのファイルをコピーする
MMS と SageMaker Inference Toolkit をインストールしたコンテナイメージに対して 3 を ENTRYPOINT としビルドする

ステップ 1 に関しては、SageMaker Inference Toolkit のライブラリ内にデフォルトの Handler が実装されており、デプロイ時に上書き可能なので、必須ではありません。ステップ 2 に関しても、SageMaker Inference Toolkit のライブラリ内にデフォルトの HandlerService が実装されているので、必須ではありません。つまり、コンテナイメージビルド時に必須となるのはステップ 3 と 4 だけになります。

ステップ 3 のコンテナ起動時のエントリーポイントは、カスタマイズが不要な場合は次のように README とほぼ同じコードになります。このファイル名を serve.py としましょう。

from sagemaker_inference import model_server
model_server.start_model_server()

ステップ 4 は次のような Dockerfile をビルドします。Dockerfile はステップ 3 の serve.py と同じディレクトリに置きます。今回はベースイメージとして Python 3.8 の Docker Official Image を利用しています。

FROM python:3.8
WORKDIR /usr/src/app
RUN apt-get update && apt-get upgrade -y && apt-get install -y openjdk-17-jdk-headless
RUN pip install --no-cache-dir numpy multi-model-server sagemaker-inference
COPY serve.py ./
ENTRYPOINT ["python", "/usr/src/app/serve.py"]

3 行目で OpenJDK をインストールしているのは MMS が JVM に依存しているためです。4 行目で NumPy, MMS, SageMaker Inference Toolkit をインストールしています。ビルド後 ECR へプッシュすれば、コンテナイメージの準備は完了です。

手順 3

続いて推論コードです。ここも手順 2 と同様に MMS と SageMaker Inference Toolkit に対応した実装が必要です。推論コードは SageMaker Inference Toolkit の Handler クラスから呼び出されます。呼び出されるインターフェイスは決まっており、具体的には 4 つの関数 model_fn(), input_fn(), predict_fn(), output_fn() が必要です。今回の NumPy のモデルでは次のように実装します。

from pathlib import Path
import numpy as np
from sagemaker_inference import content_types, decoder, encoder

def model_fn(model_dir, context=None):
    return np.load(str(Path(model_dir) / 'model.npy'))

def input_fn(input_data, content_type, context=None):
    return decoder.decode(input_data, content_type)

def predict_fn(data, model, context=None):
    return model @ data
    
def output_fn(prediction, accept, context=None):
    return encoder.encode(prediction, accept)

それぞれの関数の役割については SageMaker Inference Toolkit の GitHub を参照ください。この推論コードのファイル名を entry_point.py として、さらに空のディレクトリを作成し、その下に置きます（ここでは code というディレクトリを作成し ./code/entry_point.py という構成とします）。

手順 4

最後にデプロイと推論呼び出しです。手順 2 と 3 を正しく行っていれば MMS と SageMaker Inference Toolkit とは独立して考えることができます。次のコードを実行することでデプロイと推論呼び出しができます。image_uri には ECR へプッシュしたコンテナイメージの URI を指定します。

from sagemaker import Model
imoprt sagemaker
import numpy as np
model = Model(
    image_uri='<CUSTOM_IMAGE_URI_HERE>',
    model_data='s3://DOC-EXAMPLE-BUCKET/model.tar.gz',
    entry_point='entry_point.py',
    source_dir='./code',
    role=sagemaker.get_execution_role(),
    predictor_cls=sagemaker.predictor.Predictor,
)
predictor = model.deploy(initial_instance_count=1, instance_type='ml.m5.xlarge')
predictor.serializer = sagemaker.serializers.NumpySerializer(dtype=np.float32)
predictor.deserializer = sagemaker.deserializers.NumpyDeserializer(dtype=np.float32)
predictor.predict(np.array([1, 2]))
# array([10., 20.])

パターン 1 では Model 作成時に entry_point だけ指定しましたが、パターン 2 ではわざわざ推論コードのディレクトリを分け、entry_point と source_dir の両方を指定しています。これはクライアント SDK である SageMaker Python SDK の現時点 (v2.109) での仕様のためです。また入出力に応じて適切に serializer/deserializer を設定する必要があることに注意してください。推論結果は想定通り 2 次元ベクトルが 10 倍され返ってきました。

パターン 2 のサンプルコードとして Amazon SageMaker Examples の Multi-Model-Endpoints using your own algorithm container が参考になります。GitHub で公開されている BYOC (bring your own container)、つまりカスタムコンテナ関連のサンプルコードはパターン 2 であることが多いので、リポジトリ内を “multi-model-server” や “sagemaker-inference” などで検索してみると多数見つけることができます。

パターン 3: 独自のモデルサービングのコンテナイメージに、SageMaker 連携機能を実装する方法（SageMaker 対応の Dockerfile を参考に実装）

パターン 3 で重要となるのは、参考となる Dockerfile を特定することです。まず持ち込みたいモデルサービングのライブラリが SageMaker に対応しているか確認する必要があります。こちらのドキュメントから辿ることができますが、ランタイムとモデルサービングが混ざっており少々わかりにくくなっています。以下に、ブログ執筆時点における SageMaker がコンテナイメージを提供している代表的なモデルサービングのライブラリをまとめます。バージョンによって変わることもあるため、あくまで参考情報として利用ください。

TorchServe (PyTorch の場合に利用される)
TensorFlow Serving (TensorFlow の場合に利用される)
MMS (Hugging Face, Apache MXNet, Scikit-learn の場合に利用される)
Triton Inference Server
MLeap + Spring Boot (SparkML の場合に利用される)

もしこれらのモデルサービングのライブラリ以外を利用したければ、パターン 4 での実装が必要です。1 つ目から 3 つ目に関しては AWS Deep Learning Containers の GitHub リポジトリから Dockerfile を特定できます。例えば PyTorch v1.12 (& TorchServe v0.6.0) の CPU 推論の Dockerfile はこちらです。4 つ目の Triton Inference Server に関しては、ブログ執筆時点では Dockerfile は公開していないため、パターン 4 での実装が必要です。5 つ目に関しては SageMaker SparkML Serving Container の GitHub リポジトリから Dockerfile を特定できます。これらの公開されている Dockerfile を参考にしつつ、持ち込みたい既存コンテナイメージにレイヤーを加えて実装します。

注意点として、これらのコンテナイメージで利用している SageMaker 連携機能は、モデルサービングのライブラリの起動方法や呼び出し方法を隠蔽していることが多く、場合によってはスクラッチで実装した方が効率が良い場合があります。Dockerfile を確認し、もし実装工数が大きくなりそうならばパターン 4 に切り替えることを検討すると良いです。

パターン 4: 独自のモデルサービングのコンテナイメージに、SageMaker 連携機能を実装する方法（スクラッチで実装）

パターン 4 では、SageMaker Endpoint がどのように外部とデータをやり取りしているか、その仕様を理解するのが重要です。こちらのドキュメントに仕様が書かれているので、参考にしつつ実装します。またクライアントが SageMaker Python SDK を利用している場合は、起動時および推論リクエスト時に AWS の API (e.g., CreateEndpoint API, InvokeEndpoint API) に追加する各種パラメータを理解する必要があるため注意が必要です。

パターン 4 のサンプルコードとして Amazon SageMaker Examples の bring your own algorithm container が参考になります。こちらは Nginx + Gunicorn + Flask で構築されたコンテナイメージを SageMaker Endpoint にデプロイしているサンプルとなっています。

さいごに

カスタムコンテナは、SageMaker のマネージドなインスタンスによって得られる運用負荷の低減と、独自のコンテナによって得られる柔軟性の向上を兼ね備えた、非常に強力な利用方法です。これまでもドキュメントやサンプルコード、ブログなど様々な情報がありましたが、一歩引いた形で概要をまとめたものはない状況でした。このブログが皆様の理解の助けになることを願っています。

著者について

辻陽平 (Yohei Tsuji) は AWS Japan のエンタープライズソリューションアーキテクトで、製造領域や研究開発領域のお客様を中心にアーキテクチャ設計や構築をサポートしています。HPC や機械学習の領域を得意としており、特に分散深層学習は大好物です。最近はジム通いで身体を大きくすることとコーヒーに夢中です

Amazon Web Services ブログ

Amazon SageMaker におけるカスタムコンテナ実装パターン詳説 〜推論編〜