#!/usr/bin/env bash
# Purpose: run one model across this machine and the RPC hosts, with an explicit or a
#          proportional split, either as a server to measure against or as a benchmark
# Platform: all (the client is whichever machine you sit at; the hosts can be any track)
# Minimum memory: 8 GB on the client; the model has to fit the sum of every device
# Assumes: llama.cpp built with -DGGML_RPC=ON on the client; one ggml-rpc-server already
#          running on every host in RPC_HOSTS; the model file readable at $MODEL; the
#          environment from env-example.txt has been sourced
#
# Usage: . ./rpc.env && bash run-split.sh                 # start llama-server
#        . ./rpc.env && MODE=bench bash run-split.sh      # run llama-bench instead
#        . ./rpc.env && MODE=probe bash run-split.sh      # list the devices and exit
#
# Environment (all from env-example.txt):
#   MODEL         first shard of the GGUF                      (required)
#   RPC_HOSTS     host:port,host:port in device order          (required)
#   TENSOR_SPLIT  comma-separated proportions, e.g. 1,1,0.5    (default: by free memory)
#   CTX           context length to allocate                   (default: 8192)
#   NGL           layers to offload                            (default: 999)
#   LLAMA_BIN     directory holding the binaries               (default: found on PATH)
#   HOST_PORT     port llama-server listens on locally         (default: 8080)
#   MODE          server, bench or probe                       (default: server)
#
# The device order that --tensor-split addresses is: this machine's local devices first,
# then the RPC hosts in the order they appear in RPC_HOSTS. Run MODE=probe once and read
# the list before you write a split; guessing the order is how a split ends up backwards.

set -euo pipefail

MODE="${MODE:-server}"
CTX="${CTX:-8192}"
NGL="${NGL:-999}"
HOST_PORT="${HOST_PORT:-8080}"
TENSOR_SPLIT="${TENSOR_SPLIT:-}"

die() { echo "run-split: $*" >&2; exit 1; }

bin() {
  if [ -n "${LLAMA_BIN:-}" ]; then
    echo "$LLAMA_BIN/$1"
  else
    command -v "$1" || true
  fi
}


# --- 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"

# --- 1. Are the hosts actually there? ---------------------------------------------------
# A cluster that fails three minutes into loading a 125 GB model because one host is not
# listening is a cluster that wasted three minutes. Check first; it costs nothing.
IFS=',' read -r -a HOSTS <<< "$RPC_HOSTS"
for hp in "${HOSTS[@]}"; do
  host="${hp%%:*}"
  port="${hp##*:}"
  if ! (exec 3<>"/dev/tcp/$host/$port") 2>/dev/null; then
    die "no ggml-rpc-server answering at $hp. Start it there, and check CLUSTER_IFACE on that host."
  fi
  echo "    reachable: $hp"
done

# --- 2. Probe: list every device this run would see, then stop ---------------------------
if [ "$MODE" = "probe" ]; then
  BENCH="$(bin llama-bench)"
  [ -x "$BENCH" ] || die "llama-bench not found; set LLAMA_BIN"
  echo "==> devices visible to this client, local first, then the RPC hosts in order"
  exec "$BENCH" --rpc "$RPC_HOSTS" --list-devices
fi

[ -n "${MODEL:-}" ] || die "set MODEL to the first shard of the GGUF; see env-example.txt"
[ -f "$MODEL" ] || die "$MODEL does not exist on this machine"

echo "==> splitting $(basename "$MODEL")"
echo "    rpc hosts    : $RPC_HOSTS"
echo "    context      : $CTX tokens, offloading $NGL layers"
if [ -n "$TENSOR_SPLIT" ]; then
  echo "    tensor split : $TENSOR_SPLIT (local devices first, then RPC hosts in order)"
else
  echo "    tensor split : none given, so llama.cpp splits in proportion to free memory"
fi

# --- 3. Benchmark mode --------------------------------------------------------------------
if [ "$MODE" = "bench" ]; then
  BENCH="$(bin llama-bench)"
  [ -x "$BENCH" ] || die "llama-bench not found; set LLAMA_BIN"
  ARGS=(-m "$MODEL" --rpc "$RPC_HOSTS" -ngl "$NGL" -p 512 -n 128 -r 3 -o json)
  if [ -n "$TENSOR_SPLIT" ]; then
    # llama-bench separates split proportions with "/" where llama-server uses ","
    ARGS+=(-ts "${TENSOR_SPLIT//,//}")
  fi
  exec "$BENCH" "${ARGS[@]}"
fi

# --- 4. Server mode -----------------------------------------------------------------------
SERVER="$(bin llama-server)"
[ -x "$SERVER" ] || die "llama-server not found; set LLAMA_BIN"

ARGS=(-m "$MODEL" --rpc "$RPC_HOSTS" -ngl "$NGL" -c "$CTX"
      --host 127.0.0.1 --port "$HOST_PORT" --metrics)
[ -n "$TENSOR_SPLIT" ] && ARGS+=(-ts "$TENSOR_SPLIT")

echo "    listening on : 127.0.0.1 port $HOST_PORT, with the metrics endpoint enabled"
echo "    loading now; a split model loads slowly the first time and quickly after that,"
echo "    because each host caches its own tensors on local disk."
exec "$SERVER" "${ARGS[@]}"