#!/bin/bash
#
# Sample Prolog to start and quit the MPS server as needed
# NOTE: This is only a sample and may need modification for your environment
#

# Specify default locations of file with MPS device ID
MPS_DEV_ID_FILE="/var/run/mps_dev_id"

# Specify directory where MPS and Slurm commands are located (if not in search path)
#MPS_CMD_DIR="/usr/bin/"
#SLURM_CMD_DIR="/usr/bin"

# If job requires MPS then write device ID to file and start server as needed
if [ -n "${CUDA_VISIBLE_DEVICES}" ] &&
   [ -n "${CUDA_MPS_ACTIVE_THREAD_PERCENTAGE}" ]; then
	echo ${CUDA_VISIBLE_DEVICES} >${MPS_DEV_ID_FILE}
	unset CUDA_MPS_ACTIVE_THREAD_PERCENTAGE
	export CUDA_MPS_PIPE_DIRECTORY=/tmp/nvidia-mps_${CUDA_VISIBLE_DEVICES}
	export CUDA_MPS_LOG_DIRECTORY=/tmp/nvidia-log_${CUDA_VISIBLE_DEVICES}
	${MPS_CMD_DIR}nvidia-cuda-mps-control -d && echo "MPS control daemon started"
	sleep 1
	${MPS_CMD_DIR}nvidia-cuda-mps-control start_server -uid $SLURM_JOB_UID && echo "MPS server started for $SLURM_JOB_UID"
else
	# If job requires full GPU and previous use was MPS then quit MPS server
	if [ -f ${MPS_DEV_ID_FILE} ]; then
		MPS_DEV_ID=$(cat ${MPS_DEV_ID_FILE})
	else
		MPS_DEV_ID=-1
	fi
	if [ -n "${CUDA_VISIBLE_DEVICES}" ] &&
	   [ ${CUDA_VISIBLE_DEVICES} -eq ${MPS_DEV_ID} ] &&
	   [ -z "${CUDA_MPS_ACTIVE_THREAD_PERCENTAGE}" ]; then
		# Determine if MPS server is running
		ps aux | grep nvidia-cuda-mps-control | grep -v grep > /dev/null
		if [ $? -eq 0 ]; then
			echo "Stopping MPS control daemon"
			# Reset GPU mode to default
			${MPS_CMD_DIR}nvidia-smi -c ${CUDA_VISIBLE_DEVICES}
			# Quit MPS server daemon
			echo quit | nvidia-cuda-mps-control
			# Test for presence of MPS zombie process
			ps aux | grep nvidia-cuda-mps | grep -v grep > /dev/null
			if [ $? -eq 0 ]; then
				logger "`hostname` Slurm Epilog: MPS refusing to quit! Downing node"
				${SLURM_CMD_DIR}scontrol update nodename=${SLURMD_NODENAME} State=DOWN Reason="MPS not quitting"
			fi
			# Check GPU sanity, simple check
			${MPS_CMD_DIR}nvidia-smi > /dev/null
			if [ $? -ne 0 ]; then
				logger "`hostname` Slurm Epilog: GPU not operational! Downing node"
				${SLURM_CMD_DIR}scontrol update nodename=${SLURMD_NODENAME} State=DOWN Reason="GPU not operational"
			fi
		fi
	fi
fi

exit 0
