Amazon SageMakerで独自の学習/推論用コンテナイメージを作ってみる

こんにちは、大阪DI部の大澤です。

Amazon SageMakerはAWSの機械学習のフルマネージドサービスです。ノートブックインスタンスによるデータ探索から機械学習モデルの学習から推論用エンドポイントへのモデルのデプロイなどを行うことができます。

SageMakerには幾つかの組み込みアルゴリズムがありますが、それ以外にもTensorFlowやMXNetなどのその他の機械学習フレームワークを使ったSageMaker上での学習やモデルのホスティング、バッチ変換なども可能です。いくつかのフレームワークについてはSageMakerで学習と推論を行うときに必要なコンテナイメージが用意されています。しかしながら、用意されているコンテナイメージでは要件が満たせない場合には、自らでコンテナイメージを作成するかマーケットプレイスでアルゴリズム/モデルパッケージをサブスクライブする必要があります。 今回は学習/推論に使用するコンテナイメージを作ってみたいと思います。

対象

このエントリは以下のような読者を想定しています。

• SageMakerにおけるモデルの学習や推論用エンドポイントの作成などがざっくりわかる。
• コンテナの概念やDocker、ECRなどがざっくり分かる。触ったことがある。
• SageMakerで独自の学習/推論用コンテナイメージを作ってみたい。

学習/推論用コンテナイメージ

SageMakerで機械学習モデルの学習と推論を行う際にはコンテナを使用しますが、そのコンテナイメージは学習用と推論用で分けてもいいですし、同じものでも大丈夫です。 そのために、コンテナを起動する際のdocker run時にCMD引数としてtrainもしくはserveを渡して実行します。学習時にはtrainを渡すことで、学習用のスクリプトを実行します。推論時にはserveを渡すことで、推論用のスクリプトを実行します。処理実行中、標準出力へ出力される内容はジョブに対応するCloudWatch Logsのログストリームに転送されます。

それぞれの処理にはどういったものが必要かは以下の各項目をご覧ください。

学習用コンテナ

学習用コンテナでは特定のパスにあるデータとハイパーパラメータを読み込んで、モデルの学習処理を実行します。

ファイル構成

知っておくと良さそうな最低限のファイル構成は以下の通りです。

/opt
    ├──program
    │   └─ train
    └──ml
        ├── input
        │   ├── config
        │   │   ├── hyperparameters.json
        │   │   ├── inputdataconfig.json
        │   │   └── resourceconfig.json
        │   └── data
        │       └── <channel_name>
        │           └── <input data>
        ├── model
        │   └── <model files>
        └── output
            └── failure

• /opt/program: 学習用コンテナイメージを作成する際に配置する必要があるフォルダです。ノートブックの例に合わせて/opt/programとしていますが、コンテナ起動時に実行できる場所であれば特にパスの指定はありません。(＝Dockerfileでパス通す必要がある)
  • /opt/program/train: コンテナ起動時に実行されるスクリプトファイルです。学習処理などを記述します。
    • shebangなどの方法でファイル名で実行可能にする必要があります。
      • 例) スクリプトファイルの行頭に#!/usr/bin/env pythonを記載する。
  • train以外にファイルが存在していても問題ありません。
• /opt/ml: SageMakerによって作成されるファイル群です。（コンテナイメージ作成時には含める必要がないです）
  • /opt/ml/input/config: 学習を制御するための設定データと入力データが入っています。
    • /opt/ml/input/config/hyperparameters.json: 学習時などに使用するハイパーパラメータが辞書形式で保存されているJSONファイルです。辞書形式で値は全て文字列型になっているので、読み込む際には適切な型へのキャストを行う必要があります。
    • /opt/ml/input/config/resourceconfig.json: 分散学習のためのネットワーク構成などが記述されたJSONファイルです。
    • /opt/ml/input/config/inputdataconfig.json: 入力データの情報が記述されたJSONファイルです。入力されるデータチャネルごとにコンテンツタイプやインプットモードなどが記述されています。
  • /opt/ml/input/data/<channel_name>/: 対応するチャネルの入力データが入っています。学習処理が実行される前にS3からデータがコピーされます。
    • データ入力方法がFileの時に使います。
  • /opt/ml/input/data/<channel_name>_<epoch_number>: 対応するチャネルのエポックごとにデータがあります。各エポックに対応するデータを読むごとに次のエポックのデータが作られます。
    • データ入力方法がPipeの時に使います。
  • /opt/ml/model/: 学習の結果得られたモデルデータを出力します。
    • 保存するモデルデータは単体のファイルでも複数のファイルでも問題ありません。
    • SageMakerが自動的にmodelディレクトリをtarで圧縮し、S3に保存します。
  • /opt/ml/output: ジョブ実行時に処理が失敗した場合に原因などを示したエラーメッセージを出力します。このテキストファイルの中身は、DescribeTrainingJobで得られる内容のFailureReasonという項目として表示されます。

