使用 GPU

本页面介绍了如何在具有 GPU 的 Dataflow 上运行 Apache Beam 流水线。使用 GPU 的作业按 Dataflow 价格页面中指定的费用收费。

如需详细了解如何将 GPU 与 Dataflow 搭配使用，请参阅 Dataflow 对 GPU 的支持。

使用 Apache Beam 笔记本

如果您已经有流水线，并且想在 Dataflow 上使用 GPU 运行，则可以跳过本部分。

Apache Beam 笔记本提供了一种便捷的方式，让您可以在不设置开发环境的情况下，使用 GPU 对流水线进行原型设计和迭代开发。如需开始使用，请阅读使用 Apache Beam 笔记本进行开发指南，启动 Apache Beam 笔记本实例，然后按照“将 GPU 与 Apache Beam 配合使用”示例笔记本进行操作。

预配 GPU 配额

GPU 设备受 Google Cloud 项目的配额可用性的限制。请在您选择的区域中申请 GPU 配额。

安装 GPU 驱动程序

您必须通过将 install-nvidia-driver 附加到 worker_accelerator 选项，指示 Dataflow 将 NVIDIA 驱动程序安装到工作器。指定 install-nvidia-driver 选项后，Dataflow 使用 Container-Optimized OS 提供的 cos-extensions 实用程序将 NVIDIA 驱动程序安装到 Dataflow 工作器。指定 install-nvidia-driver 即表示用户同意接受 NVIDIA 许可协议。

NVIDIA 驱动程序安装程序提供的二进制文件和库会装载到运行 /usr/local/nvidia/ 下的流水线用户代码的容器中。

GPU 驱动程序版本取决于 Dataflow 使用的 Container-Optimized OS 版本。

构建自定义容器映像

要与 GPU 交互，您可能需要其他 NVIDIA 软件，例如 GPU 加速的库和 CUDA 工具包。您必须在运行用户代码的 Docker 容器中提供这些库。

通过提供满足 Apache Beam SDK 容器映像合同并具有所需 GPU 库的映像，您可以自定义容器映像。

如需提供自定义容器映像，您必须使用 Dataflow Runner v2 并使用 worker_harness_container_image 流水线选项提供容器映像。如果您使用的是 Apache Beam 2.30.0 或更高版本，为简单起见，可以改用较短的选项名称 sdk_container_image。如需了解详情，请参阅使用自定义容器。

方法 1.使用已为使用 GPU 而配置的现有映像

您可以基于为使用 GPU 预先配置的现有基础映像，构建满足 Apache Beam SDK 容器合同的 Docker 映像。例如，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 --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 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-gpu 软件包而不是 tensorflow 软件包。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 \
    # 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 Package Manager 安装特定 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 \
    && sh 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 \
    # 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.
    # This can be omitted in 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® 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" \
  --experiments "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：每个工作器虚拟机的启动磁盘大小（例如 50）
  • GPU_TYPE：可用的 GPU 类型（例如 nvidia-tesla-t4）
  • GPU_COUNT：要挂接到每个工作器虚拟机的 GPU 数量（例如 1）

使用 GPU 运行 Dataflow 作业的注意事项如下：

  • 如需为使用 GPU 的作业提供自定义容器，必须使用 Dataflow Runner v2。
  • 选择一个包含支持 GPU_TYPE 可用区的 REGION。Dataflow 会自动将工作器分配到此区域中具有 GPU 的可用区。
  • 由于 GPU 容器通常很大，因此我们建议您将默认启动磁盘大小增加到 50 GB 或更多，以避免耗尽磁盘空间。

如果您使用 TensorFlow，请考虑通过设置流水线选项 --experiments=no_use_multiple_sdk_containers 或使用具有一个 vCPU 的工作器，将工作器配置为使用单个进程。如果 n1-standard-1 无法提供足够的内存，您可以考虑使用自定义机器类型（例如 n1-custom-1-NUMBER_OF_MB 或 n1-custom-1-NUMBER_OF_MB-ext）以扩展内存。如需了解详情，请参阅 GPU 和工作器并行。

验证 Dataflow 作业

