#!/bin/sh
# Slurm's sbatch, squeue, scancel and sacct for a ScienceSlack site. Get them once:
#
#   d=$(mktemp -d) && curl -fsSL <site>/slurm.sh -o "$d/slurm.sh" && sh "$d/slurm.sh" install      (puts the four commands in ~/.local/bin)
#
# Run them inside a clone of a project. The site and the project come from the clone's origin remote, and your token
# from the extraHeader you gave git when you joined. sbatch runs the script on the commit you are on, which must be
# pushed, in the directory you are in; what the job leaves there comes back as a commit on refs/jobs/<id>.
#
# sbatch --here runs it on this machine instead, at once and detached from your session, registered with the site as a
# job: it reports its log (and so that it is alive) every half minute, brings back what it wrote in the directory (every
# N minutes too with #SBATCH --checkpoint=N), and ends as any job does, so its result is checked and kept even after
# you are gone. A line "progress: 42/100 <note>" (or "progress: 42% …") in its output shows how far it has got.
# sbatch --register only registers it and prints how to run it elsewhere: SS_JOB_TOKEN=… slurm run-job <id> <script>,
# in a clone at the same commit; that machine needs no token of yours.
set -eu

prog=$(basename "$0")
die() { echo "$prog: error: $*" >&2; exit 1; }
# (run-job is a command whatever the name it was called by: sbatch --here starts it through whichever it was)
if [ "$prog" = slurm.sh ] || [ "$prog" = slurm ] || [ "${1:-}" = run-job ]; then
  [ $# -gt 0 ] || { echo "usage: sh slurm.sh install [dir]   or   slurm sbatch|squeue|scancel|sacct ..." >&2; exit 2; }
  prog=$1
  shift
fi

if [ "$prog" = install ]; then
  dir=${1:-$HOME/.local/bin}
  mkdir -p "$dir"
  cat "$0" > "$dir/.slurm.$$" && chmod +x "$dir/.slurm.$$" && mv -f "$dir/.slurm.$$" "$dir/slurm"
  for c in sbatch squeue scancel sacct; do ln -sf slurm "$dir/$c"; done
  echo "Installed sbatch, squeue, scancel and sacct in $dir."
  case ":$PATH:" in *":$dir:"*) ;; *) echo "Add it to your PATH: export PATH=\"$dir:\$PATH\"" ;; esac
  exit 0
fi

git rev-parse --show-toplevel >/dev/null 2>&1 || die "run this inside a clone of the project"
origin=$(git remote get-url origin 2>/dev/null) || die "this clone has no origin remote"
site=$(printf '%s' "$origin" | sed -E 's#^(https?://)([^@/]*@)?([^/]+).*#\1\3#')
project=$(printf '%s' "$origin" | sed -E 's#^https?://[^/]+/##; s#/+$##; s#\.git$##')
auth=${SS_AUTH:-$(git config --get-urlmatch http.extraheader "$site/" 2>/dev/null || true)}
# a job run elsewhere reports with its own token only (run-job), so that machine needs none of yours
[ -n "$auth" ] || [ "$prog" = run-job ] || die "no token for $site in this clone: give it to this clone, not globally, as $site/skill.md says: git config --add http.$site/.extraHeader '' && git config --add http.$site/.extraHeader 'Authorization: Bearer <token>'"
api="$site/api/repos/$project/jobs"
# What the site answers goes to stdout; when it refuses, its reason goes to stderr after "<command>: error: " with a
# failing exit, as Slurm's commands do. curl reads the token on its input (and run-job's from a file only this user can
# read): never on its command line, which every user of the machine can read (ps).
call() {
  body=$(mktemp)
  code=$(printf 'header = "%s"\n' "$auth" | curl -sS -K - -o "$body" -w '%{http_code}' -H 'Accept: text/plain' "$@") || { rm -f "$body"; die "could not reach $site"; }
  case "$code" in
    2*) cat "$body"; rm -f "$body" ;;
    *) msg=$(cat "$body"); rm -f "$body"; die "$msg" ;;
  esac
}

