#!/usr/bin/env bash
# Purpose: collect, into one Markdown report, the evidence that explains a slow model: the
#          build and the devices llama.cpp can see, what the operating system reports for
#          accelerator memory, memory and swap, where llama-server put the weights and the KV
#          cache when it loaded the model and what --fit changed, and the prompt and
#          generation rates measured with that same placement while swap activity is sampled
# Platform: all (Linux, macOS and WSL2; bash 3.2 or later; uses whichever GPU tool exists)
# Minimum memory: 8 GB
# Assumes: llama.cpp v0.4.0 built as in this part's install lesson (llama-server and
#          llama-bench in $LLAMA_BIN, on PATH, or in ~/llama.cpp/build/bin); a GGUF model
#          file; python3 3.9 or later, curl and awk on PATH; nothing listening on $PORT.
#          It starts llama-server bound to 127.0.0.1 and stops it, then runs llama-bench;
#          it changes no setting on the machine.
#
# Usage: bash diagnose-slow-inference.sh <model.gguf> [report.md] [labbook.md]
#
# Environment. Each of the first six maps to llama.cpp options; unset means llama.cpp's
# own default, which is what a plain llama-server command gets.
#   NGL           --n-gpu-layers for the load     (default: unset, "auto": --fit decides)
#   CTX           --ctx-size for the load         (default: 4096). CTX=default leaves the
#                 option out, as a plain llama-server command does: the model's own context,
#                 which --fit may shrink. CTX=0 asks for the model's own context and forbids
#                 --fit to shrink it (common/arg.cpp at v0.4.0).
#   CACHE_TYPE    --cache-type-k and --cache-type-v, for example q8_0 (default: unset, f16)
#   DEVICE        --device, for example none      (default: unset, every device)
#   FIT_TARGET    --fit-target, MiB left free     (default: unset, 1024)
#   NO_KV_OFFLOAD 1 adds --no-kv-offload          (default: unset)
#   LLAMA_BIN     directory holding the two tools (default: PATH, then ~/llama.cpp/build/bin)
#   PORT          local port for the load probe   (default: 8089)
#   LOAD_TIMEOUT  seconds to wait for the load    (default: 600)
#   PROMPT_LEN, GEN_LEN, REPS   llama-bench --n-prompt, --n-gen, --repetitions
#                                                 (default: 512, 128, 3)
#   LABEL         name for the notebook line      (default: the report name without .md)
#
# How it works. llama-server is started with --verbose, because at v0.4.0 the placement lines
# ("offloaded N/M layers", buffer sizes, the KV cache size, the --fit decisions) are logged
# above the default verbosity. Once it prints "listening on", or exits, it is stopped with
# SIGTERM, which prints the memory breakdown. llama-bench then runs with --n-gpu-layers set to
# the N the server logged (or 0 when the load placed nothing on a device) and the same
# --device and --no-kv-offload values, so the rates describe the placement in the report.
# On Linux vmstat samples swap-in and swap-out every second during the benchmark; on macOS
# vm.swapusage is read before and after. Full outputs stay beside the report:
# <report>.load.log, <report>.bench.json, <report>.bench.err and, on Linux, <report>.vmstat.

set -euo pipefail

MODEL="${1:-}"
REPORT="${2:-slow-inference-report.md}"
LABBOOK="${3:-}"
CTX="${CTX:-4096}"
PORT="${PORT:-8089}"
LOAD_TIMEOUT="${LOAD_TIMEOUT:-600}"
PROMPT_LEN="${PROMPT_LEN:-512}"
GEN_LEN="${GEN_LEN:-128}"
REPS="${REPS:-3}"
LABEL="${LABEL:-$(basename "$REPORT" .md)}"

die() { echo "diagnose-slow-inference: $*" >&2; exit 1; }
have() { command -v "$1" >/dev/null 2>&1; }

[ -n "$MODEL" ] || die "usage: bash diagnose-slow-inference.sh <model.gguf> [report.md] [labbook.md]"
[ -f "$MODEL" ] || die "model file '$MODEL' does not exist (names are case-sensitive)"
case "$REPORT" in *.md) ;; *) die "the report name must end in .md, got '$REPORT'" ;; esac
for tool in python3 curl awk; do
  have "$tool" || die "$tool is not installed or not on PATH"
