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

Output:
    Outputs three columns summarising the current usage of GPUs on the cluster:
        -- The 'in_use' column shows the number of GPUs that are currently allocated to a job.
        -- The 'usable' column shows the number of GPUs that are not in an error state or otherwise unable to be
            allocated to a job.
        -- The 'total' column shows the total number of GPUs, including those that are currently unable to be allocated
            to a job.

    The number of free GPUs (GPUs not allocated to a job and not in an error state) is the difference between the
    'usable' column and the 'in_use' column.

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

# Get number of free GPUs
FREE=$(expr "${USABLE}" - "${USED}")

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