#!/bin/bash
# gpuq dispatcher: assigns class-queued jobs to idle workers, continuously.
# Run exactly ONE instance, on an always-on machine (the Spark):
#   tmux/setsid bash scripts/gpuq_dispatcher.sh
#
# Submit side drops jobs in  jspace:jspace/gpuq/_class-pending/<class>/ ;
# this loop moves each job to a concrete worker's pending/ the moment a
# worker in that class is idle (empty pending + GPU util <15% + fresh
# heartbeat). Single dispatcher == single writer == no claim races.
# Assignment is a bucket-side move, so a dispatcher crash loses nothing.

set -u
POLL=${1:-30}
echo "[gpuq-dispatch] up, polling ${POLL}s"
while true; do
  for cls in $(rclone lsf jspace:jspace/gpuq/_classes/ 2>/dev/null | sed 's/\.txt$//'); do
    jobs=$(rclone lsf "jspace:jspace/gpuq/_class-pending/$cls/" --include "*.sh" 2>/dev/null | sort)
    [ -z "$jobs" ] && continue
    for w in $(rclone cat "jspace:jspace/gpuq/_classes/$cls.txt" 2>/dev/null); do
      job=$(echo "$jobs" | head -1); [ -z "$job" ] && break
      node=${w%-gpu*}; gpu=${w##*-gpu}
      hb=$(rclone cat "jspace:jspace/gpuq/_health/$node.txt" 2>/dev/null) || continue
      # heartbeat freshness: first line is UTC date_%H:%M:%SZ
      hbts=$(echo "$hb" | head -1); hbage=$(( $(date -u +%s) - $(date -u -d "${hbts//_/ }" +%s 2>/dev/null || echo 0) ))
      [ "$hbage" -gt 180 ] && continue
      util=$(echo "$hb" | grep -E "^$gpu, " | head -1 | awk -F', ' '{print $2}' | tr -dc 0-9)
      pend=$(rclone lsf "jspace:jspace/gpuq/$w/pending/" --include "*.sh" 2>/dev/null | wc -l)
      if [ "$pend" -eq 0 ] && [ "${util:-100}" -lt 15 ]; then
        rclone moveto "jspace:jspace/gpuq/_class-pending/$cls/$job" \
                      "jspace:jspace/gpuq/$w/pending/$job" -q \
          && echo "[gpuq-dispatch] $(date +%H:%M:%S) $job -> $w"
        jobs=$(echo "$jobs" | tail -n +2)
      fi
    done
  done
  sleep "$POLL"
done