AWS Batchで速く/安くやるデータセットの前処理

OpenStreamアドベントカレンダーの一日目です。

結構前からやっている趣味DeepLearningですが、最近（実際は結構前から）次のような問題に当たり始めました。

• データセットが大きくなってきてHDDが厳しい
• データセットが大きくなってきて前処理がやばい

小さいデータセット＋Augmentationでなんとかなるものはいいんですが、現在最大のデータセットは 画像33万枚、220GB弱 あります。
んで、これを前処理したり何だりしていると、最終的に学習で利用するデータを作成するだけで、HDDが500GBくらい利用されてしまう状態です。

容量も当然厳しいんですが、一番厳しいのは処理時間です。現状の前処理を行うと、大体 12時間くらい かかります。趣味でやるので基本的に自分のPCでやっていると、HDDが悲鳴を上げる上に、実行している間はレイテンシが悪すぎて他の作業もできないって状態になってしまっていました。

前処理の中でGPUを利用しているので、GPUを使うような作業も出来ないという。

一部のデータセットで試して〜ってやれば出来ないこともないんですが、結局最後には全部やらないとならないので、この機会に AWS Batch を使って、一気に処理出来ないかどうかを試してみました。

Azure Batchとかもありますが、とりあえずはAWSで。Azure Batchでも同じようなことは出来るかと
GCPのDataflowのだとApache Beamに縛られるので、今回は対象外としました

前提

今回作成するジョブは大きく２つです。

• 画像のリサイズ
• エッジの抽出

また、今回は前処理だけ出来ればいいので、GPUインスタンスは使いません。

後述するように、基本的に AWS CloudFormation を利用します。基本的にAWS CLIは構築時点では利用しません。

では行ってみましょう。

AWS Batchについて

日本（の特にSIer的な人々）にとって、Batchと言えばJP1とかああいったフロー制御を想像してしまいますが、それとは異なります。

こちら でわかりやすくまとめられています。Batchって言われると脊髄反射的にGUIが・・・とか思ってしまうのはなんとかしたいです。

上記の記事で説明はされているので、ここではAWS Batchが何かということについては記述しません。

ジョブ実行用コンテナの作成

まずはジョブ実行用のDocker imageを作成します。今回はサイズにこだわって、次のような構成にしてみました。

• Alpine Linux 3.6
• Python 3.6.3
• OpenCV 3.2.0

で、実際に利用しているDockerfileがこんな感じになります。

FROM alpine:3.6

ENV OPENCV_VERSION 3.2.0
ENV PYTHON_VERSION 3.6.3

RUN apk add --update --no-cache \
    build-base \
    openblas \
    openblas-dev \
    unzip \
    curl \
    cmake \
    libjpeg \
    libjpeg-turbo-dev \
    libpng-dev \
    jasper-dev \
    tiff-dev \
    libwebp-dev \
    clang-dev \
    openssl \
    ncurses \
    libstdc++ \
    linux-headers && \
    mkdir -p /opt && \
    apk add --no-cache --virtual .build-python \
      openssl-dev \
      ncurses-dev \
      xz && \
    curl -LO https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tar.xz && \
    unxz Python-${PYTHON_VERSION}.tar.xz && \
    tar xf Python-${PYTHON_VERSION}.tar -C /opt && \
    cd /opt/Python-${PYTHON_VERSION} && \
    ./configure --enable-shared && make -j > /dev/null && make install > /dev/null && \
    cd / && \
    rm -rf /opt/Python-${PYTHON_VERSION} Python-${PYTHON_VERSION}.tar && \
    pip3 install numpy==1.13.3 && \
    apk del --virtual .build-python && \
    curl -LO https://github.com/opencv/opencv/archive/${OPENCV_VERSION}.zip && \
    unzip ${OPENCV_VERSION}.zip -d /opt && \
    rm -rf ${OPENCV_VERSION}.zip && \
    mkdir -p /opt/opencv-${OPENCV_VERSION}/build && \
    cd /opt/opencv-${OPENCV_VERSION}/build && \
    cmake \
      -D CMAKE_BUILD_TYPE=RELEASE \
      -D CMAKE_INSTALL_PREFIX=/usr/local \
      -D WITH_FFMPEG=NO \
      -D WITH_IPP=NO \
      -D WITH_OPENEXR=NO \
      -D WITH_TBB=YES \
      -D BUILD_TESTS=NO \
      -D BUILD_EXAMPLES=NO \
      -D BUILD_ANDROID_EXAMPLES=NO \
      -D INSTALL_PYTHON_EXAMPLES=NO \
      -D BUILD_DOCS=NO \
      -D BUILD_opencv_python2=NO \
      -D BUILD_opencv_python3=ON \
      -D PYTHON3_EXECUTABLE=/usr/local/bin/python3 \
      -D PYTHON3_INCLUDE_DIR=/usr/local/include/python3.6m/ \
      -D PYTHON3_LIBRARY=/usr/local/lib/libpython3.so \
      -D PYTHON_LIBRARY=/usr/local/lib/libpython3.so \
      -D PYTHON3_PACKAGES_PATH=/usr/local/lib/python3.6/site-packages/ \
      -D PYTHON3_NUMPY_INCLUDE_DIRS=/usr/local/lib/python3.6/site-packages/numpy/core/include/ \
      .. && \
    make VERBOSE=1 && \
    make -j2 && \
    make install && \
    cd / && \
    rm -rf /opt && \
    apk del gcc build-base openblas-dev clang-dev linux-headers cmake unzip