done
if [ -n "$LABBOOK" ] && [ ! -f "$LABBOOK" ]; then
  die "notebook '$LABBOOK' does not exist; create it or leave the third argument out"
fi

if [ -n "${LLAMA_BIN:-}" ]; then
  SERVER="$LLAMA_BIN/llama-server"; BENCH="$LLAMA_BIN/llama-bench"
elif have llama-server && have llama-bench; then
  SERVER="$(command -v llama-server)"; BENCH="$(command -v llama-bench)"
else
  SERVER="$HOME/llama.cpp/build/bin/llama-server"; BENCH="$HOME/llama.cpp/build/bin/llama-bench"
fi
[ -x "$SERVER" ] || die "llama-server not found at $SERVER; set LLAMA_BIN to the directory holding it"
[ -x "$BENCH" ] || die "llama-bench not found at $BENCH; set LLAMA_BIN to the directory holding it"
if curl -s --max-time 2 "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
  die "something is already answering on 127.0.0.1:$PORT; stop it or set PORT to a free port"
fi

BASE="${REPORT%.md}"
LOAD_LOG="$BASE.load.log"
BENCH_JSON="$BASE.bench.json"
BENCH_ERR="$BASE.bench.err"
VMSTAT_LOG="$BASE.vmstat"
BODY="$(mktemp)"
SWAP_BEFORE="$(mktemp)"
SWAP_AFTER="$(mktemp)"
SERVER_PID=""
VMSTAT_PID=""
cleanup() {
  if [ -n "$SERVER_PID" ]; then kill "$SERVER_PID" 2>/dev/null || true; fi
  if [ -n "$VMSTAT_PID" ]; then kill "$VMSTAT_PID" 2>/dev/null || true; fi
  rm -f "$BODY" "$SWAP_BEFORE" "$SWAP_AFTER"
}
trap cleanup EXIT
trap 'cleanup; exit 130' INT TERM

FENCE="$(printf '\140\140\140')"   # three backticks, kept out of the quoting
section() { printf '\n## %s\n\n' "$1" >> "$BODY"; }
block() { printf '%stext\n%s\n%s\n' "$FENCE" "${1:-(no output)}" "$FENCE" >> "$BODY"; }
note() { printf '%s\n' "$@" >> "$BODY"; }
run() { "$@" 2>&1 || echo "(command failed: $*)"; }
OS="$(uname -s)"

echo "==> Collecting evidence for $(basename "$MODEL") into $REPORT"

# --- Machine ---------------------------------------------------------------------------
section "Machine and operating system"
block "$(run uname -a)"
if [ -r /etc/os-release ]; then block "$(grep -E '^(PRETTY_NAME|VERSION_ID)=' /etc/os-release)"; fi
if [ "$OS" = "Darwin" ]; then block "$(run sw_vers)"; fi
WSL=0
if grep -qi microsoft /proc/version 2>/dev/null; then
  WSL=1
  note "Running inside WSL2. MemTotal below is the virtual machine's memory, set by memory= in" \
       "the [wsl2] section of %UserProfile%\\.wslconfig (default: half of Windows' memory)."
  block "$(grep -E '^(MemTotal|SwapTotal):' /proc/meminfo)"
fi

# --- Build and devices -------------------------------------------------------------------
section "Build, and the devices llama.cpp can see"
block "$(run "$SERVER" --version)"
DEVICES="$(run "$BENCH" --list-devices)"
block "$DEVICES"
note "\`(none)\` under Available devices means no GPU backend is compiled in or its runtime" \
     "did not load. The MiB free figure is what --fit plans with."

# --- Accelerator, as the operating system reports it -----------------------------------------
section "Accelerator memory, as the operating system reports it"
if have nvidia-smi; then
  block "$(run nvidia-smi --query-gpu=name,memory.total,memory.used,memory.free,utilization.gpu --format=csv)"
  block "$(run nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv)"