基本的な処理の流れ

ファイル構成を踏まえた上で、理解を深めるために学習ジョブ実行時の基本となる流れを書き出してみます。

• SageMakerによってインスタンスが準備されて、その中に学習用コンテナイメージが起動されます。その際にtrainコマンドが渡されます。
  • 想定: docker run イメージ名 train
• コンテナ内でtrainスクリプトが実行されます。以下はtrainスクリプトに最低限必要な処理です。
  • /opt/ml/input配下のconfigとdataから設定情報と入力データを読み込みます。
  • モデルの学習処理を実行します。
  • 学習したモデルデータを/opt/ml/modelに出力します。
  • trainスクリプト実行中に何かしらのエラーが発生した場合には/opt/ml/output/failureにエラー情報を出力します。
  • サンプルスクリプト: amazon-sagemaker-examples/train at master · awslabs/amazon-sagemaker-examples

その他詳細についてはドキュメントをご覧ください。

• Using Your Own Training Algorithms - Amazon SageMaker

推論用コンテナの要件

推論用コンテナでは推論リクエストを裁くためのサーバーを立ち上げる必要があります。サーバーでリクエストを受け取る必要があります。リクエストは8080番ポートに投げられます。

リクエストの種類

推論用エンドポイントへのリクエストの種類は以下の通りです。

• /ping(GET): ヘルスチェック用に定期的にリクエストが投げられます。200を返すことで正常だと判定されます。
  • コンテナ起動後最初に200が返ったタイミングで、推論リクエストの受付が開始します。コンテナ起動後30秒以内に200リクエストが返らなければ、コンテナ起動リクエストは失敗します。
• /invocations(POST): 推論するために入力データを投げて、推論結果を受け取ります。リクエストやレスポンスの形式はアルゴリズムによって変わります。
  • 推論リクエスト時の処理時間は60秒以内に抑える必要があります。推論処理が重い場合で60秒に収まらない場合はインスタンスタイプを変更するか、処理を改善する必要があります。

ファイル構成

知っておくと良さそうな最低限のファイル構成は以下の通りです。

/opt
    ├── program
    │   ├─ predict.py
    │   ├─ wsgi.py
    │   ├─ nginx.conf
    │   └─ serve
    └──ml
        └── model
            └── <model files>

• /opt/program: 推論用コンテナイメージを作成する際に配置する必要があるフォルダです。ノートブックの例に合わせて/opt/programとしていますが、コンテナ起動時に実行できる場所であれば特にパスの指定はありません。(≒Dockerfileでパス通しておく必要があります)
  • /opt/program/serve: コンテナ起動時に実行されるスクリプトファイルです。サーバー起動処理などを記述します。
    • shebangなどの方法でファイル名で実行可能にする必要があります。
  • /opt/program/wsgi.py,/opt/program/nginx.conf: 推論リクエストやヘルスチェックなどのリクエストを受けるために何かしらのサーバーを起動する必要があります。
    • 今回の例ではウェブサーバーとしてNginxを想定しているため、nginx.confとしています。
    • 今回の例ではアプリケーションサーバーとしてgunicornを想定し、アプリケーションへのキック用にwsgi.pyを使うとしています。
    • 今回紹介しているのはあくまで一例なので、同じファイル名である必要もないですし、その他のサーバーでも問題ありません。
  • /opt/program/predict.py:アプリケーションサーバーから呼び出される、ウェブアプリケーションの処理。
    • ヘルスチェック時や推論時などの処理を記述する。
    • 今回紹介している内容は一例なので、構成に応じて複数のファイルを配置したり、別のファイル名でも問題ありません。
• /opt/ml: SageMakerによって作成されるファイル群です。（コンテナイメージ作成時には含める必要がないです）
  • /opt/ml/model/: 推論に使用するモデルデータが保存されています。SageMakerによってコンテナ起動前にtar形式のモデルデータがS3からコピーされて、展開されます。

基本的な処理の流れ

理解を深めるために基本となる流れを書き出してみます。

• SageMakerによってインスタンスが準備されて、その中に学習用コンテナイメージが起動されます。その際にserveコマンドが渡されます。
  • 想定: docker run イメージ名 serve
• コンテナ内でserveスクリプトが実行されます。
  • Webサーバーとアプリケーションサーバーを起動します。
