使用 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_TYPEWORKER_ZONE
  • 容器 IMAGE URI 应包含一个标记。您可以使用 :latest,而不是省略标记。

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

验证 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', ...]

排查 Dataflow 作业问题

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

工作器无法启动

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

如果您遇到 ZONE_RESOURCE_POOL_EXHAUSTEDZONE_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 设备。

后续步骤