fi
AMD_FOUND=0
for dir in /sys/class/drm/card*/device; do
  if [ -r "$dir/mem_info_gtt_total" ]; then
    AMD_FOUND=1
    block "$(cd "$dir" && grep -H . mem_info_vram_total mem_info_vram_used mem_info_gtt_total mem_info_gtt_used 2>&1)"
  fi
done
if [ "$AMD_FOUND" = 1 ] && [ -r /sys/module/ttm/parameters/pages_limit ]; then
  block "pages_limit: $(cat /sys/module/ttm/parameters/pages_limit)"
  note "amdgpu files are bytes; pages_limit x 4096 is the GTT ceiling Part 5 set."
fi
if [ "$OS" = "Darwin" ]; then
  block "$(printf 'iogpu.wired_limit_mb: %s\nhw.memsize: %s' "$(sysctl -n iogpu.wired_limit_mb 2>&1)" "$(sysctl -n hw.memsize 2>&1)")"
  note "A wired limit of 0 is the system default; Part 5's value does not survive a restart."
fi
if ! have nvidia-smi && [ "$AMD_FOUND" = 0 ] && [ "$OS" != "Darwin" ]; then
  note "No nvidia-smi, no amdgpu memory files and not macOS: nothing to report here."
fi

# --- Memory and swap, before ----------------------------------------------------------------
section "Memory and swap before the load"
if [ "$OS" = "Darwin" ]; then
  sysctl -n vm.swapusage > "$SWAP_BEFORE" 2>&1 || true
  block "vm.swapusage: $(cat "$SWAP_BEFORE")"
else
  free -m > "$SWAP_BEFORE" 2>&1 || true
  block "$(cat "$SWAP_BEFORE")"
  if have swapon; then block "$(swapon --show 2>&1 || true)"; note "An empty block above means no swap device is configured."; fi
fi

# --- The load probe -----------------------------------------------------------------------
SERVER_ARGS=(--model "$MODEL" --host 127.0.0.1 --port "$PORT" --verbose)
BENCH_EXTRA=(--flash-attn on)
ASKED="(no --ctx-size)"
if [ "$CTX" != "default" ]; then SERVER_ARGS+=(--ctx-size "$CTX"); ASKED="--ctx-size $CTX"; fi
if [ -n "${CACHE_TYPE:-}" ]; then
  SERVER_ARGS+=(--cache-type-k "$CACHE_TYPE" --cache-type-v "$CACHE_TYPE")
  BENCH_EXTRA+=(--cache-type-k "$CACHE_TYPE" --cache-type-v "$CACHE_TYPE")
  ASKED="$ASKED --cache-type-k/v $CACHE_TYPE"
fi
if [ -n "${NGL:-}" ]; then SERVER_ARGS+=(--n-gpu-layers "$NGL"); ASKED="$ASKED --n-gpu-layers $NGL"; fi
if [ -n "${DEVICE:-}" ]; then
  SERVER_ARGS+=(--device "$DEVICE"); BENCH_EXTRA+=(--device "$DEVICE"); ASKED="$ASKED --device $DEVICE"
fi
if [ -n "${FIT_TARGET:-}" ]; then SERVER_ARGS+=(--fit-target "$FIT_TARGET"); ASKED="$ASKED --fit-target $FIT_TARGET"; fi
if [ "${NO_KV_OFFLOAD:-}" = 1 ]; then
  SERVER_ARGS+=(--no-kv-offload); BENCH_EXTRA+=(--no-kv-offload 1); ASKED="$ASKED --no-kv-offload"
fi

echo "==> Loading with llama-server $ASKED (this can take a minute)"
"$SERVER" "${SERVER_ARGS[@]}" > "$LOAD_LOG" 2>&1 < /dev/null &
SERVER_PID=$!
STATUS="timeout"
waited=0
while [ "$waited" -lt "$LOAD_TIMEOUT" ]; do
  if grep -q "listening on" "$LOAD_LOG"; then STATUS="loaded"; break; fi
  if ! kill -0 "$SERVER_PID" 2>/dev/null; then STATUS="failed"; break; fi
  sleep 1
  waited=$((waited + 1))
