使用 NVIDIA® L4 GPU 类型

本页面介绍如何使用 NVIDIA® L4 GPU 类型运行 Dataflow 流水线。L4 GPU 类型适合运行机器学习推理流水线。

要求

  • 使用 Apache Beam SDK 2.46 版或更高版本。建议使用 Apache Beam 2.50 或更高版本。
  • 运行作业的区域中需要有 L4 GPU 配额 (NVIDIA_L4_GPUS)。如需了解详情，请参阅 GPU 配额。
  • L4 GPU 类型仅适用于 G2 加速器优化机器类型。如需了解详情，请参阅 G2 机器系列。使用 L4 GPU 类型的流水线受到 G2 标准限制的约束。
  • NVIDIA L4 GPU 类型使用 NVIDIA 驱动程序 525.0 版或更高版本以及 CUDA 工具包 12.0 版或更高版本。您在流水线中使用的任何代码都必须与 NVIDIA 驱动程序版本和 CUDA 工具包版本兼容。例如，如果您使用 PyTorch，则需要使用 PyTorch 23.01 版或更高版本。

使用 NVIDIA® L4 GPU 类型运行流水线

如需使用 NVIDIA L4 GPU 类型，您需要在流水线代码中添加以下流水线选项和服务选项。

Java

--workerMachineType=G2_MACHINE_TYPE
--dataflowServiceOptions="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

Python

--machine_type=G2_MACHINE_TYPE
--dataflow_service_options="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

Go

--worker_machine_type=G2_MACHINE_TYPE
--dataflow_service_options="worker_accelerator=type:nvidia-l4;count:GPU_COUNT;install-nvidia-driver"

替换以下值：

  • G2_MACHINE_TYPE：要使用的 G2 机器类型
  • GPU_COUNT：要使用的 GPU 数量。每种 G2 机器类型都有固定数量的 NVIDIA L4 GPU。如需查找您的机器类型的正确 GPU 数量，请参阅 G2 标准机器类型表中的 GPU 数量列。

如需详细了解如何使用 GPU 运行流水线，请参阅使用 GPU 运行流水线。

管理依赖项

如需管理依赖项，请使用自定义容器。如需了解详情，请参阅在 Dataflow 中使用自定义容器。

以下 Dockerfile 示例包含使用 NVIDIA L4 GPU 类型的流水线的兼容依赖项。

RUN apt-get -y update
RUN apt-get install [system packages]

# Install the SDK.
RUN pip install --no-cache-dir apache-beam[gcp]==2.51.0
# Install the machine learning dependencies.
RUN pip install --no-cache-dir tensorflow[and-cuda]
RUN pip install xgboost
RUN pip install transformers accelerate
(etc…..)
# Verify that the image doesn't have conflicting dependencies.
RUN pip check

# Copy files from official SDK image, including the script and dependencies.
COPY --from=apache/beam_python3.10_sdk:2.51.0 /opt/apache/beam /opt/apache/beam

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

后续步骤