AWS Batchで固定されたEC2インスタンスのGPUを使う

# 背景
AWS Batchでは、Jobの定義とAMIの指定を行うことで、
1. 自動でEC2インスタンスを立ち上げ
2. Dockerコンテナを起動
3. コンテナ終了時に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 の作成](http://docs.aws.amazon.com/ja_jp/batch/latest/userguide/batch-gpu-ami.html) にある例では [Deep Learning AMI with Source Code (CUDA 9, Amazon Linux)](https://aws.amazon.com/marketplace/pp/B076T8RSXY) を使用していますが、今回は [Deep Learning Base AMI (Amazon Linux)](https://aws.amazon.com/marketplace/pp/B077GFM7L7) を使用することにしました。
（AWS Batchを用いる = Dockerを用いるので、Hostマシンに機械学習ライブラリは不要。 & Cuda, CuDNNのバージョンが複数あるので、動かすDockerImageが変わった時に助かりそう。）
ここは、必要に応じてaws marketplaceでいろいろ探してみるとよいかと思います。

## インスタンスを起動する
### IAM
インスタンスからECSに向けたIAM Roleが必要です。Roleの設定がない場合は、EC2コンソールから設定をしましょう。
#### UserData
``` sh:UserData
#!/bin/bash
echo "ECS_CLUSTER=<ClusterName>" >> /etc/ecs/ecs.config
```

## Dockerが使えるようにする
上記のAMIのままでは、Dockerが使えないので、[ドキュメント](http://docs.aws.amazon.com/ja_jp/batch/latest/userguide/batch-gpu-ami.html)を参考に、ECSで使える状態のインスタンスを作成していきます。

```sh: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にインスタンスを登録するには、インスタンス側での設定が必要です。インスタンス起動時のユーザデータで下記ファイルがきちんと作られていることを確認。できていなければ作りましょう。

``` sh:/etc/ecs/ecs.config
ECS_CLUSTER=<ClusterName>
```

[ドキュメント](http://docs.aws.amazon.com/ja_jp/batch/latest/userguide/batch-gpu-ami.html) の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もなければ、停止。

```python3
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にしとけばよかったんちゃうか、と思いながら。
ただ新しいものに触れて、上記以外の学びも得たので、それはまた別の時に。