done
kill -TERM "$SERVER_PID" 2>/dev/null || true
wait "$SERVER_PID" 2>/dev/null || true
SERVER_PID=""

section "Model load: $STATUS (llama-server $ASKED)"
# The evidence lines; buffers of 0.00 MiB come from --fit's trial pass and are left out.
block "$(grep -E 'no usable GPU|using device|common_params_fit_impl: (projected|will leave|cannot meet|context size|user has requested|entire model|  - )|common_fit_params: (successfully|failed)|offloaded [0-9]+/[0-9]+ layers|model buffer size|KV buffer size|llama_kv_cache: size|llama_context: n_ctx  |recommendedMaxWorkingSetSize|greater than the recommended|cudaMalloc failed|failed to allocate|unable to allocate|error loading model|Device memory allocation|out of memory' "$LOAD_LOG" | grep -v ' D ' | grep -vE 'buffer size = +0\.00 MiB' | head -n 60 || true)"
# The breakdown printed last: the one for the load that actually happened.
block "$(grep 'common_memory_breakdown_print: |' "$LOAD_LOG" | sed 's/.*common_memory_breakdown_print: //' | awk '/memory breakdown \[MiB\]/ { n = 0 } { rows[n++] = $0 } END { for (i = 0; i < n; i++) print rows[i] }' || true)"
note "The breakdown is llama.cpp's own accounting, in MiB: per device, total = free + (self =" \
     "model + context + compute) + unaccounted; the Host row is system memory. Full log: $LOAD_LOG"

# --- The rates, with the placement the load chose ------------------------------------------
OFFLOADED="$(grep -oE 'offloaded [0-9]+/[0-9]+ layers' "$LOAD_LOG" | tail -n 1 | awk '{ split($2, a, "/"); print a[1] }' || true)"
# Weights in buffers whose names are not host buffers (CPU, CPU_Mapped, CPU_REPACK, *_Host).
DEVICE_WEIGHTS="$(grep -E 'model buffer size' "$LOAD_LOG" | grep -v ' D ' | grep -vE ' (CPU[A-Za-z_]*|[A-Za-z0-9]+_Host) +model buffer' | awk -F'= *' '{ s += $2 } END { printf "%.0f", s }' || true)"
if [ -z "$OFFLOADED" ] || [ "${DEVICE_WEIGHTS:-0}" = 0 ]; then BENCH_NGL=0; else BENCH_NGL="$OFFLOADED"; fi
BENCH_ARGS=(--model "$MODEL" --n-prompt "$PROMPT_LEN" --n-gen "$GEN_LEN" --repetitions "$REPS" --n-gpu-layers "$BENCH_NGL" --output json)
rm -f "$BENCH_JSON" "$BENCH_ERR" "$VMSTAT_LOG"

section "Rates measured with that placement"
BENCH_STATUS="not run"
if [ "$STATUS" = "loaded" ]; then
  echo "==> Measuring with llama-bench --n-gpu-layers $BENCH_NGL (repetitions: $REPS)"
  if [ "$OS" = "Linux" ] && have vmstat; then
    vmstat 1 > "$VMSTAT_LOG" 2>&1 &
    VMSTAT_PID=$!
  fi
  BENCH_STATUS=0
  "$BENCH" "${BENCH_ARGS[@]}" "${BENCH_EXTRA[@]}" > "$BENCH_JSON" 2> "$BENCH_ERR" || BENCH_STATUS=$?
  if [ -n "$VMSTAT_PID" ]; then kill "$VMSTAT_PID" 2>/dev/null || true; wait "$VMSTAT_PID" 2>/dev/null || true; VMSTAT_PID=""; fi
  block "llama-bench --n-prompt $PROMPT_LEN --n-gen $GEN_LEN --repetitions $REPS --n-gpu-layers $BENCH_NGL ${BENCH_EXTRA[*]} (exit status $BENCH_STATUS)"
  if [ "$BENCH_STATUS" != 0 ]; then block "$(tail -n 15 "$BENCH_ERR")"; fi
