Shell脚本如何配置容器GPU访问权限

本文目录导读：

1. 前提条件检查
2. 使用Docker配置GPU访问
3. 使用Podman配置GPU访问
4. 使用Singularity/Apptainer配置
5. 多GPU和资源限制配置
6. 完整的自动化配置脚本
7. 故障排查脚本
8. 环境变量配置
9. 注意事项

在Shell脚本中配置容器的GPU访问权限，主要取决于你使用的容器运行时（如Docker、Podman等）以及NVIDIA驱动和工具链的安装情况,以下是几种常见场景的配置方法：

前提条件检查

在脚本中检查GPU和NVIDIA驱动是否可用：

#!/bin/bash
# 检查NVIDIA驱动
if ! command -v nvidia-smi &> /dev/null; then
    echo "错误: NVIDIA驱动未安装"
    exit 1
fi
# 检查GPU可用性
if ! nvidia-smi -L &> /dev/null; then
    echo "错误: 无可用GPU"
    exit 1
fi
echo "NVIDIA驱动和GPU检查通过"

使用Docker配置GPU访问

安装NVIDIA Container Toolkit

#!/bin/bash
# 设置仓库
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
sudo systemctl restart docker

运行带有GPU的容器

#!/bin/bash
# 运行带有所有GPU的容器
docker run --gpus all nvidia/cuda:12.0-base nvidia-smi
# 运行选择特定GPU
docker run --gpus '"device=0,1"' nvidia/cuda:12.0-base nvidia-smi
# 运行指定GPU数量
docker run --gpus 2 nvidia/cuda:12.0-base nvidia-smi
# 使用环境变量指定GPU
docker run -e NVIDIA_VISIBLE_DEVICES=0,1 nvidia/cuda:12.0-base nvidia-smi

完整脚本示例

#!/bin/bash
IMAGE="nvidia/cuda:12.0-base"
CONTAINER_NAME="gpu_container"
GPU_DEVICES="0,1"
WORK_DIR="/workspace"
# 清理旧容器
docker rm -f $CONTAINER_NAME 2>/dev/null
# 运行容器
docker run -itd \
    --name $CONTAINER_NAME \
    --gpus "device=$GPU_DEVICES" \
    -v $PWD:$WORK_DIR \
    -w $WORK_DIR \
    --shm-size=8g \
    $IMAGE
# 验证GPU访问
docker exec $CONTAINER_NAME nvidia-smi

使用Podman配置GPU访问

#!/bin/bash
# Podman需要额外配置
# 安装nvidia-container-toolkit（同上）
# 运行容器
podman run --security-opt=label=disable \
    --hooks-dir=/usr/share/containers/oci/hooks.d/ \
    nvidia/cuda:12.0-base nvidia-smi

使用Singularity/Apptainer配置

#!/bin/bash
# Singularity自动识别GPU
singularity exec --nv docker://nvidia/cuda:12.0-base nvidia-smi
# 或使用本地镜像
singularity build cuda.sif docker://nvidia/cuda:12.0-base
singularity exec --nv cuda.sif python train.py

多GPU和资源限制配置

#!/bin/bash
GPU_COUNT=2
GPU_MEMORY_LIMIT="4g"  # 每GPU内存限制
docker run --gpus "device=0,1" \
    -e NVIDIA_VISIBLE_DEVICES=0,1 \
    -e CUDA_VISIBLE_DEVICES=0,1 \
    --memory="16g" \
    --memory-swap="16g" \
    nvidia/cuda:12.0-base \
    nvidia-smi

完整的自动化配置脚本

