gpu-usage

#!/bin/bash
nodes=`sinfo --format="%N" --noheader`
header=true
pretty=false

print_usage () {
  cat << EOM
Usage: $0 [-n charles[01-19]]
    Get gpu usage for all, or specified nodes

Arguments
    -n    (optional) specify nodes, either comma separated list, or summarised.
          Mirrors slurm command -n usage in sinfo, and -w usage for squeue." 
    -h    (optional) do not print header
    -p    (optional) make output pretty, ignores -h if specified

Examples:
    $ # Get status of all nodes aggregated
    $ ./gpu-usage
    >    in_use,usable,total
    >    18,34,55

    $ # Suppress header
    $ ./gpu-usage -h
    >    18,34,55

    $ # Specify specific nodes in comma separated list
    $ ./gpu-usage -n charles01,charles03 -p
    >    in_use  usable  total
    >    1       2       4

    $ # Same as above but with summarised node format
    $ ./gpu-usage -n charles[01,03] -p
    >    in_use  usable  total
    >    1       2       4

    $ # Example with split list
    $ ./gpu-usage -p -n charles[01-06,11-19]
    >    in_use  usable  total
         14      26      47
 

EOM
}

while getopts 'n:hp' flag; do
  case "${flag}" in
    n) nodes="${OPTARG}" ;;
    h) header=false ;;
    p) pretty=true ;;
    *) print_usage
       exit ;;
  esac
done

if [ "$pretty" = true ]; then
  $0 -n ${nodes} | column -t -s','
  exit 0
fi

# Count number of GPUs requested by running jobs
USED=`squeue -t R -o %b -w ${nodes} | grep -o ':[0-9]' | awk ' {s+=substr($1,2)} END {print s}'`
if [ "$USED" == '' ]; then
  USED=0
fi

# Count number of GPUS in nodes marked idle or mix
USABLE=$(sinfo -t idle,mix,alloc -r -N -o %G -n ${nodes} | grep -o ':[0-9]' | awk ' {s+=substr($1,2)} END {print s}')
if [ "$USABLE" == '' ]; then
  USABLE=0
fi

# Count number of GPUs in all nodes
TOTAL=$(sinfo -N -o %G -n ${nodes} | grep -o ':[0-9]' | awk ' {s+=substr($1,2)} END {print s}')
if [ "$TOTAL" == '' ]; then
  TOTAL=0
fi

if [ "$header" = true ]; then 
  echo in_use,usable,total
fi
echo "${USED},${USABLE},${TOTAL}"