全部一つのRUNに押し込めているので、失敗すると悲惨ですが、これで作成したimageは 273MB 程度と、それなりに扱いやすいサイズになります。

このimageは、実際に作業をさせるimageのBase imageになるので、小さいにこしたことはないです。

実際に利用するImage

実際に利用するImageは、以下で管理しています。

単純に必要なPythonスクリプトをコピーして、pipで依存のインストールをしているだけです。

AWS Batchの環境作成

AWS Batchは、Managed環境とUnmanaged環境の二通りが選べますが、基本的にはManagedでいいかと思います。しかし、ManagedでもJobを動かすまでには以下のようなリソースが必要になります。

AWS Batchが様々なサービスを組み合わせて構築されている証左ですが、手動で作ってると色々とめんどくさいです。

• VPC/Subnet
  • ECSが使われるので、PublicかPrivate＋NAT Gatewayがセットアップされてないと厳しいです
• SecurityGroup
• ECSのService Role
  • InstanceProfileも必要です
• AWS BatchのService Role
• JobのRole
  • ECSのInstanceProfileとは別で必要です
• ECR
  • Publicなimageでよければ、Dockerhubとかでもいいんですが、大抵はECRが必要かと

ということで、CloudFormationのテンプレートを作りました。

ただ、必要なリソースを全部入れているので、結構量の多いテンプレートになっています。そこで、Batch系のリソースだけ抜粋します。

  ComputeEnvironment:
    Type: AWS::Batch::ComputeEnvironment
    Properties:
      Type: MANAGED
      ServiceRole: !Sub "arn:aws:iam::${AWS::AccountId}:role/AWSBatchServiceRole"
      ComputeEnvironmentName: C4OnDemand
      ComputeResources:
        MaxvCpus: 128
        SecurityGroupIds:
          - !Ref InstanceSecurityGroup
        Type: EC2
        Subnets:
          - !Ref PrivateSubnet
        MinvCpus: 0
        ImageId: ami-a1b2c3d4
        InstanceRole: ecsInstanceRole
        InstanceTypes:
          - c4.large
          - c4.xlarge
          - c4.2xlarge
          - c4.4xlarge
          - c4.8xlarge
        Ec2KeyPair: batch-compute
        Tags: {"Name": "Batch Instance - C4OnDemand"}
        DesiredvCpus: 48
      State: ENABLED

  MyRepository: 
    Type: "AWS::ECR::Repository"
    Properties: 
      RepositoryName: "image-convert"
      RepositoryPolicyText:
        Version: "2012-10-17"
        Statement: 
          - 
            Sid: AllowPushPull
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:user/${User}"
            Action: 
              - "ecr:GetDownloadUrlForLayer"
              - "ecr:BatchGetImage"
              - "ecr:BatchCheckLayerAvailability"
              - "ecr:PutImage"
              - "ecr:InitiateLayerUpload"
              - "ecr:UploadLayerPart"
              - "ecr:CompleteLayerUpload"
          - 
            Sid: AllowPull
            Effect: Allow
            Principal:
              AWS: !GetAtt JobRole.Arn
            Action: 
              - "ecr:GetDownloadUrlForLayer"
              - "ecr:BatchGetImage"
              - "ecr:BatchCheckLayerAvailability"

  JobDefinitionForResize:
    Type: 'AWS::Batch::JobDefinition'
    Properties:
      Type: container
      JobDefinitionName: !Sub
          - ${Service}-resize-image
          - { Service: !Ref ServiceName}
      Parameters:
        Bucket: !Ref Bucket
        ExcludeFiles: exclude.txt
        Size: 128
      ContainerProperties:
        Command:
          - python3
          - -m
          - resize_fix_size
          - -b
          - "Ref::Bucket"
          - -d
          - resized
          - -e
          - "Ref::ExcludeFiles"
          - -s
          - "Ref::Size"
          - --crop
          - "Ref::Prefix"
        Memory: 1000
        JobRoleArn: !Ref JobRole
        Vcpus: 1
        Image: !Sub
          - "${AWS::AccountId}.dkr.ecr.ap-northeast-1.amazonaws.com/${Repository}/image-converter"
          - {"Repository": !Ref MyRepository}
      RetryStrategy:
        Attempts: 1

  JobQueue:
    Type: AWS::Batch::JobQueue
    Properties:
      ComputeEnvironmentOrder:
        - Order: 1
          ComputeEnvironment: !Ref ComputeEnvironment
      State: ENABLED
      Priority: 1
      JobQueueName: HighPriority

