抢占GPU显卡跑程序--实用脚本

只需修改两行，将脚本复制到程序位置，预定义要抢占的显卡号，写入自己要启动的程序指令

#!/bin/bash
#=====================只需关注==================
# 要检查的 GPU 列表，使用逗号分隔，如 "4,5,6,7"
GPUS_TO_CHECK="4,5,6,7"
# 要启动的程序
PROGRAM="nohup sh train.sh"
#=====================只需关注==================

# 显存剩余百分比阈值
MEMORY_THRESHOLD=80

# 检查显卡显存空闲的函数
check_gpu_memory() {
  local gpu_id=$1
  # 获取指定 GPU 的总显存和已用显存
  local memory_total=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits -i $gpu_id)
  local memory_used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i $gpu_id)

  # 计算剩余显存百分比
  local memory_free_percentage=$((100 * (memory_total - memory_used) / memory_total))

  # 如果显存剩余百分比大于等于阈值，则认为显卡空闲
  if [ "$memory_free_percentage" -ge "$MEMORY_THRESHOLD" ]; then
    return 0
  else
    return 1
  fi
}

# 计数空闲 GPU 的数量
free_gpu_count=0

# 遍历 GPU 列表并检查显存状态
for gpu_id in $(echo $GPUS_TO_CHECK | tr ',' ' '); do
  if check_gpu_memory $gpu_id; then
    free_gpu_count=$((free_gpu_count + 1))
  fi
done

# 检查是否有至少两张 GPU 空闲
if [ "$free_gpu_count" -ge 2 ]; then
  echo "At least two GPUs are free. Starting the program..."
  CUDA_VISIBLE_DEVICES=$(echo $GPUS_TO_CHECK | cut -d',' -f1,2) $PROGRAM &
  exit 0
else
  echo "Less than two GPUs are free. Exiting..."
  exit 1
fi