Files
dockerfile_dp/finetune/configure_gpu.sh
2024-08-28 15:18:15 +08:00

48 lines
1.2 KiB
Bash
Executable File

#!/bin/bash
# 提取GPU UUID
GPU_UUIDS=$(nvidia-smi -a | grep 'GPU UUID' | awk '{print $4}')
# 生成node-generic-resources JSON片段
NODE_RESOURCES=$(echo "$GPU_UUIDS" | awk '{print "\"NVIDIA-GPU=" $1 "\","}' | tr -d '\n')
NODE_RESOURCES=${NODE_RESOURCES%,} # 移除最后一个逗号
# 生成完整的daemon.json内容
DAEMON_JSON=$(cat <<EOF
{
"runtimes": {
"nvidia": {
"path": "/usr/bin/nvidia-container-runtime",
"runtimeArgs": []
}
},
"default-runtime": "nvidia",
"node-generic-resources": [
$NODE_RESOURCES
]
}
EOF
)
# 备份当前的daemon.json文件
if [ -f /etc/docker/daemon.json ]; then
sudo cp /etc/docker/daemon.json /etc/docker/daemon.json.bak
fi
# 写入新的daemon.json内容
echo "$DAEMON_JSON" | sudo tee /etc/docker/daemon.json > /dev/null
# 添加swarm-resource配置
# swarm-resource = "DOCKER_RESOURCE_NVIDIAGPU"
sudo sed -i '/^#.*swarm-resource/s/^#//' /etc/nvidia-container-runtime/config.toml
sudo sed -i '/swarm-resource =/s/=.*/= "DOCKER_RESOURCE_GPU"/' /etc/nvidia-container-runtime/config.toml
# 重启Docker服务
sudo systemctl restart docker.service
# 验证配置
docker info | grep -i 'nvidia'
echo "GPU UUIDs have been configured and Docker has been restarted."