使用 GPU 运行流水线

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

如需详细了解如何将 GPU 与 Dataflow 搭配使用，请参阅 Dataflow 对 GPU 的支持。如需详细了解使用 GPU 构建流水线的开发者工作流，请参阅将 GPU 与 Dataflow 搭配使用的简介。

使用 Apache Beam 笔记本

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

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

预配 GPU 配额

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

安装 GPU 驱动程序

如需在 Dataflow 工作器上安装 NVIDIA 驱动程序，请将 install-nvidia-driver 附加到 worker_accelerator 服务选项。

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

如需自定义容器映像，请提供满足 Apache Beam SDK 容器映像合同并具有所需 GPU 库的映像。

如需提供自定义容器映像，请使用 Dataflow Runner v2 并使用 sdk_container_image 流水线选项提供容器映像。如果您使用的是 Apache Beam 2.29.0 或更低版本，请使用 worker_harness_container_image 流水线选项。如需了解详情，请参阅使用自定义容器。

如需构建自定义容器映像，请使用以下两种方法之一：

使用已为使用 GPU 而配置的现有映像
使用 Apache Beam 容器映像

使用已为使用 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 --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-gpu 软件包而不是 tensorflow 软件包。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" ]

在某些操作系统发行版上，使用操作系统软件包管理系统安装特定 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 \
    && 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 或更多。
考虑有多少进程在工作器虚拟机上同时使用同一 GPU。然后，决定您要将 GPU 限制为单个进程还是允许多个进程使用 GPU。
- 如果一个 Apache Beam SDK 进程可以使用大多数可用的 GPU 内存（例如通过将大型模型加载到 GPU 上），则建议您设置流水线选项 --experiments=no_use_multiple_sdk_containers 以将工作器配置为使用单个进程。或者，使用具有一个 vCPU 的工作器，方法是将自定义机器类型（例如 n1-custom-1-NUMBER_OF_MB 或 n1-custom-1-NUMBER_OF_MB-ext）用于扩展内存。如需了解详情，请参阅使用每个 vCPU 具有更多内存的机器类型。
- 如果 GPU 由多个进程共享，请使用 NVIDIA 多处理服务 (MPS) 在共享 GPU 上启用并发处理。
如需了解背景信息，请参阅 GPU 和工作器并行性。

如需使用 GPU 运行 Dataflow 作业，请使用以下命令。如需使用适配，请使用 accelerator 资源提示，而不是使用 worker_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 的 `REGION`。Dataflow 会自动将工作器分配到此区域中具有 GPU 的可用区。
IMAGE：Docker 映像的 Artifact Registry 路径
DISK_SIZE_GB：每个工作器虚拟机的启动磁盘大小（例如 50）
GPU_TYPE：可用的 GPU 类型（例如 nvidia-tesla-t4）。
GPU_COUNT：要挂接到每个工作器虚拟机的 GPU 数量（例如 1）

验证 Dataflow 作业

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

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

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

在 Google Cloud 控制台中，转到虚拟机实例页面。

转到虚拟机实例
点击虚拟机实例详情。
验证详情页面是否包含 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 利用率

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

在 Google Cloud 控制台中，前往 Monitoring 或使用以下按钮：

前往监控
在 Monitoring 导航窗格中，点击 Metrics Explorer。
对于资源类型，请指定 Dataflow Job。对于指标，请根据您要监控的指标指定 GPU utilization 或 GPU memory utilization。

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

启用 NVIDIA 多处理服务

对于在具有多个 vCPU 的工作器上运行的 Python 流水线，您可以通过启用 NVIDIA 多进程服务 (MPS) 来提升 GPU 操作的并发性。如需了解详情以及使用 MPS 的步骤，请参阅使用 NVIDIA MPS 提高共享 GPU 的性能。

将 GPU 与 Dataflow Prime 搭配使用

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

排查 Dataflow 作业问题

如果您在使用 GPU 运行 Dataflow 作业时遇到问题，请参阅排查 Dataflow GPU 作业问题。

后续步骤

详细了解 Dataflow 上的 GPU 支持。
使用 NVIDIA L4 GPU 类型运行机器学习推理流水线。
学习使用 GPU 处理 Landsat 卫星图像。