单机多卡空闲GPU自动检测Pytorch多卡训练脚本

本文分享了一个Bash脚本，它能够自动检测空闲的GPU，并根据用户的需求分配相应数量的GPU以运行特定的命令或脚本。

脚本代码

bash 复制代码

#!/bin/bash

# 亮红色和亮绿色文本的ANSI转义序列
light_red='\e[91m'
light_green='\e[92m'
no_color='\e[0m'

# 检查是否有参数输入
if [ -z "$1" ]; then
    echo -e "${light_red}Error: nproc_per_node is required. Please input a value.${no_color}"
    exit 1
fi

nproc_per_node=$1
shift  # 移除第一个参数，使得"$@"只包含除nproc_per_node之外的参数

# 检查nproc_per_node是否大于0
if ! [[ "$nproc_per_node" =~ ^[0-9]+$ ]] || [ "$nproc_per_node" -le 0 ]; then
    echo -e "${light_red}Error: nproc_per_node must be a positive integer.${no_color}"
    exit 1
fi

# 使用nvidia-smi查询空闲的GPU
free_gpus=($(nvidia-smi --query-gpu=index,gpu_name,utilization.gpu,utilization.memory,memory.free --format=csv,noheader,nounits | awk -F ', ' '{if($3 == 0 && $4 == 0) print $1}'))

# 检查是否有足够的空闲GPU
if [ ${#free_gpus[@]} -lt $nproc_per_node ]; then
    echo -e "${light_red}Error: Not enough free GPUs. Found ${#free_gpus[@]}, require $nproc_per_node.${no_color}"
    exit 1
fi

# 将所有空闲的GPU编号转换为以逗号分隔的字符串
all_free_gpus=$(IFS=,; echo "${free_gpus[*]}")

# 选择前nproc_per_node个空闲的GPU
selected_gpus=$(IFS=,; echo "${free_gpus[*]:0:$nproc_per_node}")

# 输出所有空闲的GPU编号
echo -e "${light_green}All free GPUs: ${all_free_gpus}${no_color}"
# 输出选择的GPU编号
echo -e "${light_green}Selected GPUs: ${selected_gpus}${no_color}"

# 构建并以绿色文本运行命令
cmd="CUDA_VISIBLE_DEVICES=$selected_gpus python -m torch.distributed.launch --nproc_per_node=$nproc_per_node --use_env $@"
echo -e "${light_green}Running command: $cmd${no_color}"
eval $cmd

功能简述

该脚本的核心功能是自动检测系统中空闲的GPU，并根据用户定义的需要的卡数（通过nproc_per_node参数）分配GPU资源。如果系统中没有足够的空闲GPU，脚本将输出错误信息并终止执行。

使用步骤

创建脚本文件 ：首先，您需要创建一个新的脚本文件。可以使用任何文本编辑器来创建文件，例如使用vim。创建文件的命令如下：
bash 复制代码
```
vim gpu_allocator.sh
```
在打开的编辑器中，复制上文中的脚本代码，并保存退出。
赋予执行权限：为了让脚本文件可执行，您需要修改其权限。这可以通过下面的命令来完成：
bash 复制代码
```
chmod +x gpu_allocator.sh
```
这个命令会赋予脚本文件执行权限。
执行脚本：现在，脚本已经准备好可以执行了。使用下面的格式运行您的脚本：
bash 复制代码
```
./gpu_allocator.sh <nproc_per_node> <其他命令>
```
- <nproc_per_node>：这是一个必需的参数，指定您希望使用的GPU数量（整数）。
- <其他命令>：这是脚本将要执行的其他命令或脚本，例如运行Python脚本或其他需要GPU的应用。

示例

假设您想要使用2个GPU来运行一个Python训练

脚本train.py，您可以这样执行上述脚本：

bash 复制代码

./gpu_allocator.sh 2 python train.py --batch-size 64

在这个示例中，脚本会自动寻找2个空闲的GPU，如果找到，就会使用这两个GPU来执行train.py脚本，其中--batch-size 64是传递给train.py的一个参数。