cudadl/slurm-wrapper.sh

#!/bin/bash
#
# Use this script if you cannot use MPI but you have a
# NFS available (or some other means of communicating 
# the ip/port between computers).
#
# Don't reuse the same "FILE_WITH_LOC" since the race
# conditions are ugly, and the files won't be
# automatically deleted after this script.
#
OUTPUT_FILE_DIR="$1"
FILE_WITH_LOC="$2"
CUDADL_PATH="$3"
#
if [ "$SLURM_LOCALID" == "0" ] ; then
	if [ "$SLURM_PROCID" == "0" ] ; then
		# adapted from https://superuser.com/questions/184307/bash-create-anonymous-fifo
		# and https://stackoverflow.com/questions/8297415/in-bash-how-to-find-the-lowest-numbered-unused-file-descriptor
		#
		PIPE=$(mktemp -u)
		# create a named pipe
		mkfifo $PIPE
		# need to get a bi-directional file descriptor first to prevent bash from blocking
		exec {FD}<>$PIPE
		exec {FD1}>$PIPE
		exec {FD2}<$PIPE
		# unlink the named pipe
		rm $PIPE
		# don't need this file descriptor anymore
		exec {FD}>&-
		#
		CONTROLLER_BOUNDCB_FD=$FD1 ${CUDADL_PATH}/controller "${@:4}" &
		CONTROLLER_PID=$!
		trap "kill $CONTROLLER_PID ; exit" SIGINT
		# the controller should close its write file descriptor, so we need to close ours
		exec {FD1}>&-
		#
		BYTE1=$(head -c 1 <&$FD2)
		BYTE2=$(head -c 1 <&$FD2)
		BYTE1=$(LC_CTYPE=C printf '%d' "'$BYTE1")
		BYTE2=$(LC_CTYPE=C printf '%d' "'$BYTE2")
		PORT=$((BYTE2*256+BYTE1))
		# is there no better way in bash?!
		#
		read IP <&$FD2
		#
		exec {FD2}<&-
		#
		sleep 10
		# give the controller time to get ready
		# not sure if needed, but hopefully it fixes a bug with workers disconnecting 
		#
		echo "$IP $PORT" > $FILE_WITH_LOC
		#
	fi
	#
	COUNTER=0
	START_DPNODE=1
	while [ ! -f $FILE_WITH_LOC ] ; do
		if [ $COUNTER -gt 60 ] ; then
			# file still doesn't exist after 60 seconds
			START_DPNODE=0
			break
		fi
		sleep 1
		COUNTER=$((COUNTER+1))
	done
	#
	if [ $START_DPNODE -eq 1 ] ; then
		sleep 1 # to make sure it is done writing
		read IP PORT < $FILE_WITH_LOC
		#
		OUTPUT_FILE="${OUTPUT_FILE_DIR}/dpnode_${SLURM_NODEID}_${SLURM_LOCALID}.out"
		TMP_OUTPUT_FILE="${SLURM_TMPDIR}/dpnode_${SLURM_NODEID}_${SLURM_LOCALID}.out"
		${CUDADL_PATH}/dpnode $IP $PORT > "${TMP_OUTPUT_FILE}" 2>&1
		mv "$TMP_OUTPUT_FILE" "$OUTPUT_FILE"
	fi
	#
	if [ ! -z "$CONTROLLER_PID" ] ; then
		wait "$CONTROLLER_PID"
	fi
else
	COUNTER=0
	START_WORKER=1
	while [ ! -f $FILE_WITH_LOC ] ; do
		if [ $COUNTER -gt 60 ] ; then
			# file still doesn't exist after 60 seconds
			START_WORKER=0
			break
		fi
		sleep 1
		COUNTER=$((COUNTER+1))
	done
	#
	if [ $START_WORKER -eq 1 ] ; then
		sleep 1 # to make sure it is done writing
		read IP PORT < $FILE_WITH_LOC
		#
		WORKER_ID=$((SLURM_LOCALID-1))
		#
		OUTPUT_FILE="${OUTPUT_FILE_DIR}/worker_${SLURM_NODEID}_${WORKER_ID}.out"
		TMP_OUTPUT_FILE="${SLURM_TMPDIR}/worker_${SLURM_NODEID}_${WORKER_ID}.out"
		${CUDADL_PATH}/worker $IP $PORT $WORKER_ID > "${TMP_OUTPUT_FILE}" 2>&1
		mv "$TMP_OUTPUT_FILE" "$OUTPUT_FILE"
	fi
fi