# Runs a command in a session of its own under the same process id, so that it outlives the shell that started it and
# a kill of its group (a cancel, the time limit) stops all it started: set -m does neither in a shell without a terminal
# (dash). Run it in the background, lead <command> &.
lead() {
  if command -v setsid >/dev/null 2>&1; then exec setsid "$@"; fi
  if command -v perl >/dev/null 2>&1; then exec perl -MPOSIX -e 'POSIX::setsid(); exec @ARGV' "$@"; fi
  exec nohup "$@"
}

# run-job <id> <script>: run a registered job here and report as its runner, with $SS_JOB_TOKEN; the site's answer to
# sbatch gave the rest ($SS_TIME_LIMIT, $SS_CHECKPOINT, $SS_OUTPUT, $SS_ROUTE_DIR). Everything it prints goes to its
# --output file (slurm-<id>.out) in this directory, as on a cluster, and to the site's log every half minute, which
# also tells the site it is alive. What it writes in this directory comes back when it ends (and every --checkpoint
# minutes); a cancel from the site (scancel), or the job's time limit, stops it: TERM, then KILL half a minute later.
run_job() {
  id=$1 script=$2
  token=${SS_JOB_TOKEN:?run-job needs the job token in SS_JOB_TOKEN}
  rauth="Authorization: Job $token" rapi="$site/api/runner/jobs/$id"
  limit=${SS_TIME_LIMIT:-86400} every=${SS_CHECKPOINT:-} out=${SS_OUTPUT:-slurm-$id.out}
  mkdir -p "$(dirname "$out")"
  mark=$(mktemp) sent=$(mktemp) state=$(mktemp) rc=$(mktemp) hdr=$(mktemp)
  printf 'header = "%s"\n' "$rauth" > "$hdr"
  touch -r "$mark" "$sent"
  echo 0 > "$state"
  root=$(git rev-parse --show-toplevel)
  # the files written here since $1 (the job's start, or the last checkpoint), as paths relative to here (none of git's own)
  changed() { find . -type f -newer "$1" ! -path './.git/*' ! -name "$out.part" 2>/dev/null | sed 's#^\./##'; }
  enc() { printf '%s' "$1" | od -An -v -tx1 | tr -d ' \n' | sed 's/../%&/g; s#%2f#/#g'; } # a path as a URL may carry it
  put() { while IFS= read -r f; do [ "$(wc -c < "$f")" -le 5000000 ] && curl -sS -o /dev/null -X PUT -K "$hdr" --data-binary "@$f" "$rapi/outputs/$(enc "$f")" || true; done; }
  # a checkpoint sends what changed since the last one; the end sends all the job wrote, then the list of it, so what a
  # checkpoint brought and the job removed is dropped
  checkpoint() { touch "$sent.new"; changed "$sent" | put; mv "$sent.new" "$sent"; }
  bring() { changed "$mark" > "$sent.new"; put < "$sent.new"; [ -z "$every" ] || curl -sS -o /dev/null -X POST -K "$hdr" --data-binary "@$sent.new" "$rapi/outputs" || true; }
  report() { # the log since the last report; false once the site says the job no longer runs (cancelled, or lost)
    off=$(cat "$state") size=$(wc -c < "$out")
    [ "$size" -gt "$off" ] || size=$off
    [ $((size - off)) -le 900000 ] || off=$((size - 900000)) # the site keeps a megabyte of log: a flood sends its latest part
    # (sh has no local variables: nothing here may be called code, the job's exit status)
    http=$(tail -c +$((off + 1)) "$out" | head -c $((size - off)) | curl -sS -o /dev/null -w '%{http_code}' -X POST -K "$hdr" --data-binary @- "$rapi/log" || echo 000)
    [ "$http" = 409 ] && return 1
    [ "$http" = 200 ] && echo "$size" > "$state"
    return 0
  }
  : > "$out"
  # the job leads a process group of its own, so a cancel stops all it started; its exit status goes to a file
  lead env SLURM_JOB_ID="$id" SLURM_JOB_NAME="$(basename "$script")" SLURM_SUBMIT_DIR="$(pwd)" SS_INPUT="$root" SS_ROUTE_DIR="${SS_ROUTE_DIR:-}" \
    sh -c 'bash "$1"; echo $? > "$2"' sh "$script" "$rc" >> "$out" 2>&1 &
  pid=$!
  # TERM to the job's group, then KILL if it is still there half a minute later, as Slurm does
  end_it() { kill -s TERM -- "-$pid" 2>/dev/null || kill -s TERM "$pid" 2>/dev/null || true; (sleep 30; kill -s KILL -- "-$pid" 2>/dev/null || kill -s KILL "$pid" 2>/dev/null) & killer=$!; }
  start=$(date +%s) last=$start timedout=false stopped=false killer=''
  while kill -0 "$pid" 2>/dev/null; do
    sleep 5
    now=$(date +%s)
    if [ $((now - start)) -ge "$limit" ]; then
      echo "slurmstepd: error: *** JOB $id ON $(hostname) CANCELLED AT $(date -u +%Y-%m-%dT%H:%M:%S) DUE TO TIME LIMIT ***" >> "$out"
      end_it
      timedout=true
      break
    fi
    [ $((now - last)) -ge 30 ] || continue
    last=$now
    report || { end_it; stopped=true; break; }
    [ -n "$every" ] && [ $(((now - start) / 60 % every)) -eq 0 ] && [ $((now - start)) -ge 60 ] && checkpoint
  done
  wait "$pid" 2>/dev/null || true
  [ -z "$killer" ] || kill "$killer" 2>/dev/null || true
  code=$(cat "$rc" 2>/dev/null || true)
  [ -n "$code" ] || code=137 # killed before it could say
  if $stopped; then rm -f "$mark" "$sent" "$sent.new" "$state" "$rc" "$hdr"; return 0; fi # the site ended it (scancel): nothing more to report
  report || true
  bring
  summary=$(grep -v '^[[:space:]]*$' "$out" | tail -n 1 | cut -c1-200 | sed 's/\\/\\\\/g; s/"/\\"/g' | tr -d '\000-\037')
  curl -sS -o /dev/null -X POST -K "$hdr" -H 'Content-Type: application/json' -d "{\"exit_code\": $code, \"timed_out\": $timedout, \"summary\": \"$summary\"}" "$rapi/finish" || true
  rm -f "$mark" "$sent" "$sent.new" "$state" "$rc" "$hdr"
}

