在超算系統上運行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網卡的地址。第二，當只申請一個節點進行訓練任務時，上述腳本會失敗並掛起。不過一般來說，分布式訓練任務都需要多個節點，所以這個問題影響比較小。目前正在解決這兩個問題。