フレクトのクラウドblog re:newal

http://blog.flect.co.jp/cloud/からさらに引っ越しています

クラウド最安 GPU で Jupyter Notebook を実行するには

技術開発室の佐藤です。こんにちは。

皆さんは GPU を使っていますか?筆者の周辺では年を追うごとに GPU の需要が高まっています。用途のほとんどはディープラーニングのモデルトレーニングです。特に画像系AI機能のモデルを作成する場合は、 GPU は必須です。

しかしこの GPU には問題があります。価格です。AWS などのクラウドベンダー各社は GPU インスタンスをラインアップしていますが、CPU インスタンスと比較するとかなりの高価格です。(設備の調達価格や電気料金を考えると仕方ないのかもしれませんが。。)最小販売単位が大きいことも悩みどころで、AWS と Azure は最低 4 CPU 構成からとなっています。本番稼働では必要でしょうが、学習用には過剰なのです。

筆者の知る限り、 GPU を一番小口で提供しているのは GCP で、1コアの n1-standard-1 に NVIDIA Tesla T4 をひとつ接続する構成が最小です(昨年までは K80 が最安だったが、最近 T4 が値下げされた)。これまで筆者はこのぐらいのインスタンスで Jupyter Notebook (以下、 notebook)を使って各種の手元検証をしていました。しかしこのような使い方では稼働率が低く、高額な GPU がしばしばアイドルになっていることを心苦しく思っていました。

この状況を改善するひとつの方法は「余剰割引」の利用です。AWS なら Spot Instance、 GCP なら Preemptible GPU になります。通常価格の半額以下で調達できます。ただしこの余剰割引にも問題があり、ベンダの都合で突然解放されてしまうのです(直前に予告通知はある)。元々待機状態の設備の格安提供なので、ベンダの都合でいつでも取り上げられる約束なのです。その時が来たら、ローカルストレージも全て解放されてしまいます。手元検証作業をしていたら突然中止、設定やり直し、というのは、心理的にもよろしくありません。

廉価な余剰割引を、安心して使う方法はないものでしょうか。

バッチ化、コンテナ化、GKE ジョブ化 で挑戦

そこで筆者が思いついた解決方法は、以下のようなものです。

  • Notebook は廉価な CPU インスタンスで基本調査した後、バッチ実行する。
  • このバッチ実行をコンテナ化する。
  • このコンテナを、余剰割引 GPU (preemptible GPU) ノードで GKE ジョブとして実行する。

図に書くと以下のような感じです。
f:id:masashi-sato-flect:20200131130140p:plain

GKE (Google Kubernetes Engine) は GCP が提供する Kubernetes ランタイムです。 今回は GCPGPUを利用しますので、KubernetesGCP を使うのが最も手間がないと考えました。

こうすれば、以下のメリットが得られるはずです。

  • GPU がアイドルでも preemptible GPU なので出費が少ない。
  • Kaggle や巷ブログに豊富に出回っている notebook をそのまま評価できる。
  • ジョブ実行中にノードが解放されたら(この動作は preemption と呼ばれる)、ノード障害として検出され、 GKE がジョブを自動で再実行する。

結果を先にお話しすると、この仕組みは一通り動作しました。筆者は今後はこの方法で GPU 技術の評価をしていきたいと考えています。もしかしたら同じような悩みをお持ちの方もいらっしゃるかもしれないと思いましたので、以下にその手法を説明させていただきます。

Notebook をバッチ実行する

これは実は楽勝です。Jupyter Notebook がインストールされた環境で、シェルコマンドで以下のように打つだけです。

jupyter nbconvert --ExecutePreprocessor.timeout=600 --to notebook --execute mynotebook.ipynb

Pythonコードで実行することもできます。

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor

with open(LOCAL_FILE_PATH) as fIn:
    nb = nbformat.read(fIn, as_version=4)
    ep = ExecutePreprocessor(timeout=600, kernel_name='python3')
    ep.preprocess(nb)
    with open(PROCESSED_FILE_PATH, 'w', encoding='utf-8') as fOut:
        nbformat.write(nb, fOut)