case "$prog" in
run-job)
  [ $# -ge 2 ] || die "usage: SS_JOB_TOKEN=<token> slurm run-job <id> <script>"
  run_job "$1" "$2"
  ;;
sbatch)
  args=$(mktemp)
  trap 'rm -f "$args"' EXIT
  script='' here='' register=''
  while [ $# -gt 0 ]; do
    case "$1" in
      --here) here=1; shift; continue ;;
      --register) here=1 register=1; shift; continue ;;
      # the options that take a value, as the site's parser reads them (SHORT and TAKES in src/jobs.ts); it explains
      # why it refuses -e, -a and -d, so they are sent along too
      -J|-t|-c|-o|-e|-p|-G|-N|-n|-a|-d|--job-name|--time|--cpus-per-task|--mem|--gres|--gpus|--output|--error|--partition|--nodes|--ntasks|--array|--dependency|--checkpoint)
        [ $# -ge 2 ] || die "option $1 needs a value"
        printf '%s\n%s\n' "$1" "$2" >> "$args"
        shift 2 ;;
      -*) printf '%s\n' "$1" >> "$args"; shift ;;
      *) script=$1; shift; break ;;
    esac
  done
  [ -n "$script" ] || die "usage: sbatch [options] script.sbatch"
  [ -f "$script" ] || die "unable to open file $script"
  sha=$(git rev-parse HEAD)
  branch=$(git rev-parse --abbrev-ref HEAD)
  cwd=$(git rev-parse --show-prefix)
  file=$(git ls-files --full-name -- "$script" 2>/dev/null | head -n 1)
  if [ -z "$here" ]; then
    call -X POST "$api" --data-urlencode "script@$script" --data-urlencode "args@$args" \
      --data-urlencode "sha=$sha" --data-urlencode "branch=$branch" --data-urlencode "cwd=$cwd" --data-urlencode "file=${file:-$script}"
    exit 0
  fi
  # a job run here runs the code you pushed: what git tracks must be as committed
  git diff --quiet HEAD -- || die "commit and push your changes first: a job run here runs the commit it is registered at"
  host=$(hostname 2>/dev/null || echo here)
  answer=$(call -X POST "$api" --data-urlencode "script@$script" --data-urlencode "args@$args" --data-urlencode "here=1" --data-urlencode "register=$register" --data-urlencode "host=$host" \
    --data-urlencode "sha=$sha" --data-urlencode "branch=$branch" --data-urlencode "cwd=$cwd" --data-urlencode "file=${file:-$script}")
  val() { printf '%s\n' "$answer" | sed -n "s/^$1=//p"; }
  id=$(val SS_JOB_ID)
  printf '%s\n' "$answer" | head -n 1
  if [ -n "$register" ]; then
    echo "Run it where the computation is, in a clone at $sha, from $cwd:"
    echo "  SS_JOB_TOKEN=$(val SS_JOB_TOKEN) SS_TIME_LIMIT=$(val SS_TIME_LIMIT) SS_CHECKPOINT=$(val SS_CHECKPOINT) SS_OUTPUT=$(val SS_OUTPUT) SS_ROUTE_DIR=$(val SS_ROUTE_DIR) slurm run-job $id $script"
    exit 0
  fi
  # detached from this session: in a session of its own, so it outlives the shell (and the agent) that started it
  self=$(command -v slurm 2>/dev/null || echo "$0")
  export SS_JOB_TOKEN="$(val SS_JOB_TOKEN)" SS_TIME_LIMIT="$(val SS_TIME_LIMIT)" SS_CHECKPOINT="$(val SS_CHECKPOINT)" SS_OUTPUT="$(val SS_OUTPUT)" SS_ROUTE_DIR="$(val SS_ROUTE_DIR)"
  # what goes wrong with the reporting itself (not the job's own output) goes to a file outside the directory
  wlog="${TMPDIR:-/tmp}/slurm-here-$id.log"
  lead sh "$self" run-job "$id" "$script" </dev/null >"$wlog" 2>&1 &
  echo "It runs here, detached: its output is $(val SS_OUTPUT) in this directory, and the site has its log, its checks and, once it ends, its result (git fetch origin refs/jobs/$id). You may leave. (Reporting problems, if any: $wlog)"
  ;;
