#!/usr/bin/env bash
# Purpose: ask every rpc-server in the cluster what device it is offering, one host at a
#          time, and write the device order the client will use into one short report
# Platform: all (run it on the client; the hosts may be any mixture of tracks)
# Minimum memory: 8 GB on the client; the hosts need only enough for their own share
# Assumes: llama.cpp built with -DGGML_RPC=ON on the client so llama-bench can register RPC
#          devices; one ggml-rpc-server running on every host in RPC_HOSTS; the environment
#          from env-example.txt has been sourced
#
# Usage: . ./rpc.env && bash probe-rpc-devices.sh [report.md]
#
# Environment:
#   RPC_HOSTS   host:port,host:port                      (required)
#   LLAMA_BIN   directory holding llama-bench            (default: found on PATH)
#
# Why one host at a time: --list-devices with every host registered prints one flat list,
# and nothing in that list says which host a device came from. Registering one host per
# run and recording the order is the only way to know what --tensor-split is addressing.

set -euo pipefail

REPORT="${1:-rpc-devices.md}"

die() { echo "probe-rpc-devices: $*" >&2; exit 1; }

if [ -n "${LLAMA_BIN:-}" ]; then
  BENCH="$LLAMA_BIN/llama-bench"
else
  BENCH="$(command -v llama-bench || true)"
fi
[ -n "$BENCH" ] && [ -x "$BENCH" ] || die "llama-bench not found; set LLAMA_BIN"

# --- 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:-}" ] || die "set RPC_HOSTS (or Part 18's CLUSTER_PEERS); see env-example.txt"

{
  echo "# RPC cluster device probe"
  echo
  echo "- client: $(hostname -s), $(uname -s) $(uname -m)"
  echo "- probed: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
  echo "- RPC_HOSTS: \`$RPC_HOSTS\`"
  echo
} > "$REPORT"

echo "==> local devices, which always come first in the split order"
{
  echo "## Local devices (position 1 onwards in --tensor-split)"
  echo
  echo '```'
} >> "$REPORT"
"$BENCH" --list-devices 2>&1 | tee -a "$REPORT"
echo '```' >> "$REPORT"
echo >> "$REPORT"

IFS=',' read -r -a HOSTS <<< "$RPC_HOSTS"
POSITION=1
for hp in "${HOSTS[@]}"; do
  host="${hp%%:*}"
  port="${hp##*:}"
  echo "==> $hp"
  {
    echo "## RPC host $POSITION: \`$hp\`"
    echo
  } >> "$REPORT"

  if (exec 3<>"/dev/tcp/$host/$port") 2>/dev/null; then
    echo "    reachable"
    echo "- TCP connect: succeeded" >> "$REPORT"
  else
    echo "    NOT reachable: nothing is listening on $hp" >&2
    echo "- TCP connect: **failed** — no server listening, or the wrong interface" >> "$REPORT"
    echo >> "$REPORT"
    POSITION=$((POSITION + 1))
    continue
  fi

  # Round-trip time on the path the RPC traffic will take. A number in the tens of
  # milliseconds here means Wi-Fi or a router, and this part's challenge page is about that.
  if command -v ping >/dev/null; then
    rtt="$(ping -c 4 "$host" 2>/dev/null | tail -n 1 || true)"
    [ -n "$rtt" ] && echo "- round trip: \`$rtt\`" >> "$REPORT"
  fi

  {
    echo
    echo '```'
  } >> "$REPORT"
  "$BENCH" --rpc "$hp" --list-devices 2>&1 | tee -a "$REPORT"
  {
    echo '```'
    echo
  } >> "$REPORT"
  POSITION=$((POSITION + 1))
done

{
  echo "## Split order"
  echo
  echo "The proportions given to \`--tensor-split\` are read in this order: every local"
  echo "device listed above first, then the RPC hosts in the order they appear in"
  echo "RPC_HOSTS. Copy that order into your notebook beside the split you chose."
} >> "$REPORT"

echo "==> written to $REPORT"