#!/usr/bin/env bash
# Purpose: start the Ray head node of a two-machine vLLM cluster inside the NGC vLLM
#          container, with every network variable pointed at the QSFP interface so the
#          collectives use the direct cable rather than the house network
# Platform: spark (the head node of the pair); works on nvidia with a generic vLLM image
# Minimum memory: 128 GB on this node for the models this lab serves
# Assumes: prepare-spark-pair.sh has passed on both nodes and written cluster-env.sh;
#          vLLM's examples/ray_serving/run_cluster.sh has been downloaded beside this
#          script; docker works without sudo; a .env filled in from env-example.txt.
#          This script does not return: it holds the container open, so run it inside
#          tmux or screen. Closing the shell tears the cluster down.
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
if [ -f "${HERE}/cluster-env.sh" ]; then
    # shellcheck disable=SC1091  # written by prepare-spark-pair.sh
    . "${HERE}/cluster-env.sh"
fi

RUN_CLUSTER="${RUN_CLUSTER:-${HERE}/run-cluster.sh}"
VLLM_IMAGE="${VLLM_IMAGE:-}"
NODE_SELF_ADDR="${NODE_SELF_ADDR:-}"
NODE_HEAD_ADDR="${NODE_HEAD_ADDR:-$NODE_SELF_ADDR}"
CX7_IFNAME="${CX7_IFNAME:-}"
HF_HOME="${HF_HOME:-$HOME/.cache/huggingface}"

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

command -v docker >/dev/null 2>&1 || fail "docker is not on PATH."
docker ps >/dev/null 2>&1 || fail "docker ps failed. Add yourself to the docker group ('sudo usermod -aG docker \$USER' then 'newgrp docker') on both nodes."
[ -f "$RUN_CLUSTER" ] || fail "Ray cluster helper not found at ${RUN_CLUSTER}. Download vLLM's examples/ray_serving/run_cluster.sh to that path, or set RUN_CLUSTER."
[ -n "$VLLM_IMAGE" ] || fail "VLLM_IMAGE is not set. Take the tag for your hardware from the DGX Spark vLLM playbook and put it in .env; both nodes must use the same tag."
[ -n "$NODE_SELF_ADDR" ] || fail "NODE_SELF_ADDR is not set."
[ -n "$CX7_IFNAME" ] || fail "CX7_IFNAME is not set."

if [ "$NODE_SELF_ADDR" != "$NODE_HEAD_ADDR" ]; then
    fail "This node's NODE_SELF_ADDR is not NODE_HEAD_ADDR, so it is a worker. Run start-ray-worker.sh here instead."
fi

PASSTHROUGH=(
    -e "VLLM_HOST_IP=${NODE_SELF_ADDR}"
    -e "MASTER_ADDR=${NODE_HEAD_ADDR}"
    -e "NCCL_SOCKET_IFNAME=${CX7_IFNAME}"
    -e "GLOO_SOCKET_IFNAME=${CX7_IFNAME}"
    -e "TP_SOCKET_IFNAME=${CX7_IFNAME}"
    -e "UCX_NET_DEVICES=${CX7_IFNAME}"
    -e "OMPI_MCA_btl_tcp_if_include=${CX7_IFNAME}"
    -e "RAY_memory_monitor_refresh_ms=0"
)
if [ -n "${NCCL_IB_HCA:-}" ]; then
    PASSTHROUGH+=(-e "NCCL_IB_HCA=${NCCL_IB_HCA}")
fi
if [ -n "${NCCL_DEBUG:-}" ]; then
    PASSTHROUGH+=(-e "NCCL_DEBUG=${NCCL_DEBUG}")
fi

cat <<INFO
==> Ray head node
    image             ${VLLM_IMAGE}
    this node         ${NODE_SELF_ADDR} on ${CX7_IFNAME}
    model cache       ${HF_HOME}
    passthrough       ${#PASSTHROUGH[@]} arguments, all naming ${CX7_IFNAME}

    This shell now holds the cluster open. Leave it running, start the worker on the
    other node, and do everything else from a second terminal. Confirm the pool with:

        docker exec \$(docker ps --format '{{.Names}}' | grep -E '^node-[0-9]+\$') ray status

INFO

exec bash "$RUN_CLUSTER" "$VLLM_IMAGE" "$NODE_HEAD_ADDR" --head "$HF_HOME" "${PASSTHROUGH[@]}"