GPU の使用

このページでは、GPU を使用して Dataflow で Apache Beam パイプラインを実行する方法について説明します。GPU を使用するジョブでは、Dataflow の料金ページで定められた料金が発生します。

Dataflow で GPU を使用する方法については、Dataflow での GPU のサポートをご覧ください。GPU を使用したパイプラインのビルドに関するデベロッパーワークフローの詳細については、GPU を使用した開発をご覧ください。

Apache Beam ノートブックの使用

Dataflow で GPU を使用して動作するパイプラインがすでに存在する場合は、このセクションをスキップできます。

Apache Beam ノートブックを使用すると、開発環境を設定することなく、GPU を使用してパイプラインのプロトタイプを作成と反復開発を行えます。使用を始めるには、Apache Beam ノートブックを使用した開発ガイドに沿って Apache Beam ノートブックインスタンスを起動し、サンプルノートブックの Apache Beam で GPU を使用するに沿って操作します。

GPU 割り当てのプロビジョニング

GPU デバイスは、Google Cloud プロジェクトの割り当て枠の対象になります。選択したリージョンで GPU 割り当てをリクエストします。

GPU ドライバのインストール

install-nvidia-driver を worker_accelerator オプションに追加することで、Dataflow に、NVIDIA ドライバをワーカーにインストールするように指示します。install-nvidia-driver オプションを指定すると、Container-Optimized OS によって提供される cos-extensions ユーティリティを使用して、Dataflow が Dataflow ワーカーに NVIDIA ドライバをインストールします。install-nvidia-driver を指定すると、NVIDIA ライセンス契約に同意したことになります。

NVIDIA ドライバインストーラによって提供されるバイナリとライブラリは、/usr/local/nvidia/ のパイプラインユーザーコードを実行しているコンテナにマウントされます。

GPU ドライバのバージョンは、Dataflow で現在使用されている Container-Optimized OS のバージョンによって異なります。特定の Dataflow ジョブの GPU ドライバのバージョンを確認するには、ジョブの Dataflow ステップのログで GPU driver を検索します。

カスタムコンテナイメージのビルド

GPU とやり取りするには、GPU アクセラレーションによるライブラリや CUDA ツールキットなど、NVIDIA ソフトウェアが追加で必要になる場合があります。ユーザーコードを実行する Docker コンテナにこれらのライブラリを指定します。

コンテナイメージは、Apache Beam SDK コンテナイメージの契約を満たし、必要な GPU ライブラリが入ったイメージを指定することでカスタマイズできます。

カスタムコンテナイメージを用意するには、Dataflow Runner v2 を使用し、worker_harness_container_image パイプラインオプションでコンテナイメージを指定します。Apache Beam 2.30.0 以降を使用する場合は、簡略化のために短縮されたオプション名 sdk_container_image を使用できます。詳細については、カスタムコンテナの使用をご覧ください。

方法 1:GPU を使用するために構成された既存のイメージの使用

Apache Beam SDK コンテナ契約を満たす Docker イメージは、GPU の使用が事前に構成された既存のベースイメージからビルドできます。たとえば、TensorFlow の Docker イメージや NVIDIA コンテナイメージは、GPU を使用向けに事前に構成されています。

Python 3.6 を含む TensorFlow Docker イメージをベースに構築されるサンプル Dockerfile は、次のようになります。

ARG BASE=tensorflow/tensorflow:2.5.0-gpu
FROM $BASE

# Check that the chosen base image provides the expected version of Python interpreter.
ARG PY_VERSION=3.6
RUN [[ $PY_VERSION == `python -c 'import sys; print("%s.%s" % sys.version_info[0:2])'` ]] \
   || { echo "Could not find Python interpreter or Python version is different from ${PY_VERSION}"; exit 1; }

RUN pip install --upgrade pip \
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 \
    # Verify that there are no conflicting dependencies.
    && pip check

# Copy the Apache Beam worker dependencies from the Beam Python 3.6 SDK image.
COPY --from=apache/beam_python3.6_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Apache Beam worker expects pip at /usr/local/bin/pip by default.
# Some images have pip in a different location. If necessary, make a symlink.
# This line can be omitted in Beam 2.30.0 and later versions.
RUN [[ `which pip` == "/usr/local/bin/pip" ]] || ln -s `which pip` /usr/local/bin/pip

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

TensorFlow Docker イメージを使用する場合は、TensorFlow 2.5.0 以降を使用します。以前の TensorFlow Docker イメージでは、tensorflow パッケージではなく tensorflow-gpu パッケージがインストールされます。この違いは、TensorFlow 2.1.0 リリース以降では重要ではありませんが、tfx など、ダウンストリームのパッケージには、tensorflow パッケージを必要とするものが複数あります。

