Google Vertex AI自定义Docker镜像无法识别GPU问题求助

核心原因 #

nvidia/cuda:xx-base镜像仅包含CUDA运行时基础依赖，没有预装NVIDIA容器工具包（nvidia-container-toolkit）和设备驱动相关组件，GCP的自定义任务环境需要这些组件来暴露GPU设备到容器内。

具体修复步骤 #

1. 修改Dockerfile，添加NVIDIA工具依赖 #

在你的Dockerfile中补充以下配置，确保容器能识别GPU：

# 保留原有基础镜像
FROM nvidia/cuda:12.4.1-base-ubuntu20.04

# 安装NVIDIA容器工具包依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
    gnupg2 \
    curl \
    && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | apt-key add - \
    && curl -s -L https://nvidia.github.io/libnvidia-container/ubuntu20.04/libnvidia-container.list | tee /etc/apt/sources.list.d/nvidia-container-toolkit.list \
    && apt-get update && apt-get install -y --no-install-recommends \
    nvidia-container-toolkit \
    && rm -rf /var/lib/apt/lists/*

# 配置CUDA环境变量
ENV PATH=/usr/local/cuda/bin:$PATH
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

2. 本地验证镜像GPU识别能力 #

构建镜像后，用以下命令测试容器内是否能识别GPU：

docker run --gpus all --rm 你的镜像名称 nvidia-smi

若能正常输出GPU型号、显存等信息，说明镜像配置有效。

3. 调整GCP任务提交命令 #

显式指定加速器数量（默认1，可按需调整），确保参数完整：

gcloud ai custom-jobs create \
  --region=xxxxx \
  --display-name=xxx123 \
  --worker-pool-spec=machine-type=n1-standard-4,accelerator-type=NVIDIA_TESLA_T4,accelerator-count=1,container-image-uri=xxxxx \
  --args=xx,xxx-job-train,xxxxx

4. 在训练脚本中验证GPU可用性 #

在训练代码开头添加验证逻辑（以Python为例）：

import torch
print("GPU可用状态:", torch.cuda.is_available())
print("GPU数量:", torch.cuda.device_count())
if torch.cuda.is_available():
    print("当前GPU:", torch.cuda.get_device_name(0))

若输出GPU相关信息，说明问题已解决。

额外注意事项 #

• 确保GCP项目已启用Compute Engine API和AI Platform API
• 确认n1-standard-4机器类型支持NVIDIA_TESLA_T4加速器（GCP多数n1系列实例均支持）
• 若使用TensorFlow/PyTorch等框架，需保证框架版本与CUDA 12.4.1兼容

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