• SageMakerが/pingにGETリクエストを定期的に投げます。その際に200が返ってきたら準備完了したとみなして、推論エンドポイントへの推論リクエストの受け付けを開始します。
• ユーザーが/invocationsにPOSTリクエストを投げることで推論処理をリクエストし、推論結果がレスポンスとして返されます。

その他詳細についてはドキュメントをご覧ください。

• Using Your Own Inference Code - Amazon SageMaker

やってみる

SageMakerのノートブックインスタンス上で以下のノートブックの内容を試してみます。

• amazon-sagemaker-examples/scikit_bring_your_own.ipynb at master · awslabs/amazon-sagemaker-examples

概要

• 学習/推論兼用のコンテナイメージを作成し、ECRのリポジトリへプッシュします。
• ランダムフォレスト分類モデルの学習します。
  • データセット: Iris
• 学習させたもモデルを推論エンドポイントにホストさせて、推論リクエストを投げて確認します。

学習/推論用コンテナイメージ

まずは今回作成するコンテナイメージに関連するcontainerディレクトリを見ていきます。

container
├── Dockerfile
├── build_and_push.sh
└── decision_trees
    ├── nginx.conf
    ├── predictor.py
    ├── serve
    ├── train
    └── wsgi.py

• Dockerfile: コンテナイメージをビルドする際の設定ファイル
  • Ubuntuのイメージをベースに以下のような処理を行います。
  • pythonやnginxなどといったパッケージやpandasやscikit-learnなどのpythonパッケージのインストールを行います。
  • スクリプトを配置する/opt/programにパスを通したりといった、環境変数の設定を行います。
  • ローカルからコンテナ内にスクリプト群をコピーし、作業ディレクトリを設定します。
• build_and_push.sh: コンテナイメージのビルドとECRのリポジトリへのプッシュを行うスクリプト
  • この後、実行するスクリプトとほぼ同じです。
• decition_trees: コンテナ内の/opt/programに配置するファイルを格納しています。
  • decision_trees/train: 学習時に実行されるスクリプト
    • 配置された入力データとハイパーパラメータを読み込み、ランダムカットフォレストの分類モデルを学習する。
  • decision_trees/serve: 推論エンドポイント作成時に実行されるスクリプト
    • サーバー(Nginxとgunicorn)を起動させます。
  • decision_trees/wsgi.py: gunicornがリクエストを受けた際にpredictor.pyに定義されたアプリケーションを呼び出す
  • decision_trees/nginx.conf: Nginx用の設定ファイル。
  • decision_trees/predictor.py: gunicornから呼び出されるアプリケーションが記述されたスクリプトです。
    • 推論やヘルスチェックリクエストに対応する処理が書かれています。

自分で学習/推論用コンテナを作る際にはtrainとpredictor.py、Dockerfile、build_and_push.shを必要に応じて書き換えるだけで良さそうです。もちろん要件によってはそれだけで済まない場合もあるので、確認/検討は必要です。

では、このファイル構成を踏まえて、以下のスクリプトでコンテナイメージを作成し、ECRのリポジトリへプッシュします。

※以下のスクリプトはbuild_and_push.shとほぼ同じ内容です。違う点はalgorithm_nameを実行時引数をもとに定義するかどうかです。

%%sh

# アルゴリズム名
algorithm_name=decision-trees-sample

# コンテナデータのあるパスまで移動
cd container

# コンテナ起動時に実行するファイルを実行可能にする
chmod +x decision_trees/train
chmod +x decision_trees/serve

# アカウントIDを取得する
account=$(aws sts get-caller-identity --query Account --output text)

# リージョンを取得する。リージョンが取得できなければ、us-west-2とする
region=$(aws configure get region)
region=${region:-us-west-2}

# コンテナイメージを保存する、ECRのリポジトリ名
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# リポジトリが存在しない場合には新たに作成する
aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1

if [ $? -ne 0 ]
then
    aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null
fi

# DockerでECRへログインするためのコマンドを取得し、そのまま実行する
$(aws ecr get-login --region ${region} --no-include-email)


# 学習/推論用コンテナイメージを作成(ビルド)する
docker build  -t ${algorithm_name} .

# 作成したコンテナイメージにECR用のタグを設定する
docker tag ${algorithm_name} ${fullname}

# ECRのリポジトリへコンテナイメージをプッシュ(保存)する
docker push ${fullname}

数分ほどで処理が完了します。

動作確認

