#!/usr/bin/env bash
# Purpose: join the second DGX Spark to the Ray cluster started by start-ray-head.sh,
#          inside the same container image and with the same interface variables, so both
#          ranks agree about which cable carries the collectives
# Platform: spark (the worker 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: start-ray-head.sh is already running on the head node and still has its shell
#          open; prepare-spark-pair.sh has passed here; the SAME container image tag as
#          the head; vLLM's run_cluster.sh downloaded beside this script; docker without
#          sudo. This script does not return, so run it inside tmux or screen.
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:-}"
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 on this node too."
[ -f "$RUN_CLUSTER" ] || fail "Ray cluster helper not found at ${RUN_CLUSTER}. Download the same run_cluster.sh the head node uses."
[ -n "$VLLM_IMAGE" ] || fail "VLLM_IMAGE is not set. It must be the same tag the head node pulled."
[ -n "$NODE_SELF_ADDR" ] || fail "NODE_SELF_ADDR is not set. On a worker it is THIS machine's address, not the head's."
[ -n "$NODE_HEAD_ADDR" ] || fail "NODE_HEAD_ADDR is not set. It is the head node's address on the QSFP link."
[ -n "$CX7_IFNAME" ] || fail "CX7_IFNAME is not set."

if [ "$NODE_SELF_ADDR" = "$NODE_HEAD_ADDR" ]; then
    fail "NODE_SELF_ADDR equals NODE_HEAD_ADDR, so this is the head node. Run start-ray-head.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 worker node
    image             ${VLLM_IMAGE}
    this node         ${NODE_SELF_ADDR} on ${CX7_IFNAME}
    joining head at   ${NODE_HEAD_ADDR}
    model cache       ${HF_HOME}

    VLLM_HOST_IP is this node's own address, not the head's: vLLM's documentation says it
    is unique per worker, and getting it wrong is the commonest reason a worker appears to
    join and then produces nothing.

    Leave this shell running. If Ray reports "Unable to connect to GCS", the head is not
    running, the address is wrong, the two nodes have different Ray versions, or a
    firewall is in the way.

INFO

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