igneum/relay/clients/agent.sh

#!/usr/bin/env bash
# Igneum relay agent for the Mac (or Linux/WSL): the bash twin of igneum-agent.ps1 for `run` tasks whose body is a bash script.
#   agent.sh                 register this machine and poll every 20 s; run each `run` task, post a result, mark it done
#   agent.sh once            one pass (used by the tests)
# Env: RELAY_MACHINE overrides the name (default: what the relay returns for this hostname, else `hostname -s`).
# State: ~/.local/state/igneum-relay (state.json, tasks/, logs/). Needs curl, python3 (jq optional).
set -uo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
RELAY_URL='__RELAY_URL__'; RELAY_KEY='__RELAY_KEY__'; RELAY_TOKEN='__RELAY_TOKEN__'
BASE="$RELAY_URL/r/$RELAY_TOKEN/api"
STATE="${XDG_STATE_HOME:-$HOME/.local/state}/igneum-relay"; mkdir -p "$STATE/tasks" "$STATE/logs"
POLL="${RELAY_POLL:-20}"; TAIL=65536
log() { echo "[$(date +%H:%M:%S)] $*"; }
get() { curl -sS --max-time 60 "$BASE/$1" -H "x-igneum-key: $RELAY_KEY"; }
post() { curl -sS --max-time 120 -X POST "$BASE/$1" -H 'Content-Type: application/json' -H "x-igneum-key: $RELAY_KEY" --data-binary "$2"; }
py() { python3 -c "$@"; }
register() {
  local info
  info="$(py 'import json,platform,shutil,subprocess,os
gpus=[]
try:
    out=subprocess.run(["system_profiler","SPDisplaysDataType"],capture_output=True,text=True,timeout=20).stdout
    gpus=[l.split(":",1)[1].strip() for l in out.splitlines() if "Chipset Model" in l]
except Exception: pass
if not gpus and shutil.which("nvidia-smi"):
    try: gpus=[l.strip() for l in subprocess.run(["nvidia-smi","--query-gpu=name","--format=csv,noheader"],capture_output=True,text=True,timeout=10).stdout.splitlines() if l.strip()]
    except Exception: pass
print(json.dumps({"hostname":platform.node().split(".")[0],"info":{"os":platform.platform(),"user":os.environ.get("USER",""),"gpus":gpus,"nvcc":bool(shutil.which("nvcc")),"agent":"agent.sh v1","dir":os.getcwd()}}))')"
  local r; r="$(post register "$info")" || { log "register failed"; return 1; }
  MACHINE="${RELAY_MACHINE:-$(printf '%s' "$r" | py 'import json,sys; print(json.load(sys.stdin)["name"])')}"
  ROLE="$(printf '%s' "$r" | py 'import json,sys; print(json.load(sys.stdin).get("role",""))')"
  printf '%s\n' "$MACHINE" > "$STATE/machine.txt"
  log "registered as $MACHINE (role ${ROLE:-unset})"
}
post_result() { # id title code log note
  local id="$1" title="$2" code="$3" logf="$4" note="${5:-}" body
  body="$(tail -c "$TAIL" "$logf" 2>/dev/null | py 'import sys,json; print(json.dumps(sys.stdin.read()))')"
  local t; t="$(printf '%s' "$title: exit $code${note:+ ($note)}" | py 'import sys,json; print(json.dumps(sys.stdin.read()))')"
  post drop "{\"kind\":\"result\",\"from\":$(printf '%s' "$MACHINE" | py 'import sys,json; print(json.dumps(sys.stdin.read()))'),\"to\":\"all\",\"task_id\":$id,\"title\":$t,\"body\":$body,\"flags\":{\"exit_code\":$code,\"pass\":${RELAY_PASS:-1}}}" >/dev/null && log "result posted for #$id"
  if [ "$(stat -f%z "$logf" 2>/dev/null || stat -c%s "$logf")" -gt "$TAIL" ]; then RELAY_MACHINE="$MACHINE" RELAY_TITLE="$title full log" bash "$HERE/send.sh" "$logf" >/dev/null || true; fi
}
run_task() { # json-of-one-item pass
  local it="$1" pass="${2:-1}" id title body elevated rc
  id="$(printf '%s' "$it" | py 'import json,sys; print(json.load(sys.stdin)["id"])')"
  title="$(printf '%s' "$it" | py 'import json,sys; print(json.load(sys.stdin)["title"])')"
  elevated="$(printf '%s' "$it" | py 'import json,sys; print("1" if json.load(sys.stdin)["flags"].get("elevated") else "")')"
  local rebootc; rebootc="$(printf '%s' "$it" | py 'import json,sys; print("1" if json.load(sys.stdin)["flags"].get("reboot_continue") else "")')"
  printf '%s' "$it" | py 'import json,sys; sys.stdout.write(json.load(sys.stdin)["body"])' > "$STATE/tasks/task-$id.sh"
  local logf="$STATE/logs/task-$id-pass$pass-$(date +%Y%m%d-%H%M%S).log"
  log "task #$id '$title' pass $pass${elevated:+ elevated}${rebootc:+ reboot_continue}"
  if [ -n "$elevated" ]; then
    RELAY_PASS="$pass" RELAY_TASK_ID="$id" RELAY_MACHINE="$MACHINE" RELAY_SEND="$HERE/send.sh" sudo -n -E bash "$STATE/tasks/task-$id.sh" > >(tee -a "$logf") 2>&1; rc=$?
  else
    RELAY_PASS="$pass" RELAY_TASK_ID="$id" RELAY_MACHINE="$MACHINE" RELAY_SEND="$HERE/send.sh" bash "$STATE/tasks/task-$id.sh" > >(tee -a "$logf") 2>&1; rc=$?
  fi
  wait 2>/dev/null; sleep 0.2
  echo "__RELAY_EXIT__=$rc" >> "$logf"
  if grep -q 'RELAY-REBOOT' "$logf"; then
    if [ -n "$rebootc" ]; then
      post_result "$id" "$title" "$rc" "$logf" "rebooting, resumes as pass $((pass+1))"
      printf '{"pending":%s,"pass":%s}\n' "$id" "$((pass+1))" > "$STATE/state.json"
      log "task asked for a reboot; re-run agent.sh after the restart to resume (no auto-restart on the Mac)"; return 0
    fi
  fi
  post_result "$id" "$title" "$rc" "$logf"
  post done "{\"id\":$id,\"exit_code\":$rc}" >/dev/null
}
one_pass() {
  if [ -f "$STATE/state.json" ]; then
    local pend pass; pend="$(py 'import json; d=json.load(open("'"$STATE/state.json"'")); print(d["pending"])')"; pass="$(py 'import json; d=json.load(open("'"$STATE/state.json"'")); print(d["pass"])')"
    rm -f "$STATE/state.json"; run_task "$(get "item?id=$pend" | py 'import json,sys; print(json.dumps(json.load(sys.stdin)["item"]))')" "$pass"
  fi
  local j; j="$(get "inbox?machine=$MACHINE&kind=run&ack=1")" || { log "poll failed"; return 1; }
  local n; n="$(printf '%s' "$j" | py 'import json,sys; print(len(json.load(sys.stdin)["items"]))')"
  local i=0; while [ "$i" -lt "$n" ]; do run_task "$(printf '%s' "$j" | py 'import json,sys; print(json.dumps(json.load(sys.stdin)["items"]['"$i"']))')" 1; i=$((i+1)); done
  [ "$n" = 0 ] && printf '\r[%s] idle as %s   ' "$(date +%H:%M:%S)" "$MACHINE"
  return 0
}
register || exit 1
if [ "${1:-}" = "once" ]; then one_pass; exit $?; fi
while true; do one_pass; sleep "$POLL"; done