#!/bin/bash
# Per-node gpuq supervisor: keeps one worker per GPU alive and publishes a
# heartbeat to the bucket. Run ONE per node, in its own tmux window:
#   tmux new-window -t ssh_tmux -n gpuq_sup \
#     'bash ~/jspace/scripts/gpuq_supervisor.sh node1 5'
#
# Usage: gpuq_supervisor.sh <node-id> <num-gpus> [tmux-session]
# Worker ids are <node-id>-gpu<i>; heartbeat lands at
#   jspace:jspace/gpuq/_health/<node-id>.txt   (age > ~2 min == node in trouble)
# Check fleet health from anywhere:
#   rclone cat jspace:jspace/gpuq/_health/node1.txt

set -u
NODE=$1
NGPU=$2
SESSION=${3:-ssh_tmux}
DIR="$(cd "$(dirname "$0")" && pwd)"

echo "[gpuq-sup $NODE] supervising $NGPU workers in session $SESSION"
while true; do
  for i in $(seq 0 $((NGPU - 1))); do
    if ! pgrep -f "gpuq_worker.sh $NODE-gpu$i " >/dev/null 2>&1; then
      echo "[gpuq-sup $NODE] $(date +%F_%H:%M:%S) worker gpu$i dead — respawning"
      tmux new-window -d -t "$SESSION" -n "gpuq$i" \
        "bash $DIR/gpuq_worker.sh $NODE-gpu$i $i" 2>/dev/null \
        || nohup bash "$DIR/gpuq_worker.sh" "$NODE-gpu$i" "$i" \
             >> "$HOME/gpuq/$NODE-gpu$i.worker.log" 2>&1 &
    fi
  done
  {
    date -u +%F_%H:%M:%SZ
    uptime
    nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.total \
      --format=csv,noheader 2>/dev/null
    df -h / /dev/shm 2>/dev/null | tail -2
    echo "workers: $(pgrep -fc "gpuq_worker.sh $NODE-gpu" || echo 0)/$NGPU"
    echo "pending: $(ls "$HOME"/gpuq/$NODE-gpu*/pending/*.sh 2>/dev/null | wc -l) local"
  } > /tmp/gpuq_health_$NODE.txt
  rclone copyto /tmp/gpuq_health_$NODE.txt \
    "jspace:jspace/gpuq/_health/$NODE.txt" -q 2>/dev/null
  sleep 60
done