#!/bin/bash
  
  # This script functions as a wrapper of a bash command that uses GPUs.
  #
  # It sets the CUDA_VISIBLE_DEVICES variable so that it limits the number of GPUs
  # used for programs. It is neccesary for running a job on the grid if the job
  # would automatically grabs all resources available on the system, e.g. a
  # TensorFlow program.
  
  num_gpus=1 # this variable indicates how many GPUs we will allow the command
             # passed to this script will run on. We achieve this by setting the
             # CUDA_VISIBLE_DEVICES variable
  set -e
  
  if [ "$1" == "--num-gpus" ]; then
    num_gpus=$2
    shift
    shift
  fi
  
  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
  
  if [ $# -eq 0 ]; then
    echo "Usage:  $0 [--num-gpus <num-gpus>] <command> [<arg1>...]"
    echo "Runs <command> with args after setting CUDA_VISIBLE_DEVICES to "
    echo "make sure exactly <num-gpus> GPUs are visible (default: 1)."
    exit 1
  fi
  
  CUDA_VISIBLE_DEVICES=
  num_total_gpus=`nvidia-smi -L | wc -l`
  num_gpus_assigned=0
  
  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
  
      export CUDA_VISIBLE_DEVICES=$(echo $CUDA_VISIBLE_DEVICES | sed "s=,$==g")
  
      echo "$0: Running the job on GPU(s) $CUDA_VISIBLE_DEVICES"
  fi
  
  "$@"