在超算系统上使用sbatch提交MXNet分布式训练任务

在超算系统上运行MXNet分布式训练任务时，面临着一个IP地址相关的问题。我们在提交MXNet的分布式任务时，需要知道各个GPU节点的IP地址，把这些IP地址放到一个hosts文件中，以供分布式训练使用。因此，一种常用的方式是先使用salloc或yhalloc申请若干节点，然后依次登录这些节点，查询它们的IP地址，手动写入到一个hosts文件中，再使用MXNet提供的脚本提交分布式训练任务。显然，这种方法具有很多劣势。首先，当集群资源不足时，我们需要人工守在电脑前，等待有空闲资源时再手动申请节点，费时费力；其次，这种方式需要在.bashrc文件中修改全局环境变量，导致同一时间段只能有1位同学做实验，不能利用好剩余的资源；最后，在整个任务运行期间，我们也需要守在电脑前，等待任务结束后手动释放资源，否则也会造成机器空转，浪费机时。因此，我们需要寻找其他的方式，每当集群中有空闲资源时，都能够自动申请资源，运行任务，然后释放掉相关资源，而且能够使多为同学一起做实验，互不影响。这样一来，我们就不需要守在电脑屏幕旁边，既方便了自己，又节省了机时，一举两得。

当前，超算集群一般都使用Slurm进行管理。Slurm支持利用sbatch命令采用批处理方式运行作业，sbatch在脚本正确传递给作业调度系统后立即退出，同时获取到一个作业号，作业会等待所需资源满足后开始运行。sbatch会提交一个批处理作业脚本到Slurm，批处理脚本名可以在命令行上通过传递给sbatch，如没有指定文件名，则sbatch从标准输入中获取脚本内容。关于sbatch命令的详细内容请参考链接1以及链接2。

sbatch脚本文件的基本格式包括以下几点：

• 第一行以#!/bin/sh等指定该脚本的解释程序，/bin/sh可以变为/bin/bash、/bin/csh等。
• 在可执行命令之前的每行#SBATCH前缀后跟的参数作为作业调度系统参数。
• 在任何非注释及空白之后的#SBATCH将不再作为Slurm参数处理。

例如，下面的脚本申请了4个GPU节点，并在这些节点上运行python test.py命令。

#!/bin/bash
#SBATCH -N 4
#SBATCH -p gpu_v100
python test.py

看起来，sbatch命令完美地契合了我们的需求，能够自动运行任务并管理资源，只需要我们写一个脚本，在脚本里放上我们需要运行的分布式训练任务相关的代码，然后使用sbatch提交就可以了。但是，一个关键的问题是，在通过sbatch申请到节点资源之前，我们并不能预知所申请到的节点是哪些，也就不知道它们的IP地址；然而，我们写在脚本里的分布式训练命令，又要求我们提供节点的IP地址。一种解决方法是在提供给sbatch的脚本里，指定要申请的节点列表，根据这个列表可以推算出相应的IP地址。下面的代码表示的就是这种情况，我们在脚本里指定申请gpu1到gpu4这4个节点，然后提前在hosts文件里把这4个节点的IP地址填进去，这样就可以完成自动运行训练任务且释放资源了。

#!/bin/bash
#SBATCH -N 2
#SBATCH -p gpu_v100
#SBATCH -w gpu[1-4]

python ../../tools/launcher.py --launcher ssh -H hosts -n 4 python train_imagenet.py

不幸的是，上面的方法仍然具有很强的局限性。当集群中有空闲节点而我们又没有在申请列表中指定它们时，就会导致我们的任务在等待，不利于我们进行实验。因此，想要完美地发挥出sbatch命令的优势，就必须能够在脚本中获取所申请节点的IP地址。已知在脚本中，我们可以通过srun -l /bin/hostname获取到每个节点的主机名，那么问题就变成如何根据主机名获取到对应的IP地址。所幸，我们可以通过getent hosts hostname来从主机名获取到其对应的IP地址，如下图所示：

那么，我们只要通过一些编程技术，在申请资源后，从主机名获取到对应的IP地址并写入到hosts文件中，然后再运行分布式训练命令，就可以发挥出sbatch的优势，解放自己的双手~

在sbatch脚本中，获取到主机名后，通过xargs向getent hosts命令传输参数，然后利用awk命令抽取相应的IP地址并写入到hosts文件中，就达成了我们的目的。关于xargs和awk命令，请参考xargs 命令教程和awk 命令教程。完整的脚本文件如下所示，其中环境变量可以根据自己的实验需求进行配置。

#!/bin/bash
#SBATCH -N 4
#SBATCH -p gpu_v100

export PYTHONEPATH=CUSTOM_PYTHON_PATH
export OTHER_ENV_VAR=VALUES

MACHINEFILE="nodes.$SLURM_JOB_ID"
srun -l /bin/hostname | sort -n | awk '{print $2}' | xargs -L 1 getent ahosts | awk '{match($0,/([0-9]{1,3}\.){3}[0-9]{1,3}/,a); print a[0]}' > $MACHINEFILE
sort $MACHINEFILE | uniq > hosts
rm $MACHINEFILE

python ../../tools/launcher.py --launcher ssh -H hosts -n 4 python train_imagenet.py

使用以上方式获取IP地址，仍然有两个局限性：第一，只能获取到默认网卡的IP地址，不能获取其他的IP地址。例如在广州V100集群上，获取到的是IB网卡的地址。第二，当只申请一个节点进行训练任务时，上述脚本会失败并挂起。不过一般来说，分布式训练任务都需要多个节点，所以这个问题影响比较小。目前正在解决这两个问题。