要确认作业使用具有 GPU 的工作器虚拟机，请按以下步骤操作：

  1. 验证作业的 Dataflow 工作器是否已启动。
  2. 在作业运行期间，查找与该作业关联的工作器虚拟机。
    1. 在搜索产品和资源提示框中粘贴作业 ID。
    2. 选择与作业关联的 Compute Engine 虚拟机实例。

您还可以在 Compute Engine 控制台中找到所有正在运行的实例的列表。

  1. 在 Google Cloud Console 中，转到虚拟机实例页面。

    转到虚拟机实例

  2. 点击虚拟机实例详情。

  3. 验证详情页面是否包含 GPU 部分，以及是否已附加 GPU。

如果作业未使用 GPU 启动，请检查 --worker_accelerator 实验是否配置正确并且显示在 Dataflow 监控界面的 experiments 中。加速器元数据中的令牌顺序非常重要。

例如，Dataflow 监控界面中的“实验”流水线选项可能如下所示：

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

查看 GPU 利用率

如需查看工作器虚拟机上的 GPU 利用率，请执行以下操作：

  1. 在 Google Cloud Console 中，转到 Monitoring 或使用以下按钮：

    转至 Monitoring

  2. 在 Monitoring 导航窗格中，点击 Metrics Explorer。

  3. 将 Dataflow Job 指定为资源类型，并将 GPU utilization 或 GPU memory utilization 指定为指标，具体取决于您要监控的指标。

如需了解详情，请参阅 Metrics Explorer 指南。

将 GPU 与 Dataflow Prime 搭配使用

利用 Dataflow Prime，您可以为流水线的特定步骤请求加速器。如需搭配使用 GPU 和 Dataflow Prime，请勿使用 --experiments=worker_accelerator 流水线选项。请改为使用 accelerator 资源提示请求 GPU。如需了解详情，请参阅使用资源提示。

排查 Dataflow 作业问题

如果您在使用 GPU 运行 Dataflow 作业时遇到问题，请按照下面可能解决问题的问题排查步骤操作。

工作器无法启动

如果您的作业卡住，并且 Dataflow 工作器一直不开始处理数据，则可能是使用自定义容器和 Dataflow 时出现问题。如需了解详情，请参阅自定义容器问题排查指南。

如果您是 Python 用户，请验证是否满足以下条件：

  • 容器映像中的 Python 解释器次要版本与您启动流水线时使用的版本相同。如果不匹配，您可能会看到类似 SystemError: unknown opcode 的错误，并且堆栈轨迹涉及 apache_beam/internal/pickler.py。
  • 如果您使用的是 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，请验证只使用一个 TensorFlow 进程访问一个 GPU 设备。如需了解详情，请参阅 GPU 和工作器并行。

未使用 GPU

如果流水线运行成功，但未使用 GPU，请验证以下内容：

  • 容器映像中安装的 NVIDIA 库符合流水线用户代码和它使用的库的要求。
  • 容器映像中已安装的 NVIDIA 库可作为共享库访问。

如果这些设备不可用，则您使用的软件配置可能不兼容。例如，如果您使用 TensorFlow，请验证您是否使用了 TensorFlow、cuDNN 版本和 CUDA 工具包版本的兼容组合。

如需验证映像配置，建议运行一个简单的流水线，该流水线只会检查 GPU 是否可用并可供工作器访问。

使用独立虚拟机进行调试

在设计和迭代适合您的容器映像时，通过在独立虚拟机上试用容器映像可以更快地缩短反馈环。

您可以在具有 GPU 的独立虚拟机上调试自定义容器，方法是，创建在 Container-Optimized OS 上运行 GPU 的 Compute Engine 虚拟机、安装驱动程序并如下所示启动容器。

  1. 创建虚拟机实例。

    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

  2. 使用 ssh 连接到虚拟机。

    gcloud compute ssh INSTANCE_NAME --project "PROJECT"

  3. 安装 GPU 驱动程序。通过 ssh 连接到虚拟机后，请在虚拟机上运行以下命令：

    # 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

  4. 启动自定义容器。

    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 路径。

  5. 验证容器中安装的 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 上的流水线。

如需了解详情，请参阅：

后续步骤