既定のタイムアウトは短い(60秒)ので、指定するようにします。

バッチ実行をコンテナ化する

Jupyter Notebook のコンテナは各種ありますが、今回は以下の TensorFlow 公式コンテナを使用しました。

  • tensorflow/tensorflow:latest-gpu-py3-jupyter

この上で、以下のような Python スクリプトを実行します。環境変数で Cloud Storage の所在を受け取ってダウンロード、 Notebook を実行してアップロードして返すだけの簡単な内容です。

import os
from google.cloud import storage
# https://nbconvert.readthedocs.io/en/latest/execute_api.html#executing-notebooks-using-the-python-api-interface
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor

def download_blob(bucket_name, source_blob_name, destination_file_name):
    # 中略
    # https://cloud.google.com/storage/docs/downloading-objects
# end of download_blob

def upload_blob(bucket_name, source_file_name, destination_blob_name):
    # 中略
    # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
# end of upload_blob

def main():
    LOCAL_FILE_PATH     = 'test.ipynb'
    PROCESSED_FILE_PATH = 'test_processed.ipynb'
    BUCKET_NAME         = os.environ['BUCKET_NAME']
    DOWNLOAD_FILE_PATH  = os.environ['DOWNLOAD_FILE_PATH']
    UPLOAD_FILE_PATH    = os.environ['UPLOAD_FILE_PATH']

    download_blob(
        BUCKET_NAME,
        DOWNLOAD_FILE_PATH,
        LOCAL_FILE_PATH
    )

    with open(LOCAL_FILE_PATH) as fIn:
        nb = nbformat.read(fIn, as_version=4)
        ep = ExecutePreprocessor(timeout=600, kernel_name='python3')
        ep.preprocess(nb)
        with open(PROCESSED_FILE_PATH, 'w', encoding='utf-8') as fOut:
            nbformat.write(nb, fOut)
            upload_blob(
                BUCKET_NAME,
                PROCESSED_FILE_PATH,
                UPLOAD_FILE_PATH
            )
        # end of with fOut
    # end of with fIn
# end of main

if __name__ == "__main__":
    # execute only if run as a script
    main()

コンテナの Dockerfile は以下のようになりました。こちらも上記の Python スクリプトを動かすだけの簡単な内容です。

FROM tensorflow/tensorflow:latest-gpu-py3-jupyter
COPY requirements.txt /tf
COPY tf_exec.py /tf
RUN pip install --upgrade pip
RUN pip install -r /tf/requirements.txt
ENTRYPOINT python /tf/tf_exec.py

requirements.txt は、Cloud Storage の SDK だけです。

google-cloud-storage

以下コマンドでコンテナを作って Container Registry に上げます。

docker build --force-rm --tag=test:latest.
gcloud auth configure-docker
docker tag test:latest gcr.io/${PROJECT_ID}/${REPOSITORY_NAME}:latest
docker push gcr.io/${PROJECT_ID}/${REPOSITORY_NAME}:latest

コンテナを、preemptible GPU ノードで GKE ジョブとして実行する

このステップは少々苦労しました。以下のような課題がありました。

  • ジョブに GCP の権限を付与する必要があった。
  • CPU ノードも用意する必要があった。
  • ジョブが preemptible GPU ノードに配置されるように指定する必要があった。

ジョブに GCP の権限を付与する必要があった

今回実行するコンテナは Cloud Storage をアクセスしますので、専用のサービスアカウントを作成し、 Storage Object Admin のロールを付与しました。Container Registry もアクセスしますが、こちらは Cloud Storage の権限がそのまま適用されるとありますので、追加作業はありません。

このサービスアカウントを、GKE ノードを作成するときに指定します(後述)。

CPU ノードも用意する必要があった

Kubernetesではシステム管理のコンテナが常時活動しており、これらを実行しているインスタンスで preemption が発生すると、回復までにより多くの時間(5~10分ぐらい?)を要します。ダウンタイムを減らすには、常時起動の管理ノードと preemptible なノードの両方を用意します。管理ノードには小さな廉価なインスタンスを使用しました。

