GPU を使用してパイプラインを実行する

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

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

Apache Beam ノートブックを使用する

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

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

GPU 割り当てをプロビジョニングする

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

GPU ドライバをインストールする

Dataflow ワーカーに NVIDIA ドライバをインストールするには、worker_accelerator サービスオプションに install-nvidia-driver を追加します。

install-nvidia-driver オプションを指定すると、Dataflow は Container-Optimized OS によって提供される cos-extensions ユーティリティを使用して、NVIDIA ドライバを Dataflow ワーカーにインストールします。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 を使用し、sdk_container_image パイプラインオプションでコンテナイメージを指定します。Apache Beam バージョン 2.29.0 以前を使用している場合は、worker_harness_container_image パイプラインオプションを使用します。詳細については、カスタムコンテナを使用するをご覧ください。

カスタムコンテナイメージをビルドするには、次のいずれかの方法を使用します。

GPU 用に構成された既存のイメージを使用する
Apache Beam コンテナイメージを使用する

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 インタープリタをインストールできます。

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

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

# Select an NVIDIA base image with needed 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 バージョンをインストールすることが難しい場合があります。その場合は、Minconda や 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 needed 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" ]

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 ワーカーに接続する GPU のタイプと数を構成するには、worker_accelerator サービスオプションを使用します。ユースケースとパイプラインでの GPU 利用方法に基づいて、GPU のタイプと数を選択します。

Dataflow でサポートされている GPU タイプのリストについては、Dataflow での GPU のサポートをご覧ください。

GPU を使用してジョブを実行する

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

通常、GPU コンテナは大きいため、ディスク容量が不足しないように、次の操作を行います。
- デフォルトのブートディスクサイズを 50 GB 以上に増やします。
ワーカー VM で同じ GPU を同時に使用するプロセスの数を検討します。次に、GPU を単一のプロセスに制限するか、複数のプロセスで GPU を使用できるようにするかを決定します。
- 1 つの Apache Beam SDK プロセスで、使用可能な GPU メモリのほとんどを使用できる場合（たとえば、GPU に大きなモデルを読み込む場合など）、パイプラインオプション --experiments=no_use_multiple_sdk_containers を設定して単一のプロセスを使用するようにワーカーを構成することをおすすめします。または、拡張メモリに対応した、n1-custom-1-NUMBER_OF_MB または n1-custom-1-NUMBER_OF_MB-ext などのカスタムマシンタイプを使って、1 つの vCPU を持つワーカーを使用します。詳細については、vCPU あたりのメモリ容量が大きいマシンタイプを使用するをご覧ください。
- GPU が複数のプロセスで共有されている場合は、NVIDIA Multi-Processing Service（MPS）を使用して、共有 GPU で同時処理を有効にします。
背景情報については、GPU とワーカーの並列処理をご覧ください。

GPU を使用して Dataflow ジョブを実行するには、次のコマンドを使用します。Right Fitting を使用するには、worker_accelerator サービスオプションではなく、accelerator リソースヒントを使用します。

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: Dataflow リージョン（例: us-central1）。GPU_TYPE をサポートするゾーンが存在する「リージョン」を選択します。Dataflow は、このリージョン内の GPU があるゾーンにワーカーを自動的に割り当てます。
IMAGE: Docker イメージの Artifact Registry パス
DISK_SIZE_GB: 各ワーカー VM のブートディスクのサイズ（例: 50）
GPU_TYPE: 利用可能な GPU のタイプ（例: nvidia-tesla-t4）
GPU_COUNT: 各ワーカー VM に接続する GPU の数（例: 1）

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] をクリックします。
リソースの種類には Dataflow Job を指定します。指標には、モニタリングする指標に応じて GPU utilization または GPU memory utilization を指定します。

詳しくは、Metrics Explorer をご確認ください。

NVIDIA Multi-Processing Service を有効にする

複数の vCPU を持つワーカーで実行される Python パイプラインでは、NVIDIA Multi-Processing Service（MPS）を有効にすることで、GPU オペレーションの同時実行を改善できます。MPS の詳細と設定手順については、NVIDIA MPS を使用して共有 GPU のパフォーマンスを改善するをご覧ください。

Dataflow Prime で GPU を使用する

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

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

GPU で Dataflow ジョブを実行する際に問題が発生した場合は、Dataflow GPU ジョブのトラブルシューティングをご覧ください。

次のステップ

Dataflow での GPU サポートについて確認する。
NVIDIA L4 GPU タイプで ML 推論パイプラインを実行する。
GPU を使用した Landsat 衛星画像の処理に取り組む。

GPU を使用してパイプラインを実行する

Apache Beam ノートブックを使用する

GPU 割り当てをプロビジョニングする

GPU ドライバをインストールする

カスタム コンテナ イメージをビルドする

GPU 用に構成された既存のイメージを使用する

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

Apache Beam コンテナ イメージを使用する

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

GPU を使用してジョブを実行する

Python

Dataflow ジョブを確認する

GPU 使用率を表示する

NVIDIA Multi-Processing Service を有効にする

Dataflow Prime で GPU を使用する

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

次のステップ

カスタムコンテナイメージをビルドする

Apache Beam コンテナイメージを使用する