使用 Dataflow GPU 时的最佳做法

本页面介绍了使用 GPU 构建流水线的最佳做法。

如需了解如何在 Dataflow 作业中启用 GPU 的相关信息和示例，请参阅使用 GPU 运行流水线和使用 GPU 处理 Landsat 卫星图像。

在 Dataflow 中使用 GPU 的前提条件

• 如需在 Dataflow 作业中使用 GPU，您必须使用 Runner v2。
• Dataflow 会在 Docker 容器内的工作器虚拟机中运行用户代码。这些工作器虚拟机运行 Container-Optimized OS。 为了让 Dataflow 作业使用 GPU，您需要满足以下前提条件：
  • GPU 驱动程序安装在工作器虚拟机上，并可供 Docker 容器访问。如需了解详情，请参阅安装 GPU 驱动程序。
  • 您的流水线所需的 GPU 库（例如 NVIDIA CUDA-X 库或 NVIDIA CUDA 工具包）安装在自定义容器映像中。如需了解详情，请参阅配置容器映像。
• 由于 GPU 容器通常很大，因此为了避免耗尽磁盘空间，请将默认启动磁盘大小增加到 50 GB 或更多。

注意事项

在设计测试环境和生产环境时，请考虑以下因素。

本地开发

通过将 Apache Beam 与 NVIDIA GPU 结合使用，您可以创建用于进行预处理和推理的大规模数据处理流水线。使用 GPU 进行本地开发时，请考虑以下信息：

• 通常，数据处理工作流会使用额外的库，这些库需要安装在 Dataflow 工作器上的启动环境和执行环境中。此配置向开发工作流添加了用于配置流水线要求或在 Dataflow 中使用自定义容器的步骤。最好是具有尽可能模拟生产环境的本地开发环境。

• 如果您的工作流同时满足以下两个条件，则无需构建自定义容器或更改开发工作流来配置流水线要求：

  • 您使用的库隐式使用了 NVIDIA GPU。
  • 您的代码无需进行任何更改即可支持 GPU。

• 某些库不能以透明方式在 CPU 和 GPU 使用之间切换，因此需要特定的构建和不同的代码路径。如需复制此场景的“代码-运行-代码”开发生命周期，您将需要执行一些额外的步骤。

• 运行本地实验时，请尽可能近地复制 Dataflow 工作器的环境。根据具体库，您可能需要使用安装有 GPU 和所需 GPU 库的机器。此机器类型可能不会在您的本地环境中提供。您可以使用在配备有 GPU 的 Google Cloud 虚拟机上运行的容器来模拟 Dataflow Runner 环境。

机器类型规范

如需详细了解每种 GPU 型号支持的机器类型，请参阅 GPU 平台。支持 N1 机器类型的 GPU 同时也支持自定义 N1 机器类型。

GPU 的类型和数量决定了工作器可以拥有的 vCPU 和内存数量的上限。如需了解相应的限制，请参阅可用性。

指定更多 CPU 或内存时，您可能需要指定更多 GPU。

如需了解详情，请参阅 Compute Engine 上的 GPU。

优化资源使用情况

大多数流水线并不完全由需要 GPU 的转换组成。典型的流水线具有注入阶段，该阶段使用 Apache Beam 提供的众多来源之一。该阶段之后是一些数据操纵或构形转换，它们之后会馈送到 GPU 转换。

适配功能使用 Apache Beam 资源提示为批处理流水线自定义工作器资源。启用适配后，Dataflow 仅会在需要 GPU 的流水线阶段使用 GPU。因此，此功能可提高流水线的灵活性和功能，同时可能会降低费用。

如需了解详情，请参阅适配。

GPU 和工作器并行性

对于使用 Dataflow Runner v2 架构的 Python 流水线，Dataflow 会为每个虚拟机核心启动一个 Apache Beam SDK 进程。每个 SDK 进程都在其各自的 Docker 容器中运行，从而产生许多线程，每个线程都会处理传入的数据。

GPU 使用多个进程架构，并且 Dataflow 工作器中的 GPU 对所有进程和线程都可见。

