#!/bin/bash

# Launches the Nvidia CUDA MPS daemon which allows processes to share cuda
# contexts when executing CUDA kernels/inference engines


set -euo pipefail
LOG=/home/zoox/.log/glog/nvidia-cuda-mps.log
if [ $1 == "--dryrun" ]; then
  LOG=/tmp/nvidia-cuda-mps.log
fi
echo "starting shell script" > $LOG
FILE=/usr/bin/nvidia-cuda-mps-control

# Kill the existing mps node if there is one before starting a new one.
if [ -z $(pidof ${FILE}) ]; then
    # Do nothing if there is not existing mps node.
    :
else
    pid=$(pidof ${FILE})
    echo $(sudo kill $pid) >> $LOG
    timer=0
    while pgrep -f ${FILE} &> /dev/null; do
        echo 'waiting for mps node to terminate...' >> $LOG
        sleep 1
        # Timeout and exit if the mps node is not terminated after ten seconds.
        ((++timer))
        if [ $timer -gt 10 ]; then
            echo 'timeout waiting for mps node to terminate, will exit now' >> $LOG
            exit 1
        fi
    done
    echo "existing mps node with pid ${pid} has been terminated" >> $LOG
fi

if hash $FILE > /dev/null 2>&1; then
    echo "$FILE found" >> $LOG
    if [ $1 != "--dryrun" ]; then
            export CUDA_VISIBLE_DEVICES=0; # Select GPU 0.
            export CUDA_MPS_PIPE_DIRECTORY=/tmp/nvidia-mps; # Select a location thats accessible
            export CUDA_MPS_LOG_DIRECTORY=/var/log; # Select a location thats accessible
            exec sudo $FILE -f >> $LOG
        else
            echo "nvidia_cuda_mps_launch dryrun successful" | tee -a $LOG
        fi
else
    echo "$FILE does not exist"  >> $LOG
    exit 1
fi
