概要

Amazon SageMaker で、独自コンテナを含む推論パイプラインエンドポイントを構築する 方法について書きます

SageMaker で提供されるコンテナで所望の前処理が実現できず、独自コンテナによる前処理をしたい場合、カスタム前処理コンテナはどのようにつくるか、どのように組み込むか調べた際のメモです

組み込みアルゴリズムと推論パイプライン

Amazon SageMaker では組み込みアルゴリズムが提供されており、解きたい問題にフィットするものがあれば、MLアルゴリズムに関するコードを書かずに手軽にMLモデルを作れるようになっています

組み込みアルゴリズムに引き渡す入力形式はアルゴリズムごとに決まっており、基本的には、保有するデータを整形する前処理が必要です
こうした 前処理は学習／推論で同じ処理なため、共通化されていることが望ましい ということで、SageMaker では前処理アルゴリズムをコンテナ化して学習／推論時に実行させる仕組みが提供されています

前処理コンテナでは SageMaker でのコンテナの挙動を踏襲し、fit でインスタンス上に前処理コンテナを展開してデータ整形／S3バケットへ出力したり、deploy でインスタンス上に前処理コンテナを展開してリクエストデータ整形／次処理へ渡す、といった挙動をさせます
とくに推論処理では、エンドポイントへ届いたリクエストを前処理して後続の推論アルゴリズムへ渡すために、「前処理＋推論処理（＋後処理等）」をセットにした 推論パイプラインエンドポイント が構築できるようになっています

前処理で独自コンテナが必要になるケース

推論パイプラインのサンプルには、SageMaker で MLフレームワークを使うときと同様 sagemaker.sklearn.estimator.SKLearn など Estimator を利用して、任意のスクリプトを所定のコンテナ環境下で実行させる例が載っています
提供されているコンテナで所望の処理が実現できれば、カスタムコンテナは不要です

from sagemaker.sklearn.estimator import SKLearn

script_path = 'sklearn_abalone_featurizer.py'
sklearn_preprocessor = SKLearn(
    entry_point=script_path,
    role=role,
    train_instance_type="ml.c4.xlarge",
    sagemaker_session=sagemaker_session)

しかし提供されるコンテナにないライブラリや、インストールされていないミドルウェアを使いたい場合には、独自の環境をコンテナ化して実行させる必要があります

たとえば自然言語を扱う課題で MeCab で形態素解析したいといった場合には、MeCab が動作してかつ Amazon SageMaker の仕様に従ったコンテナが必要になります

前処理コンテナのつくりかた

Amazon SageMaker で提供されているコンテナ群は GitHub で公開されているので、ライブラリを追加する程度であれば、既存コンテナを拡張するのが楽かもしれません
用意されているコンテナを FROM 文で呼び出した上に、追加のインストールやセットアップを施します

そうした提供済みコンテナで都合が悪い（ミドルウェアのバージョンが合わないとか、既存の資産を使いたいとか）場合には独自コンテナを作ります

前処理コンテナの要件

基本的な挙動は SageMaker の独自のトレーニングイメージの仕様にあわせる必要があります
fit や deploy といった命令をうけて所定の処理が RUN される必要があるため、 scikit_bring_your_own などをベースに、枠組みそのままで中身の処理を書き換えるのがよさそうです
また、パイプラインエンドポイントのための独自の設定がいくつか必要です

構成要素
- train：トレーニング（fit）時に呼び出される処理
- serve：エンドポイント構築（deploy）時に呼び出される処理
- predictor.py：推論エンドポイントで呼び出される処理
その他の要件
- Dockerfile に次のコードが必要
  - LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
- エンドポイントのリクエスト受け付けポートが次の環境変数に依存する
  - SAGEMAKER_BIND_TO_PORT
- Dockerコンテナリポジトリに、次のポリシー付与が必要
  - イメージをプルするための Amazon SageMaker のアクセス許可

前処理コンテナに実装する内容

scikit_bring_your_own 等をベースに、下記の機能を実装します

train

生データを受け取って整形加工し、MLモデルトレーニングに使用する訓練データをつくり、S3へ配置します
訓練データ生成の際に作成される「辞書データ」や「ベクトライザ」は推論時にも同じものを使うため、推論エンドポイントへ引き継ぐ必要がありますが、SageMaker のアルゴリズムで /opt/ml/model 以下を共有する仕組みにのせて、推論エンドポイントの同一ディレクトリへ展開させるようにします

Estimator.fit で呼び出される
- Amazon SageMaker がトレーニングイメージを実行する方法
指定されたタイプのインスタンスに、inputs={'raw':'s3://'} で引き渡した S3 パスのデータが /opt/ml/input/data/{channel} 以下へ 自動的に 展開される
/opt/ml/input/data/{channel} からデータを参照して一括整形する
MLトレーニングで使用する、加工整形した訓練データを S3 上の任意のパスへアップロードする
推論エンドポイントでデータの整形処理に必要になるモデルデータ（辞書やベクトライザなど）を /opt/ml/model へ出力する
- joblib.dump など
/opt/ml/model 以下のファイル群が tar.gz 形式でアーカイブされて output_path で指定された S3 パスへ 自動的に アップロードされる