如果您在共享 GPU 上运行多个 SDK 进程，可以通过启用 NVIDIA 多进程服务 (MPS) 来提高 GPU 效率和利用率。MPS 可提高 GPU 流水线的工作器并行性和总体吞吐量，尤其是对于 GPU 资源用量较低的工作负载。如需了解详情，请参阅使用 NVIDIA MPS 提高共享 GPU 的性能。

为避免 GPU 内存超额订阅，您可能需要管理 GPU 访问权限。如果您使用的是 TensorFlow，则以下任何建议都可以帮助您避免 GPU 内存超额订阅：

• 无论工作器 vCPU 数量是多少，请将 Dataflow 工作器配置为仅启动一个容器化的 Python 进程。如需进行此配置，请在启动作业时使用以下流水线选项：

  • --experiments=no_use_multiple_sdk_containers
  • --number_of_worker_harness_threads

  如需详细了解要使用多少个线程，请参阅减少线程数。

• 启用 MPS。

推理工作负载

当您使用机器学习 (ML) 模型进行本地和远程推理时，请使用内置的 Apache Beam RunInference 转换。RunInference API 是针对机器学习推断进行了优化的 PTransform。当您在流水线中使用机器学习模型时，使用 RunInference 转换可以提高效率。

工作流程

下面的两阶段工作流展示了如何使用 GPU 构建流水线。该工作流会分别处理 GPU 相关和非 GPU 相关的问题，并缩短反馈环。

1. 创建流水线

   创建可在 Dataflow 上运行的流水线。将需要 GPU 的转换替换为不使用 GPU 但提供相同功能的转换：

   1. 创建与 GPU 使用相关的所有转换，例如数据注入和操纵。

   2. 通过直通模式或架构更改为 GPU 转换创建存根。

2. 在本地测试

   在模拟 Dataflow 工作器执行环境的环境中测试流水线代码的 GPU 部分。以下步骤介绍了运行此测试的方法之一：

   1. 创建一个包含所有必要库的 Docker 映像。

   2. 开始开发 GPU 代码。

   3. 使用包含该 Docker 映像的 Google Cloud 虚拟机开始“代码-运行-代码”周期。为了排除库不兼容的问题，请在独立于 Apache Beam 流水线的本地 Python 进程中运行 GPU 代码。然后，在 Direct Runner 上运行整个流水线或在 Dataflow 上启动流水线。

使用运行容器优化型操作系统的虚拟机

对于满足最低要求的环境，请使用容器优化型虚拟机 (VM)。如需了解详情，请参阅创建挂接 GPU 的虚拟机。

常规工作流如下：

1. 创建一个虚拟机。

2. 连接到虚拟机并运行以下命令：

   sudo cos-extensions install gpu -- -version latest
   sudo mount --bind /var/lib/nvidia /var/lib/nvidia
   sudo mount -o remount,exec /var/lib/nvidia
   

3. 确认 GPU 可用：

   ./nvidia-smi
   

4. 从作为卷装载的虚拟机启动包含 GPU 驱动程序的 Docker 容器。例如：

   sudo 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 gcr.io/bigdatapivot/image_process_example:latest
   

如需查看 Dockerfile 示例，请参阅构建自定义容器映像。将流水线所需的所有依赖项都添加到 Dockerfile 中。

如需详细了解如何使用预先配置使用 GPU 的 Docker 映像，请参阅使用已为使用 GPU 而配置的现有映像。

用于使用容器优化型系统的工具

• 如需将 Docker CLI 配置为使用 docker-credential-gcr 作为默认 Google Container Registry (GCR) 注册表集的凭据帮助程序，请使用以下命令：

  sudo docker-credential-gcr configure-docker
  

  如需详细了解如何设置 Docker 凭据，请参阅 docker-credential-gcr。

• 如需向虚拟机或从虚拟机复制文件（例如流水线代码），请使用 toolbox。此方法在使用自定义优化映像时很有用。 例如：

  toolbox /google-cloud-sdk/bin/gsutil cp gs://bucket/gpu/image_process/* /media/root/home/<userid>/opencv/
  

  如需了解详情，请参阅使用工具箱调试节点问题。

后续步骤

• 详细了解 Dataflow 对 GPU 的支持。
• 详细了解如何使用 GPU 运行流水线。
• 学习使用 GPU 处理 Landsat 卫星图像。