#!/bin/sh -e

##########################################################################
#   Script description:
#       Power down idle compute nodes to save power.
#       PowerEdge R415 with dual Opteron 4386, dual power supplies and
#       iDRAC Enterprise draws 14W powered off, 100W on and idle.
#       (300W under full CPU load in case you were curious)
#
#       Nodes can be powered on remotely with something like the following
#       if remote IPMI is enabled:
#
#           ipmitool -v -I lanplus -H IP-address [-P password] \
#                    -U root chassis power on
#
#       To enable, run the following on the node while it's up:
#
#           ipmitool -I open lan set 1 access on
#       
#   History:
#   Date        Name        Modification
#   2020-02-23  Jason Bacon Begin
##########################################################################

auto-root-check $0 "Needs to set drain state."

case `auto-ostype` in
FreeBSD)
    power_off=-p
    ;;

RHEL)
    power_off=-h
    ;;

*)
    printf "$0: Not supported on $(auto-ostype).\n"
    exit 1
    ;;

esac

idle_nodes=$(sinfo -h -o '%T %n' | awk '$1 ~ "idle" { print $2 }')
idle_nodes="$idle_nodes"
printf "Idle nodes:\n$idle_nodes\n"
file_servers=$(cluster-file-servers)

for node in $idle_nodes; do
    # Exclude file servers
    if ! awk '{ print $1 }' /etc/fstab | grep -q "^${node}.*:"; then
	compute_only_nodes="$compute_only_nodes $node"
    else
	file_servers="$file_servers $node"
    fi
done
printf "Compute-only nodes:\n$compute_only_nodes\n"
printf "File servers:\n$file_servers\n"

for node in $compute_only_nodes; do
    printf "Draining idle node $node...\n"
    scontrol update nodename=$node state=drain reason=power-save
done

# Strip spaces
temp=$(echo "$compute_only_nodes" | sed -e 's| ||g')
if [ ! -z $temp ]; then
    printf "\nGiving SLURM time to set drain state...\n"
    sleep 5
fi

for node in $compute_only_nodes; do
    state=$(sinfo -h -n $node -o '%T')
    reason=$(sinfo -h -n $node -o '%E')
    case $state in
    "drained"|"drained*")
	if [ 0"$reason" = 0power-save ]; then
	    printf "Shutting down idle node $node...\n"
	    ssh -o ConnectTimeout=5 $node shutdown -p now || true
	fi
	;;
    
     *)
	;;
    esac
done

temp=$(echo "$file_servers" | sed -e 's| ||g')
if [ ! -z $temp ]; then
    if [ "0$(squeue -h)" = 0 ]; then
	printf "No jobs running.  Shutting down file servers.\n"
	for node in $file_servers; do
	    if [ $(node-type $node) = compute ]; then
		printf "Draining $node...\n"
		scontrol update nodename=$node state=drain reason=power-save || true
	    fi
	done
    
	# Strip spaces
	printf "\nGiving SLURM time to set drain state...\n"
	sleep 5
	
	for node in $file_servers; do
	    state=$(sinfo -h -n $node -o '%T')
	    reason=$(sinfo -h -n $node -o '%E')
	    case $state in
	    "drained"|"drained*")
		if [ 0"$reason" = 0power-save ]; then
		    printf "Unmounting $node on head node...\n"
		    for mount_point in $(awk -v node=$node '$1 ~ node { print $2 }' /etc/fstab); do
			if [ "0$mount_point" != 0 ]; then
			    printf "Unmounting $mount_point...\n"
			    umount -f $mount_point
			    ssh -o ConnectTimeout=10 $node shutdown -p now
			fi
		    done
		fi
		;;
	    
	    *)
		;;
	    
	    esac
	done
    else
	printf "Jobs running, leaving file servers up.\n"
    fi
fi
