本页面介绍了如何使用 GPU 运行 Dataflow 作业。在预览版期间,使用 GPU 的作业按 Dataflow 价格页面中指定的费用收费。
如需详细了解如何将 GPU 与 Dataflow 配合使用,请参阅 Dataflow 对 GPU 的支持。
预配 GPU 配额
GPU 设备受 Google Cloud 项目的配额可用性约束。在您选择的区域中申请 GPU 配额。
安装 GPU 驱动程序
您必须通过将 install-nvidia-driver
附加到 worker_accelerator
配置,指示 Dataflow 将 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 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 工作器是否已启动。
- 在作业运行时,查找与该作业关联的工作器虚拟机。
- 在搜索产品和资源提示框中粘贴作业 ID。
- 选择与作业关联的 Compute Engine 虚拟机实例。
您还可以在 Compute Engine 控制台中找到所有正在运行的实例的列表。
在 Google Cloud Console 中,转到虚拟机实例页面。
点击虚拟机实例详情。
验证详情页面是否具有 GPU 部分,以及是否挂接了 GPU。
如果您的作业没有使用 GPU 启动,请检查 experiments
中的 Dataflow 监控界面是否正确配置并显示 --worker_accelerator
实验。加速器元数据中令牌的顺序很重要。
例如,Dataflow 监控界面中的“实验”流水线选项可能如下所示:
['use_runner_v2','worker_accelerator=type:nvidia-tesla-t4;count:1;install-nvidia-driver', ...]
对 Dataflow 作业进行问题排查
如果您在使用 GPU 运行 Dataflow 作业时遇到问题,请了解可能解决问题的问题排查步骤。
工作器无法启动
如果您的作业卡住,并且 Dataflow 工作器无法启动,请验证您使用的是否是兼容的机器类型:
- GPU 仅受 N1 机器或自定义 N1 机器类型支持。
- 如果您使用的工作器具有较多的 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 映像和 AI Platform 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 卫星图像。