概要

AWS の Sagemakerの推論エンドポイントで、
デフォルトでは入っていないライブラリを追加する方法の説明です。
(多分推論エンドポイント以外でも使えます）

AWS公式の機械学習フレームワーク用Dockerイメージに
任意のライブラリを追加でインストールし、新しいDockerイメージを作ります。

方法

1. Sagemakerのノートブックインスタンスを起動する

前準備として、作業をするためのノートブックインスタンスを起動します。
ローカルで実行することも可能かもしれませんが、imageを取得したりプッシュする時の読み込みスピードが速いので、ノートブックインスタンスを利用することをおすすめします。

以下の内容でノートブックインスタンスを起動してください。

リージョン：
これから作成するimageを使用したいリージョン

IAMロール：
「AmazonSageMaker-ExecutionRole」（「新しいロールを作成する」で作成されるもの）に
「EC2InstanceProfileForImageBuilderECRContainerBuilds」ポリシーをアタッチしたもの
（これをアタッチすることによって、作成したimageをECRにプッシュすることができるようになります。）

2. 拡張したいコンテナのURLを探す

以下のサイトから自分が拡張したいコンテナのURLを選びます。
deep-learning-containers/available_images.md at master · aws/deep-learning-containers · GitHub

大体の場合はGeneral Framework Containersの中から選べばいいと思います。
Elastic Inferenceを使いたい場合は多分Elastic Inference Containersから選べばOKです。
学習時に使うコンテナの場合はJob Typeがtrainingのものを、推論時に使うならinferenceのものを選んでください。
また、自分の使うpythonのバージョンに合ったものを選んでください。

ページ上部にあるリージョンの表から、作ったコンテナを使いたいリージョンを選択し、
URLの前半部分は該当するURLに差し替えてください。

例）東京リージョンで使う予定で、CPUインスタンスでTensorflow2.3.1を使って推論を実行したいなら
763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/tensorflow-inference:2.3.1-cpu-py37-ubuntu18.04

3. Dockerfileの作成

ノートブックインスタンスを起動したら、コンテナを作るためのDockerfileを作成します。
※Dockerfileと同じディレクトリにはDocker関係のファイル以外は置かない方がいいです。（docker buildを実行する際、ディレクトリ内にある全てのファイルの情報がイメージに送られるため、同じディレクトリに余計なファイルがあるとその分buildに時間がかかります。）

Dockerfile

FROM <2.で選んだURL>

RUN pip install <入れたいパッケージ>
RUN apt-get install -y <入れたいもの>

上記のように、
拡張したいdocker imageをFROMで取得してから
RUNを使って追加したいパッケージを入れるようにします。

4. コンテナの構築とECRへのアップロード

Dockerfileと同じディレクトリに
以下のシェルスクリプトファイルを作成します。
（image_nameには自分が分かりやすい名前を指定してください。（4.でimage_uriを指定する時に使います））

build_docker_image.sh

# sagemakerの公式コンテナを拡張したDockerイメージをノートブックインスタンス上でビルドするためスクリプト。
# 主にイメージの動作確認のために用いるものです。
# (Dockerfileがあるディレクトリに移動した状態で実行してください。)

image_name=tensorflow231cpu-cv2-numpy-boto3  # 新しく作るイメージの名前(任意)
region=ap-northeast-1  # DockerfileでFROMで取得するイメージのリージョン

# ECRにログイン（AWS公式のECRレジストリにログインします）
official_registry=763104351884.dkr.ecr.${region}.amazonaws.com
aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${official_registry}

# 拡張されたイメージをノートブックインスタンス上のDockerでビルドします
docker build -t ${image_name} .

# docker loginによってノートブックインスタンス内のディレクトリに保存された認証情報を削除します（インスタンス停止時には削除されますが念のため）
docker logout ${official_registry}

build_push_docker_image.sh

# sagemakerの公式コンテナを拡張したDockerイメージを作成し、ECRにプッシュするためのスクリプト。
# イメージを利用したいアカウント&リージョンでこのスクリプトを実行することで、ECRにイメージをプッシュして利用可能にできます。
# (Dockerfileがあるディレクトリに移動した状態で実行してください。)


image_name=tensorflow231cpu-cv2-numpy-boto3  # 作成するイメージの名前（任意）

region=$(aws configure get region)  # 現在のリージョン (DockerfileでFROMで取得するイメージのリージョンと同じである必要があります)
account=$(aws sts get-caller-identity --query Account --output text)  # 現在使用しているアカウントのID


# *************** 1. 拡張イメージのビルド ***************
# ECRにログイン（AWS公式のECRレジストリにログインします）
official_registry=763104351884.dkr.ecr.${region}.amazonaws.com
aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${official_registry}

# 拡張されたイメージをノートブックインスタンス上のDockerでビルドします
docker build -t ${image_name} .

# docker loginによってノートブックインスタンス内のディレクトリに保存された認証情報を削除します（インスタンス停止時には削除されますが念のため）
docker logout ${official_registry}


# *************** 2. ビルドしたイメージをECRにプッシュ ***************
# ECRにログイン（作成したイメージを利用したいアカウントのECRレジストリにログインします）
my_registry=${account}.dkr.ecr.${region}.amazonaws.com
aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${my_registry}

# push先のリポジトリがECRにない場合は新しく作成します
aws ecr describe-repositories --repository-names "${image_name}" > /dev/null 2>&1
if [ $? -ne 0 ]
then
aws ecr create-repository --repository-name "${image_name}" > /dev/null
fi

# 作成したイメージをfullnameに命名しなおします
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${image_name}:latest"  # プッシュするイメージの最終的な名前
docker tag ${image_name} ${fullname}

# 作成したイメージをECRにプッシュします
docker push ${fullname}

# docker loginによってノートブックインスタンス内のディレクトリに保存された認証情報を削除します（インスタンス停止時には削除されますが念のため）
docker logout ${my_registry}

build_docker_image.sh の方を実行すると、
Dockerfileに基づいて拡張されたDocker imageがノートブックインスタンス内のDockerで作成されます。
（pushする前にimageの内容を確認したい場合などに使います）

build_push_docker_image.sh の方を実行すると、
ノートブックインスタンス内のDockerでimageをビルドした後、そのimageをそのままECRにプッシュします。
これを実行することによってsagemakerから拡張imageを使えるようになります。

上記のスクリプトを実行するには、
まずターミナルを起動してDockerfileや上記のシェルスクリプトがあるディレクトリに移動します。

※この時、Jupyterlabでデフォルトで表示されているディレクトリは
ターミナルで最初にいるディレクトリとは異なるので注意してください。

cd SageMaker

でSageMakerディレクトリに移動するとjupyterlabのデフォルトと同じディレクトリに移動できます。

移動したら、

sh build_push_docker_image.sh

でシェルスクリプトを実行してください。

4. 作ったコンテナをSagemakerで使用する

ECRにimageをプッシュしたら、今後はそのリージョンでSagemakerのサービスを使う時に利用できるようになります。
例えば、Tensorflowの学習済みモデルを使ってSagemakerのモデルを作りたい場合、
モデルのオブジェクトを作る時のimage_uri引数として
さっきのfullnameにあたるURLを指定してやれば使えます。
（ECRコンソールのプライベートリポジトリからもプッシュしたimageのURIを確認できます。）

from sagemaker.tensorflow.model import TensorFlowModel
model = TensorFlowModel(model_data="hoge.tar.gz",
                        role=role,
                        image_uri="<自分のアカウント番号>.dkr.ecr.<2.で選んだリージョン>.amazonaws.com/<4.のimage_name>")

参考ページ

SageMaker関連

コンテナの拡張方法を説明したAWS公式のガイド
構築済みのコンテナの拡張 - Amazon SageMaker
公式コンテナのURLのリスト（2.で使ったものです）
deep-learning-containers/available_images.md at master · aws/deep-learning-containers · GitHub

以下のサイトでは、
TensorFlowModelクラスなど、
Sagemakerのモデルを作る時に使うクラスの仕様（引数などについて）が解説されてます。

全てのフレームワークで共通の仕様
Model — sagemaker 2.26.0 documentation
各フレームワークごとの仕様
Frameworks — sagemaker 2.26.0 documentation
- 例えば、Tensorflowを使うなら以下のページ（それ以外のフレームワークにも同様のページがあります）
  [TensorFlow — sagemaker 2.26.0 documentation]
  (https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/sagemaker.tensorflow.html)

Docker関連

Dockerfileの記述方法について
Dockerfileの書き方, 利用する命令, 作成手順 - わくわくBank
docker build時に同じディレクトリに余計はファイルを置かない方が良い理由について
(「Dockerデーモンに送るtarファイルが大きい」を参照してください）
docker buildを速くするコツ - keroxpのScrapbox
docker loginやDockerhub、AWS ECRについての説明
さわって理解するDocker入門第3回 | オブジェクトの広場