#!/bin/bash
# Per-GPU work queue worker, fed from the S3 bucket (rclone remote "jspace").
#
# Usage:  gpuq_worker.sh <worker-id> <gpu-index>
# e.g.:   gpuq_worker.sh node1-gpu3 3
# Run ONE per GPU, each in its own tmux window:
#   tmux new-window -t ssh_tmux -n gpuq3 'bash ~/jspace/scripts/gpuq_worker.sh node1-gpu3 3'
#
# Queue layout (per worker, race-free by construction — no work stealing):
#   jspace:jspace/gpuq/<worker-id>/pending/<job>.sh   submitted jobs
#   jspace:jspace/gpuq/<worker-id>/done/<job>.sh|.log finished jobs + logs
# Jobs run oldest-first by name — prefix with 00_, 01_, ... to order them.
#
# Job contract: plain bash, repo at $GPUQ_REPO (default ~/jspace),
# CUDA_VISIBLE_DEVICES is set by the worker.
#
# DATA CONTRACT (selective sync — declare in job-file comments):
#   # gpuq-in:  results-loop/mbpp_data.json results-loop/adapter_code.pt
#   # gpuq-out: results-loop/eval_code_gate*.json results/jbar_31b.pt
# Paths are repo-relative and mirror to jspace:jspace/<same path> in the
# bucket. Inputs are pulled before the job starts; outputs are pushed
# every SYNC_EVERY seconds while the job runs AND once at exit — a dying
# node loses at most one sync interval. Globs allowed in the basename.
# Repeat the comment lines to declare more paths. No declaration = no sync
# (old jobs keep working; they can still rclone by hand).
#
# GRACEFUL RESTART: touch $HOME/gpuq/<worker-id>/STOP — the worker exits
# between jobs, never mid-job. NEVER kill the tmux session while a job
# runs (that killed a 3.5h scan once).

set -u
ID=$1
GPU=$2
Q="jspace:jspace/gpuq/$ID"
LOCAL="$HOME/gpuq/$ID"
REPO="${GPUQ_REPO:-$HOME/jspace}"
SYNC_EVERY="${GPUQ_SYNC_EVERY:-120}"
mkdir -p "$LOCAL/pending"

sync_paths() {  # $1 = direction (in|out), $2 = job file
  local dir line p
  while IFS= read -r line; do
    for p in $line; do
      local d b
      d=$(dirname "$p"); b=$(basename "$p")
      mkdir -p "$REPO/$d"
      if [ "$1" = in ]; then
        rclone copy "jspace:jspace/$d" "$REPO/$d" --include "$b" -q 2>/dev/null
      else
        rclone copy "$REPO/$d" "jspace:jspace/$d" --include "$b" -q 2>/dev/null
      fi
    done
  done < <(sed -n "s/^# gpuq-$1:[[:space:]]*//p" "$2")
}

echo "[gpuq $ID] worker up on GPU $GPU, polling $Q/pending (repo $REPO)"
while true; do
  if [ -e "$LOCAL/STOP" ]; then
    rm -f "$LOCAL/STOP"
    echo "[gpuq $ID] $(date +%F_%H:%M:%S) STOP sentinel — draining (no job running)"
    break
  fi
  rclone move "$Q/pending" "$LOCAL/pending" --include "*.sh" -q 2>/dev/null
  job=$(ls "$LOCAL"/pending/*.sh 2>/dev/null | sort | head -1)
  if [ -z "$job" ]; then
    sleep 20
    continue
  fi
  name=$(basename "$job" .sh)
  echo "[gpuq $ID] $(date +%F_%H:%M:%S) START $name"
  sync_paths in "$job"
  CUDA_VISIBLE_DEVICES=$GPU bash "$job" > "$LOCAL/$name.log" 2>&1 &
  pid=$!
  while kill -0 $pid 2>/dev/null; do
    sleep "$SYNC_EVERY"
    sync_paths out "$job"
    rclone copy "$LOCAL/$name.log" "$Q/running/" -q 2>/dev/null
  done
  wait $pid
  rc=$?
  echo "[gpuq $ID] $(date +%F_%H:%M:%S) DONE  $name rc=$rc"
  echo "rc=$rc" >> "$LOCAL/$name.log"
  sync_paths out "$job"
  rclone copy "$LOCAL/$name.log" "$Q/done/" -q
  rclone deletefile "$Q/running/$name.log" -q 2>/dev/null
  rclone moveto "$job" "$Q/done/$name.sh" -q
done