#!/usr/bin/env bash
# Purpose: collect, into one report, every piece of evidence needed to explain a slow
#          llama.cpp RPC cluster: which interface carries the traffic, its MTU and whether
#          it is wireless, the route and round trip to each host, whether each rpc-server
#          answers, whether RDMA is available, the split in use, and whether the model
#          would have fitted on this machine alone
# Platform: all (Linux and macOS paths; each section says when it could not run)
# Minimum memory: 8 GB
# Assumes: run on the client; the environment from env-example.txt has been sourced; it
#          reads and measures only, changes nothing, and needs no privileges
#
# Usage: . ./rpc.env && bash diagnose-rpc-cluster.sh [report.md]
#
# Environment:
#   RPC_HOSTS      host:port,host:port                    (required)
#   CLUSTER_IFACE  the interface you believe is in use    (recommended)
#   MODEL          the model file being split             (optional but useful)
#   TENSOR_SPLIT   the proportions in use                 (optional)
#   LLAMA_BIN      directory holding the binaries         (default: found on PATH)
#
# Read the report top to bottom before forming a theory. Almost every slow cluster is
# explained by one of the first three sections, and guessing before reading them is how an
# afternoon disappears.

set -euo pipefail

REPORT="${1:-rpc-cluster-report.md}"
OS="$(uname -s)"

say() { printf '%s\n' "$*" >> "$REPORT"; }
section() { printf '\n## %s\n\n' "$*" >> "$REPORT"; }
code_start() { printf '```\n' >> "$REPORT"; }
code_end() { printf '```\n' >> "$REPORT"; }


# --- Part 18 compatibility --------------------------------------------------------------
# Part 18's .env lists every machine in CLUSTER_PEERS, space separated, as names without
# ports. When RPC_HOSTS is not set, build it from those names in the order they appear,
# appending RPC_PORT to each. Set RPC_HOSTS yourself whenever the order matters, when only
# some peers take part, or when a peer listens on a different port: the order is what
# --tensor-split addresses.
RPC_PORT="${RPC_PORT:-50052}"
if [ -z "${RPC_HOSTS:-}" ] && [ -n "${CLUSTER_PEERS:-}" ]; then
  read -r -a PART18_PEERS <<< "$CLUSTER_PEERS"
  for peer in "${PART18_PEERS[@]}"; do
    RPC_HOSTS="${RPC_HOSTS:+$RPC_HOSTS,}${peer}:${RPC_PORT}"
  done
  echo "    RPC_HOSTS built from Part 18's CLUSTER_PEERS: $RPC_HOSTS"
fi

[ -n "${RPC_HOSTS:-}" ] || { echo "diagnose-rpc-cluster: set RPC_HOSTS, or Part 18's CLUSTER_PEERS; see env-example.txt" >&2; exit 1; }

: > "$REPORT"
say "# RPC cluster diagnostic"
say ""
say "- client: $(hostname -s), $OS $(uname -m)"
say "- collected: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
say "- RPC_HOSTS: \`$RPC_HOSTS\`"
say "- TENSOR_SPLIT: \`${TENSOR_SPLIT:-not set, so llama.cpp splits by free memory}\`"

# --- 1. What is this client running? ----------------------------------------------------
section "1. Build"
code_start
if [ -n "${LLAMA_BIN:-}" ] && [ -x "$LLAMA_BIN/llama-cli" ]; then
  "$LLAMA_BIN/llama-cli" --version >> "$REPORT" 2>&1 || true
elif command -v llama-cli >/dev/null; then
  llama-cli --version >> "$REPORT" 2>&1 || true
else
  say "llama-cli not found; set LLAMA_BIN"
fi
code_end
say ""
say "A build without RPC compiled in cannot use \`--rpc\` at all, so if the client refused"
say "the option, rebuild it with \`-DGGML_RPC=ON\` before reading any further."

# --- 2. Interfaces, addresses and MTU ---------------------------------------------------
section "2. Interfaces on this machine"
code_start
case "$OS" in
  Darwin)
    ifconfig 2>/dev/null | grep -E '^[a-z0-9]+:|inet |media:' >> "$REPORT" || true
    ;;
  *)
    ip -br addr >> "$REPORT" 2>&1 || true
    echo "---" >> "$REPORT"
    ip -br link >> "$REPORT" 2>&1 || true
    ;;