先ほどECRのリポジトリへコンテナイメージをプッシュしてしまいましたが、念のためローカルで動作確認を行います。動作確認用のスクリプトがcontainer/local_test/に入っているので、それらを実行することで正常に動くことを確認します。 もしここでエラーが出た場合には対応箇所を修正し、改めてコンテナイメージのビルドとプッシュを行う必要があります。

まずはモデルの学習処理を動作確認してみます。

cd container/local_test/
./train_local.sh イメージ名

上手くいったようです。

次は推論処理の動作確認をしてみます。そのために、先に推論用のコンテナを起動させます。

./serve_local.sh イメージ名

別でターミナルを開いて推論確認用スクリプトを実行し、推論処理の確認を行います。

./predict.sh ./payload.csv

※　学習処理実行時に以下のようなエラーが出ることがあります。

ImportError: libopenblasp-r0-8dca6697.3.0.dev.so: cannot open shared object file: No such file or directory

私の場合はDockerfileのnumpyのバージョン固定を無くすことでエラーが出ずに動作するようになりました。よければ試してみてください。

pip install numpy==1.14.5 scipy scikit-learn pandas flask gevent gunicorn && \
↓
pip install numpy scipy scikit-learn pandas flask gevent gunicorn && \

それでも動かないことがあったので、以下のようにscipy周りのファイル削減用処理をコメントアウトすることで動作するようになりました。

         (cd /usr/local/lib/python2.7/dist-packages/scipy/.libs; rm *; ln ../../numpy/.libs/* .) && \
↓
#         (cd /usr/local/lib/python2.7/dist-packages/scipy/.libs; rm *; ln ../../numpy/.libs/* .) && \

学習

先ほど作成したコンテナイメージを使ってSageMaker上でモデルの学習ジョブを実行します。 今回は使いやすいデータがすでに用意されているので、処理がほとんどありません。データをアップロードして、学習をセットアップし、実行するという流れです。

# 学習用データやモデルデータを保存するS3のパスの接頭辞(バケットはSageMakerのデフォルトバケットを使用)
prefix = 'DEMO-scikit-byo-iris'

# 必要なモジュールの読み込み
import boto3
import re
import os
import numpy as np
import pandas as pd
from sagemaker import get_execution_role
import sagemaker as sage
from time import gmtime, strftime

# 学習やエンドポイントの作成時に使用するIAMロールを指定
role = get_execution_role()

# セッション取得
sess = sage.Session()

#データがあるパスを指定
WORK_DIRECTORY = 'data'

# データをS3へアップロード
data_location = sess.upload_data(WORK_DIRECTORY, key_prefix=prefix)

# アカウントIDとリージョン名を元に学習用コンテナイメージのarnを指定する
# リポジトリ名を変更している場合は"decision-tree-sample"のところを変更する必要があります
account = sess.boto_session.client('sts').get_caller_identity()['Account']
region = sess.boto_session.region_name
image = '{}.dkr.ecr.{}.amazonaws.com/decision-trees-sample:latest'.format(account, region)

# 学習をハンドルするestimatorを作成します。
# 引数などは組み込みアルゴリズムなどを使用する場合とほとんど変わりません。
tree = sage.estimator.Estimator(image,
                       role, 1, 'ml.c4.2xlarge',
                       output_path="s3://{}/output".format(sess.default_bucket()),
                       sagemaker_session=sess)

# 学習ジョブ開始
tree.fit(data_location)

数分でジョブは完了します。

推論

先ほど学習させた分類モデルをエンドポイントに配置して、推論処理が上手くいくか確認します。

まずはモデルをデプロイします。コンテナイメージなどは学習時にestimatorで設定してあるので、ここでの設定は不要です。

from sagemaker.predictor import csv_serializer
predictor = tree.deploy(1, 'ml.m4.xlarge', serializer=csv_serializer)

エンドポイントが起動したら、irisデータを読み込んで推論リクエストを投げてみます。

shape=pd.read_csv("data/iris.csv", header=None)

import itertools

a = [50*i for i in range(3)]
b = [40+i for i in range(10)]
indices = [i+j for i,j in itertools.product(a,b)]

test_data=shape.iloc[indices[:-1]]
test_X=test_data.iloc[:,1:]
test_y=test_data.iloc[:,0]

print(predictor.predict(test_X.values).decode('utf-8'))

推論できているようです。

確認が済んだので、エンドポイントを削除しておきます。

sess.delete_endpoint(predictor.endpoint)

さいごに

Amazon SageMakerでの学習/推論用のコンテナの概要と、コンテナイメージを作って推論するまでの流れについてご紹介しました。サンプルノートブックで作成するコンテナイメージを基にすれば割と簡単に独自の学習/推論用コンテナを作成することができそうです。

お読みいただきありがとうございましたー!