これでも全体の1/3くらいです。AWS Batchは、Compute Environmentでリソースの総量を管理し、Job Queueで流量を制御し、Job Definitionで流れるJobの内容を制御します。

実はCloudFormationのドキュメントにある AWS::Batch::JobDefinition のサンプルには罠があり、RetryStrategyが無いため、真似するともれなくエラーになります。（記述時点:2017/12）

大抵コピーで済ます人（自分）は間違いなく踏む地雷かと。

AWS BatchのServiceRoleについて

AWS BatchのServiceRoleは、恐らくCLI/CloudFormationで作成する人だけが引っかかる地雷があります。ドキュメントのTroubleShootingにわざわざ書いてあるくらいなので、よくあったんでしょう・・・私も見事に踏みました。

単純に言うとARNの指定間違いなのですが、これをやってしまうと1時間程度消すことも消す準備をすることもできず、もれなく待ち状態になります。

Important
Do not attempt to delete a compute environment that is in an INVALID state due to a misconfigured AWS Batch service role. This could cause your environment to get stuck in a DELETING state for up to an hour, and you cannot update the compute environment until the operation times out and fails back to INVALID.
http://docs.aws.amazon.com/batch/latest/userguide/troubleshooting.html#invalid_service_role_arn から引用

Management ConsoleからRoleを作成していて、それをそのまま利用していれば基本的には起こりません。AWS Batchを手動で作るのめんどいなーって人や自動化バンザイって人はご注意を。

Jobを実行する

今回対象にしているテータセットは、全部sha256 hashを名前にしています。それを利用して、prefixをJobのパラメータに渡すようにしてみました。

実際にprefixの分け方＝並列度で、どれくらい実行時間が変わるかを試してみました。試すのは画像のリサイズ処理です。

ちなみにローカルでは、8並列で平均して100枚/5秒です。無視するデータを抜いた25万枚だと250分＝4時間くらいかかります。

• prefixを0000〜ffffにした時（65535個のJob）
  • 1jobあたり平均して5から6枚を処理
  • 途中で断念
    • Jobのsubmitだけで4時間くらいかかる（4job/秒）
    • Jobあたりの枚数が少なすぎて、起動→終了のオーバーヘッドの方が大きい
• prefixを000〜fffにした時（4096個のJob）
  • 1jobあたり平均して80枚を処理
  • 実行時間は40分ほど
    • 25万枚/40分＝100枚/秒
  • vCPUの利用率は大体15〜20％くらい
    • Single threadだったので、Networkが明らかにボトルネックになっている
  • 金額はだいたい以下の通り

    c4.8xlarge * 3 = 0.6 * 3 = $1.8/h
    c4.4xlarge * 2 = 0.18 * 2 = $0.36/h
    $2.16/h * 45/60 = $1.54
    

• prefixを00〜ffにした時（256個のJob）
  • 1jobあたり平均して1300枚を処理
    • 25万枚/7分＝600枚/秒
  • 実行時間は7分ちょっと（！）
    • ただし、各コンテナで8threadで処理するように
  • vCPUの利用率は劇的に改善。大体70〜80％くらい
    • ネットワークの利用度合いが明らかによくなっている
  • 金額はだいたい以下の通り

    c4.8xlarge * 3 = 0.6 * 3 = $1.8/h
    c4.4xlarge * 2 = 0.18 * 2 = $0.36/h
    $2.16/h * 7/60 = $0.24
    

S3→Instance→S3という経路を辿るため、Networkがボトルネックになります。そのため、1vCPUしかわたしていないコンテナでも、threadを利用したほうがいいと判断した所、大当たりでした。

まさか一回あたり30円くらいまで下げられるとは思いませんでした。

