CUDA activation in shel script

activate_cuda.sh

#!/bin/bash
trap "exit" INT TERM ERR
trap "kill 0" EXIT

HOME_DIR=$(getent passwd "$USER" | cut -d: -f6)
CURRENT_DIR=${PWD##*/}
WORKSPACE_PATH=$(pwd)

CUDA_VERSION=10.0
export CUDA_VISIBLE_DEVICES=0,1

# Activate CUDA
SET_CUDA=true
cuda_check () {

  if [[ ":$PATH:" == *":/usr/local/cuda"* ]]; then
    CURRENT_CUDA_PATH=$(readlink /usr/local/cuda)
    CURRENT_CUDA=${CURRENT_CUDA_PATH##*-}

    if [[ $CURRENT_CUDA == $CUDA_VERSION ]]; then
      echo ""
      echo "CUDA already set to $CURRENT_CUDA."
      echo ""
      SET_CUDA=false
    else
      echo ""
      echo "CUDA already set to $CURRENT_CUDA in PATH. CUDA $CURRENT_CUDA will be replaced with $CUDA_VERSION"
      echo ""
      read -p "Press enter to continue or ctrl-c to exit"
      echo ""
      echo "Deactivating $CURRENT_CUDA before activating another version."
      echo ""
      cuda_deactivate

    fi
  fi

}

cuda_activate () {

  DEFAULT_CUDA_PATH=$(readlink /usr/local/cuda)
  DEFAULT_CUDA=${DEFAULT_CUDA_PATH##*-}

  echo ""
  echo "Installed CUDA versions:"

  cudas=()
  for cuda_install_path in `ls -d /usr/local/cuda-*`
  do
    cudas+=("${cuda_install_path##*-}")
    echo ${cuda_install_path##*-}
  done
  echo ""

  # If argument provided for CUDA version use it if valid
  [[ ${cudas[@]} =~ (^|[[:space:]])"$1"($|[[:space:]]) ]] && CUDA_SET=$1 || read -p "Cuda version [default=$DEFAULT_CUDA]: " CUDA_SET
  CUDA_SET_VERSION=$CUDA_SET


  [[ ${cudas[@]} =~ (^|[[:space:]])"$CUDA_SET"($|[[:space:]]) ]] && echo 'Setting CUDA to version:' || echo 'Not a valid CUDA installation'
  [[ ${cudas[@]} =~ (^|[[:space:]])"$CUDA_SET"($|[[:space:]]) ]] && echo $CUDA_SET_VERSION || return

  export ORIGINAL_PATH=$PATH
  export PATH=/usr/local/cuda-$CUDA_SET/bin:$PATH

  export ORIGINAL_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
  export LD_LIBRARY_PATH=/usr/local/cuda-$CUDA_SET/lib64:/usr/local/cuda-$CUDA_SET/extras/CUPTI/lib64:/lib/nccl/cuda-$CUDA_SET:$LD_LIBRARY_PATH

  sudo rm /usr/local/cuda
  sudo ln -s /usr/local/cuda-$CUDA_SET /usr/local/cuda

  echo ""
  echo "Version info for CUDA $CUDA_SET:"
  echo ""
  echo "$(nvcc -V)"
}

# Deactivate CUDA
cuda_deactivate () {

  if [[ -z "${ORIGINAL_PATH}" ]]; then
    echo "CUDA is not activated. Aborting."
    return 1
  fi

  export PATH=$ORIGINAL_PATH
  unset ORIGINAL_PATH

  export LD_LIBRARY_PATH=$ORIGINAL_LD_LIBRARY_PATH
  unset ORIGINAL_LD_LIBRARY_PATH

  echo "CUDA deactivated"
}

cuda_check
if [[ $SET_CUDA == true ]]; then
  cuda_activate $CUDA_VERSION
fi

wait