Nvidia MPS activate script for slurm

activate-mps.sh

#!/bin/bash

env | grep CUDA_

export CUDA_MPS_PIPE_DIRECTORY=$HOME/.nv_mps/$SLURM_JOB_ID/pipe
export CUDA_MPS_LOG_DIRECTORY=$HOME/.nv_mps/$SLURM_JOB_ID/log
mkdir -p $CUDA_MPS_PIPE_DIRECTORY $CUDA_MPS_LOG_DIRECTORY

cleanup_mps() {
    echo "shutdown MPS server ..."
    echo "quit" | nvidia-cuda-mps-control
}

trap cleanup_mps EXIT SIGINT SIGTERM SIGKILL

echo "start MPS server ..."
nvidia-cuda-mps-control -d
sleep 5··
MPS_PROCS=$(ps aux | grep -E "nvidia-cuda-mps-(control|server)" | grep -v grep | wc -l)
if [ $MPS_PROCS -eq 0 ]; then
    echo "ERROR: MPS server not found"
    exit 1
fi
echo "MPS server start successfully, nprocs: $MPS_PROCS"