使用Vertex AI Pipelines运行带GPU的自定义Docker容器：如何安装NVIDIA驱动？

核心认知：无需手动安装CUDA驱动 #

Vertex AI Pipelines的GPU执行环境（基于GCE或GKE的运行时），底层节点已由Google预置适配的NVIDIA CUDA驱动，你不需要在容器内或运行时手动安装驱动。只需保证容器内的CUDA Toolkit版本与底层驱动版本兼容即可。

实现步骤 #

1. 构建兼容的自定义Docker镜像 #

使用NVIDIA官方基础镜像作为底座，安装PyTorch及应用依赖：

# 选择与Vertex AI底层驱动兼容的CUDA版本（推荐11.x系列，适配多数GCP GPU节点）
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

# 安装匹配CUDA版本的PyTorch
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 复制应用代码并设置工作目录
COPY your_pytorch_app.py /app/
WORKDIR /app

# 定义启动命令
CMD ["python", "your_pytorch_app.py"]

2. 配置Vertex AI Pipeline的GPU运行时 #

用KFP代码定义组件

from kfp.v2 import dsl
from kfp.v2.dsl import component

@component(base_image="gcr.io/your-project/your-custom-gpu-image:latest")
def pytorch_gpu_task():
    import your_pytorch_app
    your_pytorch_app.run()

@dsl.pipeline(
    name="gpu-custom-container-pipeline",
    pipeline_root="gs://your-bucket/pipeline-root"
)
def pipeline():
    gpu_task = pytorch_gpu_task()
    # 指定GPU资源
    gpu_task.set_accelerator_type('NVIDIA_TESLA_T4')
    gpu_task.set_accelerator_count(1)

用Vertex AI UI配置

创建任务时选择自定义镜像，在「资源配置」中指定GPU类型和数量即可。

3. 版本兼容性验证 #

运行时可通过以下命令查看底层驱动版本，确保容器内CUDA Toolkit版本不超过驱动支持的最高版本：

nvidia-smi

例：驱动版本525.x支持最高CUDA 11.7，容器内就不要用高于11.7的CUDA Toolkit。

灵活扩展方案 #

若不想依赖GCP深度学习容器，仍可基于NVIDIA官方镜像自由扩展，只需注意：

• 容器内安装的依赖不要破坏CUDA运行时环境
• 禁止在容器内安装NVIDIA驱动（会与底层节点驱动冲突）

内容的提问来源于stack exchange，提问作者skeller88