serve

train 時に生成したモデルファイルをインスタンス上に配置し、 /invocations のリクエストを受け付ける web サーバを起動します

Estimator.deploy で呼び出される
- Amazon SageMaker が推論イメージを実行する方法
output_path に指定された S3 パスから、Training 時に生成済みのモデル群をインスタンスの /opt/ml/model 以下へ 自動的に 展開する
- Amazon SageMaker がモデルアーティファクトをロードする方法
/invocations と /ping に応答するウェブサーバーを起動する
- パイプラインモデルのエンドポイントは、リクエストを受け付けるポートを環境変数 SAGEMAKER_BIND_TO_PORT で指定されるため、web サーバの設定でポートを環境変数を参照するようにする（環境変数がなければ 8080 を使う）

predict

serve で起動したエンドポイントで、推論を行いたい生データを所定の書式で受け取り、後続の推論アルゴリズムが要求する書式に整形して返します

serve の際に配置したモデルファイル群を /opt/ml/model 以下から取得し、コード上へ読み出す
- joblib.load など
推論リクエストに含まれる生データを取得し、文字列の正規化を行う
正規化したデータについて、先に取得したベクトライザ等を使用して推論アルゴリズムが要求する形式のデータへ変換する
変換した形式のデータをレスポンスとして返す

コンテナのビルド、ECRへのPUSH

上記実装したコード群をコンテナ化し、ECR のリポジトリへ push します

push した ECR のリポジトリには、SageMaker のパイプラインモデルからの呼び出しを許可するポリシーを付与する必要があります（推論パイプラインの Amazon ECR アクセス許可のトラブルシューティング）

ECR コンソールから Repositories > 該当リポジトリリンク > Permissions と進んで、アクセス許可ポリシーを下記 JSON の通り追加します


{
    "Version": "2008-10-17",
    "Statement": [
        {
            "Sid": "allowSageMakerToPull",
            "Effect": "Allow",
            "Principal": {
                "Service": "sagemaker.amazonaws.com"
            },
            "Action": [
                "ecr:GetDownloadUrlForLayer",
                "ecr:BatchGetImage",
                "ecr:BatchCheckLayerAvailability"
            ]
        }
    ]
}

ここまででパイプライン用の前処理コンテナができました

MLプロセスへの組み込み

組み込みMLアルゴリズムと、独自コンテナによる前処理を組み合わせ、MLフローを実行します

前処理（訓練データの準備）

生データを用意してS3へアップロードし、以下のようなコードによって、訓練データへ変換します
＆前処理で生成したベクトライザ等の生成物を S3 へ格納します


import sagemaker
sess = sagemaker.Session()
role = sagemaker.get_execution_role()

preprocess_image_name = 'preprocess_ecr_container_image'
preprocess_job_name = '{preprocess_id}'
rawdata_input_path = 's3_rawdata_path'
train_dataset_path = 's3_preprocessed_data_path'

preprocessor = sagemaker.estimator.Estimator(
            image_name           = preprocess_image_name,
            role                 = role,
            train_instance_count = preprocess_instance_count, 
            train_instance_type  = preprocess_instance_type,
            train_volume_size    = preprocess_volume_size,
            output_path          = train_dataset_path,
            sagemaker_session    = sess )
preprocessor.fit(
    inputs={'rawdata': rawdata_input_path},
    job_name = preprocess_job_name )

学習

前処理で訓練に必要なデータはすべて S3 上の所定の場所へアップロードされている状態です
下記のコードによって、組み込みアルゴリズムによるトレーニングを実行します

import sagemaker

algorithm_name = 'ml_hoge'
ml_container_name = get_image_uri(boto3.Session().region_name, algorithm_name)
ml_model_path = 's3_ml_model_path'
ml_training_job_name = '{algorithm_name}-{preprocess_id}-{timestamp}'

ml = sagemaker.estimator.Estimator(
    image_name           = ml_container_name,
    role                 = role, 
    sagemaker_session    = sess,
    train_instance_count = train_instance_count, 
    train_instance_type  = train_instance_type,
    output_path          = ml_model_path )
ml.fit(
    inputs={
        'train':      train_dataset_path_train, 
        'validation': train_dataset_path_valid, 
        'auxiliary':  train_dataset_path_aux, 
        'test':       train_dataset_path_test },
    job_name = ml_training_job_name )

エンドポイント構築

前処理で生成したベクトライザ等を利用するデータ前処理コンテナと、組み込みアルゴリズムの学習済モデルを参照する推論コンテナをセットにして、推論パイプラインエンドポイントを構築します
（ sagemaker.pipeline.PipelineModel の models 変数へ渡した配列のモデル順に、リクエストに応答するエンドポイントが構築されます）


import sagemaker

