nodeinfo

#!/usr/bin/env bash

# nodeinfo

set -eu -o pipefail

help() {
    # Display help information
    printf "%s\n" "nodeinfo / Display node information by partition, CPU/GPU models, and state"
    printf "%s\n" ""
    printf "%s\n" "Usage: nodeinfo [-h] [-p <partition>] [-s <state>]"
    printf "%s\n" ""
    printf "%s\n" "Options:"
    printf "%s\n" ""
    printf "%s\n" "-h  Display help"
    printf "%s\n" ""
    printf "%s\n" "-p  Filter partitions"
    printf "%s\n" ""
    printf "%s\n" "    * Default value shows all partitions"
    printf "%s\n" "    * Multiple partitions can be specified with comma separation"
    printf "%s\n" ""
    printf "%s\n" "-s  Filter node states"
    printf "%s\n" ""
    printf "%s\n" "    * Default value shows all nodes (in any state)"
    printf "%s\n" "    * Common node state options include the following:"
    printf "%s\n" "        alloc,down,drain,fail,idle,maint,mix"
    printf "%s\n" "    * Multiple options can be specified with comma separation"
    printf "%s\n" ""
    printf "%s\n" "Examples:"
    printf "%s\n" ""
    printf "%s\n" "  nodeinfo"
    printf "%s\n" "  nodeinfo -h"
    printf "%s\n" "  nodeinfo -p epyc-64"
    printf "%s\n" "  nodeinfo -s idle"
    printf "%s\n" "  nodeinfo -s idle,mix"
    printf "%s\n" "  nodeinfo -p largemem -s idle"
    printf "%s\n" "  nodeinfo -p main,gpu -s idle,mix"
}

summarize() {
    # Summarize partition status
    printf "%s\n" "--------------------------------------------------------------------"
    printf "%s\n" "Partition    Timelimit   Status Nodes by state  CPUs by state       "
    printf "%s\n" "                                A/I/O/Total     A/I/O/Total         "
    printf "%s\n" "------------ ----------- ------ --------------- --------------------"
    sinfo -h -o "%12P %11l %6a %15F %20C"
    printf "%s\n" "------------"
    printf "%s\n" "A = allocated or mixed; I = idle; O = other"
}

# Set default values for option arguments
states=all
partitions=
# There is no all value for sinfo --partition option, so leave blank
# See if statement below

# Define options
while getopts ":hzp:s:" option; do
    case $option in
      h) # display help
         help
         exit 0
         ;;
      z) # summarize partition status
         summarize
         exit 0
         ;;
      p) # select partitions
         partitions="$OPTARG"
         if [[ $partitions =~ ^- ]]; then
             printf "%s\n" "Error: Invalid argument to option -p"
             exit 1
         fi
         ;;
      s) # select node states
         states="$OPTARG"
         if [[ $states =~ ^- ]]; then
             printf "%s\n" "Error: Invalid argument to option -s"
             exit 1
         fi
         ;;
      \?) # invalid argument
         printf "%s\n" "Error: Invalid option -$OPTARG" >&2
         exit 1
         ;;
    esac
done

printf "%s\n" "-----------------------------------------------------------------------------------------------------------------------"
printf "%s\n" "Partition    Timelimit   CPU model     CPUs/ Memory(MB)/ GPU model          State  Nodes Nodelist                      "
printf "%s\n" "                                       node  node                                                                      "
printf "%s\n" "------------ ----------- ------------- ----- ----------- ------------------ ------ ----- ------------------------------"

if [[ $partitions = "" ]]; then
    # List all partitions
    sinfo -h -o "%12P %11l %13b %5c %11m %18G %6t %5D %30N" -t $states
else
    # List specified partitions
    sinfo -h -o "%12P %11l %13b %5c %11m %18G %6t %5D %30N" -t $states -p $partitions
fi