#!/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 ~/jspace, venv at ~/jspace/.venv (or
# /venv/main on vast images — job script picks), CUDA_VISIBLE_DEVICES is set
# by the worker. Job handles its own LOOP_OUT / result rclone if needed.

set -u
ID=$1
GPU=$2
Q="jspace:jspace/gpuq/$ID"
LOCAL="$HOME/gpuq/$ID"
mkdir -p "$LOCAL/pending"

echo "[gpuq $ID] worker up on GPU $GPU, polling $Q/pending"
while true; do
  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"
  CUDA_VISIBLE_DEVICES=$GPU bash "$job" 2>&1 | tee "$LOCAL/$name.log"
  rc=${PIPESTATUS[0]}
  echo "[gpuq $ID] $(date +%F_%H:%M:%S) DONE  $name rc=$rc"
  echo "rc=$rc" >> "$LOCAL/$name.log"
  rclone copy "$LOCAL/$name.log" "$Q/done/" -q
  rclone moveto "$job" "$Q/done/$name.sh" -q
done