preprocessor     = sagemaker.estimator.Estimator.attach(preprocess_job_name)
preprocess_model = preprocessor.create_model()
ml               = sagemaker.estimator.Estimator.attach(ml_training_job_name)
ml_model         = ml.create_model()

pipeline_endpoint_name = '{algorithm_name}-pipeline-endpoint-{version}'
pipeline_model = sagemaker.pipeline.PipelineModel(
    name   = pipeline_model_name,
    role   = role,
    models = [ preprocess_model, ml_model ] )
pipeline_model.deploy(
    initial_instance_count = pipeline_instance_count, 
    instance_type          = pipeline_instance_type, 
    endpoint_name          = pipeline_endpoint_name )

ここまでで、独自コンテナによる前処理を組み込んだ推論エンドポイントが起動します

推論リクエストへの応答

生データを所定の形式で引き渡せば、前処理コンテナで整形 -> 組み込みアルゴリズムの推論コンテナで推論 した結果が返ります

import sagemaker 
from sagemaker.content_types import CONTENT_TYPE_JSON

ml_predictor = sagemaker.predictor.RealTimePredictor(
            endpoint    = pipeline_endpoint_name,
            serializer  = sagemaker.predictor.json_serializer,
            deserializer= sagemaker.predictor.json_deserializer,
            content_type= CONTENT_TYPE_JSON,
            accept      = CONTENT_TYPE_JSON )

payload = {'target':'日本語文字列'}
pprint(ml_predictor.predict(payload))

トラブルシューティング

エンドポイント構築時に ping health check に失敗する

エンドポイント構築時に ping に失敗して構築できないというエラーが発生する場合

ValueError: Error hosting endpoint pipeline-endpoint-vXX: Failed Reason:  The container-1 for production variant AllTraffic did not pass the ping health check. Please check CloudWatch logs for this endpoint.

独自コンテナで起動させたWebサーバが、環境変数 SAGEMAKER_BIND_TO_PORT ポートでリクエストを受け付けていない可能性があります（CloudWatchLogs にも何の情報も出ず途方に暮れるのですが、落ち着いて）
Webサーバ（ nginx など）の設定を見直し、SAGEMAKER_BIND_TO_PORT 環境変数が存在する場合には、環境変数の指定するポートで受け付けるようにします

ECR コンテナリポジトリに必要な権限が付与されていない

パイプラインエンドポイントを構築しようとする際、権限が足りないといったエラーが出る場合

ValueError: Error hosting endpoint pipeline-endpoint-vXX: Failed Reason:  The repository of your image 111111111.dkr.ecr.xxxxx.amazonaws.com/XXXXXXXXXXXXXXXXXXXX does not grant ecr:GetDownloadUrlForLayer, ecr:BatchGetImage, ecr:BatchCheckLayerAvailability permission to sagemaker.amazonaws.com service principal.

推論パイプラインエンドポイントを構築する際には、パイプラインに組み込む独自コンテナのリポジトリに SageMaker からの読み出しを許可するポリシー付与が必要です
推論パイプラインの Amazon ECR アクセス許可を参考に、ECR コンソールから必要なポリシーを付与します

{
    "Version": "2008-10-17",
    "Statement": [
        {
            "Sid": "allowSageMakerToPull",
            "Effect": "Allow",
            "Principal": {
                "Service": "sagemaker.amazonaws.com"
            },
            "Action": [
                "ecr:GetDownloadUrlForLayer",
                "ecr:BatchGetImage",
                "ecr:BatchCheckLayerAvailability"
            ]
        }
    ]
}

その他

推論パイプラインのトラブルシューティング参照

まとめ

独自コンテナを含む推論パイプラインエンドポイントを構築する際、次の点に注意してください

基本的な挙動は SageMaker のカスタムコンテナと同じ
- Amazon SageMaker がトレーニングイメージを実行する方法
- Amazon SageMaker ホスティングサービスでの独自の推論コードの使用
エンドポイントのリクエスト受け付けポートは次の環境変数に依存する
- SAGEMAKER_BIND_TO_PORT
- 推論パイプラインでリアルタイム予測を実行
Dockerコンテナリポジトリに、次のポリシー付与が必要
- イメージをプルするための Amazon SageMaker のアクセス許可

日本語処理など独自コンテナを必要とするケースで、ぜひ参考にしていただければと思います

参考資料

パイプラインのコンテナは、(8080 ではなく) SAGEMAKER_BIND_TO_PORT 環境変数で指定されたポートでリッスンします。

コンテナがこの要件に準拠していることを示すには、次のコマンドを使用して Dockerfile にラベルを追加します
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true

Amazon SageMaker 組み込みアルゴリズムを含むパイプラインでカスタム Docker イメージを使用する場合は、Amazon ECR ポリシーが必要です。ポリシーは、イメージをプルするために Amazon SageMaker のアクセス許可を Amazon ECR レポジトリに与えます。

Amazon SageMaker の推論パイプラインで、独自コンテナを組み合わせる方法

概要