本文目录导读:

在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)
注意事项
- 权限问题:确保运行脚本的用户有docker组权限
- 版本兼容:NVIDIA驱动、CUDA版本和容器镜像需要兼容
- 资源限制:合理设置内存和GPU限制,避免资源争抢
- 安全考虑:不要在生产环境中公开暴露GPU设备
选择合适的配置方式取决于你的具体需求和基础设施环境,建议先在测试环境中验证配置的正确性。