else
  echo "==> The model did not load ($STATUS); skipping the benchmark"
  note "Not run: the load ended with status $STATUS, and the load section above holds the reason."
fi
if [ "$OS" = "Darwin" ]; then sysctl -n vm.swapusage > "$SWAP_AFTER" 2>&1 || true; else free -m > "$SWAP_AFTER" 2>&1 || true; fi
section "Memory and swap during and after the benchmark"
if [ -s "$VMSTAT_LOG" ]; then
  block "$(awk 'NR <= 2 { print; next } { rows[n++] = $0 } END { for (i = (n > 6 ? n - 6 : 1); i < n; i++) print rows[i] }' "$VMSTAT_LOG")"
  note "si and so are pages swapped in and out per second; the first sample after the headers, an average since boot, is left out."
fi
block "$(cat "$SWAP_AFTER")"

# --- Signals: the lines to read first, and the notebook record --------------------------------
SIGNALS="$(python3 - "$LOAD_LOG" "$BENCH_JSON" "$VMSTAT_LOG" "$SWAP_BEFORE" "$SWAP_AFTER" "$DEVICES" \
  "$STATUS" "$ASKED" "$BENCH_NGL" "$MODEL" "$LABEL" "$LABBOOK" "$WSL" <<'PY'
import json, re, statistics, sys
from datetime import datetime, timezone
(load_log, bench_json, vmstat_log, swap_before, swap_after, devices, status, asked,
 bench_ngl, model, label, labbook, wsl) = sys.argv[1:14]
T95 = {2: 12.706, 3: 4.303, 4: 3.182, 5: 2.776, 6: 2.571, 7: 2.447, 8: 2.365, 9: 2.306, 10: 2.262}

def read(path):
    try:
        return open(path, encoding="utf-8", errors="replace").read()
    except OSError:
        return ""

log = "\n".join(l for l in read(load_log).splitlines() if " D " not in l)
last = lambda pattern: (re.findall(pattern, log) or [None])[-1]
dev_lines = re.findall(r"^\s+(\S+): .*\((\d+) MiB, (\d+) MiB free\)$", devices, re.M)
buffers = re.findall(r"load_tensors:\s+(\S+) model buffer size =\s+([\d.]+) MiB", log)
is_host = lambda name: name.startswith("CPU") or name.endswith("_Host")
dev_w = sum(float(v) for n, v in buffers if not is_host(n))
host_w = sum(float(v) for n, v in buffers if is_host(n))
kv_bufs = re.findall(r"llama_kv_cache:\s+(\S+) KV buffer size =\s+([\d.]+) MiB", log)
kv_where = ", ".join(f"{n} {float(v):.0f}" for n, v in kv_bufs if float(v) > 0) or "-"
fit = [m.split("common_params_fit_impl: ", 1)[1] for m in re.findall(
    r".*common_params_fit_impl: (?:will leave|cannot meet|context size|user has requested|  - ).*", log)]
fit_failed = last(r"common_fit_params: failed to fit params[^\n]*")
errors = re.findall(r".*(?:cudaMalloc failed|failed to allocate|unable to allocate|error loading model|"
                    r"Device memory allocation|greater than the recommended max working set).*", log)

def swap_used(text):
    m = re.search(r"^Swap:\s+\d+\s+(\d+)", text, re.M) or re.search(r"used = ([\d.]+)M", text)
    return float(m.group(1)) if m else None

si = so = None
rows = [l.split() for l in read(vmstat_log).splitlines()]
head = next((r for r in rows if "si" in r and "so" in r), None)
if head:
    data = [r for r in rows if len(r) == len(head) and r[0].isdigit()][1:]  # row 1 is since boot
    if data:
        si = max(int(r[head.index("si")]) for r in data)
        so = max(int(r[head.index("so")]) for r in data)

