#!/bin/bash
# One-screen status: progress, speed, ETA, loss, checkpoints, disk, helper scripts.
cd /fsx/ajay-stu/AR
STOP=$(grep -oP '^STOP_AT=\K[0-9]+' run_args.sh)
J=$(squeue -u $USER -h -n ar-real -t RUNNING -o %i | head -1)
if [ -z "$J" ]; then
  echo "NO training job running"; squeue -u $USER -o "%.8i %.8T %.4D %.10M %j"
else
  NN=$(squeue -j $J -h -o %D)
  L=$(grep "^step" logs/real_$J.log 2>/dev/null | tail -1)
  if [ -z "$L" ]; then
    echo "job $J on $NN nodes | starting up: loading checkpoint / compiling (first step in ~3-5 min)"
  else
    # tokens per step depend on node count (same table as the job script)
    case $NN in 3|6) TPSTEP=393216;; 5) TPSTEP=327680;; 7) TPSTEP=458752;; *) TPSTEP=524288;; esac
    S=$(echo "$L" | awk '{print $2}')
    TPS=$(echo "$L" | awk '{gsub("k", "", $10); print $10 * 1000}')
    SEC=$(awk -v l=$((STOP - S)) -v t=$TPSTEP -v r=$TPS 'BEGIN {printf "%d", (r > 0) ? l * t / r : 0}')
    echo "job $J on $NN nodes | $L"
    awk -v s=$S -v stop=$STOP -v sec=$SEC 'BEGIN {
      printf "progress: %d / %d (%.1f%%)   left: %d steps, ~%.1f days\n", s, stop, 100*s/stop, stop-s, sec/86400 }'
    echo "finish ETA: $(date -u -d "+$SEC sec" '+%a %d %b %H:%M UTC')  ($(TZ=Asia/Kolkata date -d "+$SEC sec" '+%H:%M IST'))"
    grep "^step" logs/real_$J.log | tail -100 | awk '{s+=$5} END {printf "avg loss (last 100 prints): %.3f\n", s/NR}'
  fi
  echo "last checkpoint: $(grep -E '^\[ckpt\]' logs/real_$J.log 2>/dev/null | tail -1)"
fi
# disk now, lowest in the last 24 h, and helper scripts
echo "disk now:   $(tail -1 logs/disk.log)"
echo "lowest 24h: $(grep free= logs/disk.log | tail -288 | sort -t= -k2 -n | head -1)"
pgrep -f watchdog.sh   >/dev/null && echo "watchdog:   running" || echo "watchdog:   NOT running"
pgrep -f diskwatch.sh  >/dev/null && echo "diskwatch:  running" || echo "diskwatch:  NOT running"
pgrep -f autoreload.sh >/dev/null && echo "autoreload: running" || echo "autoreload: NOT running"
echo "tokenizing jobs: $(squeue -u $USER -h -n ar-tok | wc -l)"