[scripts] Extend limit_num_gpus.sh to support --num-gpus 0. (kaldi-asr#3027)

Author: oplatek

egs/wsj/s5/utils/parallel/limit_num_gpus.sh

@@ -18,8 +18,8 @@ if [ "$1" == "--num-gpus" ]; then
   shift
 fi
 
-if ! printf "%d" "$num_gpus" >/dev/null || [ $num_gpus -le 0 ]; then
-  echo $0: Must pass a positive interger after --num-gpus
+if ! printf "%d" "$num_gpus" >/dev/null || [ $num_gpus -le -1 ]; then
+  echo $0: Must pass a positive interger or 0 after --num-gpus
   echo e.g. $0 --num-gpus 2 local/tfrnnlm/run_lstm.sh
   exit 1
 fi
@@ -35,18 +35,24 @@ CUDA_VISIBLE_DEVICES=
 num_total_gpus=`nvidia-smi -L | wc -l`
 num_gpus_assigned=0
 
-for i in `seq 0 $[$num_total_gpus-1]`; do
-# going over all GPUs and check if it is idle, and add to the list if yes
-  if nvidia-smi -i $i | grep "No running processes found" >/dev/null; then
-    CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}$i, && num_gpus_assigned=$[$num_gpus_assigned+1]
-  fi
-# once we have enough GPUs, break out of the loop
-  [ $num_gpus_assigned -eq $num_gpus ] && break
-done
+if [ $num_gpus -eq 0 ] ; then
+    echo "$0: Running the job on CPU. Disabling submitting to gpu"
+    export CUDA_VISIBLE_DEVICES=""
+else
+    for i in `seq 0 $[$num_total_gpus-1]`; do
+    # going over all GPUs and check if it is idle, and add to the list if yes
+      if nvidia-smi -i $i | grep "No running processes found" >/dev/null; then
+        CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}$i, && num_gpus_assigned=$[$num_gpus_assigned+1]
+      fi
+    # once we have enough GPUs, break out of the loop
+      [ $num_gpus_assigned -eq $num_gpus ] && break
+    done
 
-[ $num_gpus_assigned -ne $num_gpus ] && echo Could not find enough idle GPUs && exit 1
+    [ $num_gpus_assigned -ne $num_gpus ] && echo Could not find enough idle GPUs && exit 1
 
-export CUDA_VISIBLE_DEVICES=$(echo $CUDA_VISIBLE_DEVICES | sed "s=,$==g")
+    export CUDA_VISIBLE_DEVICES=$(echo $CUDA_VISIBLE_DEVICES | sed "s=,$==g")
+
+    echo "$0: Running the job on GPU(s) $CUDA_VISIBLE_DEVICES"
+fi
 
-echo "$0: Running the job on GPU(s) $CUDA_VISIBLE_DEVICES"
 "$@"