記事一覧

Amazon SageMaker Training で Bring Your Own Container する

今まで、 AWS が用意していたコンテナイメージを利用して Training Job を起動していました。
具体的には TensorFlow, PyTorch, MXNet, Scikit-Learn のコンテナを利用する例を紹介しました。

では、独自コンテナを使う場合はどうなんでしょうか？というのを試してみます。

例にもれず、使用しているコードは GitHub にありますので clone してお使いください。

Hello My Container

早速 Hello World 相当のことをしてみましょう。
このようなコードと Dockerfile を用意してみました。

./container/3-1/Dockerfile

FROM python:3.9.10-slim-buster
COPY ./hello_my_container_image.py /
ENTRYPOINT ["python","/hello_my_container_image.py"]

./container/3-1/hello_my_container_image.py

print('Hello my container .')
exit()

python の 3.9 slim-buster のコンテナイメージで、 Hello my container image. と出力するだけの簡単なコードをコンテナの中に入れて、ENTRYPOINT で実行しなさい、と指定しているだけです。

さて、このコンテナを使って SageMaker Training を実行してみたいのですが、SageMaker Training で自分のコンテナイメージを使うには、
まず Amazon Elastic Container Registry (ECR) にプッシュする必要があります。ビルドしてプッシュしましょう。

ビルド

./3_byoc.ipynb

IMAGE_NAME='sagemaker-slim-buster'
TAG = ':python39'

%cd container/3-1
!docker build -t {IMAGE_NAME}{TAG} . 
%cd ../../

プッシュ

./3_byoc.ipynb

# boto3の機能を使ってリポジトリ名に必要な情報を取得する
import boto3
MY_ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account')
REGION = boto3.session.Session().region_name
# ECR へ Docker Image の push
MY_ECR_ENDPOINT = f'{MY_ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/'

MY_REPOSITORY_URI = f'{MY_ECR_ENDPOINT}{IMAGE_NAME}'
MY_IMAGE_URI = f'{MY_REPOSITORY_URI}{TAG}'

!$(aws ecr get-login --region {REGION} --registry-ids {MY_ACCOUNT_ID} --no-include-email)
 
# リポジトリの作成
!aws ecr delete-repository --repository-name {IMAGE_NAME} --force # 同名のリポジトリがあった場合削除
!aws ecr create-repository --repository-name {IMAGE_NAME}
 
!docker tag {IMAGE_NAME}{TAG} {MY_IMAGE_URI}
!docker push {MY_IMAGE_URI}

print(f'コンテナイメージは {MY_IMAGE_URI} へ登録されています。')

これらを実行することによって ECR へコンテナイメージをプッシュできました。
最後に実行してみます。

./3_byoc.ipynb

from sagemaker.estimator import Estimator

estimator = Estimator(
    image_uri = MY_IMAGE_URI,
    role=sagemaker.get_execution_role(),
    instance_count=1,
    instance_type='ml.m5.xlarge',
)
estimator.fit()

出力結果

(略)hello my container .
(略)

無事 Hello my container image. という出力を得ることが出来ました。
こんな感じで、エンジニア目線で始める Amazon SageMaker Training ①機械学習を使わないはじめての Training Job の表現に戻りますが、
「Amazon SageMaker Training とは、①用意したコードを ~~②用意したデータと~~ ③用意した環境で実行してくれ ~~、 ④結果を自動で保存してくれ~~ る、バッチ処理サービスです。」
が体感できました。
環境を準備し、コンテナ内に処理したい内容をコードに記載して入れてしまえば実行ができることがわかりました。

独自コンテナで SageMaker の機能を使うためには

Hello my container の例では自前で作ったコンテナイメージを SageMaker Training ジョブで動かせましたが、SageMaker Training では、コンテナの外(S3)からコードとデータを与え、成果物を S3 に出力する、という機能を使いたいですが、このままではその機能を利用できません。
やり方は公式 Doc に記載があります。

Adapting Your Own Training Container

詳細は Doc を読んでいただきたいのですが、肝は Dockerfile の以下の部分です。

RUN pip3 install sagemaker-training

sagemaker-training をインストールすることで、記事①で紹介した機能を利用できるようになります。
さて、実際にやってみましょう。

./container/3-2/Dockerfile

FROM python:3.9.10-slim-buster
RUN apt update -y && \
    apt install gcc -y && \
    pip3 install sagemaker-training

apt で gcc をインストールしてますが、sagemaker-training のインストールに gcc が必要だったためインストールしておりますが、ご覧の通り先程利用した python:3.9.10-slim-buster に sagemaker-training をインストールしているだけです。

先程同様ビルドとプッシュをしておきます。

./3_byoc.ipynb

IMAGE_NAME='sagemaker-slim-buster'
TAG = ':python39-sagemaker-training'

%cd container/3-2
!docker build -t {IMAGE_NAME}{TAG} . 
%cd ../../

# boto3の機能を使ってリポジトリ名に必要な情報を取得する
import boto3
MY_ACCOUNT_ID = boto3.client('sts').get_caller_identity().get('Account')
REGION = boto3.session.Session().region_name
# ECR へ Docker Image の push
MY_ECR_ENDPOINT = f'{MY_ACCOUNT_ID}.dkr.ecr.{REGION}.amazonaws.com/'

MY_REPOSITORY_URI = f'{MY_ECR_ENDPOINT}{IMAGE_NAME}'
MY_IMAGE_URI = f'{MY_REPOSITORY_URI}{TAG}'

!$(aws ecr get-login --region {REGION} --registry-ids {MY_ACCOUNT_ID} --no-include-email)

!docker tag {IMAGE_NAME}{TAG} {MY_IMAGE_URI}
!docker push {MY_IMAGE_URI}

print(f'コンテナイメージは {MY_IMAGE_URI} へ登録されています。')

さて、このコンテナイメージで、記事①の 1-5 で行ったコードとデータを持ち込んで成果物を S3 に保存を実行してみます。

実行

./3_byoc.ipynb

input_s3_uri = sagemaker.session.Session().upload_data(path='./data/1-5/', bucket=sagemaker.session.Session().default_bucket(), key_prefix='training/1-5')
from sagemaker.tensorflow import TensorFlow
estimator = TensorFlow(
    entry_point='./src/1-5/concat.py',
    image_uri=MY_IMAGE_URI,
    instance_count=1,
    instance_type='ml.m5.xlarge',
    role=sagemaker.get_execution_role()
)
estimator.fit(input_s3_uri)

結果確認

./3_byoc.ipynb

import tarfile
prefix = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts'].replace(f's3://{sagemaker.session.Session().default_bucket()}/','')
sagemaker.session.Session().download_data('./', sagemaker.session.Session().default_bucket(), key_prefix=prefix)
with tarfile.open('./model.tar.gz', 'r') as f:
    f.extractall()
!cat output.csv

出力結果

1,2
3,4
5,6
7,8

AWS の SageMaker Training のコンテナを拡張する場合

先程の例では slim-buster をベースイメージとして実行した例ですが、SageMaker Training のコンテナを拡張したいケースもあるかと思います。
その場合は、sagemaker-training 相当のものはすでに入っているのでインストール不要ですので、必要なモジュールを追加を Dockerfile に記載するだけで OK です。

詳細はこちらをご覧ください。(コードを紹介するまでもないので…)