背景

AWS Batchでは、Jobの定義とAMIの指定を行うことで、

自動でEC2インスタンスを立ち上げ
Dockerコンテナを起動
コンテナ終了時にEC2インスタンスの自動削除
を行ってくれます。

しかし、EC2インスタンスの料金体系では、毎回インスタンスを立ち上げ・削除するのはコスパが悪いため、Jobが実行されるEC2インスタンスを固定して、実行中以外は削除ではなく、「停止」にしたい！

ということで、その手順をまとめました。
AWS Batch、あるいはAmazon ECS（Elastic Container Service）でGPUを使うための設定も含まれていますので、インスタンスの固定はいいや、という方にも参考になれば、と思います。

AWS Batchの設定

Compute environments（=ECS Cluster）

AWS Batch > Compute environments > Create environment から、AWS BatchのJobを実行する計算環境を作成します。この計算環境は、ECSのClusterとして保存されます。（ECS Clusterが自動で生成され、Compute environmentsに紐付けされます。）

Managed:
AMIを指定しておくことで、Jobの実行時にそのAMIで新たにEC2インスタンスを立ち上げてくれます。
Unmanaged:
Compute environmentsに紐付いたECS Clusterを生成します。EC2インスタンスは、自分でそのClusterに登録する必要があります。

作成されたClusterをECSコンソールから確認しましょう。後ほど使います。
<EnvironmentName>_Batch_<ランダムっぽい文字列> のような名前担っているはずです。

Job Queue

Job Queueは、上記で作成したEnvironmentを選択してください。

##Job Definition

Volumes

Name	Source path
nvidia	/var/lib/nvidia-docker/volumes/nvidia_driver/latest

Mount points

Container path	Source volume	Read only
/usr/local/nvidia	nvidia	false

ECSで使える（Dockerが動かせる）EC2インスタンスを用意する

ECS用に最適化されたAMIはAWSから提供されていますが、GPUが載ったものでECS用に最適化されたAMIは提供されていません。

用いるAMIを決める

AWS Batchのドキュメント GPU ワークロードの AMI の作成にある例では Deep Learning AMI with Source Code (CUDA 9, Amazon Linux) を使用していますが、今回は Deep Learning Base AMI (Amazon Linux) を使用することにしました。
（AWS Batchを用いる = Dockerを用いるので、Hostマシンに機械学習ライブラリは不要。 & Cuda, CuDNNのバージョンが複数あるので、動かすDockerImageが変わった時に助かりそう。）
ここは、必要に応じてaws marketplaceでいろいろ探してみるとよいかと思います。

インスタンスを起動する

IAM

インスタンスからECSに向けたIAM Roleが必要です。Roleの設定がない場合は、EC2コンソールから設定をしましょう。

UserData

UserData

#!/bin/bash
echo "ECS_CLUSTER=<ClusterName>" >> /etc/ecs/ecs.config

Dockerが使えるようにする

上記のAMIのままでは、Dockerが使えないので、ドキュメントを参考に、ECSで使える状態のインスタンスを作成していきます。

configure-gpu.sh

#!/bin/bash
# Install ecs-init, start docker, and install nvidia-docker
sudo yum install -y ecs-init
sudo service docker start
wget https://github.com/NVIDIA/nvidia-docker/releases/download/v1.0.1/nvidia-docker-1.0.1-1.x86_64.rpm
sudo rpm -ivh --nodeps nvidia-docker-1.0.1-1.x86_64.rpm

# Validate installation
rpm -ql nvidia-docker
rm nvidia-docker-1.0.1-1.x86_64.rpm

# Make sure the NVIDIA kernel modules and driver files are bootstraped
# Otherwise running a GPU job inside a container will fail with "cuda: unknown exception"
echo '#!/bin/bash' | sudo tee /var/lib/cloud/scripts/per-boot/00_nvidia-modprobe > /dev/null
echo 'nvidia-modprobe -u -c=0' | sudo tee --append /var/lib/cloud/scripts/per-boot/00_nvidia-modprobe > /dev/null
sudo chmod +x /var/lib/cloud/scripts/per-boot/00_nvidia-modprobe
sudo /var/lib/cloud/scripts/per-boot/00_nvidia-modprobe