コンテナサイズが大きくなると、ワーカーの起動時間が遅くなります。このパフォーマンスの変化は、Deep Learning Containers などのコンテナを使用している場合に発生することがあります。

特定の Python バージョンのインストール

Python のバージョンに対する要件が厳しい場合は、必要な GPU ライブラリがある NVIDIA ベースイメージからイメージをビルドした後、Python インタープリタをインストールできます。

次の例では、Python インタープリタを含まない CUDA コンテナイメージカタログから NVIDIA イメージを選択します。目的の Python 3 と pip は、このサンプルを微修正することでインストールできます。この例では TensorFlow を使用します。したがって、イメージを選択すると、ベースイメージの CUDA バージョンと cuDNN バージョンが TensorFlow バージョンの要件を満たします。

サンプル Dockerfile は次のようになります。

# Select an NVIDIA base image with desired GPU stack from https://ngc.nvidia.com/catalog/containers/nvidia:cuda

FROM nvidia/cuda:11.0.3-cudnn8-runtime-ubuntu20.04

RUN \
    # Add Deadsnakes repository that has a variety of Python packages for Ubuntu.
    # See: https://launchpad.net/~deadsnakes/+archive/ubuntu/ppa
    apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F23C5A6CF475977595C89F51BA6932366A755776 \
    && echo "deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu focal main" >> /etc/apt/sources.list.d/custom.list \
    && echo "deb-src http://ppa.launchpad.net/deadsnakes/ppa/ubuntu focal main" >> /etc/apt/sources.list.d/custom.list \
    && apt-get update \
    && apt-get install -y curl \
        python3.8 \
        # With python3.8 package, distutils need to be installed separately.
        python3-distutils \
    && rm -rf /var/lib/apt/lists/* \
    && update-alternatives --install /usr/bin/python python /usr/bin/python3.8 10 \
    && curl https://bootstrap.pypa.io/get-pip.py | python \
    && pip install --upgrade pip \
    # Install Apache Beam and Python packages that will interact with GPUs.
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 tensorflow==2.4.0 \
    # Verify that there are no conflicting dependencies.
    && pip check

# Copy the Apache Beam worker dependencies from the Beam Python 3.8 SDK image.
COPY --from=apache/beam_python3.8_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

OS ディストリビューションによっては、OS パッケージマネージャーを使用して特定の Python バージョンをインストールすることが難しい場合があります。このような場合は、Miniconda や pyenv などのツールを使用して Python インタープリタをインストールできます。

サンプル Dockerfile は次のようになります。

FROM nvidia/cuda:11.0.3-cudnn8-runtime-ubuntu20.04

# The Python version of the Dockerfile must match the Python version you use
# to launch the Dataflow job.

ARG PYTHON_VERSION=3.8

# Update PATH so we find our new Conda and Python installations.
ENV PATH=/opt/python/bin:/opt/conda/bin:$PATH

RUN apt-get update \
    && apt-get install -y wget \
    && rm -rf /var/lib/apt/lists/* \
    # The NVIDIA image doesn't come with Python pre-installed.
    # We use Miniconda to install the Python version of our choice.
    && wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
    && bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \
    && rm Miniconda3-latest-Linux-x86_64.sh \
    # Create a new Python environment with desired version, and install pip.
    && conda create -y -p /opt/python python=$PYTHON_VERSION pip \
    # Remove unused Conda packages, install necessary Python packages via pip
    # to avoid mixing packages from pip and Conda.
    && conda clean -y --all --force-pkgs-dirs \
    && pip install --upgrade pip \
    # Install Apache Beam and Python packages that will interact with GPUs.
    && pip install --no-cache-dir apache-beam[gcp]==2.29.0 tensorflow==2.4.0 \
    # Verify that there are no conflicting dependencies.
    && pip check \
    # Apache Beam worker expects pip at /usr/local/bin/pip by default.
    # You can omit this line when using Beam 2.30.0 and later versions.
    && ln -s $(which pip) /usr/local/bin/pip

# Copy the Apache Beam worker dependencies from the Apache Beam SDK for Python 3.8 image.
COPY --from=apache/beam_python3.8_sdk:2.29.0 /opt/apache/beam /opt/apache/beam

# Set the entrypoint to Apache Beam SDK worker launcher.
ENTRYPOINT [ "/opt/apache/beam/boot" ]

方法 2:Apache Beam コンテナイメージの使用

GPU を使用するためのコンテナイメージは、事前構成されたイメージを使用しなくても構成できます。この方法は、事前構成されたイメージが機能しない場合に使用します。独自のコンテナイメージを設定するには、互換性のあるライブラリを選択し、その実行環境を構成する必要があります。

サンプル Dockerfile は次のようになります。

FROM apache/beam_python3.7_sdk:2.24.0
ENV INSTALLER_DIR="/tmp/installer_dir"

# The base image has TensorFlow 2.2.0, which requires CUDA 10.1 and cuDNN 7.6.
# You can download cuDNN from NVIDIA website
# https://developer.nvidia.com/cudnn
COPY cudnn-10.1-linux-x64-v7.6.0.64.tgz $INSTALLER_DIR/cudnn.tgz
RUN \
    # Download CUDA toolkit.
    wget -q -O $INSTALLER_DIR/cuda.run https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda_10.1.243_418.87.00_linux.run && \

    # Install CUDA toolkit. Print logs upon failure.
    sh $INSTALLER_DIR/cuda.run --toolkit --silent || (egrep '^\[ERROR\]' /var/log/cuda-installer.log && exit 1) && \
    # Install cuDNN.
    mkdir $INSTALLER_DIR/cudnn && \
    tar xvfz $INSTALLER_DIR/cudnn.tgz -C $INSTALLER_DIR/cudnn && \

    cp $INSTALLER_DIR/cudnn/cuda/include/cudnn*.h /usr/local/cuda/include && \
    cp $INSTALLER_DIR/cudnn/cuda/lib64/libcudnn* /usr/local/cuda/lib64 && \
    chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn* && \
    rm -rf $INSTALLER_DIR

# A volume with GPU drivers will be mounted at runtime at /usr/local/nvidia.
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/nvidia/lib64:/usr/local/cuda/lib64

コンテナ内で /usr/local/nvidia/lib64 のドライバライブラリを共有ライブラリとして検出できるように、LD_LIBRARY_PATH 環境変数を構成する必要があります。

TensorFlow を使用する場合は、互換性のある CUDA ツールキットと cuDNN バージョンの組み合わせを選択する必要があります。詳細については、ソフトウェア要件とテスト済みのビルド構成をご覧ください。

Dataflow ワーカーの GPU のタイプと数の選択

Dataflow では、worker_accelerator パラメータを使用して Dataflow ワーカーに接続する GPU のタイプと数を構成できます。ユースケースとパイプラインでの GPU 利用方法に基づいて、GPU のタイプと数を選択できます。

Dataflow では、次の GPU タイプがサポートされています。

NVIDIA® A100
NVIDIA® Tesla® T4
NVIDIA® Tesla® P4
NVIDIA® Tesla® V100
NVIDIA® Tesla® P100
NVIDIA® Tesla® K80

パフォーマンスデータなど、各 GPU タイプの詳細については、GPU の比較表をご覧ください。

GPU を使用したジョブの実行

GPU を使用して Dataflow ジョブを実行するには、次のコマンドを使用します。

Python

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --worker_harness_container_image "IMAGE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --dataflow_service_options "worker_accelerator=type:GPU_TYPE;count:GPU_COUNT;install-nvidia-driver" \
  --experiments "use_runner_v2"

次のように置き換えます。

PIPELINE: パイプラインのソースコードファイル
PROJECT: Google Cloud プロジェクト名
BUCKET: Cloud Storage バケット
REGION: リージョンエンドポイント（例: us-central1）
IMAGE: Docker イメージの Container Registry パス
DISK_SIZE_GB: 各ワーカー VM のブートディスクのサイズ（例: 50）
GPU_TYPE: 利用可能な GPU のタイプ（例: nvidia-tesla-t4）
GPU_COUNT: 各ワーカー VM に接続する GPU の数（例: 1）

GPU を使用した Dataflow ジョブを実行する際の考慮事項は次のとおりです。

GPU を使用してカスタムコンテナをジョブに提供するには、Dataflow Runner v2 を使用する必要があります。
REGION には、GPU_TYPE をサポートするゾーンがあるリージョンを選択します。Dataflow は、このリージョン内の GPU があるゾーンにワーカーを自動的に割り当てます。
通常、GPU コンテナは大きいため、ディスク容量が不足しないように、デフォルトのブートディスクサイズを 50 GB 以上に増やします。

TensorFlow を使用する場合は、パイプラインオプション --experiments=no_use_multiple_sdk_containers を設定するか、1 つの vCPU でワーカーを使用することで、ワーカーが単一プロセスを使用するように構成することを検討してください。n1-standard-1 に十分なメモリがない場合は、n1-custom-1-NUMBER_OF_MB や n1-custom-1-NUMBER_OF_MB-ext（拡張メモリ）などのカスタムマシンタイプを検討してください。詳細については、GPU とワーカーの並列処理をご覧ください。

Dataflow ジョブの確認

ジョブが GPU を持つワーカー VM を使用していることを確認する手順は次のとおりです。

ジョブの Dataflow ワーカーが開始されたことを確認します。
ジョブの実行中に、ジョブに関連付けられているワーカー VM を見つけます。
1. [プロダクトとリソースの検索] プロンプトでジョブ ID を貼り付けます。
2. ジョブに関連付けられた Compute Engine VM インスタンスを選択します。

Compute Engine コンソールで実行中のインスタンスのリストを確認することもできます。

Google Cloud コンソールで、[VM インスタンス] ページに移動します。

[VM インスタンス] に移動
[VM インスタンスの詳細] をクリックします。
詳細ページに [GPU] セクションがあり、GPU が接続されていることを確認します。

ジョブが GPU で開始されなかった場合は、worker_accelerator サービスオプションが正しく構成され、Dataflow モニタリングインターフェースの dataflow_service_options に表示されていることを確認します。アクセラレータメタデータ内のトークンの順序は重要です。

たとえば、Dataflow モニタリングインターフェースの dataflow_service_options パイプラインオプションは次のようになります。

['worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver', ...]

GPU 使用率の表示

ワーカー VM の GPU 使用率を表示する手順は次のとおりです。

Google Cloud コンソールで、[Monitoring] に移動するか、次のボタンを使用します。

[Monitoring] に移動
[Monitoring] のナビゲーションパネルで、[Metrics Explorer] をクリックします。
モニタリングする指標に応じて、[Resource type] に [Dataflow Job] を指定し、[Metric] に [GPU utilization] または [GPU memory utilization] を指定します。

詳細については、Metrics Explorer のガイドをご覧ください。

Dataflow Prime での GPU の使用

Dataflow Prime では、パイプラインの特定のステップにアクセラレータをリクエストできます。GPU を Dataflow Prime で使用する場合、--dataflow-service_options=worker_accelerator パイプラインオプションを使用しないでください。代わりに、accelerator リソースヒントを指定して GPU をリクエストします。詳細については、リソースヒントの使用をご覧ください。

Dataflow ジョブのトラブルシューティング

GPU を使用した Dataflow ジョブの実行で問題が発生した場合は、次の手順でトラブルシューティングを行ってください。

ワーカーが起動しない

ジョブが動かなくなり、Dataflow ワーカーがデータの処理をまったく始めない場合は、Dataflow でカスタムコンテナを使うことに関連する問題が発生した可能性があります。詳細については、カスタムコンテナのトラブルシューティングガイドをご覧ください。

Python を使用している場合は、次の条件が満たされていることを確認します。

コンテナイメージの Python インタープリタのマイナーバージョンは、パイプラインの起動時に使用するバージョンと同じであること。不一致があると、apache_beam/internal/pickler.py を含むスタックトレースで SystemError: unknown opcode のようなエラーが発生しています。
Apache Beam SDK 2.29.0 以前を使用している場合は、イメージ上で /usr/local/bin/pip の pip にアクセスできる必要があります。

初めてカスタムイメージを使用する場合は、カスタマイズを最小限の動作構成に抑えることをおすすめします。このページの例で示されているサンプルのカスタムコンテナイメージを使用します。GPU をリクエストせずに、このコンテナイメージで単純な Dataflow パイプラインを実行できることを確認します。その後、解決策を繰り返します。

コンテナイメージのダウンロードに十分なディスク容量がワーカーにあることを確認します。必要に応じてディスクサイズを調整します。イメージが大きいとダウンロードに時間がかかり、ワーカーの起動時間が長くなります。

起動時にジョブがすぐに失敗する

ZONE_RESOURCE_POOL_EXHAUSTED または ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS エラーが発生する場合は、次の手順を行います。

Dataflow が最適なゾーンを選択するように、ワーカーゾーンを指定しないようにします。
パイプラインを異なるゾーンで起動するか、別のアクセラレータタイプで起動します。

実行時にジョブが失敗する

ジョブが実行時に失敗した場合は、ワーカーマシンと GPU のメモリ不足（OOM）エラーを確認します。GPU の OOM エラーは、ワーカーのログで cudaErrorMemoryAllocation out of memory エラーとして現れる場合があります。TensorFlow を使用している場合は、1 つの GPU デバイスへのアクセスに 1 つの TensorFlow プロセスのみ使用していることを確認します。詳細については、GPU とワーカーの並列処理をご覧ください。

GPU の使用状況がない

パイプラインは問題なく実行されていても、GPU が使用されていない場合は、次の点を確認します。

コンテナイメージにインストールされている NVIDIA ライブラリが、パイプラインのユーザーコードと使用するライブラリの要件と一致する。
コンテナイメージにインストールされた NVIDIA ライブラリが、共有ライブラリとしてアクセスできる。

使用可能なデバイスがない場合は、互換性のないソフトウェア構成を使用している可能性があります。たとえば、TensorFlow を使用している場合は、TensorFlow、cuDNN のバージョン、CUDA ツールキットのバージョンが互換性のある組み合わせになっていることを確認します。

イメージ構成を検証するには、GPU が利用可能でワーカーにアクセス可能であることを確認するだけの単純なパイプラインを実行することを検討してください。

スタンドアロン VM のデバッグ

コンテナイメージを設計して繰り返し作業を行う間は、スタンドアロン VM でコンテナイメージを試用することで、フィードバックループを短縮して高速化できます。

スタンドアロン VM でカスタムコンテナをデバッグするには、以下に示すように Container-Optimized OS で GPU を使用する Compute Engine VM を作成し、ドライバをインストールしてコンテナを起動します。

VM インスタンスを作成します。

gcloud compute instances create INSTANCE_NAME \
  --project "PROJECT" \
  --image-family cos-stable \
  --image-project=cos-cloud  \
  --zone=us-central1-f \
  --accelerator type=nvidia-tesla-t4,count=1 \
  --maintenance-policy TERMINATE \
  --restart-on-failure  \
  --boot-disk-size=200G \
  --scopes=cloud-platform

ssh を使用して VM に接続します。

gcloud compute ssh INSTANCE_NAME --project "PROJECT"

GPU ドライバをインストールします。ssh を使用して VM に接続した後、VM で次のコマンドを実行します。

# Run these commands on the virtual machine
cos-extensions install gpu
sudo mount --bind /var/lib/nvidia /var/lib/nvidia
sudo mount -o remount,exec /var/lib/nvidia
/var/lib/nvidia/bin/nvidia-smi

カスタムコンテナを起動します。

Apache Beam SDK コンテナは、/opt/apache/beam/boot エントリポイントを使用します。デバッグでは、別のエントリポイントを使用してコンテナを手動で起動できます。
```
docker-credential-gcr configure-docker
docker run --rm \
  -it \
  --entrypoint=/bin/bash \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin \
  --privileged \
  IMAGE
```
IMAGE は、Docker イメージの Container Registry パスに置き換えます。
コンテナにインストールされている GPU ライブラリが、GPU デバイスにアクセスできることを確認します。

TensorFlow を使用している場合は、次のように、Python インタープリタで利用可能なデバイスを出力できます。
```
>>> import tensorflow as tf
>>> print(tf.config.list_physical_devices("GPU"))
```
PyTorch を使用している場合は、次のように Python インタープリタで利用可能なデバイスを調べることができます。
```
>>> import torch
>>> print(torch.cuda.is_available())
>>> print(torch.cuda.device_count())
>>> print(torch.cuda.get_device_name(0))
```

パイプラインで反復処理を行うには、Direct Runner でパイプラインを起動します。また、この環境から Dataflow Runner でパイプラインを起動することもできます。

詳しくは以下をご覧ください。

次のステップ

Dataflow での GPU サポートについて確認する。
GPU を使用した Landsat 衛星画像の処理に取り組む。

GPU の使用

Apache Beam ノートブックの使用

GPU 割り当てのプロビジョニング

GPU ドライバのインストール

カスタム コンテナ イメージのビルド

方法 1:GPU を使用するために構成された既存のイメージの使用

特定の Python バージョンのインストール

方法 2:Apache Beam コンテナ イメージの使用

Dataflow ワーカーの GPU のタイプと数の選択

GPU を使用したジョブの実行

Python

Dataflow ジョブの確認

GPU 使用率の表示

Dataflow Prime での GPU の使用

Dataflow ジョブのトラブルシューティング

ワーカーが起動しない

起動時にジョブがすぐに失敗する

実行時にジョブが失敗する

GPU の使用状況がない

スタンドアロン VM のデバッグ

次のステップ

カスタムコンテナイメージのビルド

方法 2:Apache Beam コンテナイメージの使用