#!/bin/sh
# fillrate-task: gpu-map-rebuild
# Rebuild the Vast daemon's GPU UUID -> index map after a hardware change
# (board swap, GPUs moved slots, GPU replaced).
#
# Background (149175, 2026-10-16): kaalia's docker shim builds per-container
# CDI specs from /var/lib/vastai_kaalia/data/uuid-idx-mapping/HEAD. HEAD is a
# symlink (absolute path) to a file named <unix-ns>-<pid> holding compact JSON
# {"GPU-<uuid>":<index>,...} in index order, no spaces, no trailing newline,
# owned vastai_kaalia:docker mode 644. Kaalia wrote it once at install and did
# not rewrite it after the swap, so new containers got 6 of 8 GPUs and failed
# with "unresolvable CDI devices". Pointing HEAD at a map built from
# nvidia-smi fixed it.
#
# Steps (ticket mode dry-run stops after step 3):
#   1. Precheck: refuse if nvidia-smi fails, the ticket expired, or a container
#      up OLD_SEC (120 s) or more is running. Younger running containers are
#      kaalia's create/retry loop.
#   2. Build the new map from nvidia-smi --query-gpu=index,uuid. Indices must
#      be 0..N-1 and UUIDs well formed.
#   3. Show the current HEAD target, its entries, and the diff.
#   4. Apply: stop vastai, wait for retry containers to exit (any container
#      that stays running or reaches OLD_SEC: start vastai again and refuse).
#   5. Back up the whole uuid-idx-mapping dir and every /etc/cdi/D.*.yaml to
#      /var/lib/fillrate/backup/gpu-map-rebuild-<stamp>.
#   6. Write <unix-ns>-fillrate with the kaalia owner and mode, then swap HEAD
#      to it atomically (old map files are kept). Verify HEAD index by index
#      against nvidia-smi; on mismatch put the old HEAD back.
#   7. Delete D.*.yaml whose container no longer exists (specs of existing
#      containers, including a stopped rental's, are kept).
#   8. Start vastai. Verify kaalia.log reports N GPUs and nvidia_smi.json
#      uuid_idx_mapping (rewritten by kaalia) matches HEAD index by index.
#
# Never stops, starts, or removes containers or volumes, never restarts
# docker, no driver or package changes, no reboot, no Vast API calls.
# FILLRATE_* overrides are for tests; a real run uses the defaults.

set -u

KAALIA_DIR=${FILLRATE_KAALIA_DIR:-/var/lib/vastai_kaalia}
MAP_DIR=$KAALIA_DIR/data/uuid-idx-mapping
SMI_JSON=$KAALIA_DIR/data/nvidia_smi.json
KAALIA_LOG=$KAALIA_DIR/kaalia.log
CDI_DIR=${FILLRATE_CDI_DIR:-/etc/cdi}
BACKUP_ROOT=${FILLRATE_GPU_MAP_BACKUP_ROOT:-/var/lib/fillrate/backup}
STATE=${FILLRATE_GPU_MAP_STATE:-/var/lib/fillrate-runner/gpu-map-rebuild}
OLD_SEC=${FILLRATE_GPU_MAP_OLD_SEC:-120}
SETTLE_SEC=${FILLRATE_GPU_MAP_SETTLE_SEC:-15}
SETTLE_MAX_SEC=${FILLRATE_GPU_MAP_SETTLE_MAX_SEC:-90}
WAIT_MIN=${FILLRATE_GPU_MAP_WAIT_MIN:-60}
WAIT_MAX=${FILLRATE_GPU_MAP_WAIT_MAX:-240}
EXPIRES=${FILLRATE_EXPIRES:-0}
MODE=${FILLRATE_MODE:-}

STAMP=$(date +%Y%m%d-%H%M%S)
mkdir -p "$STATE" && chmod 700 "$STATE"
LOG="$STATE/run-$STAMP.log"
: > "$LOG"
VASTAI_STOPPED=0

say() {
  printf '%s\n' "$*"
  printf '%s %s\n' "$(date +%H:%M:%S)" "$*" >> "$LOG"
}

