AI Platform(GCP)でGPU 100個同時に使いテンションあがった

<img width="600" alt="aiplatform_jobs.gif" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/47605/b4052100-ec93-d74f-0914-a7676a46afcb.gif">

ひゃっはー

# 概要

機械学習基盤系のサービスとしてAWSのSageMakerが有名かと思いますが
本投稿ではGCPの類似サービスAI Platformを使ってGPU 100並列同時に動かした話をします。

# AI Platformでできること

SageMakerと似た感じで、AI Platformでは以下のようにうたっています。
（[公式ドキュメント](https://cloud.google.com/ml-engine/docs/technical-overview?hl=ja)より引用）

> AI Platform を使用して、ML ワークフローの次の段階を管理できます。
![ml-workflow.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/47605/e9e98fae-203a-67cd-e5f2-f1917d6d1af2.png)

> * データで ML モデルをトレーニングする。
    - モデルをトレーニング
    - モデルの精度を評価する
    - ハイパーパラメータを調整する
* トレーニング済みモデルをデプロイする。
* モデルに予測リクエストを送信する。
    - オンライン予測
    - バッチ予測（TensorFlow のみ）
* 継続的に予測をモニタリングする。
* モデルとモデル バージョンを管理する。

いわゆる機械学習基盤として使えますよ便利ですね、ってサービスなんですが
単に時間のかかるバッチジョブ（MLの場合、特に学習ジョブ）をサーバレスで走行する環境として使うだけでも割とその恩恵を受けられます。
（ ~~というかマジメに機械学習基盤として使うにはちょっとまだ機能連携が足りていないような...~~ ）

GPUをGCEなどのIaaSで利用するとインスタンスの起動・停止がめんどいので
学習が走行されている時間だけ課金されるマネージドなジョブ実行ができるのはありがたいです。

バッチを走行するだけなら他にも方法あるやん、って思われるかもしれないですが
意外と「GPUが使える」かつ「時間制限がない」ってサービスはあまりなくて
他の選択肢だとGKEでk8s jobくらいなのかな、と思います（他にもあったらごめんなさい...）

なので、GKEよりは手軽にML学習のジョブを走らせられるのは、それだけでもメリットと感じます。

# GPUの同時利用上限

AI Platformには同時に利用できるGPU数に上限が設けられています。

> * Tesla K80 GPUs: 30
* Tesla P100 GPUs: 30
* Tesla V100 GPUs: 8

SageMakerのデフォルト上限は下記のようです。（最新版は[公式ドキュメント](https://docs.aws.amazon.com/ja_jp/general/latest/gr/sagemaker.html)を参照してください）

| インスタンスタイプ            | GPU     | 同時利用インスタンス数上限 |
| :--:| :--:|:--:|
| ml.p2.xlarge      | K80 * 1 | 1  |
| ml.p2.8xlarge     | K80 * 8 | 1  |
| ml.p2.16xlarge    | K80 * 16 | 1 |
| ml.p3.2xlarge     | V100 * 1 | 2 |
| ml.p3.8xlarge     | V100 * 4 | 2 |
| ml.p3.16xlarge    | V100 * 8 | 2 |

合計のGPU数はさほどAI Platformと変わらないように思えますが
1つのインスタンスにGPUが複数搭載されていても、GPU分散で学習させないと意味がないです。
そういう意味だと、1GPUのジョブを何個も同時に実行したい場合にはSageMakerのデフォルトではちょっと足りないです。
（もちろんSageMakerでも上限を増やす申請はできます）

# 上限を増やす

通常、GCPの上限を増やす際の申請は、GCPコンソール上の「IAMと管理」=>「割り当て」から行うのですが
AI PlatformのGPU上限はなぜか別管理されており、[ここのGoogleフォーム](https://docs.google.com/a/google.com/forms/d/1f-tmYYl8uiV6KcWFLsNKwWnTvk4KwQYsCzZHcb3g2YM/viewform?hl=ja)から行います。
（[公式ドキュメント](https://cloud.google.com/ml-engine/docs/quotas#requesting_a_quota_increase)参照）

で、私はここで `"70 P100s", "30 V100s"` と申請し、無事通りました。
（申請から承認まで大体1週間くらいかかりました）

これでGPUめっちゃ同時に使う環境が手に入りました。


# ユースケース

ところでGPUを同時にたくさん利用したいケースってどんなんでしょうか。

* モデル実装、パラメータ、データ、等のバリエーションを短期間で試したいとき
* GPU分散で学習させたいとき
* 上記のようなタスクが複数人同時に発生するとき
* ~~Qiitaで釣りっぽいタイトルを書きたいとき~~

とかですかね。用法・用量を守って正しく使いましょう。

ちなみに私はデータのバリエーションがたくさんあって、一気に複数のモデルを作って比較したい
というちゃんとした理由は一応ありました。

# 学習ジョブの実行

AI Platformでの学習ジョブ実行手順はざっくり以下です。

1. 「pythonのコマンドとして実行できる形式でパッケージ」 or 「Dockerイメージを作成」
2. 「パッケージをGCSにアップロード」 or 「DockerイメージをGCRにpush」
3. 学習データをGCSにアップロード
4. gcloudコマンドでAI Platformジョブを開始

以前はDockerコンテナでの実行はベータ版だったのですが
今はGAになっていますし、Dockerがやりやすいでしょう。

以下ではDockerコンテナを用いた手順を記載
と思ったのですが、 ~~時間がないので~~ [公式チュートリアル](https://cloud.google.com/ml-engine/docs/custom-containers-training?hl=ja)が充実しているので
そちらにリンクを貼っておき、「？」となりそうなところだけ補足します。（みんな大好きMNISTです）
基本的にチュートリアルのコピペで動くはずです。

また、単にGPUを使ったジョブ実行を体験したいだけなら、↑のチュートリアルすべてを実施する必要はありません。

具体的には

* [始める前に](https://cloud.google.com/ml-engine/docs/custom-containers-training?hl=ja#before_you_begin)
* [カスタム コンテナでの GPU の使用](https://cloud.google.com/ml-engine/docs/custom-containers-training?hl=ja#using_gpus_with_custom_containers)

だけで可能です。

## Dockerfile

チュートリアルなんで素直にそのまま使えばいいのですが、ちょっとだけ補足。

```Docker
FROM nvidia/cuda:9.0-cudnn7-runtime
```

古い。けどまぁ無視していいかと。
例えばTensorflow 1.12以上を使いたいときとかはCUDA 10のイメージを指定しましょう。

```Docker
RUN wget -nv \
    https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz && \
    mkdir /root/tools && \
    tar xvzf google-cloud-sdk.tar.gz -C /root/tools && \
    rm google-cloud-sdk.tar.gz && \
    /root/tools/google-cloud-sdk/install.sh --usage-reporting=false \
        --path-update=false --bash-completion=false \
        --disable-installation-options && \
    rm -rf /root/.config/* && \
    ln -s /root/.config /config && \
    # Remove the backup directory that gcloud creates
    rm -rf /root/tools/google-cloud-sdk/.install/.backup

# Path configuration
ENV PATH $PATH:/root/tools/google-cloud-sdk/bin
# Make sure gsutil will use the default service account
RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
```

Dockerコンテナではなく、Pythonパッケージをアップロードする方法でやる場合は
すでに色々とインストールされているインスタンスで実行されるのですが
Dockerコンテナ内では当然必要なツール等は自分でインストールしておく必要があります。

AI PlatformはGCSを入出力ファイル置き場として使うので、
ここではそれに必要なgcloud SDKのインストール、およびgsutilコマンドの認証設定を行っています。


## ジョブの実行

DockerイメージをGCRにpushしたら、あとはgcloudコマンドでAI Platformのジョブを実行するだけです。かんたんですね。

```bash
$ export MODEL_DIR=pytorch_model_gpu_$(date +%Y%m%d_%H%M%S)
$ export REGION=us-central1
$ export JOB_NAME=custom_container_job_gpu_$(date +%Y%m%d_%H%M%S)
```

と変数を定義しておいて、下記を実行。

```bash

$ gcloud ai-platform jobs submit training $JOB_NAME \
  --scale-tier BASIC_GPU \
  --region $REGION \
  --master-image-uri $IMAGE_URI \
  -- \
  --epochs=5 \
  --model-dir=gs://$BUCKET_NAME/$MODEL_DIR
```

* `--scale-tier`
  ジョブを実行するインスタンスのスペックを指定します。（スケール階層と呼ばれてます）
  `BASIC_GPU` はk80が1個アサインされたインスタンスが使われます。
  指定できるスケール階層は [こちら](https://cloud.google.com/ml-engine/docs/machine-types?hl=ja#scale_tiers)。

* `--region`
  ジョブを実行するリージョンを選びます。
  アクセスするGCSと同一リージョンを指定しなければならない、という謎制約があります。（ファイル転送の速度を考慮してくれてる、というポジティブ解釈）
  また、リージョンによっては使えるGPUの種類が違ったりします。

* `--` 以降の引数は、gcloudコマンドの引数ではなく、実行するpythonコマンドの引数として与えられます。
  [チュートリアルのソースコード](https://github.com/GoogleCloudPlatform/cloudml-samples/blob/master/pytorch/containers/quickstart/mnist/trainer/mnist.py)を見ればわかるのですが
  `--epochs` と `--model-dir` という引数がコマンドに渡されて実行されることになります。
自作コマンドを作るときも、こんな感じで可変なパラメータを引数化します。

上記のgcloudコマンドを実行して

```
jobId: custom_container_job_gpu_20191219_142126
state: QUEUED
```

と出力されれば、ジョブ発行は成功です。

チュートリアルでは

```bash
$ gcloud ai-platform jobs describe $JOB_NAME
$ gcloud ai-platform jobs stream-logs $JOB_NAME
```

でジョブの進捗を確認しよう、となっていますが
GCPコンソール上で [AI Platformのジョブ一覧](https://console.cloud.google.com/ai-platform/jobs)を見たほうが簡単かもです。

<img width="1028" alt="スクリーンショット 2019-12-19 23.27.50.png" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/47605/b9d68dfe-0761-3302-2b65-c1577c0bb137.png">

あとはGPUの数だけジョブを発行しまくれば、ひゃっはーできます。

# 最後に

タイトルから釣りっぽい投稿になってしまい、反省しています。

そもそもGPUを100個同時に使うケースはあまりないにしても、
GPUを使いたいとき使いたいだけ利用できるサーバレスタイプの環境として、
という観点だけでもAI Platformを使ってみる価値はあるかもです。

また、今回取り上げませんでしたが[ハイパーパラメータの自動調整機能](https://cloud.google.com/ml-engine/docs/using-hyperparameter-tuning?hl=ja)を使うと、複数GPUを一気に使えるメリットもあるってもんです。
MNISTを100並列で動かしてどうすんねん、という感じなので、もっと有効に使っていきましょう。
（一応私ももっとマジメに使ってます）