ジョブが preemptible GPU ノードに配置されるように指定する必要があった

Kubernetes は taint (忌避要件?)と toleration (忌避要件許容?)という2つの概念を用いてジョブのアサイン先を条件設定することができます。今回の場合は、以下の作戦を取ります。

  • Preemptible ノードに taint 「cloud.google.com/gke-preemptible="true":NoSchedule」を設定し、システム管理コンテナがスケジュールされないようにする。
  • ジョブスケジュールに以下の条件設定をする。

以上の諸要件を勘案し、まず GKE クラスタと管理ノードを作成します。

gcloud container clusters create @{TEST_ID} \
  --zone=us-west1-b \
  --machine-type=e2-small \
  --num-nodes=1 \
  --disk-size=10GB

次にこのクラスタに preemptible GPU ノードを追加します。

gcloud container node-pools create ppool \
  --zone=us-west1-b \
  --cluster=@{TEST_ID} \
  --num-nodes=1 \
  --machine-type=n1-standard-1 \
  --disk-size=50GB \
  --preemptible \
  --accelerator=type=nvidia-tesla-k80,count=1 \
  --node-taints=cloud.google.com/gke-preemptible="true":NoSchedule \
  --service-account=@{TEST_ID}@${PROJECT_ID}.iam.gserviceaccount.com

最後の4つのコマンドオプションに注意してください。上からそれぞれ以下のような意味があります。

  • Preemptible GPU インスタンス
  • NVIDIA Tesla K80 を1つ使用
  • Taint の設定
  • Cloud Storage へのアクセス権を付与したサービスアカウントをノードに設定

コンテナが起動したら、ノードインスタンスに NVIDIA ドライバをインストールする daemonSet を仕掛けます

kubectl apply -f daemonset-preloaded.yaml

以上で準備は完了です。以下のジョブを実行します。

apiVersion: batch/v1
kind: Job
metadata:
  name: ${TEST_ID}gpu
spec:
  template:
    spec:
      containers:
      - name: ${TEST_ID}
        image: gcr.io/${PROJECT_ID}/${REPOSITORY_NAME}:latest
        resources:
          limits:
            nvidia.com/gpu: 1
        env:
        - name: BUCKET_NAME
          value: ${BUCKET_NAME}
        - name: DOWNLOAD_FILE_PATH
          value: ${SOURCE_NOTEBOOK_PATH}
        - name: UPLOAD_FILE_PATH
          value: ${DESTINATION_NOTEBOOK_PATH}
      restartPolicy: Never
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
            - matchExpressions:
              - key: cloud.google.com/gke-preemptible
                operator: In
                values: 
                - "true"
      tolerations:
      - key: cloud.google.com/gke-preemptible
        operator: Equal
        value: "true"
        effect: NoSchedule
      - key: nvidia.com/gpu
        operator: Exists
        effect: NoSchedule
  backoffLimit: 0

以下の部分に注意してください。

resources:
    limits:
        nvidia.com/gpu: 1

この指定をしないと、コンテナから GPU が見えなくなってしまいます。

このバッチで実行した notebook の冒頭では、 TensorFlow から GPU が使用できる状態になっていることが確認できました。

f:id:masashi-sato-flect:20200131130247p:plain

ようやく当初目論見だった preemptible GPU ノードで Kubernetes ジョブ実行ができました。

Preemption からの復帰時間はどのぐらいか

Preemption は GCP 側が起こしてくるので意図して発生させることはできませんが、同等のイベントを発生させてテストすることができます

gcloud compute instances simulate-maintenance-event gke-${TEST_ID}-ppool-xxxx

実際のダウンタイムはどのくらいになるのでしょうか?筆者が今回調査の際に経験したダウンタイムは、最低2分、長いものでは10分ほどでした。結構長いですね・・・。運用については順次調整していきたいと思います。

最後までお読み下さり、ありがとうございました。