refuse() {
  if [ "$VASTAI_STOPPED" = 1 ]; then
    say "run: systemctl start vastai (abort)"
    systemctl start vastai
  fi
  say "refusing: $1"
  say "summary: gpu-map-rebuild result=refused reason=$1 log=$LOG"
  exit 1
}

window_open() {
  now=$(date +%s) || return 1
  [ "$now" -le "$EXPIRES" ]
}

running_ids() {
  docker ps -q --no-trunc
}

running_ages() {
  now=$(date +%s)
  for cid in $(running_ids); do
    started=$(docker inspect --format '{{.State.StartedAt}}' "$cid" 2>/dev/null) || started=""
    image=$(docker inspect --format '{{.Config.Image}}' "$cid" 2>/dev/null) || image=unknown
    st=$(date -d "$started" +%s 2>/dev/null) || st=0
    if [ "$st" -gt 0 ] 2>/dev/null; then age=$((now - st)); else age=999999; fi
    printf '%s %s %s\n' "$cid" "$age" "$image"
  done
}

only_young_running() {
  rows=$(running_ages) || return 1
  [ -n "$rows" ] || return 0
  old=0
  while read -r cid age image; do
    [ -n "$cid" ] || continue
    if [ "$age" -ge "$OLD_SEC" ]; then
      say "  running $cid age=${age}s image=$image (long-running: treated as a real rental)"
      old=1
    else
      say "  running $cid age=${age}s image=$image (young: treated as a kaalia retry)"
    fi
  done << ROWS
$rows
ROWS
  [ "$old" = 0 ]
}

# "index uuid" lines from nvidia-smi, sorted by index.
smi_pairs() {
  nvidia-smi --query-gpu=index,uuid --format=csv,noheader | tr -d ' ' | awk -F, 'NF == 2 { print $1, $2 }' | sort -n
}

# "index uuid" lines from a map file or a JSON object string on stdin.
map_pairs() {
  tr -d ' \n\r' | tr '{},' '\n\n\n' | awk -F'":' '/^"GPU-/ { u = $1; sub(/^"/, "", u); print $2, u }' | sort -n
}

# The uuid_idx_mapping object from nvidia_smi.json.
smi_json_map() {
  [ -r "$SMI_JSON" ] || return 1
  sed -n 's/.*"uuid_idx_mapping":\({[^}]*}\).*/\1/p' "$SMI_JSON"
}

case "$MODE" in
  apply|dry-run) ;;
  *) refuse "mode must be apply or dry-run (got ${MODE:-none})" ;;
esac
say "gpu-map-rebuild start mode=$MODE machine=${FILLRATE_MACHINE_ID:-unknown} stamp=$STAMP"

# ---- 1. precheck ----
window_open || refuse "ticket expired"
for c in docker nvidia-smi systemctl; do
  command -v "$c" > /dev/null 2>&1 || refuse "$c not found"
done
systemctl cat vastai.service > /dev/null 2>&1 || refuse "vastai.service not found"
[ -d "$MAP_DIR" ] || refuse "missing $MAP_DIR"
running_ids > /dev/null || refuse "docker ps failed"
only_young_running || refuse "a long-running container is running (up ${OLD_SEC}s or more)"

# ---- 2. new map from nvidia-smi ----
pairs=$(smi_pairs) || refuse "nvidia-smi query failed"
ngpu=$(printf '%s\n' "$pairs" | awk 'NF == 2' | wc -l | tr -d ' ')
[ "$ngpu" -gt 0 ] || refuse "nvidia-smi shows no GPUs"
printf '%s\n' "$pairs" | awk -v n="$ngpu" '
  $1 != NR - 1 { bad = 1 }
  $2 !~ /^GPU-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/ { bad = 1 }
  END { exit bad }' || refuse "nvidia-smi indices are not 0..N-1 or a UUID is malformed"
NEW=$(printf '%s\n' "$pairs" | awk 'BEGIN { printf "{" } { if (NR > 1) printf ","; printf "\"%s\":%s", $2, $1 } END { printf "}" }')
say "nvidia-smi gpus=$ngpu"
printf '%s\n' "$pairs" | while read -r i u; do say "  new $i $u"; done

