#!/bin/bash
# Submit jobs to gpuq workers.
#
#   gpuq_submit.sh <worker-id> <jobfile.sh> [...]     # pin to a worker
#   gpuq_submit.sh --class <class> <jobfile.sh> [...] # pick idlest in class
#   gpuq_submit.sh --status
#
# Class rosters live in the bucket: gpuq/_classes/<class>.txt (one worker id
# per line). Scheduling: prefer workers with 0 pending AND <15% GPU util
# (from the node heartbeat); tie-break by fewest pending, then random.

set -eu
if [ "$1" = "--status" ]; then
  echo "== pending =="
  rclone lsf -R jspace:jspace/gpuq --include "pending/*.sh" 2>/dev/null
  echo "== done (last 20) =="
  rclone lsl -R jspace:jspace/gpuq --include "done/*.sh" 2>/dev/null | sort -k2,3 | tail -20
  exit 0
fi

pick_worker() {  # $1 = class
  local best="" best_score=999999
  for w in $(rclone cat "jspace:jspace/gpuq/_classes/$1.txt" 2>/dev/null); do
    local node=${w%-gpu*} gpu=${w##*-gpu}
    local pend util score
    pend=$(rclone lsf "jspace:jspace/gpuq/$w/pending/" --include "*.sh" 2>/dev/null | wc -l)
    util=$(rclone cat "jspace:jspace/gpuq/_health/$node.txt" 2>/dev/null \
           | grep -E "^$gpu, " | head -1 | awk -F', ' '{print $2}' | tr -dc 0-9)
    util=${util:-100}
    score=$((pend * 100 + util))
    if [ "$score" -lt "$best_score" ]; then best=$w; best_score=$score; fi
  done
  echo "$best"
}

if [ "$1" = "--class" ]; then
  CLASS=$2; shift 2
  for f in "$@"; do
    rclone copy "$f" "jspace:jspace/gpuq/_class-pending/$CLASS/" -q
    echo "queued $(basename $f) -> class $CLASS (dispatcher assigns on idle)"
  done
  exit 0
fi

ID=$1; shift
for f in "$@"; do
  rclone copy "$f" "jspace:jspace/gpuq/$ID/pending/" -q
  echo "queued $(basename $f) -> $ID"
done