Deep Learning AMIをECSで使えるようにしてみた

ECSやAWS Batchで機械学習などのGPUを使った処理を実行できるように、Deep Learning AMIをECSで使えるようにしたAMIを作成します。

基本的にはAWSのドキュメントに従って進めていきます ^*1。

• Creating a GPU Workload AMI - AWS Batch

やってみた

インスタンスを作成

Deepe Learning AMIでEC2インスタンスを作成します。

• スポットインスタンス(費用節約のためスポットインスタンスを選択していますが、オンデマンドでも問題ありません。)
• インスタンスタイプ: p2.xlarge(GPUがあるインスタンスタイプであればOK)
• AMI: Deep Learning AMI
• IPv4 パブリック IP の自動割り当て: 有効化

セットアップ

先ほど立ち上げたインスタンスにsshで接続して、作業を行なっていきます。

まずは、次のようなスクリプトファイルconfigure-gpu.shを作成します。 ecs-initとNVIDIA Dockerをインストールして、DockerのランタイムのデフォルトをNVIDIA Dockerに設定するという内容です。

#!/bin/bash

# ecs-initをインストール
sudo yum install -y ecs-init

# NVIDIA Dockerをインストール
sudo service docker start
DOCKER_VERSION=$(docker -v | awk '{ print $3 }' | cut -f1 -d"-")
DISTRIBUTION=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/$DISTRIBUTION/nvidia-docker.repo | \
  sudo tee /etc/yum.repos.d/nvidia-docker.repo
PACKAGES=$(sudo yum search -y --showduplicates nvidia-docker2 nvidia-container-runtime | grep $DOCKER_VERSION | awk '{ print $1 }')
sudo yum install -y $PACKAGES
sudo pkill -SIGHUP dockerd

# CUDA versionを取得
CUDA_VERSION=$(cat /usr/local/cuda/version.txt | awk '{ print $3 }' | cut -f1-2 -d".")

# コンテナを起動して動作確認
sudo docker run --privileged --runtime=nvidia --rm nvidia/cuda:$CUDA_VERSION-base nvidia-smi

# デフォルトでnvidia-container-runtimeを使用するようにDocker daemon.jsonを更新
sudo tee /etc/docker/daemon.json <<EOF
{
    "runtimes": {
        "nvidia": {
            "path": "/usr/bin/nvidia-container-runtime",
            "runtimeArgs": []
        }
    },
    "default-runtime": "nvidia"
}
EOF

# Dockerを再起動
sudo service docker restart

作成したconfigure-gpu.shを実行します。

bash ./configure-gpu.sh

動作確認のために、対応したバージョンのCUDAのコンテナを特権モードで起動して、nvidia-smiを実行します。

CUDA_VERSION=$(cat /usr/local/cuda/version.txt | awk '{ print $3 }' | cut -f1-2 -d".")
sudo docker run --privileged --rm nvidia/cuda:$CUDA_VERSION-base nvidia-smi

下のような感じでGPU情報が表示されればOKです。

AMIの容量を減らすためにDockerコンテナとイメージを削除します。 これまでに書いてあった内容以外にコンテナを起動させていない場合は、Dockerコンテナの削除コマンド実行時にエラーが出ますが、問題ありません。

sudo docker rm $(sudo docker ps -aq)
sudo docker rmi $(sudo docker images -q)

ECSエージェントを動かしている場合には、AMIを作る前にECSエージェントを停止させて永続的なデータチェックポイントを削除する必要があります。データチェックポイントが残った状態で作成したAMIを使って、インスタンスを作成してもECSエージェントが動きません。 ※これまでに書いてあった内容以外の操作をしていない場合（ecsエージェントを動かしていない場合）は不要な作業です。

sudo stop ecs
sudo rm -rf /var/lib/ecs/data/ecs_agent_data.json

これでセットアップは終了です。

AMIを作成する

マネジメントコンソールから作業していたインスタンスのAMIを作成すれば、Deep Learning AMIを基にしたECSに対応したAMIの作成は完了です！ 作成したAMIをインスタンスのカスタムイメージとして使用することで、ECSやAWS BatchでGPUを使用できるようになります。

• Amazon ECS コンテナインスタンスの起動 - Amazon Elastic Container Service
• コンピューティング環境の作成 - AWS Batch

インスタンスの削除

無駄なお金が費用がかからないように、作業用インスタンスは忘れずに削除しておきましょう。

さいごに

Deep Learning AMIを元にECSに対応した機械学習用のAMIを作成しました。ドキュメントに従うことで簡単にできました。 このAMIを使ってAWS Batchで機械学習ジョブを回す方法についてもいずれご紹介したいと思います。

お読みくださり有難うございました〜！

脚注