#!/bin/bash
set -e
# 配置参数
DOCKER_IMAGE="pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel"
CONTAINER_NAME="ml_container"
GPU_IDS="0,1,2"
WORKSPACE_DIR="/data/workspace"
SHARED_MEMORY="32g"
DATA_DIRS=("/data/datasets" "/models")
# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
NC='\033[0m'
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}
log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}
# 检查依赖
check_dependencies() {
    if ! command -v docker &> /dev/null; then
        log_error "Docker未安装"
        exit 1
    fi
    if ! command -v nvidia-smi &> /dev/null; then
        log_error "NVIDIA驱动未安装"
        exit 1
    fi
    log_info "所有依赖检查通过"
}
# 设置NVIDIA容器运行时
setup_nvidia_runtime() {
    if ! docker info | grep -q "nvidia"; then
        log_info "正在安装NVIDIA Container Toolkit..."
        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
        sudo systemctl restart docker
        log_info "NVIDIA Container Toolkit安装完成"
    else
        log_info "NVIDIA Container Toolkit已安装"
    fi
}
# 创建必要的目录
create_directories() {
    for dir in "${DATA_DIRS[@]}"; do
        mkdir -p "$dir"
    done
}
# 运行容器
run_container() {
    # 清理旧容器
    docker rm -f $CONTAINER_NAME 2>/dev/null || true
    # 构建卷挂载参数
    VOLUME_ARGS=""
    for dir in "${DATA_DIRS[@]}"; do
        VOLUME_ARGS+=" -v $dir:$dir"
    done
    # 运行容器
    docker run -itd \
        --name $CONTAINER_NAME \
        --gpus "device=$GPU_IDS" \
        -e NVIDIA_VISIBLE_DEVICES=$GPU_IDS \
        -e CUDA_VISIBLE_DEVICES=$GPU_IDS \
        $VOLUME_ARGS \
        -v $WORKSPACE_DIR:/workspace \
        -w /workspace \
        --shm-size=$SHARED_MEMORY \
        --ipc=host \
        --ulimit memlock=-1 \
        --ulimit stack=67108864 \
        $DOCKER_IMAGE \
        /bin/bash
    if [ $? -eq 0 ]; then
        log_info "容器 $CONTAINER_NAME 启动成功"
        # 验证GPU访问
        docker exec $CONTAINER_NAME nvidia-smi
        log_info "GPU访问验证通过"
    else
        log_error "容器启动失败"
        exit 1
    fi
}
# 主函数
main() {
    check_dependencies
    setup_nvidia_runtime
    create_directories
    run_container
}
# 执行主函数
main

故障排查脚本

#!/bin/bash
# GPU访问诊断脚本
echo "=== GPU访问诊断 ==="
# 检查驱动
echo "1. NVIDIA驱动版本:"
nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null || echo "未检测到驱动"
# 检查CUDA版本
echo "2. CUDA版本:"
nvidia-smi | grep "CUDA Version" 2>/dev/null || echo "未检测到CUDA"
# 检查Docker NVIDIA支持
echo "3. Docker NVIDIA支持:"
docker info 2>/dev/null | grep -i nvidia || echo "Docker未配置NVIDIA支持"
# 测试GPU访问
echo "4. 测试GPU访问:"
docker run --rm --gpus all nvidia/cuda:12.0-base nvidia-smi 2>/dev/null && \
    echo "GPU访问正常" || echo "GPU访问失败"

环境变量配置

#!/bin/bash
# 常用GPU相关环境变量
export NVIDIA_VISIBLE_DEVICES=all          # 使所有GPU可见
export CUDA_VISIBLE_DEVICES=0,1            # CUDA可见的GPU
export NVIDIA_DRIVER_CAPABILITIES=compute,utility  # 驱动能力
export NVIDIA_REQUIRE_CUDA="cuda>=12.0"    # CUDA版本要求
# 或基于文件配置
echo "0,1" > /tmp/gpu_list.txt
export CUDA_VISIBLE_DEVICES=$(cat /tmp/gpu_list.txt)

注意事项

1. 权限问题：确保运行脚本的用户有docker组权限
2. 版本兼容：NVIDIA驱动、CUDA版本和容器镜像需要兼容
3. 资源限制：合理设置内存和GPU限制，避免资源争抢
4. 安全考虑：不要在生产环境中公开暴露GPU设备

选择合适的配置方式取决于你的具体需求和基础设施环境,建议先在测试环境中验证配置的正确性。