# ---- 3. current HEAD ----
OLD_TARGET=$(readlink "$MAP_DIR/HEAD" 2>/dev/null || true)
if [ -n "$OLD_TARGET" ]; then
  say "current HEAD -> $OLD_TARGET"
elif [ -f "$MAP_DIR/HEAD" ]; then
  say "current HEAD is a regular file"
else
  say "current HEAD missing"
fi
cur=""
[ -r "$MAP_DIR/HEAD" ] && cur=$(map_pairs < "$MAP_DIR/HEAD")
say "current HEAD entries=$(printf '%s\n' "$cur" | awk 'NF == 2' | wc -l | tr -d ' ')"
printf '%s\n' "$cur" | while read -r i u; do [ -n "$i" ] && say "  cur $i $u"; done
same=0
[ "$cur" = "$pairs" ] && same=1
if [ "$same" = 1 ]; then say "HEAD already matches nvidia-smi"; else say "HEAD differs from nvidia-smi: rebuild needed"; fi

if [ "$MODE" = "dry-run" ]; then
  for f in "$CDI_DIR"/D.*.yaml; do [ -f "$f" ] && say "dry-run: D spec present $(basename "$f")"; done
  say "dry-run: would stop vastai, back up $MAP_DIR and D specs to $BACKUP_ROOT/gpu-map-rebuild-$STAMP, write the new map, swap HEAD, delete D specs of removed containers, start vastai"
  say "summary: gpu-map-rebuild mode=dry-run result=ok gpus=$ngpu head_current=$same log=$LOG"
  exit 0
fi

# ---- 4. stop vastai, wait out retry containers ----
window_open || refuse "ticket expired"
only_young_running || refuse "a long-running container is running (up ${OLD_SEC}s or more)"
say "run: systemctl stop vastai"
systemctl stop vastai || refuse "systemctl stop vastai failed"
VASTAI_STOPPED=1
sleep "$SETTLE_SEC"
settled=$SETTLE_SEC
while :; do
  ids=$(running_ids) || refuse "docker ps failed after stopping vastai"
  [ -n "$ids" ] || break
  only_young_running || refuse "a container kept running after vastai stopped and is now ${OLD_SEC}s or older"
  [ "$settled" -lt "$SETTLE_MAX_SEC" ] || refuse "containers still running ${SETTLE_MAX_SEC}s after vastai stopped"
  sleep 5
  settled=$((settled + 5))
done
say "no containers running ${settled}s after vastai stopped"

# ---- 5. backup ----
BACKUP="$BACKUP_ROOT/gpu-map-rebuild-$STAMP"
mkdir -p "$BACKUP" || refuse "backup dir failed"
chmod 700 "$BACKUP"
cp -a "$MAP_DIR" "$BACKUP/uuid-idx-mapping" || refuse "map backup failed"
mkdir -p "$BACKUP/cdi"
for f in "$CDI_DIR"/D.*.yaml; do
  [ -f "$f" ] || continue
  cp -p "$f" "$BACKUP/cdi/" || refuse "spec backup failed: $f"
done
say "backed up map dir and $(find "$BACKUP/cdi" -type f | wc -l | tr -d ' ') D specs to $BACKUP"

# ---- 6. write map, swap HEAD, verify ----
owner=$(stat -c '%U:%G' "$MAP_DIR") || owner=root:root
if [ "$same" = 0 ]; then
  name="$(date +%s%N)-fillrate"
  tmp="$MAP_DIR/.$name.tmp"
  printf '%s' "$NEW" > "$tmp" || refuse "write map failed"
  chown "$owner" "$tmp" 2>/dev/null
  chmod 644 "$tmp"
  mv -f "$tmp" "$MAP_DIR/$name" || refuse "rename map failed"
  ln -s "$MAP_DIR/$name" "$MAP_DIR/.HEAD.tmp" || refuse "symlink failed"
  chown -h "$owner" "$MAP_DIR/.HEAD.tmp" 2>/dev/null
  mv -T -f "$MAP_DIR/.HEAD.tmp" "$MAP_DIR/HEAD" || refuse "HEAD swap failed"
  say "wrote $MAP_DIR/$name ($(wc -c < "$MAP_DIR/$name" | tr -d ' ') bytes, $owner 644); HEAD -> $MAP_DIR/$name"