rates = {}
try:
    for row in json.load(open(bench_json, encoding="utf-8")):
        kind = "pp" if row["n_prompt"] and not row["n_gen"] else "tg"
        ts = [(row["n_prompt"] + row["n_gen"]) * 1e9 / ns for ns in row["samples_ns"]]
        sd = statistics.stdev(ts) if len(ts) > 1 else 0.0
        half = T95.get(len(ts), 2.0) * sd / len(ts) ** 0.5
        rates[kind] = {"test": f"{kind}{row['n_prompt'] or row['n_gen']}", "mean": round(row["avg_ts"], 2),
                       "sd": round(row["stddev_ts"], 2), "ci95": [round(row["avg_ts"] - half, 2),
                       round(row["avg_ts"] + half, 2)], "backends": row["backends"],
                       "devices": row["devices"], "ngl": row["n_gpu_layers"]}
except (OSError, ValueError, KeyError, TypeError):
    pass

offload = last(r"offloaded (\d+/\d+) layers")
out = [f"load status        {status}   (asked: {asked})",
       "devices            " + (", ".join(f"{n} {f} MiB free of {t}" for n, t, f in dev_lines) or "(none)"),
       "no-GPU warning     " + ("yes" if "no usable GPU found" in log else "no"),
       "offload line       " + (offload or "not printed"),
       f"weights, MiB       on devices {dev_w:.0f}, on host {host_w:.0f}",
       "context            " + (last(r"llama_context: n_ctx\s+= (\d+)") or "-"),
       "KV cache, MiB      " + (last(r"llama_kv_cache: size =\s+([\d.]+) MiB") or "-") + f"   (buffers: {kv_where})",
       "fit decision       " + ("; ".join(f.strip() for f in fit) if fit else "-")]
if fit_failed:
    out.append("fit failed         " + fit_failed.split(": ", 2)[-1])
for e in errors[:3]:
    out.append("load error         " + e.strip()[-110:])
b, a = swap_used(read(swap_before)), swap_used(read(swap_after))
out.append(f"swap used, MiB     before {b if b is not None else '-'}, after {a if a is not None else '-'}"
           + (f"; max si {si}, max so {so} per s during the benchmark" if si is not None else ""))
if wsl == "1":
    out.append("WSL2               " + (re.search(r"MemTotal:\s+\d+ kB", read("/proc/meminfo")) or [""])[0])
for kind in ("pp", "tg"):
    r = rates.get(kind)
    out.append(f"{kind} rate            " + (f"{r['test']} {r['mean']:.2f} ± {r['sd']:.2f} t/s, 95% interval "
               f"{r['ci95'][0]:.2f} to {r['ci95'][1]:.2f} (backend {r['backends']}, ngl {r['ngl']}, "
               f"dev {r['devices']})" if r else "not measured"))
print("\n".join(out))

if labbook:
    record = {"lab": "part-06/challenge-two-tokens-per-second", "label": label,
              "date": datetime.now(timezone.utc).isoformat(timespec="seconds"), "model_path": model,
              "asked": asked, "load_status": status, "offloaded": offload,
              "device_weights_mib": round(dev_w, 2), "host_weights_mib": round(host_w, 2),
              "kv_mib": float(last(r"llama_kv_cache: size =\s+([\d.]+) MiB") or 0) or None,
              "n_ctx": int(last(r"llama_context: n_ctx\s+= (\d+)") or 0) or None,
              "fit": fit, "fit_failed": fit_failed, "no_gpu_warning": "no usable GPU found" in log,
              "swap_used_mib_before": b, "swap_used_mib_after": a, "max_si": si, "max_so": so,
              "bench_ngl": int(bench_ngl), "rates": rates}
    with open(labbook, "a", encoding="utf-8") as fh:
        fh.write(json.dumps(record) + "\n")
PY
)"

{
  printf '# Slow inference report: %s\n\n' "$LABEL"
  printf -- '- Collected: %s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  printf -- '- Model file: %s (%s bytes)\n' "$MODEL" "$(wc -c < "$MODEL" | tr -d ' ')"
  printf '\n## Signals\n\n%stext\n%s\n%s\n' "$FENCE" "$SIGNALS" "$FENCE"
  cat "$BODY"
} > "$REPORT"

printf '%s\n' "$SIGNALS"
echo "==> Wrote $REPORT${LABBOOK:+ and one part-06/challenge-two-tokens-per-second line to $LABBOOK}"