#!/bin/bash

if [ -z "${ANDURIL_NODELIST}" ]
then NODELIST=( vm3 vm4 vm5 vm6 vm7 vm8 )
else read -a NODELIST <<< "$ANDURIL_NODELIST"
fi

if [ "$1" == "-h" ]
then echo -ne '
This tools selects the lowest cpu allocated node for slurm run
export ANDURIL_NODELIST="vm3 vm4 vm5"
 ^ to have your own preferred list of nodes
Current nodelist: "'${NODELIST[@]}'"'"( ${#NODELIST[@]} nodes)\n"
   exit
fi


NODERUNS=( )
# find node with max free CPUs
for e in ${NODELIST[@]}
do NODERUNS+=( 0 )
done
while [ true ]
do
    MAX=0
    INDEX=0
    for (( i=0; i<${#NODERUNS[@]}; i++ ))
    do  NODERUNS[$i]=$(( $( scontrol -o show node ${NODELIST[$i]} | sed 's,.*CPUAlloc=\([0-9]\+\).*CPUTot=\([0-9]\+\).*,\2-\1 ,' ) ))
        scontrol -o show node ${NODELIST[$i]} | grep State=DOWN > /dev/null && NODERUNS[$i]=-1
        echo -n "${NODELIST[$i]}:${NODERUNS[$i]} "
        [ ${NODERUNS[$i]} -gt $MAX ] && {
            MAX=${NODERUNS[$i]}
            INDEX=$i
        }    
    done
    [ $MAX -gt 0 ] && {
        break
    } || {
        SECONDS=$(( ( RANDOM % 30 ) ))
        echo "waiting for free sockets for $SECONDS s."
        sleep $SECONDS
    }
done

echo "srun: Node ${NODELIST[$INDEX]}, Free sockets: ${NODERUNS[$INDEX]}" >&2

srun --nodelist=${NODELIST[$INDEX]} "$@"