#!/usr/bin/env bash
# Purpose: check one node of a DGX Spark pair before any engine is started - the QSFP
#          interface is up, it has the address you configured, the peer answers, SSH works
#          without a password - and write the interface environment the vLLM and
#          TensorRT-LLM paths both need into cluster-env.sh for the other scripts to source
# Platform: spark (run once on each of the two nodes); the checks are read-only
# Minimum memory: 128 GB per node for the models this lab serves; this script needs none
# Assumes: the NVIDIA connect-two-Sparks playbook has already assigned addresses to the
#          QSFP interfaces and distributed SSH keys; ibdev2netdev, ip and ssh on PATH; a
#          .env file beside this script, copied from env-example.txt and filled in.
#          Changes nothing on the machine: it prints the netplan stanza it would write
#          rather than writing it.
set -euo pipefail

HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "${HERE}/.env" ]; then
    set -a
    # shellcheck disable=SC1091  # written by the reader from env-example.txt
    . "${HERE}/.env"
    set +a
fi

CX7_IFNAME="${CX7_IFNAME:-}"
CX7_ROCE_DEV="${CX7_ROCE_DEV:-}"
NODE_SELF_ADDR="${NODE_SELF_ADDR:-}"
NODE_PEER_ADDR="${NODE_PEER_ADDR:-}"
CLUSTER_USER="${CLUSTER_USER:-$(id -un)}"
ENV_OUT="${ENV_OUT:-${HERE}/cluster-env.sh}"

fail() { printf '%s\n' "$*" >&2; exit 1; }
note() { printf '    %s\n' "$*"; }

[ -n "$CX7_IFNAME" ] || fail "CX7_IFNAME is not set. Run ibdev2netdev, pick the interface that reports (Up), and put it in .env."
[ -n "$NODE_SELF_ADDR" ] || fail "NODE_SELF_ADDR is not set. It is this machine's address on the QSFP link, from 'ip addr show ${CX7_IFNAME}'."
[ -n "$NODE_PEER_ADDR" ] || fail "NODE_PEER_ADDR is not set. It is the other machine's address on the same link."

for tool in ip ssh; do
    command -v "$tool" >/dev/null 2>&1 || fail "${tool} is not on PATH."
done

printf '==> 1. QSFP ports and their RoCE devices\n'
if command -v ibdev2netdev >/dev/null 2>&1; then
    ibdev2netdev | sed 's/^/    /'
    if ! ibdev2netdev | grep -q "${CX7_IFNAME}.*(Up)"; then
        fail "${CX7_IFNAME} is not reported as (Up). Check the cable, confirm both nodes use the same physical port, and reboot before looking at software."
    fi
    note "OK: ${CX7_IFNAME} is up."
else
    note "ibdev2netdev is not on PATH; skipping the RoCE pairing check."
    note "On DGX OS it ships with the ConnectX-7 tooling. Without it you cannot see which"
    note "RoCE device belongs to which Ethernet interface, which NCCL_IB_HCA needs."
fi

printf '\n==> 2. The address on this node\n'
if ip -4 addr show "$CX7_IFNAME" | grep -qw "$NODE_SELF_ADDR"; then
    note "OK: ${CX7_IFNAME} carries the address in NODE_SELF_ADDR."
else
    printf '    %s\n' "${CX7_IFNAME} does not carry NODE_SELF_ADDR." >&2
    printf '    %s\n' "Either .env is wrong, or the address was assigned with 'ip addr add' and lost" >&2
    printf '    %s\n' "on the last reboot. The playbook's netplan option is the one that persists." >&2
    printf '\n    A netplan stanza for this node would look like this. Review it, write it to\n'
    printf '    /etc/netplan/40-cx7.yaml yourself, chmod 600 it, and run "sudo netplan apply":\n\n'
    cat <<STANZA
    network:
      version: 2
      ethernets:
        ${CX7_IFNAME}:
          dhcp4: no
          addresses:
            - ${NODE_SELF_ADDR}/24
STANZA
    exit 1
fi

printf '\n==> 3. The peer answers on the link\n'
if command -v ping >/dev/null 2>&1 && ping -c 2 -W 2 "$NODE_PEER_ADDR" >/dev/null 2>&1; then
    note "OK: the peer replies on the QSFP subnet."
else
    fail "The peer does not reply on the QSFP subnet. Check the cable and both netplan files before anything else."
fi

printf '\n==> 4. Passwordless SSH, this node to the peer\n'
if ssh -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o ConnectTimeout=5 \
       "${CLUSTER_USER}"@"${NODE_PEER_ADDR}" hostname >/dev/null 2>&1; then
    note "OK: SSH to the peer as ${CLUSTER_USER} needs no password."
else
    printf '    %s\n' "SSH to the peer failed or asked for a password." >&2
    printf '    %s\n' "Run 'mkdir -p ~/.ssh && chmod 700 ~/.ssh' on both nodes, then NVIDIA's" >&2
    printf '    %s\n' "discover-sparks script, or ssh-copy-id from each node to the other." >&2
    printf '    %s\n' "Both machines must also share the same username; this node is using ${CLUSTER_USER}." >&2
    exit 1
fi

printf '\n==> 5. Writing the interface environment to %s\n' "$ENV_OUT"
{
    printf '# Generated by prepare-spark-pair.sh. Source it before launching an engine.\n'
    printf '# Every variable here names the same interface, which is what stops a collective\n'
    printf '# from quietly running over the house network instead of the QSFP cable.\n'
    printf 'export VLLM_HOST_IP="%s"\n' "$NODE_SELF_ADDR"
    printf 'export MASTER_ADDR="%s"\n' "${NODE_HEAD_ADDR:-$NODE_SELF_ADDR}"
    printf 'export NCCL_SOCKET_IFNAME="%s"\n' "$CX7_IFNAME"
    printf 'export GLOO_SOCKET_IFNAME="%s"\n' "$CX7_IFNAME"
    printf 'export TP_SOCKET_IFNAME="%s"\n' "$CX7_IFNAME"
    printf 'export UCX_NET_DEVICES="%s"\n' "$CX7_IFNAME"
    printf 'export OMPI_MCA_btl_tcp_if_include="%s"\n' "$CX7_IFNAME"
    if [ -n "$CX7_ROCE_DEV" ]; then
        printf 'export NCCL_IB_HCA="=%s"\n' "$CX7_ROCE_DEV"
    else
        printf '# NCCL_IB_HCA is unset: NCCL will pick a RoCE device itself. Set CX7_ROCE_DEV\n'
        printf '# in .env to the roce* name ibdev2netdev pairs with %s to pin it.\n' "$CX7_IFNAME"
    fi
} > "$ENV_OUT"
sed 's/^/    /' "$ENV_OUT"

cat <<'DONE'

==> This node is ready.
    Run this script on the other node too, then start the Ray head on the node whose
    NODE_SELF_ADDR equals NODE_HEAD_ADDR.

    To see which transport NCCL actually chose once an engine is running, set
    NCCL_DEBUG=TRACE and look for "NET/IB/GDRDMA" (RDMA, efficient) rather than
    "NET/Socket" (plain TCP, not efficient for cross-node tensor parallelism).
DONE