Docker容器GPU支持完全指南：从原理到实践验证 🚀⚡

# 添加NVIDIA容器工具包仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

# 安装工具包
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

# 配置Docker使用NVIDIA运行时
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# 检查NVIDIA驱动
nvidia-smi

# 检查CUDA工具包
nvcc --version

# 检查Docker安装
docker --version

# 检查NVIDIA容器运行时
docker info | grep -i runtime

# 基础验证：运行nvidia-smi
docker run --rm --gpus all nvidia/cuda:12.1.1-base-ubuntu20.04 nvidia-smi

# 进阶验证：运行CUDA样例
docker run --rm --gpus all nvidia/cuda:12.1.1-base-ubuntu20.04 nvidia/cuda-sample:vectorAdd

# 实际应用验证：运行PyTorch测试
docker run --rm --gpus all -it pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel \
python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}'); print(f'GPU count: {torch.cuda.device_count()}')"

# 指定特定GPU
docker run --rm --gpus '"device=0,1"' nvidia/cuda:12.1.1-base-ubuntu20.04 nvidia-smi

# 基础ROCm验证
docker run --rm --device=/dev/kfd --device=/dev/dri \
rocm/rocm-terminal:latest rocm-smi

# 运行HIP示例
docker run --rm --device=/dev/kfd --device=/dev/dri \
rocm/rocm-terminal:latest /opt/rocm/hip/bin/hipinfo

# 使用PyTorch ROCm版本验证
docker run --rm --device=/dev/kfd --device=/dev/dri \
rocm/pytorch:latest python -c "import torch; print(f'ROCm available: {torch.cuda.is_available()}')"

# Dockerfile.gpu-test
FROM nvidia/cuda:12.1.1-base-ubuntu20.04

RUN apt-get update && apt-get install -y \
    python3 \
    python3-pip \
    && rm -rf /var/lib/apt/lists/*

RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

COPY gpu_test.py /app/gpu_test.py

CMD ["python3", "/app/gpu_test.py"]

# gpu_test.py - 全面的GPU功能测试
import torch
import subprocess
import sys

def test_basic_gpu():
    """测试基础GPU功能"""
    print("=== GPU基础功能测试 ===")
    print(f"CUDA可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU数量: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            print(f"  显存: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f} GB")

def test_computation():
    """测试GPU计算能力"""
    print("\n=== GPU计算能力测试 ===")
    if torch.cuda.is_available():
        device = torch.device('cuda')
        # 创建张量并转移到GPU
        x = torch.randn(10000, 10000).to(device)
        y = torch.randn(10000, 10000).to(device)
        
        # 矩阵乘法测试
        import time
        start_time = time.time()
        z = torch.matmul(x, y)
        end_time = time.time()
        
        print(f"矩阵乘法耗时: {end_time - start_time:.2f} 秒")
        print(f"结果张量形状: {z.shape}")
        
        # 清理显存
        del x, y, z
        torch.cuda.empty_cache()

def test_nvidia_smi():
    """测试nvidia-smi命令"""
    print("\n=== nvidia-smi测试 ===")
    try:
        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
        if result.returncode == 0:
            print("nvidia-smi命令执行成功")
            # 提取关键信息
            lines = result.stdout.split('\n')
            for line in lines[:10]:  # 只显示前10行
                print(line)
        else:
            print(f"nvidia-smi执行失败: {result.stderr}")
    except Exception as e:
        print(f"执行nvidia-smi时出错: {e}")

if __name__ == "__main__":
    test_basic_gpu()
    test_computation()
    test_nvidia_smi()

# 错误：docker: Error response from daemon: could not select device driver...
# 解决方案：重新安装NVIDIA容器工具包
sudo apt-get install --reinstall nvidia-container-toolkit

# 错误：Failed to initialize NVML: Driver/library version mismatch
# 解决方案：重启宿主机或重新加载NVIDIA模块
sudo rmmod nvidia_uvm nvidia_drm nvidia_modeset nvidia
sudo modprobe nvidia nvidia_modeset nvidia_drm nvidia_uvm

# 错误：CUDA error: out of memory
# 解决方案：监控显存使用，合理分配资源
docker run --rm --gpus all nvidia/cuda:12.1.1-base-ubuntu20.04 nvidia-smi --query-gpu=memory.used,memory.total --format=csv

# Dockerfile.ai-training
FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04

# 设置Python环境
ENV PYTHONUNBUFFERED=1
RUN apt-get update && apt-get install -y python3-pip

# 安装PyTorch和其他依赖
RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
RUN pip3 install transformers datasets accelerate

# 设置工作目录
WORKDIR /workspace

# 多GPU训练启动脚本
COPY start_training.sh /usr/local/bin/
RUN chmod +x /usr/local/bin/start_training.sh

CMD ["start_training.sh"]

# start_training.sh - 多GPU训练启动脚本
#!/bin/bash

# 自动检测可用GPU数量
NUM_GPUS=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)

echo "检测到 $NUM_GPUS 个GPU"

# 根据GPU数量设置并行训练参数
if [ $NUM_GPUS -gt 1 ]; then
    echo "使用多GPU训练模式"
    export CUDA_VISIBLE_DEVICES=0,1,2,3
    python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS \
        train.py --multi_gpu
else
    echo "使用单GPU训练模式"
    python train.py
fi

# docker-compose.gpu.yml
version: '3.8'

services:
  cuda-11.7:
    image: pytorch/pytorch:1.13.1-cuda11.7-cudnn8-devel
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./workspace:/workspace
    working_dir: /workspace
    command: sleep infinity

  cuda-12.1:
    image: pytorch/pytorch:2.0.1-cuda12.1-cudnn8-devel
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      - ./workspace:/workspace
    working_dir: /workspace
    command: sleep infinity

  rocm-env:
    image: rocm/pytorch:latest
    devices:
      - /dev/kfd
      - /dev/dri
    volumes:
      - ./workspace:/workspace
    working_dir: /workspace
    command: sleep infinity