#!/usr/bin/env bash
# 29_dspark_isolate.sh — Minimal test: does DSpark spec work on K3 at all?
# Try WITHOUT -fa on (FA is the most likely graph-breaker). 64in/64out, single probe.
set -euo pipefail

cd "$(dirname "$0")"
source lib/common.sh
source lib/server.sh

BIN="/root/llama.cpp/build/bin/llama-server"
MODEL="/root/models/Kimi-K3-GGUF/UD-Q4_K_XL/Kimi-K3-UD-Q4_K_XL-00001-of-00032.gguf"
DRAFT="/root/models/k3-draft/draft.gguf"
PROMPT_FILE="/root/k3-test/prompts/gen/coding_64tok.txt"
PORT=8899
THREADS="${THREADS:-112}"
PRED=64
LOG_ROOT="$(pwd)/logs"
PROGRESS="$LOG_ROOT/progress.log"

mkdir -p "$LOG_ROOT"
log() { echo "[$(date +%H:%M:%S)] $*" | tee -a "$PROGRESS"; }

PROMPT=$(cat "$PROMPT_FILE")
stop_server "$PORT" 2>/dev/null || true
sleep 2

# NO -fa on: let it auto (likely off). Draft on CPU. Default batch 512.
log "### dspark-isolate: cmoe + dspark + FA-auto + t${THREADS} (NO explicit FA)"
LLAMA_MMAP_NO_PREFETCH=1 \
"$BIN" \
  --host 127.0.0.1 \
  --port "$PORT" \
  -m "$MODEL" \
  -ngl 999 \
  --tensor-split "0.3,1,1,1,1,1,1,1" \
  --cpu-moe \
  -t "$THREADS" \
  -b 512 \
  -ub 512 \
  -np 1 \
  -md "$DRAFT" \
  -ngld 0 \
  --spec-type draft-dspark \
  > "$LOG_ROOT/server_iso.log" 2>&1 &

SRV_PID=$!
log "server[ISO]: pid=$SRV_PID, waiting for health…"

READY=0
for i in $(seq 1 480); do
  if ! kill -0 "$SRV_PID" 2>/dev/null; then
    log "server[ISO]: DIED during load"
    tail -30 "$LOG_ROOT/server_iso.log" | tee -a "$PROGRESS"
    exit 1
  fi
  HEALTH=$(curl -s -o /dev/null -w '%{http_code}' "http://127.0.0.1:$PORT/health" 2>/dev/null || echo "000")
  if [[ "$HEALTH" == "200" ]]; then
    log "server[ISO]: READY after ${i}s (pid $SRV_PID)"
    READY=1
    break
  fi
  sleep 1
done

if [[ "$READY" != "1" ]]; then
  log "server[ISO]: TIMEOUT"
  tail -20 "$LOG_ROOT/server_iso.log" | tee -a "$PROGRESS"
  exit 1
fi

# Single probe; capture crash if any
log "[iso_cold] starting…"
T0=$(date +%s%N)
RESULT=$(python3 -c "
import urllib.request, json
data = json.dumps({
    'prompt': $(python3 -c "import json; print(json.dumps(open('$PROMPT_FILE').read()))"),
    'n_predict': $PRED,
    'temperature': 0,
    'top_k': 1,
    'seed': 42
}).encode()
req = urllib.request.Request('http://127.0.0.1:$PORT/completion', data=data, headers={'Content-Type': 'application/json'})
try:
    resp = urllib.request.urlopen(req, timeout=900)
    r = json.loads(resp.read())
    t = r.get('timings', {})
    print(f\"{t.get('prompt_per_second', 0)} {t.get('predicted_per_second', 0)}\")
except Exception as e:
    print(f'ERROR {e}')
" 2>&1)
T1=$(date +%s%N)
WALL=$(( (T1 - T0) / 1000000000 ))
log "[iso_cold] wall=${WALL}s result=${RESULT}"

# Check if server still alive
if kill -0 "$SRV_PID" 2>/dev/null; then
  log "server[ISO]: still alive after probe"
else
  log "server[ISO]: CRASHED during probe — last log:"
  tail -25 "$LOG_ROOT/server_iso.log" | grep -vE "New LWP|Thread debugging|host libthread" | tee -a "$PROGRESS"
fi

echo "$RESULT" | grep -q ERROR || {
  PP=$(echo "$RESULT" | awk '{print $1}')
  TG=$(echo "$RESULT" | awk '{print $2}')
  DISK=$(io_read_bytes 2>/dev/null || echo 0)
  echo -e "iso_cold\t0\t${WALL}\t0\t${PP}\t${TG}\t${DISK}" >> "$LOG_ROOT/metrics.tsv"
}

log "server[ISO]: stopping…"
kill "$SRV_PID" 2>/dev/null || true
sleep 3
kill -9 "$SRV_PID" 2>/dev/null || true
log "### dspark-isolate DONE"
date +%s > "$LOG_ROOT/k3_suite_status.txt"