#!/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."