#!/bin/bash
# Test the given nodes in parallel, pick the 8 best for training, list the rest to hand over.
# Usage: ./testnodes.sh ip-a,ip-b,...        (about 1-2 minutes)
cd /fsx/ajay-stu/AR
NODES=$(echo "$1" | tr ',' ' ')
COLLEAGUE="ip-10-0-136-14 ip-10-0-136-94 ip-10-0-136-253 ip-10-0-138-178 ip-10-0-139-113 ip-10-0-140-247"
OUT=logs/nodetest; rm -rf $OUT; mkdir -p $OUT

# 1) one short whole-node job per node, all at once; give up on a node Slurm can't give within 60 s
for n in $NODES; do
  srun -p dev -N1 -n1 -w $n --exclusive --gres=gpu:8 --immediate=60 --time=5 \
    bash /fsx/ajay-stu/AR/nodecheck.sh > $OUT/$n.txt 2> $OUT/$n.err &
done
wait

# 2) table and verdict per node
printf "%-17s %4s %7s %6s %5s %6s %6s %6s  %s\n" node gpus gpuMiB ramGB swap load users fsxMB verdict
for n in $NODES; do
  unset h g gu ram sw ld us mb
  read h g gu ram sw ld us mb < <(cat $OUT/$n.txt 2>/dev/null)
  if [ -z "$h" ]; then v="NOT AVAILABLE in Slurm"
  elif echo " $COLLEAGUE " | grep -q " $n "; then v="colleague node"
  elif grep -qx "$n" logs/bad_nodes_perm.txt 2>/dev/null; then v="known bad"
  elif [ "$g" -ne 8 ] || [ "$gu" -ge 2000 ] || [ "$ram" -lt 1500 ] || [ "$mb" -lt 100 ] \
       || awk -v l="$ld" 'BEGIN {exit !(l > 16)}'; then v="BAD"
  else v="ok"; echo "$us $ld $n" >> $OUT/ok.txt; fi
  printf "%-17s %4s %7s %6s %5s %6s %6s %6s  %s\n" $n "${g:--}" "${gu:--}" "${ram:--}" "${sw:--}" "${ld:--}" "${us:--}" "${mb:--}" "$v"
done

# 3) best 8 ok nodes (fewest other users, then lowest load); everything else is handed over
PICK=$(sort -k1,1n -k2,2g $OUT/ok.txt 2>/dev/null | head -8 | awk '{print $3}')
echo "$PICK" | grep . > logs/picked_nodes.txt
REST=$(for n in $NODES; do grep -qx $n logs/picked_nodes.txt || echo $n; done)
echo
echo "USE for training ($(echo "$PICK" | grep -c .)): $(echo $PICK | tr ' ' ',')"
echo "GIVE AWAY ($(echo "$REST" | grep -c .)): $(echo $REST | tr ' ' ',')"