多显卡的实验自动挂载脚本

如果你在一台拥有多张显卡的服务器上运行深度学习实验，那么你可能会遇到这样的问题：

1. 一张显卡上运行一个实验。实验结束之后，再手动挂下一个实验。这需要有人一直关注实验运行进程。
2. 为了节约中间挂载时间，可能会一张显卡上同时挂多个实验。虽然会运行更久，但是短暂地解放了操作者。不需要有人一直盯着实验进程。当然，有时候可能会因此爆显存，而当事人不知道，导致显卡白白空闲好几天。

为了解决以上问题，我写了一个实验自动挂载的shell脚本。

核心思路如下：

1. makeExperiments.sh生成需要运行的命令。可以通过shell语法循环生成。
2. autoRun.sh轮训显卡使用情况，选择空闲显卡挂载实验。为了避免多个实验挂在同一张显卡，需要设置大一点的显存阈值，比如当显存占用小于200MB时，认为显卡空闲，同时设置文件锁。

makeExperiments.sh

这里主要依靠命令行传入不同参数实现运行不同实验。因此，需要对原实验代码的参数处理部分进行修改。下面展示大致思路：

#!/bin/bash
# 生成实验列表

OUTPUT_FILE="" # 这里换成自己的输出目录

# 先清空，冒号:是空命令，用于占位
: > "$OUTPUT_FILE"
# 循环生成
models=("resnet101" "resnet152")
OHDomains=("art" "clipart" "product" "realworld")
PassArt=0
for model in "${models[@]}";do
  for idx in "${!OHDomains[@]}"; do # 遍历下标
    for seed in 1 2; do
      echo "python main.py --path yourPath --seen_index $idx --model $model --randomSeed" >> "$OUTPUT_FILE" # 这里是核心命令部分
    done
  done
done
# 手动输入
{
echo "python main.py --path yourPath --seen_index 3 --model WideResNet --randomSeed";
echo "python main.py --path yourPath --seen_index 3 --model WideResNet --randomSeed";
echo "python main.py --path yourPath --seen_index 3 --model WideResNet --seed 0";
} >> "$OUTPUT_FILE"

这里展示了两种生成方式：

1. 通过循环生成实验命令
2. 手动输入实验命令

[!NOTE]

>：将输出重定向到XX文件，会覆盖之前的内容。

>>：将输出重定向到XX文件，追加输出，不会覆盖之前的内容。

autoRun.sh

#!/bin/bash
# 自动挂在实验。
# 检测空闲GPU，从experiment.txt中逐行取出实验并分配执行

EXPERIMENTS_FILE="yourPath/experiments.txt" # 实验文件，每行一个完整实验命令
CHECKINTERVAL=30 # GPU 状态检查间隔（秒）
MEM_THRESHOLD=200 # 显存小于阈值（MB），认为为空闲。不能太小，不然短时间内，多个实验都会认为同一个卡是空闲
LOCK_FILE="yourPath/exp_queue.lock" # 锁文件（自动创建），避免多张卡跑同一个实验。劝告锁，并不是强制锁

# 冷启动。检测流程，不会真的执行实验
DRY_RUN=false

[[ ! -f "${EXPERIMENTS_FILE}" ]] && {
  echo "错误：实验文件 ${EXPERIMENTS_FILE} 不存在"
  exit 1
}

while true; do
  # 获得GPU的信息，下标、uuid、显存使用情况。但是index可能不是按顺序排列的，不过不影响
  gpu_info=$(nvidia-smi --query-gpu=index,uuid,memory.used --format=csv,noheader,nounits 2>/dev/null)
  # 获取当前GPU的计算进程对应的uuid。如果张卡上有多个实验，那么就会多次出现这张卡的uuid
  busy_info=$(nvidia-smi --query-compute-apps=gpu_uuid --format=csv,noheader,nounits 2>/dev/null) # 执行结果是字符串
  while IFS="," read -r gpu uuid mem;do
    gpu=$(echo "${gpu}" | xargs) # 去掉多余空格。本来的作用是来接收参数
    uuid=$(echo "${uuid}" | xargs)
    mem=$(echo "${mem}" | xargs)

    # 显存占用超过阈值，不认为是空闲
    (( mem > MEM_THRESHOLD)) && continue
    # 这张卡上没有计算进程。为了进一步避免同一张卡上挂多个实验。 -q参数是quiet，就是没有输出。-x是整行匹配。匹配到了返回1，那么就跳过。没有匹配到就返回0。<<< 表示把后面的字符串按行读入
    grep -qx "${uuid}" <<< "${busy_info}" && continue

    (
      # 使用文件锁，从实验文件中提取第一行，执行后，删除。只包含了取实验，不包含执行实验。
      TASK=$(flock "${LOCK_FILE}" -c "head -n 1 '${EXPERIMENTS_FILE}' && sed -i '1d' '${EXPERIMENTS_FILE}'")

      # 如果取出来的任务为空，那么直接结束子进程。有可能会取出来为空行
      # 为什么每次都是从GPU 1开始发布任务，因为makeExperiments.sh第一个是空行，所以GPU 0在第一次发布任务的时候被跳过了。
      [[ -z "${TASK}" ]] && exit 0

      if "$DRY_RUN"; then
        echo "[$(date '+%m-%d %H:%M:%S')] [DRY RUN] GPU ${gpu} 将会执行: ${TASK}"
      else
        echo "[$(date '+%m-%d %H:%M:%S')] GPU ${gpu} 挂载实验: ${TASK}"
        # nohup -> no hang up， 退出终端后，也不会结束进程.
        # > /dev/null 2>&1 &：将实验日志，标准输出重定向到/dev/null，其实就是丢掉。2>&1把错误输出重定向到标准输出
        # & 后台执行。
  #        LOGFILE="yourPath/gpu${gpu}_$(date +%Y%m%d_%H%M%S).log"
        CUDA_VISIBLE_DEVICES="${gpu}" nohup bash -c "${TASK}" > /dev/null 2>&1 &
      fi
    ) &
  done <<< "${gpu_info}"
# 等子进程完毕，即分配任务。因为真正运行的实验进程是子进程的子进程，不归wait管。wait只等当前脚本下的子进程结束。
  wait
  # 实验任务为空，则跑完了
  if [[ ! -s "${EXPERIMENTS_FILE}" ]]; then
     if $DRY_RUN; then
            echo "[DRY RUN] 所有实验已分配完毕（未实际运行），脚本退出。"
        else
            echo "所有实验已挂载完成，脚本退出。"
      fi
      break
  fi

  sleep "${CHECKINTERVAL}"
done

虽然轮训查询显卡状态的方式会给服务器带来一定开销，但是这点开销对于跑深度学习实验来说是九牛一毛，不必担心。而且间隔时间为30秒，所以开销更加微不足道。

为了进一步避免同一张显卡上挂多个实验，进行了显卡计算进程判断。如果一张显卡上有多个计算进程，那么就不挂在这张显卡上。

进一步优化。利用并发的思想，将分配任务作为一个后台执行的子进程，这样可以瞬间占满所有资源，而不是一次挂在一个实验。