使用 GPU

本页面介绍了如何使用 GPU 运行 Dataflow 作业。在预览版期间，使用 GPU 的作业按 Dataflow 价格页面中指定的费用收费。

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

预配 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 交互，您可能需要其他 NVIDIA 软件，例如 GPU 加速的库和 CUDA 工具包。您必须在运行用户代码的 Docker 容器中提供这些库。

通过提供满足 Apache Beam SDK 容器映像合同并具有所需库的映像，或者通过基于随 Apache Beam SDK 版本发布的映像构建映像，您可以自定义容器映像。

如需提供自定义容器映像，您必须使用 Dataflow Runner v2 并使用 worker_harness_container_image 流水线选项提供容器映像。

如需了解详情，请参阅使用自定义容器。

如果您想在 GPU 上对流水线进行原型设计作为概念验证，但又不想费心设置开发环境，请参阅使用 Apache Beam 笔记本进行开发，启动 Apache Beam 笔记本实例，并按照将 GPU 与 Apache Beam 搭配使用示例操作。

使用 GPU 软件通过预先存在的基础映像构建自定义容器映像

您可以基于现有基础映像构建满足 Apache Beam SDK 容器合同的 Docker 映像。例如，TensorFlow Docker 映像和 AI Platform Deep Learning Containers 已预先配置使用 GPU。

示例 Dockerfile 如下所示：

# Use a GPU-enabled Tensorflow Docker image. The image has Python 3.6.
FROM tensorflow/tensorflow:2.4.0-gpu
RUN pip install --no-cache-dir apache-beam[gcp]==2.26.0

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

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

验证映像中的 Apache Beam 版本和 Python 解释器次要版本是否与用于启动流水线的版本匹配。为获得最佳结果，请考虑使用所选基础映像和 Apache Beam 的最新版本。

基于 Apache Beam SDK 映像构建自定义容器映像

如需基于 Apache Beam SDK 版本发布的 Docker 映像构建映像，您必须将所需的 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_zone "WORKER_ZONE" \
  --worker_harness_container_image "IMAGE" \
  --experiment "worker_accelerator=type:GPU_TYPE;count:GPU_COUNT;install-nvidia-driver" \
  --experiment "use_runner_v2"

请替换以下内容：

PIPELINE：流水线源代码文件
PROJECT：Google Cloud 项目名称
BUCKET：Cloud Storage 存储分区
REGION：区域端点
WORKER_ZONE：用于启动工作器实例的 Compute Engine 可用区
IMAGE：Docker 映像的 Container Registry 路径
GPU_TYPE：可用的 GPU 类型
GPU_COUNT：要附加到每个工作器虚拟机的 GPU 数量

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

如需为使用 GPU 的作业提供自定义容器，必须使用 Dataflow Runner v2。
选择支持 GPU_TYPE 的 WORKER_ZONE。
容器 IMAGE URI 应包含一个标记。您可以使用 :latest，而不是省略标记。

如果您使用 TensorFlow，请考虑选择具有 1 个 vCPU 的机器类型。如果 n1-standard-1 无法提供足够的内存，您可以考虑使用自定义机器类型（例如 n1-custom-1-NUMBER_OF_MB 或 n1-custom-1-NUMBER_OF_MB-ext）以扩展内存。指定此机器类型时，NUMBER_OF_MB 必须是 256 的倍数。

验证 Dataflow 作业

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

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

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

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

转到虚拟机实例
点击虚拟机实例详情。
验证详情页面是否包含 GPU 部分，以及是否已附加 GPU。

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

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

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

排查 Dataflow 作业问题

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

工作器无法启动

如果您的作业卡住，并且 Dataflow 工作器无法启动，请验证您使用的是否是兼容的机器类型：

仅 N1 机器或自定义 N1 机器类型支持 GPU。
如果您使用的工作器具有较多的 vCPU，则可能需要更多 GPU。
如需排除与传递自定义容器相关的问题，请参阅自定义容器问题排查指南。

如果您遇到 ZONE_RESOURCE_POOL_EXHAUSTED 或 ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS 错误，请尝试在其他可用区或使用其他加速器类型启动流水线。

未使用 GPU

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

安装在工作器容器中的 NVIDIA 库符合流水线用户代码的要求。
已安装的 NVIDIA 库可作为共享库访问。

如果您使用的是支持 GPU 设备的框架（例如 TensorFlow），请确保这些框架可以访问所挂接的设备。例如，如果您使用 TensorFlow，则可以使用以下代码输出可用的设备：

import logging
import tensorflow as tf
gpu_devices = tf.config.list_physical_devices("GPU")
logging.info("GPU devices: {}".format(gpu_devices))
if len(gpu_devices) == 0:
  logging.warning("No GPUs found, defaulting to CPU")

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

如果您无法找出不匹配的情况，请尝试从有效的已知配置开始并进行迭代。为此，您可以查看为使用 GPU 预先配置的示例教程和 Docker 映像，例如 TensorFlow Docker 映像和 Deep Learning Containers 映像中的映像。

使用独立虚拟机进行调试

您可以在具有 GPU 的独立虚拟机上调试自定义容器，方法是，创建在 Container-Optimized OS 上运行 GPU 的 Compute Engine 虚拟机、安装驱动程序并启动容器。如需详细了解这些步骤，请参阅《使用入门：在 Container-Optimized OS 上运行 GPU》。

Apache Beam SDK 容器使用 /opt/apache/beam/boot 入口点。出于调试目的，您可以使用其他入口点手动启动容器，如以下示例所示：

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 设备。

后续步骤

详细了解 Dataflow 上的 GPU 支持。
学习使用 GPU 处理 Landsat 卫星图像。