使用 Vertex AI Pipelines 通過 GPU 運行自定義 Docker 容器 - 如何安裝 NVIDIA 驅動程序？

Question

我需要使用 Vertex AI Pipelines 運行具有 GPU 支持的自定義 Docker 容器，但我沒有看到明確的方法來做到這一點。 這需要幾個組件：

1. 應用程序（在我的例子中是pytorch）
2. CUDA 工具包
3. CUDA驅動程序
4. 英偉達 GPU

我可以為 #1 和 #2 使用NVIDIA Docker基礎映像，為 #4 使用 GCP 加速器，但如何在 Vertex AI 管道中安裝 CUDA 驅動程序？ 有關於如何在 GCE 實例和GKE 節點上安裝 NVIDIA 驅動程序的文檔，但沒有關於 Vertex AI 的文檔。

一種選擇是基於 GCP Deep Learning Container 創建一個衍生容器，但我必須使用 GCP 容器並且對環境沒有那么多的控制權。

Answer 1

在我看來，您可以使用其中一個預構建的 GPU docker 容器：

https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers