#!/bin/sh
# fillrate-task: cdi-spec-refresh
# Clear stale per-container CDI specs that the Vast daemon (kaalia docker shim)
# wrote from an old GPU list, typically after a board swap moves GPUs to new
# PCI slots. Symptom: "failed to inject CDI devices: unresolvable CDI devices"
# while /etc/cdi/nvidia.yaml and nvidia-smi are both correct.
#
# Steps (ticket mode apply; dry-run does steps 1, 3 as a plan, 5, and 6 only):
#   1. Precheck: refuse if nvidia-smi fails, the ticket has expired, or a
#      long-running container (up OLD_SEC or more, default 120 s) is running.
#      Younger containers are taken to be kaalia's create/retry loop, which
#      starts a short-lived container about once a minute while CDI injection
#      keeps failing. Record every existing container ID.
#   2. systemctl stop vastai (this stops the retries), wait SETTLE_SEC, then
#      wait up to SETTLE_MAX_SEC for running containers to exit. If any running
#      container reaches OLD_SEC, or one is still running at the deadline,
#      start vastai again and refuse. Containers are never stopped by this task.
#   3. Copy /etc/cdi/D.*.yaml to /var/lib/fillrate/backup/cdi-spec-refresh-<stamp>,
#      then delete each D.<id>.yaml whose <id> is not an existing container.
#      The spec of an existing (stopped) container is kept.
#   4. systemctl start vastai, wait at least 60 s.
#   5. Verify: kaalia.log written since the start says it found N nvidia gpus,
#      where N is the nvidia-smi count; every D spec created since the start
#      lists all N UUIDs, each on the /dev/nvidiaN that nvidia-smi gives as its
#      minor number; nvidia-ctk cdi list shows nvidia.com/gpu=0..N-1 and all.
#   6. Read-only: ls -la /var/lib/vastai_kaalia; every file there (depth 3,
#      skipping .ssh, .docker, .cache, tls, data, docker_tmp, metrics_spool)
#      whose name looks like a GPU cache or that contains GPU UUIDs, with
#      size, mtime, and the UUIDs / PCI bus ids it lists (never other content);
#      the last kaalia.log lines that name GPU UUIDs; tail of
#      purge_stale_cdi.log. This is to find where the daemon gets a 6-GPU list.
#
# Never touches containers or volumes (no docker rm/stop/start/volume), never
# restarts docker, never changes drivers or packages, never reboots, makes no
# Vast API calls. The runner starts it with a clean environment. The FILLRATE_*
# path overrides below are for tests; a real run uses the defaults.

set -u

CDI_DIR=${FILLRATE_CDI_DIR:-/etc/cdi}
BACKUP_ROOT=${FILLRATE_CDI_BACKUP_ROOT:-/var/lib/fillrate/backup}
STATE=${FILLRATE_CDI_STATE:-/var/lib/fillrate-runner/cdi-spec-refresh}
KAALIA_DIR=${FILLRATE_KAALIA_DIR:-/var/lib/vastai_kaalia}
KAALIA_LOG=$KAALIA_DIR/kaalia.log
OLD_SEC=${FILLRATE_CDI_OLD_SEC:-120}
SETTLE_SEC=${FILLRATE_CDI_SETTLE_SEC:-15}
SETTLE_MAX_SEC=${FILLRATE_CDI_SETTLE_MAX_SEC:-90}
WAIT_MIN=${FILLRATE_CDI_WAIT_MIN:-60}
WAIT_MAX=${FILLRATE_CDI_WAIT_MAX:-180}
EXPIRES=${FILLRATE_EXPIRES:-0}
MODE=${FILLRATE_MODE:-apply}
[ -n "$MODE" ] || MODE=apply

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

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

refuse() {
  say "refusing: $1"
  say "summary: cdi-spec-refresh 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
}

# Print "<id> <age seconds> <image>" for each running container.
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
    [ "$st" -gt 0 ] || st=0
    if [ "$st" = 0 ]; then age=999999; else age=$((now - st)); fi
    printf '%s %s %s\n' "$cid" "$age" "$image"
  done
}

# 0 when every running container is younger than OLD_SEC; logs each one.
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 ]
}

# uuid minor pairs from nvidia-smi -q ("GPU UUID" and "Minor Number" lines).
gpu_map() {
  nvidia-smi -q | awk -F': *' '
    /^[[:space:]]*GPU UUID[[:space:]]*:/ { u = $2; sub(/[[:space:]]+$/, "", u) }
    /^[[:space:]]*Minor Number[[:space:]]*:/ { m = $2; sub(/[[:space:]]+$/, "", m); if (u != "") print u, m; u = "" }
  '
}

# uuid minor pairs from a CDI spec: a "- name: GPU-..." entry followed by its
# first /dev/nvidiaN device node path.
spec_map() {
  awk '
    /^[[:space:]]*-[[:space:]]*name:/ {
      n = $0; sub(/^[^:]*:[[:space:]]*/, "", n); gsub(/["\047[:space:]]/, "", n)
      cur = (n ~ /^GPU-/) ? n : ""
      next
    }
    cur != "" && /path:[[:space:]]*\/dev\/nvidia[0-9]+[[:space:]]*$/ {
      p = $0; sub(/.*\/dev\/nvidia/, "", p); gsub(/[[:space:]]/, "", p)
      print cur, p; cur = ""
    }
  ' "$1"
}

case "$MODE" in
  apply|dry-run) ;;
  *) refuse "mode must be apply or dry-run (got $MODE)" ;;
esac

say "cdi-spec-refresh start mode=$MODE machine=${FILLRATE_MACHINE_ID:-unknown} stamp=$STAMP"

# ---- 1. precheck ----
window_open || refuse "ticket expired"
command -v docker > /dev/null 2>&1 || refuse "docker not found"
command -v nvidia-smi > /dev/null 2>&1 || refuse "nvidia-smi not found"
command -v nvidia-ctk > /dev/null 2>&1 || refuse "nvidia-ctk not found"
systemctl cat vastai.service > /dev/null 2>&1 || refuse "vastai.service not found"

running_ids > /dev/null || refuse "docker ps failed"
only_young_running || refuse "a long-running container is running (up ${OLD_SEC}s or more)"

gpus=$(gpu_map) || refuse "nvidia-smi -q failed"
ngpu=$(printf '%s\n' "$gpus" | awk 'NF == 2' | wc -l | tr -d ' ')
[ "$ngpu" -gt 0 ] || refuse "nvidia-smi shows no GPUs"
say "nvidia-smi gpus=$ngpu"
printf '%s\n' "$gpus" | while read -r u m; do say "  gpu uuid=$u dev=/dev/nvidia$m"; done

all_ids=$(docker ps -a -q --no-trunc) || refuse "docker ps -a failed"
say "existing containers (kept; the task never stops or removes containers):"
docker ps -a --no-trunc --format '{{.ID}} {{.Names}} {{.Status}}' 2>/dev/null | while read -r line; do say "  keep $line"; done
[ -n "$all_ids" ] || say "  none"

keep=""
remove=""
for f in "$CDI_DIR"/D.*.yaml; do
  [ -f "$f" ] || continue
  [ -L "$f" ] && refuse "symlinked spec $f"
  base=$(basename "$f" .yaml)
  id=${base#D.}
  if printf '%s\n' "$all_ids" | grep -F -x -q -- "$id"; then
    keep="$keep $f"
    say "plan keep $f (existing container)"
  else
    remove="$remove $f"
    say "plan delete $f (no such container)"
  fi
done
[ -n "$keep$remove" ] || say "no D.*.yaml specs present"

if [ "$MODE" = "apply" ]; then
  window_open || refuse "ticket expired"
  only_young_running || refuse "a long-running container is running (up ${OLD_SEC}s or more)"

  # ---- 2. stop vastai, then wait for retry containers to exit ----
  say "run: systemctl stop vastai"
  systemctl stop vastai || refuse "systemctl stop vastai failed"
  sleep "$SETTLE_SEC"
  settled=$SETTLE_SEC
  while :; do
    ids=$(running_ids) || { systemctl start vastai; refuse "docker ps failed after stopping vastai"; }
    [ -n "$ids" ] || break
    if ! only_young_running; then
      say "run: systemctl start vastai (abort)"
      systemctl start vastai
      refuse "a container kept running after vastai stopped and is now ${OLD_SEC}s or older"
    fi
    if [ "$settled" -ge "$SETTLE_MAX_SEC" ]; then
      say "run: systemctl start vastai (abort)"
      systemctl start vastai
      refuse "containers still running ${SETTLE_MAX_SEC}s after vastai stopped"
    fi
    sleep 5
    settled=$((settled + 5))
  done
  say "no containers running ${settled}s after vastai stopped"
  all_ids=$(docker ps -a -q --no-trunc) || { systemctl start vastai; refuse "docker ps -a failed"; }
  keep=""
  remove=""
  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 keep="$keep $f"; else remove="$remove $f"; fi
  done
  say "final plan: keep $(echo "$keep" | wc -w | tr -d ' '), delete $(echo "$remove" | wc -w | tr -d ' ')"

  # ---- 3. backup then delete ----
  BACKUP="$BACKUP_ROOT/cdi-spec-refresh-$STAMP"
  if ! mkdir -p "$BACKUP" || ! chmod 700 "$BACKUP"; then
    systemctl start vastai
    refuse "backup dir failed"
  fi
  for f in $keep $remove; do
    cp -p "$f" "$BACKUP/" || { systemctl start vastai; refuse "backup of $f failed"; }
  done
  say "backed up to $BACKUP: $(ls "$BACKUP" | wc -l | tr -d ' ') files"
  for f in $remove; do
    say "run: rm -f $f"
    rm -f -- "$f"
  done

  # ---- 4. start vastai ----
  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"
  waited=0
  found=""
  while [ "$waited" -lt "$WAIT_MAX" ]; do
    sleep 5
    waited=$((waited + 5))
    if [ -f "$KAALIA_LOG" ]; then
      found=$(tail -c +"$((log_off + 1))" "$KAALIA_LOG" | grep -E "found [0-9]+ nvidia gpus" | tail -n 1)
    fi
    [ "$waited" -ge "$WAIT_MIN" ] && [ -n "$found" ] && break
  done
  say "waited ${waited}s; vastai $(systemctl is-active vastai 2>/dev/null)"
else
  say "dry-run: would stop vastai, back up all listed specs to $BACKUP_ROOT/cdi-spec-refresh-$STAMP, delete the planned specs, start vastai"
  log_off=0
  found=""
  [ -f "$KAALIA_LOG" ] && found=$(grep -E "found [0-9]+ nvidia gpus" "$KAALIA_LOG" | tail -n 1)
  marker=""
fi

# ---- 5. verify ----
fail=0
if [ -n "$found" ]; then
  say "kaalia: $found"
  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, nvidia-smi $ngpu)"; fail=1; fi
else
  say "verify_kaalia=fail (no 'found N nvidia gpus' line since start)"
  fail=1
fi

check_spec() {
  s=$(spec_map "$1")
  bad=0
  printf '%s\n' "$gpus" | while read -r u m; do
    got=$(printf '%s\n' "$s" | awk -v u="$u" '$1 == u { print $2; exit }')
    if [ -z "$got" ]; then
      say "  $1: missing $u" >&2
      echo MISMATCH
    elif [ "$got" != "$m" ]; then
      say "  $1: $u on /dev/nvidia$got, nvidia-smi says /dev/nvidia$m" >&2
      echo MISMATCH
    fi
  done | grep -x -q MISMATCH && bad=1
  extra=$(printf '%s\n' "$s" | awk 'NF == 2' | wc -l | tr -d ' ')
  [ "$extra" = "$ngpu" ] || bad=1
  return $bad
}

new_specs=0
for f in "$CDI_DIR"/D.*.yaml; do
  [ -f "$f" ] || continue
  if [ -n "$marker" ] && [ "$f" -nt "$marker" ]; then
    new_specs=$((new_specs + 1))
    if check_spec "$f"; then say "verify_dspec=ok $f"; else say "verify_dspec=fail $f"; fail=1; fi
  else
    if check_spec "$f"; then say "existing spec ok: $f"; else say "existing spec stale (kept, not changed): $f"; fi
  fi
done
[ "$new_specs" -gt 0 ] || say "verify_dspec=none-yet (kaalia writes a D spec when it creates a container; recheck after the next self-test)"

list=$(nvidia-ctk cdi list 2>/dev/null)
cdi_ok=1
i=0
while [ "$i" -lt "$ngpu" ]; do
  printf '%s\n' "$list" | grep -F -x -q "nvidia.com/gpu=$i" || cdi_ok=0
  i=$((i + 1))
done
printf '%s\n' "$list" | grep -F -x -q "nvidia.com/gpu=all" || cdi_ok=0
if [ "$cdi_ok" = 1 ]; then say "verify_cdi_list=ok (nvidia.com/gpu=0..$((ngpu - 1)) and all)"; else say "verify_cdi_list=fail"; printf '%s\n' "$list" | grep '^nvidia.com/' | while read -r l; do say "  $l"; done; fail=1; fi

# ---- 6. read-only listing and GPU-list diagnostics ----
say "ls -la $KAALIA_DIR:"
ls -la "$KAALIA_DIR" 2>&1 | while read -r l; do say "  $l"; done

say "gpu list candidates under $KAALIA_DIR (read-only; only UUIDs and PCI bus ids shown):"
cands=$(find "$KAALIA_DIR" -maxdepth 3 \
  \( -name .ssh -o -name .docker -o -name .cache -o -name tls -o -name data -o -name docker_tmp -o -name metrics_spool \) -prune -o \
  -type f -size -20M ! -name 'kaalia*.log*' -print 2>/dev/null | while read -r f; do
    case "$(basename "$f")" in
      *gpu*|*GPU*|*cache*|*machine_info*|*mach_info*|*.json) printf '%s\n' "$f"; continue ;;
    esac
    grep -q -E 'GPU-[0-9a-f]{8}-' "$f" 2>/dev/null && printf '%s\n' "$f"
  done)
if [ -z "$cands" ]; then
  say "  none found"
else
  printf '%s\n' "$cands" | while read -r f; do
    [ -n "$f" ] || continue
    n=$(grep -o -E 'GPU-[0-9a-f-]{36}' "$f" 2>/dev/null | sort -u | wc -l | tr -d ' ')
    say "  $(ls -l --time-style=+%Y-%m-%dT%H:%M "$f" 2>/dev/null | awk '{ print $5, $6 }') $f uuids=$n"
    grep -o -E 'GPU-[0-9a-f-]{36}' "$f" 2>/dev/null | sort | uniq -c | while read -r c u; do say "    $u x$c"; done
    grep -o -E '[0-9a-fA-F]{4,8}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-9]' "$f" 2>/dev/null | sort -u | head -n 16 | tr '\n' ' ' | while read -r b; do say "    bus: $b"; done
  done
fi
if [ -f "$KAALIA_LOG" ]; then
  say "last kaalia.log lines naming GPU UUIDs (first 300 chars):"
  grep -E 'GPU-[0-9a-f]{8}-' "$KAALIA_LOG" 2>/dev/null | tail -n 5 | cut -c1-300 | while read -r l; do say "  $l"; done
fi
if [ -f "$KAALIA_DIR/purge_stale_cdi.log" ]; then
  say "tail -n 40 purge_stale_cdi.log:"
  tail -n 40 "$KAALIA_DIR/purge_stale_cdi.log" | cut -c1-300 | while read -r l; do say "  $l"; done
else
  say "purge_stale_cdi.log not present"
fi

result=ok
deleted=0
[ "$MODE" = apply ] && deleted=$(echo "$remove" | wc -w | tr -d ' ')
[ "$fail" = 0 ] || result=verify-failed
say "summary: cdi-spec-refresh mode=$MODE result=$result gpus=$ngpu kept=$(echo "$keep" | wc -w | tr -d ' ') deleted=$deleted new_specs=$new_specs log=$LOG"
[ "$fail" = 0 ] || exit 2
exit 0