Step Functions で自作 Dockerfile を使って SageMaker の GPU マシンを動かす方法

Step Functions で SageMaker のリソースを特にカスタムコンテナイメージで使う場合，単純に InstanceType として "ml.g4dn.xlarge" などの GPU マシンを設定するだけでは GPU を使った学習はできなくて，使いたい Dockerfile に少し手を加える必要がある．

今回は備忘録も兼ねて Dockerfile の中身を紹介しながら，GPU 環境での動作確認を行っていく．

以下の手順で動作確認を行う．

1. ローカルで Dockerfile と動作確認用の Python スクリプトを作成する
2. ローカルで build し，AWS ECR に build したイメージを push する
3. push したイメージの URI を SageMaker Training Job の TrainingImage に指定する
4. Step Functions を実行する
5. CloudWatch Logs を確認する

今回の動作確認フローは下図のようなイメージになっている．

自作の「Dockerfile.gpu」ファイル

今回は Tensorflow-gpu のベースイメージを使っている．そのイメージに nvidia-docker を追加でインストールすることで Step Functions で GPU 用のカスタムコンテナイメージを使って SageMaker を動作させることができる．

TensorFlow Docker Imagesから好きな GPU 用のイメージを選択する．今回は「tensorflow/tensorflow:2.6.1-gpu」を使用することにした．Optional Features にも記載されているが，nvidia-docker が必要だよとのことなので，この通りにする．

-gpu tags are based on Nvidia CUDA. You need nvidia-docker to run them. NOTE: GPU versions of TensorFlow 1.13 and above (this includes the latest- tags) require an NVidia driver that supports CUDA 10. See NVidia's support matrix.

大事な部分は以下の4行．以前はこちらの記事にも書かれていますが，nvidia-container-toolkit をインストールする必要があったみたいだが，今は nvidia-docker2 をインストールすることで，nvidia-container-toolkit も一緒にインストールされて，よりシンプルになっている．

RUN distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
    && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
RUN sudo apt-get update && sudo apt-get install -y nvidia-docker2

今回使用した Dockerfile 以下の通りです．

# Dockerfile.gpu
FROM tensorflow/tensorflow:2.6.1-gpu

# Set some environment variables.
# PYTHONUNBUFFERED keeps Python from buffering our standard
# output stream, which means that logs can be delivered to the user quickly.

ENV PYTHONUNBUFFERED=TRUE

# PYTHONDONTWRITEBYTECODE keeps Python from writing the .pyc files which
# are unnecessary in this case.
ENV PYTHONDONTWRITEBYTECODE=TRUE

# DEBIAN_FRONTEND prevent from stoping docker build with tzdata
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get -y update && apt-get install -y --no-install-recommends \
        curl \
        sudo \
        libmecab-dev \
        python3.8 \
        python3-distutils \
        python3-six \
        git \
        file \
        wget \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# ここが今回のポイントで，nvidia-docker をインストールする
RUN distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
    && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
RUN sudo apt-get update && sudo apt-get install -y nvidia-docker2

# 必要なライブラリーを requirements.lock を使ってインストールする
COPY requirements.lock /tmp/requirements.lock
RUN python3 -m pip install -U pip \
    && python3 -m pip install -r /tmp/requirements.lock \
    && python3 -m pip install sagemaker \
    && python3 -m pip install sagemaker-training \
    && rm /tmp/requirements.lock \
    && rm -rf /root/.cache

# Timezone jst
RUN ln -sf /usr/share/zoneinfo/Asia/Tokyo /etc/localtime

# Locale Japanese
ENV LC_ALL=ja_JP.UTF-8

# Set up the program in the image
ENV PROGRAM_DIR=/opt/program
COPY app $PROGRAM_DIR
WORKDIR $PROGRAM_DIR
ENV PATH="/opt/program:${PATH}"

RUN chmod +x $PROGRAM_DIR/hello_gpu.py

CMD ["python3"]

また，SageMaker で GPU を認識しているかを確認するための Python スクリプトは以下になる．

# hello_gpu.py
import tensorflow as tf
from tensorflow.python.client import device_lib


def main():
    # TensorflowのGPU確認
    print(f'GPUs Available: {tf.test.is_gpu_available()}')
    print(f"Num GPUs Available: {len(tf.config.list_physical_devices('GPU'))}")
    print(f'{device_lib.list_local_devices()}')


if __name__ == "__main__":
    main()

これらのファイルを用意して，ローカルで build を行い，build した後はそのイメージを ECR に push することで，Step Functions で使用できるようになる．

※ ECR への push 方法や設定は今回割愛する．

Step Functions の設定

Step Functions とは，AWS が提供する各種サービスを組み合わせたパイプラインを構築するためのワークフローサービスである．機械学習向けのアクションも用意されていて，今回使用する SageMaker Training Job もその一つ．

設定は yaml ファイルのような形式で AWS のコンソール画面上で打ち込んでいく．今回実施する内容の記述は以下の通り．

{
  "Comment": "Check GPU env",
  "StartAt": "Hello-GPU",
  "States": {
    "Hello-GPU": {
      "Comment": "GPUの動作確認",
      "Type": "Task",
      "Resource": "arn:aws:states:::sagemaker:createTrainingJob.sync",
      "Parameters": {
        "RoleArn": "<SageMaker Training Jobの実行権限がアタッチされているロール>",
        "TrainingJobName": "sample-training-job",
        "AlgorithmSpecification": {
          "EnableSageMakerMetricsTimeSeries": true,
          "TrainingImage": "<アカウントID>.dkr.ecr.ap-northeast-1.amazonaws.com/sample:latest-gpu",
          "TrainingInputMode": "File"
        },
        "EnableInterContainerTrafficEncryption": true,
        "EnableManagedSpotTraining": true,
        "Environment": {
          "SAGEMAKER_PROGRAM": "/opt/program/hello_gpu.py"
        },
        "ResourceConfig": {
          "InstanceCount": 1,
          "InstanceType": "ml.g4dn.xlarge",
          "VolumeSizeInGB": 20
        },
        "StoppingCondition": {
          "MaxRuntimeInSeconds": 12345,
          "MaxWaitTimeInSeconds": 12345
        }
      },
      "End": true
    }
  }
}

細かい設定内容に関しては，CreateTrainingJob というドキュメントを参考にすると良い．

ここで，Environment（環境変数）の SAGEMAKER_PROGRAM について説明しておく．この変数に指定したプログラムは Training Job のエントリーポイントにすることができる．

元々は以下のコマンドが実行されるのだが（train.py があればそれが対象となる），実行したいプログラムのパスを指定することで任意のプログラムを実行することができる．

docker run <イメージ> train

ただし，実行権限を与えておく必要があるので，Dockerfile 内で RUN chmod +x $PROGRAM_DIR/hello_gpu.py としている．

あとは，実行結果を CloudWatch Logs で確認して，以下の内容がログに出力されていれば大丈夫．

GPUs Available: True
Num GPUs Available: 1

今回は，Step Functions で SageMaker の Training Job をカスタムコンテナイメージを使って GPU 環境で動かす方法を紹介した．この方法を使えば，深層学習などの GPU 環境を必要とした学習もパイプラインに組み込むことが可能になる．また，Training Job を使った学習ができれば，実験結果は SageMaker Experiments に保存されるので，再現性を担保することもできる．

Step Functions でカスタムコンテナイメージを使って GPU 環境で学習させたい場合には，参考にして頂ければと！

参考

• Available SageMaker Studio Instance Types
• NVIDIA Dockerって今どうなってるの？ (20.09 版)