esac
code_end
say ""
if [ -n "${CLUSTER_IFACE:-}" ]; then
  say "You said the cluster link is \`$CLUSTER_IFACE\`. Its MTU and kind:"
  code_start
  case "$OS" in
    Darwin)
      ifconfig "$CLUSTER_IFACE" >> "$REPORT" 2>&1 || true
      networksetup -listallhardwareports 2>/dev/null \
        | grep -B 1 "Device: $CLUSTER_IFACE\$" >> "$REPORT" || true
      ;;
    *)
      ip -d link show dev "$CLUSTER_IFACE" >> "$REPORT" 2>&1 || true
      if [ -e "/sys/class/net/$CLUSTER_IFACE/wireless" ] \
         || [ -d "/sys/class/net/$CLUSTER_IFACE/phy80211" ]; then
        echo "THIS INTERFACE IS WIRELESS" >> "$REPORT"
      fi
      ;;
  esac
  code_end
  say ""
  say "An MTU of 1500 on a link you configured for jumbo frames means the change did not"
  say "take, or a switch in the path did not take it. A wireless interface here is the"
  say "single most common reason a cluster is slower than one machine."
else
  say "CLUSTER_IFACE is not set, so this report cannot say which link you meant to use."
fi

# --- 3. Route and round trip to each host ------------------------------------------------
section "3. Route, reachability and round trip"
IFS=',' read -r -a HOSTS <<< "$RPC_HOSTS"
for hp in "${HOSTS[@]}"; do
  host="${hp%%:*}"
  port="${hp##*:}"
  say "### \`$hp\`"
  say ""
  code_start
  case "$OS" in
    Darwin) route get "$host" >> "$REPORT" 2>&1 || true ;;
    *)      ip route get "$(getent ahostsv4 "$host" 2>/dev/null | awk 'NR==1 {print $1}')" \
              >> "$REPORT" 2>&1 || echo "could not resolve $host" >> "$REPORT" ;;
  esac
  code_end
  say ""
  say "The interface in that route is the one that will actually carry the traffic. If it"
  say "is not \`${CLUSTER_IFACE:-your cluster link}\`, you have found the fault."
  say ""
  code_start
  if (exec 3<>"/dev/tcp/$host/$port") 2>/dev/null; then
    echo "TCP connect to $hp: succeeded" >> "$REPORT"
  else
    echo "TCP connect to $hp: FAILED - no rpc-server listening, or it bound elsewhere" >> "$REPORT"
  fi
  ping -c 4 "$host" >> "$REPORT" 2>&1 || echo "ping failed or is blocked" >> "$REPORT"
  code_end
  say ""
  say "Path MTU, by sending a packet that must not be fragmented:"
  say ""
  code_start
  case "$OS" in
    Darwin) ping -c 2 -D -s 1472 "$host" >> "$REPORT" 2>&1 || echo "1500-byte path failed" >> "$REPORT" ;;
    *)      ping -c 2 -M "do" -s 1472 "$host" >> "$REPORT" 2>&1 || echo "1500-byte path failed" >> "$REPORT" ;;
  esac
  case "$OS" in
    Darwin) ping -c 2 -D -s 8972 "$host" >> "$REPORT" 2>&1 || echo "9000-byte path failed (no jumbo frames)" >> "$REPORT" ;;
    *)      ping -c 2 -M "do" -s 8972 "$host" >> "$REPORT" 2>&1 || echo "9000-byte path failed (no jumbo frames)" >> "$REPORT" ;;
  esac
  code_end
  say ""
done

# --- 4. Is RDMA available at all? ----------------------------------------------------------
section "4. RDMA"
code_start
if command -v ibv_devices >/dev/null; then
  ibv_devices >> "$REPORT" 2>&1 || true
else
  echo "ibv_devices not found: no verbs library on this machine, so this client is on TCP" >> "$REPORT"
fi
if [ -n "${GGML_RPC_NO_RDMA:-}" ]; then
  echo "GGML_RPC_NO_RDMA is set in this environment: RDMA is switched off deliberately" >> "$REPORT"
fi
code_end
say ""
say "RDMA is negotiated per connection and it is point-to-point: a connection made over a"
say "different interface from the RDMA-capable one stays on TCP however good the hardware is."

# --- 5. Would it have fitted on this machine? -----------------------------------------------
section "5. The question the challenge is really about"
if [ -n "${MODEL:-}" ] && [ -f "$MODEL" ]; then
  bytes="$(python3 -c 'import os,sys; print(os.path.getsize(sys.argv[1]))' "$MODEL" 2>/dev/null || echo 0)"
  say "- model file: \`$(basename "$MODEL")\`"
  say "- first shard: $((bytes / 1000000)) MB (later shards are not counted here)"
else
  say "- MODEL is not set or does not exist here, so the fit question cannot be answered."
fi
code_start
case "$OS" in
  Darwin) sysctl -n hw.memsize 2>/dev/null | awk '{ printf "physical memory: %d MB\n", $1/1000000 }' >> "$REPORT" || true ;;
  *)      grep -E '^(MemTotal|MemAvailable|SwapTotal)' /proc/meminfo >> "$REPORT" 2>&1 || true ;;
esac
code_end
say ""
say "If every shard of the model would have fitted in this machine's memory with room for"
say "the KV cache, the cluster was never going to win, and the fix is to stop splitting."

echo "==> written to $REPORT"
echo "    Read sections 2 and 3 first: interface, MTU, route."