この辺はバッチコンピューティングの経験が問われそうです・・・

依存関係のあるJobを実行する

実際には、リサイズ以外にも処理が必要になります。そこで、こんなスクリプトを作ってみました。

# submit-all-jobs.sh
#!/bin/bash
JOB_QUEUE=xxxxxx

RESIZE_JOB_ARN=$(aws batch describe-job-definitions --job-definition-name image-converter-resize-image --status ACTIVE | jq -r '.jobDefinitions | max_by(.revision).jobDefinitionArn')
EDGE_JOB_ARN=$(aws batch describe-job-definitions --job-definition-name image-converter-extract-edge --status ACTIVE | jq -r '.jobDefinitions | max_by(.revision).jobDefinitionArn')

MAKE_SEQ=$(cat <<EOF
for i in range(0, 0xff):
    print("{:0>2x}".format(i + 1))
EOF
        )

SEQ=$(python -c "$MAKE_SEQ")
echo "$SEQ" | xargs -P 8 -I {} -n 1 \
                    ./submit-jobs.sh $JOB_QUEUE $RESIZE_JOB_ARN $EDGE_JOB_ARN {}

# submit-jobs.sh
#!/bin/bash
JOB_QUEUE=$1
RESIZE_JOB_ARN=$2
EDGE_JOB_ARN=$3
SEQ=$4

resize_job=$(aws batch submit-job --job-name "resize-$SEQ-$now" --job-queue $JOB_QUEUE \
    --job-definition $RESIZE_JOB_ARN \
    --parameters Prefix=full/$SEQ
          )

resize_job_id=$(jq '.jobId' "$resize_job")

aws batch submit-job --job-name "edge-$SEQ-$now" --job-queue $JOB_QUEUE \
    --job-definition $RESIZE_JOB_ARN \
    --parameters Prefix=resized/$SEQ \
    --depends-on jobId=$resize_job_id

xargsでお手軽にparallel実行するためのスクリプトです。8並列で行うと、256×2のジョブ登録に数分といったとこ
ろです。

AWS Batchでの依存関係は、submit-jobの depends-on パラメータで指定します。実際にはListなので、複数のJobが終わってから実行する、みたいなことも出来ます。

ここでは、リサイズが終わってからエッジ抽出をしたいので、そういうふうに指定しています。

なお、リサイズ＋エッジ抽出を全部やっても 10分 くらいで終わります。今までの苦労は何だったんや・・・。

前処理の要求は止まらない

30万枚の画像が30円くらいでリサイズ出来るようになり、前処理が捗るようになりました。とりあえず現状は概ね満足しています。もっと前処理が必要となってもなんとかなりそうって感覚を得られましたし。

ただ、まだ色々と課題はあります。

• TFRecordへの変換がめんどくさい
  • 現状はローカルにダウンロードしてきて変換スクリプトかましてます
  • TFRecordにするまでが前処理なので、そのへんが難しいところです
• 前処理で別のNNを利用する場合どうするか
  • GPUを使うと10倍以上速いので利用したいところですが、個人でやるレベルかって・・・

TFRecordまで一連のJobで出来ると学習が捗るんですが・・・。できそうなので、後で試してみようかと思います。

AWS Batchを今回初めてまともに触ってみましたが、なんとなく既存の仕組みの延長にある感じなので、理解はしやすいように思いました。
EC2が秒単位課金になったこともあり、一気にInstanceを立ち上げてさっさと終わらせるってことができるようになったのも大きいですね。

EMRを使うほどじゃないとか、EMRだと使いにくい言語で大量のデータ処理をしないといけないってなった時には検討してみてはどうでしょうか。

（番外）AWS Batchの注意点

今回試してみて、いくつか？ってなったりハマった点がありました。いずれも2017/12時点です。

• DashBoardのJob Queueには1000以上は表示できない
  • 1000を超えると 100+ って表示になります。1000超えてるのに何故100？ってなることうけあいです
• Jobの一覧で検索できない
  • 失敗したJobが100を超えると悲惨です
• Jobの一覧で並べ替えが出来ない
  • Started atで並べ替えしたい・・・
• まれにJobがRUNNINGでStuckする
  • TerminateもCancelも効かず、最終的にはJob Queueを削除しました
  • 多分30分くらいすると解消される感じです
• vCPU/Memoryのバランスに注意
  • ECSと一緒ですね
• Instanceのスケールダウンが結構早い
  • 5分とかjobが投入されないと全部消す勢いです

GAになってからまだ一年経っていないので、これからもっと使いやすくなっていくんじゃないかと思います。

明日は @pegass85 さんです。