如何使用 Bash 脚本保持所有 GPU 设备运行任务？

Question

I have an 8-GPU server and I'd like to train one neural network on each of them simultaneously. I have several tens of such networks to train and I would like to schedule the training task. Currently I'm writing my own bash script for this scheduling task.

for l1 in {1e-4,2e-4,5e-4,1e-3}; do

      python train.py --lr $l1 --attr 0 --device 0 &
      python train.py --lr $l1 --attr 1 --device 1 &
      python train.py --lr $l1 --attr 2 --device 2 &
      python train.py --lr $l1 --attr 3 --device 3 &
      python train.py --lr $l1 --attr 4 --device 4 &
      python train.py --lr $l1 --attr 5 --device 5 &
      python train.py --lr $l1 --attr 6 --device 6 &
      python train.py --lr $l1 --attr 7 --device 7

      sleep 1
      wait 
done 

In the above script, --device flag chooses the GPU to use while other flags just determine the hyper-parameters of my deep neural networks. What this script does is that, for each iteration of the for-loop, it launches one training task on each GPU and wait for all of them to finish before starting the next iteration. The issue is that, each of the training task may take different time to run, thus there will be a significant amount of time that I'm using less than 8 GPUs simultaneously, which lengthens the time for whole task to finish.

I'm wondering whether there is some way for me to detect which GPU has finished its task and launch a new task on it, so that I can always have 8 GPUs running.

Thanks a lot!

Answer 1

I saw you are not using a cluster, that means the GPUs are on your local machine.

In this case, you can use this library: https://pypi.org/project/simple-gpu-scheduler/

Hope this helps.

Answer 2

Just found that Ray is a great package for managing your experiment. ( https://github.com/ray-project/ray )