Bash-скрипт, который ждет, пока GPU станет свободным

Если мой скрипт обучения на PyTorch запустить дважды одновременно, второй запуск завершится с ошибкой из-за нехватки памяти на GPU.

Поэтому я хочу скрипт bash, который можно было бы вызывать как ./scripts/wait.sh && ./scripts/train.sh, чтобы запуск обучения начинался только в том случае, если, согласно nvidia-smi, GPU не используется.

Я хотел бы, чтобы скрипт был устойчив к условиям гонки, когда несколько терминалов запускают скрипт wait.sh.

В качестве альтернативы, можно проверять состояние процесса вместо проверки nvidia-smi, но я думаю, что это может быть сложнее. Может быть, в Linux есть способ управления файловыми блокировками?

Как мне разобрать вывод nvidia-smi в bash?

Используйте flock для создания файла блокировки. Обычный способ — запускать основную программу через flock:

flock /tmp/gpu.lock ./roomheater.sh

Используйте цикл while или until, чтобы ждать, пока GPU будет свободен (если в данный момент он используется чем-то, что не захватило блокировку).

until <некоторая команда для проверки, что gpu свободен>; do
    sleep 10
done

Что касается разбора nvidia-smi, вы можете использовать

ngpu=$(nvidia-smi -L | wc -l);
idx=$(seq 0 $(($ngpu-1)) | tr "\n" , | sed 's/,$//');
frees=$(nvidia-smi pmon -c 1 -i $idx | awk -vn=$ngpu 'BEGIN {for(i=0;i<n; i++) { gfree[i] = 1 } } $3 == "C" {gfree[$1] = 0}  END { for (key in gfree) { if (gfree[key] == 1) { print key } } }')

что даст вам индивидуальные номера свободных GPU в $frees, а echo $frees | wc -w вернет количество свободных. Если у вас есть только один GPU, вы можете просто использовать nvidia-smi pmon -c 1 | awk '{print $3}' | grep C, чтобы проверить, использует ли какой-либо процесс GPU для вычислений. Я использую вышеуказанное, чтобы установить CUDA_VISIBLE_DEVICES на совместимых много-GPU узлах, чтобы избежать переподписки.

Для дополнительной защиты от условий гонки сочетайте с решением с flock от grawity.

Это кажется полезной вещью и для меня. Поэтому я написал скрипт, который я бы использовал.

#!/bin/bash
SCRIPT_NAME="${0##*/}"
USAGE="Использование: $SCRIPT_NAME необходимая_vram ожидаемая_используемая_vram команда [аргументы]"
WAIT_VRAM="$1"
EXPECT_VRAM="$2"

function getnvfree() {
  nvstatus="$(nvidia-smi | grep -o "[0-9]\+MiB / *[0-9]\+MiB" | head -1 | tr -d '[MiB ]')"
  IFS="https://superuser.com/" read -r -a nvvalues <<< "$nvstatus"
  echo $(( ${nvvalues[1]} - ${nvvalues[0]} ))
}

if [[ "$3" == "" ]] ; then
  echo $USAGE
  exit 1
fi
if (( $WAIT_VRAM <= 0 || $EXPECT_VRAM <= 0 )) ; then
  echo $USAGE
  exit 1
fi

shift
shift

set -m
while true ; do
  exec {lock_fd}>"/var/lock/${SCRIPT_NAME}.lock" || exit 1
  flock -n "$lock_fd" || { echo "ОШИБКА: flock() не удалась." >&2; exit 1; }
  # ... команды выполняются под блокировкой ...

  # Ждите, пока первый GPU не будет иметь [аргумент 1] MiB свободной VRAM.
  nvfree=$(getnvfree)
  if (( $nvfree > $WAIT_VRAM )) ; then
    # Затем запустите программу, указанную в аргументах 3...
    "$@" &
    # Затем ждите, пока она не использует по крайней мере [аргумент 2] больше VRAM.
    while true ; do
      sleep 1
      nvfreenow=$(getnvfree)
      if (( ( $nvfree - $nvfreenow ) >= $EXPECT_VRAM )) ; then
        flock -u "$lock_fd"
        fg %1
        exit 0
      fi
    done
  fi
  flock -u "$lock_fd"
  sleep 10
done

Использование — передать желаемую свободную VRAM, за ней следует объем VRAM, ожидаемого для использования после запуска приложения. Таким образом, файл блокировки не поддерживается навсегда; другие экземпляры блокируются только за счет использования VRAM. Таким образом, если вы укажете точное требование к RAM, меньше чем половина свободной VRAM, более одного процесса сможет запуститься. (Вам это может быть не полезно, но я думаю, что мне будет.)

В зависимости от вашего использования, для команды, требующей 10000 MiB, но которая стартует, используя только 5000, ваша команда может выглядеть так:

./scripts/wait.sh 10000 5000 ./scripts/train.sh

Функция getnvfree — это часть, которая разбирает использование VRAM от nvidia-smi для первого GPU и выполняет математические расчеты для нахождения свободной VRAM.