使用 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"]

後續步驟