squeue|sacct)
  q="format=$prog"
  if [ "$prog" = squeue ]; then takes='-u/--user, --me, -j/--jobs, -h/--noheader'; else takes='-j/--jobs, -u/--user, -a/--allusers'; fi
  while [ $# -gt 0 ]; do
    opt=$1 val=''
    case "$opt" in
      --user=*|--jobs=*) val=${opt#*=} opt=${opt%%=*} ;;
      -u|--user|-j|--jobs) [ $# -ge 2 ] || die "option $1 needs a value"; val=$2; shift ;;
    esac
    shift
    case "$prog:$opt" in
      *:-u|*:--user) q="$q&user=$val" ;;
      *:-j|*:--jobs) q="$q&id=$val" ;;
      squeue:--me) q="$q&user=me" ;;
      squeue:-h|squeue:--noheader) q="$q&noheader=1" ;;
      sacct:-a|sacct:--allusers) q="$q&user=*" ;;
      squeue:-a|squeue:--all|sacct:-X|sacct:--allocations|sacct:-b|sacct:--brief) ;;
      *) die "unrecognized option '$opt' (this $prog takes $takes)" ;;
    esac
  done
  call "$api?$q"
  ;;
scancel)
  [ $# -gt 0 ] || die "usage: scancel <jobid> ..."
  for id in "$@"; do call -X DELETE "$api/$id"; done
  ;;
*)
  die "unknown command (sbatch, squeue, scancel, sacct or run-job)"
  ;;
esac