fi
got=$(map_pairs < "$MAP_DIR/HEAD")
if [ "$got" = "$pairs" ]; then
  say "verify_head=ok (index by index vs nvidia-smi, $ngpu gpus)"
else
  say "verify_head=fail"
  if [ -n "$OLD_TARGET" ]; then
    ln -s "$OLD_TARGET" "$MAP_DIR/.HEAD.tmp" && chown -h "$owner" "$MAP_DIR/.HEAD.tmp" 2>/dev/null
    mv -T -f "$MAP_DIR/.HEAD.tmp" "$MAP_DIR/HEAD" && say "restored HEAD -> $OLD_TARGET"
  fi
  refuse "new HEAD does not match nvidia-smi"
fi

# ---- 7. stale D specs ----
all_ids=$(docker ps -a -q --no-trunc) || refuse "docker ps -a failed"
deleted=0
kept=0
for f in "$CDI_DIR"/D.*.yaml; do
  [ -f "$f" ] || continue
  id=$(basename "$f" .yaml)
  id=${id#D.}
  if printf '%s\n' "$all_ids" | grep -F -x -q -- "$id"; then
    kept=$((kept + 1))
    say "keep $f (existing container)"
  else
    say "run: rm -f $f"
    rm -f -- "$f"
    deleted=$((deleted + 1))
  fi
done

# ---- 8. start vastai and verify ----
log_off=0
[ -f "$KAALIA_LOG" ] && log_off=$(wc -c < "$KAALIA_LOG" | tr -d ' ')
marker="$STATE/start-$STAMP"
: > "$marker"
say "run: systemctl start vastai"
systemctl start vastai || say "warning: systemctl start vastai returned non-zero"
VASTAI_STOPPED=0
waited=0
found=""
fresh_json=0
while [ "$waited" -lt "$WAIT_MAX" ]; do
  sleep 5
  waited=$((waited + 5))
  [ -f "$KAALIA_LOG" ] && found=$(tail -c +"$((log_off + 1))" "$KAALIA_LOG" | grep -E "found [0-9]+ nvidia gpus" | tail -n 1)
  [ -f "$SMI_JSON" ] && [ "$SMI_JSON" -nt "$marker" ] && fresh_json=1
  [ "$waited" -ge "$WAIT_MIN" ] && [ -n "$found" ] && [ "$fresh_json" = 1 ] && break
done
say "waited ${waited}s; vastai $(systemctl is-active vastai 2>/dev/null)"

fail=0
n=$(printf '%s\n' "$found" | sed -n 's/.*found \([0-9][0-9]*\) nvidia gpus.*/\1/p')
if [ "$n" = "$ngpu" ]; then say "verify_kaalia=ok ($n gpus)"; else say "verify_kaalia=fail (kaalia '${n:-none}', nvidia-smi $ngpu)"; fail=1; fi

js=$(smi_json_map | map_pairs)
if [ "$fresh_json" = 0 ]; then
  say "verify_smi_json=stale (nvidia_smi.json not rewritten since start; comparing anyway)"
fi
if [ "$js" = "$pairs" ]; then
  say "verify_smi_json=ok (uuid_idx_mapping matches HEAD index by index)"
else
  say "verify_smi_json=fail"
  printf '%s\n' "$js" | while read -r i u; do [ -n "$i" ] && say "  json $i $u"; done
  fail=1
fi
got=$(map_pairs < "$MAP_DIR/HEAD")
if [ "$got" = "$pairs" ]; then say "verify_head_after_start=ok"; else say "verify_head_after_start=fail (kaalia changed HEAD)"; fail=1; fi

result=ok
[ "$fail" = 0 ] || result=verify-failed
say "summary: gpu-map-rebuild mode=apply result=$result gpus=$ngpu head_rebuilt=$((1 - same)) specs_kept=$kept specs_deleted=$deleted backup=$BACKUP log=$LOG"
[ "$fail" = 0 ] || exit 2
exit 0