# Start the nvidia-docker-plugin and run a container with 
# nvidia-docker (retry up to 4 times if it fails initially)
sudo -b nohup nvidia-docker-plugin > /tmp/nvidia-docker.log
sudo docker pull nvidia/cuda:9.0-cudnn7-devel
COMMAND="sudo nvidia-docker run nvidia/cuda:9.0-cudnn7-devel nvidia-smi"
for i in {1..5}; do $COMMAND && break || sleep 15; done

# Create symlink to latest nvidia-driver version
nvidia_base=/var/lib/nvidia-docker/volumes/nvidia_driver
sudo ln -s $nvidia_base/$(ls $nvidia_base | sort -n  | tail -1) $nvidia_base/latest

スクリプトの実行（Step 4.）

bash ./configure-gpu.sh

これで環境は完成していますが、念のため動作確認（Step 5. ）

sudo docker run --privileged -v /var/lib/nvidia-docker/volumes/nvidia_driver/latest:/usr/local/nvidia nvidia/cuda:9.0-cudnn7-devel nvidia-smi

DockerからGPUが使えることを確認できたら（以下のような $ nvidia-smi の結果が見られれば） OKです。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 384.81                 Driver Version: 384.81                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:00:17.0 Off |                    0 |
| N/A   43C    P0    42W / 300W |     10MiB / 16152MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

ECS Clusterにインスタンスを登録する

Check 1
インスタンスからECSに向けたIAM Roleが必要です。Roleの設定がない場合は、EC2コンソールから設定をしましょう。

Check 2
Clusterにインスタンスを登録するには、インスタンス側での設定が必要です。インスタンス起動時のユーザデータで下記ファイルがきちんと作られていることを確認。できていなければ作りましょう。

/etc/ecs/ecs.config

ECS_CLUSTER=<ClusterName>

ドキュメントのStep 6., 7.的なことをしてあげると、Clusterにインスタンスが登録されます。（しばらく時間がかかります）

sudo docker rm $(sudo docker ps -aq)

sudo docker rmi $(sudo docker images -q)

sudo restart ecs

Submit Job

お好きなフックでどうぞ

EC2インスタンスを停止する

Managedであれば自動でインスタンスを削除してくれますが、Unmanagedの場合、インスタンスの停止も自分で行わなければなりません。
今回は、Jobの成功、失敗にかかわらず停止したいので、以下のpythonスクリプトを自作し、AWS Lambda x CloudWatchEventsで走らせることにしました。

['SUBMITTED', 'PENDING', 'RUNNABLE', 'STARTING', 'RUNNING'] のいずれかの状態のJobがあればそのまま。
いずれの状態のJobもなければ、停止。

import os
import boto3

JOB_NAME = os.environ['JOB_NAME']
JOB_QUEUE = os.environ['JOB_QUEUE']
INSTANCE_ID = os.environ['INSTANCE_ID']


def is_running():
    batch = boto3.client('batch')

    check_status = ['SUBMITTED', 'PENDING', 'RUNNABLE', 'STARTING', 'RUNNING']
    for s in check_status:
        list = batch.list_jobs(
            jobQueue=JOB_QUEUE,
            jobStatus=s
        )['jobSummaryList']
        if list:
            print(f"{[j['jobId'] for j in list]} are still {s}.")
            return True
    else:
        return False


def ec2_stop():
    ec2 = boto3.client('ec2')

    response = ec2.describe_instances(
        InstanceIds=[INSTANCE_ID]
    )
    if response['Reservations'] and response['Reservations'][0]['Instances']:
        state = response['Reservations'][0]['Instances'][0]['State']['Name']

        if state in {'pending', 'running'}:
            print(f"Start to stop the instance: {INSTANCE_ID}")
            response = ec2.stop_instances(
                InstanceIds=[INSTANCE_ID]
            )
        else:
            print(f"The instance {INSTANCE_ID} is not running.")
    else:
        raise ValueError(f'No Such Instances: [{INSTANCE_ID}]')

def lambda_handler(event, context):
    if not is_running():
        ec2_stop()

最後に

こんなことするくらいなら、最初からBatchじゃなくて、ECSにしとけばよかったんちゃうか、と思いながら。
ただ新しいものに触れて、上記以外の学びも得たので、それはまた別の時に。