#!/bin/bash
# ONE-COMMAND bootstrap for rented GPU nodes (vast.ai and similar).
#
#   bash node_bootstrap.sh <node-id> [num-gpus] [--dry-run]
#
# Encodes every lesson this project paid for (LESSONS.md, memory notes):
#   - storage: /workspace (survives instance stop) — NEVER /dev/shm (RAM,
#     wiped on stop). Aborts if /workspace is missing rather than degrade.
#   - repo: init/fetch/checkout pattern (safe when the dir already exists;
#     plain `git clone` fails silently into non-empty dirs).
#   - models: pinned by revision, loaded by SNAPSHOT PATH (never refs/main;
#     no HF_HUB_OFFLINE — it breaks dataset streaming). Use pin_model().
#   - sync: if ~/.config/rclone/rclone.conf exists (owner places it — never
#     copied by automation), workers sync results directly to S3 via the
#     gpuq data contract (# gpuq-in/out job comments) and a 120s sidecar
#     mirrors logs. Without credentials the node still works; use presigned
#     PUT URLs (scripts/gpuq_presign.py, run at home) embedded in jobs.
#   - workers: one gpuq worker per GPU, graceful STOP sentinel
#     (touch ~/gpuq/<id>/STOP), never kill mid-job.
set -u
NODE_ID=${1:?usage: node_bootstrap.sh <node-id> [num-gpus] [--dry-run]}
NGPU=${2:-1}
DRY=${3:-}

say() { echo "[bootstrap $NODE_ID] $*"; }

# --- storage ---------------------------------------------------------
if [ ! -d /workspace ]; then
  say "FATAL: /workspace missing — rent instances with a disk allocation"
  exit 1
fi
export GPUQ_REPO=/workspace/jspace
export HF_HOME=/workspace/hf
say "storage: repo=$GPUQ_REPO hf=$HF_HOME ($(df -h /workspace | awk 'NR==2{print $4}') free)"

# --- credentials / sync mode ----------------------------------------
if [ -f "$HOME/.config/rclone/rclone.conf" ]; then
  SYNC_MODE=rclone
else
  SYNC_MODE=presign
fi
say "sync mode: $SYNC_MODE"

[ "$DRY" = "--dry-run" ] && { say "dry-run OK"; exit 0; }

# --- python + deps ---------------------------------------------------
source /venv/main/bin/activate 2>/dev/null || true
uv pip install -q "transformers==5.13.*" datasets accelerate

# --- repo ------------------------------------------------------------
mkdir -p "$GPUQ_REPO" && cd "$GPUQ_REPO"
git init -q 2>/dev/null
git remote add origin https://git.draic.info/nils/jspace.git 2>/dev/null
git fetch -q origin main && git checkout -q -f FETCH_HEAD || { say "FATAL: repo fetch failed"; exit 1; }
mkdir -p "$GPUQ_REPO/results" "$GPUQ_REPO/results-loop"

# --- model pinning helper (source this file to use it in job scripts) -
pin_model() {  # pin_model <repo-id> <revision>; echoes snapshot path
  local M=$1 REV=$2
  for i in 1 2 3 4 5; do
    timeout 3600 hf download "$M" --revision "$REV" >&2 && break
    echo "RETRY $i: $M" >&2; sleep 5
  done
  mkdir -p "$HF_HOME/hub/models--${M//\//--}/refs"
  echo "$REV" > "$HF_HOME/hub/models--${M//\//--}/refs/main"
  echo "$HF_HOME/hub/models--${M//\//--}/snapshots/$REV"
}
export -f pin_model

# --- log sidecar (rclone mode only) ----------------------------------
if [ "$SYNC_MODE" = rclone ] && ! pgrep -f "[n]ode_logsync" > /dev/null; then
  cat > /root/node_logsync.sh <<EOF
#!/bin/bash
while true; do
  rclone copy \$HOME/gpuq jspace:jspace/gpuq-nodelogs/$NODE_ID/ --include "*.log" -q 2>/dev/null
  sleep 120
done
EOF
  chmod +x /root/node_logsync.sh
  setsid nohup /root/node_logsync.sh > /dev/null 2>&1 < /dev/null &
  say "log sidecar started"
fi

# --- workers ---------------------------------------------------------
for g in $(seq 0 $((NGPU - 1))); do
  W="$NODE_ID-gpu$g"
  if tmux has-session -t "gpuq_$W" 2>/dev/null; then
    say "worker $W already running"
  else
    tmux new-session -d -s "gpuq_$W" \
      "GPUQ_REPO=$GPUQ_REPO bash $GPUQ_REPO/scripts/gpuq_worker.sh $W $g"
    say "worker $W started (queue: jspace:jspace/gpuq/$W/pending)"
  fi
done
say "NODE-READY  (submit: gpuq_submit.sh $NODE_ID-gpu0 job.sh; drain: touch ~/gpuq/<id>/STOP)"