Automate shutdown when no GPU processes are running and no users are logged in

shutdown_idle_gpu_machine.sh

#!/bin/bash

# A script for shutting down an idle VM with GPUs attached when they are idle and there
# are no users logged in. Add the below crontab via `crontab -e` to run this script every
# 10 minutes and append its output to a log:
#
#   */10 * * * * ~/shutdown_idle_gpu_machine.sh >> ~/shutdown_idle_gpu_machine.log 2>&1
#
# This script should live in your HOME directory. Your user should have sudoer privileges.

function log() {
  echo "$(date -u): $1"
}

function user_logged_in() {
  # Ignore tmux sessions, reporting only live ssh connections
  test "$(who | grep -v tmux | wc -l)" -gt 0
}

function nvidia_process_running() {
  PROCESSES=$(nvidia-smi -q | perl -n -e'/Processes\s+: (.+)/ && print $1')
  if [[ "$PROCESSES" =~ ^(None)+$ ]] ; then
    return 1
  else
    return 0
  fi
}

function uptime_more_than_ten_minutes() {
  UPTIME_SECONDS=$(cat /proc/uptime | cut -d' ' -f1 | awk '{print int($0)}')
  test $UPTIME_SECONDS -lt 600
}


uptime_more_than_ten_minutes
if [[ $? -eq 0 ]] ; then
  log "Uptime is less than 10 minutes, exiting without shutdown."
  exit 1
fi

user_logged_in
if [[  $? -eq 0 ]] ; then
  log "A user is currently logged in, exiting without shutdown."
  exit 1
fi

nvidia_process_running
if [[ $? -eq 0 ]] ; then
  log "A process is running an NVIDIA GPU, exiting without shutdown."
  exit 1
fi

log "All conditions met, shutting down machine."
sudo shutdown
exit 0