#!/usr/bin/env bash
# Multi-node SGLang deployment for GLM-5.2-FP8 on 2x RTX 6000D nodes (16 GPUs).
#
# Topology:
#   node0 = pro6000D.1 (gpu1, 174.1.51.5, 10.101.0.11 / 10.102.0.11)
#   node1 = pro6000D.3 (gpu3, 174.1.51.7, 10.101.0.13 / 10.102.0.13)
# Each node has 8x RTX 6000D (85.6GB) = 685GB; GLM-5.2-FP8 is ~700GB, so a
# single node cannot hold the weights -> 2-node TP=16 is mandatory.
#
# SGLang multi-node: every node runs its own `python3 -m sglang.launch_server`
# and they connect via NCCL using --dist-init-addr/--nnodes/--node-rank.
# node0 exposes the HTTP API; node1 is compute-only. The benchmark client only
# talks to node0.

# --- Path bootstrap -------------------------------------------------------
# run_bench.sh sources platform.sh (which sets ROOT_DIR) before this file, but
# the start_* scripts source this file directly. Provide a fallback so
# ${ROOT_DIR} / ${SCRIPT_DIR} expand correctly in both paths.
SCRIPT_DIR="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")" && pwd)}"
ROOT_DIR="${ROOT_DIR:-$(cd "${SCRIPT_DIR}/../../.." 2>/dev/null && pwd)}"
# If the relative climb failed (e.g. dir moved), fall back to the script dir.
[[ -z "$ROOT_DIR" || ! -d "$ROOT_DIR/platforms" ]] && ROOT_DIR="$SCRIPT_DIR"

EXPERIMENT="glm52_pro6000_sglang_multinode_tp16"
MODEL_NAME="GLM-5.2-FP8"
MODEL_PATH="/data/hf_models/GLM-5.2-FP8"
SERVED_MODEL_NAME="GLM-5.2-FP8"

SGLANG_PORT="${SGLANG_PORT:-30031}"

# Python interpreter for orchestration scripts (parse_backend.py, compare.py,
# etc.) and the benchmark client. Defaults to the system python3 if the sglang
# venv does not exist on the host.
VENV_CLIENT="${VENV_CLIENT:-/root/.miniconda3/envs/sglang}"

# Run the benchmark client natively (0) or inside Docker (1).
USE_DOCKER_CLIENT="${USE_DOCKER_CLIENT:-1}"

# All 8 GPUs per node participate in TP=16 (8 per node x 2 nodes).
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"

# Runtime working directory for logs, pid files, and tmp. Defaults to a local
# directory under this experiment so the benchmark is self-contained.
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"

# ---------------------------------------------------------------------------
# Multi-node topology
# ---------------------------------------------------------------------------
# Management network is used for the NCCL bootstrap (dist-init-addr). The data
# plane (NCCL/RoCE) is selected automatically by NCCL inside the container via
# the NCCL_* env vars below.
NODE0_HOST="${NODE0_HOST:-pro6000D.1}"
NODE1_HOST="${NODE1_HOST:-pro6000D.3}"
NODE0_IP="${NODE0_IP:-174.1.51.5}"
NODE1_IP="${NODE1_IP:-174.1.51.7}"
NNODES="${NNODES:-2}"
# Port for NCCL bootstrap (dist-init-addr). Must be free on node0.
DIST_INIT_PORT="${DIST_INIT_PORT:-50000}"
# SSH alias for reaching node1 from the control machine (node0).
NODE1_SSH="${NODE1_SSH:-root@174.1.51.7}"

# ---------------------------------------------------------------------------
# Parallel configuration. TP=16 across 2 nodes (8 GPUs/node).
# Constraint (sglang): tp_size * pp_size must be divisible by nnodes.
#   16 * 1 = 16, 16 % 2 == 0  OK
# ---------------------------------------------------------------------------
TP_SIZE="${TP_SIZE:-16}"
DP_SIZE="${DP_SIZE:-1}"
# PARALLEL_CONFIGS is kept for compatibility with the matrix runner; the single
# multi-node config is "16 1".
declare -a PARALLEL_CONFIGS=(
  "16 1"
)

# ---------------------------------------------------------------------------
# SGLang server settings (aligned with dsv4_pro6000_sglang_tp_dp_matrix).
# ---------------------------------------------------------------------------
MEM_FRACTION_STATIC="${MEM_FRACTION_STATIC:-0.80}"
MOE_RUNNER_BACKEND="${MOE_RUNNER_BACKEND:-auto}"
# TileLang DSA on SM120 currently requires BF16 KV cache. Keep the smoke-test
# context at 128K so the doubled KV-cache element size fits beside FP8 weights.
CONTEXT_LENGTH="${CONTEXT_LENGTH:-131072}"
MAX_RUNNING_REQUESTS="${MAX_RUNNING_REQUESTS:-64}"
# GLM-5.2 is shipped as FP8 weights; tell SGLang explicitly.
QUANTIZATION="${QUANTIZATION:-fp8}"
KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-bfloat16}"

# ---------------------------------------------------------------------------
# NCCL tuning for cross-node RoCE (mandatory, otherwise bandwidth collapses).
# Verified in the bandwidth report: dual-NIC tuned all_reduce = 41.6 GB/s,
# vs 21 GB/s without these knobs.
# ---------------------------------------------------------------------------
export NCCL_IB_HCA="${NCCL_IB_HCA:-mlx5_0,mlx5_3}"
export NCCL_MIN_NCHANNELS="${NCCL_MIN_NCHANNELS:-8}"
export NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-4}"
export NCCL_NET_GDR_LEVEL="${NCCL_NET_GDR_LEVEL:-PHB}"
# Bind the NCCL socket (bootstrap/OOB) to the management NIC eth1 (174.1.51.x).
# eth0/eth3 are the RoCE data-plane NICs; eth1 carries the mgmt IP.
export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-eth1}"
export NCCL_IB_RETRY_CNT="${NCCL_IB_RETRY_CNT:-14}"
export NCCL_IB_TIMEOUT="${NCCL_IB_TIMEOUT:-16}"

# ---------------------------------------------------------------------------
# Deployment switch. 0 = native sglang venv, 1 = Docker.
# ---------------------------------------------------------------------------
USE_DOCKER="${USE_DOCKER:-1}"
DOCKER_IMAGE="${DOCKER_IMAGE:-lmsysorg/sglang:nightly-dev-cu13-20260720-b3570a45}"

# To use ShareGPT, set BENCH_DATASET_NAME=random and DATASET_PATH explicitly.
BENCH_DATASET_NAME="${BENCH_DATASET_NAME:-random}"
DATASET_PATH="${DATASET_PATH:-${ROOT_DIR}/dataset/ShareGPT_V3_unfiltered_cleaned_split.json}"
SGLANG_BENCH_MODULE="${SGLANG_BENCH_MODULE:-sglang.benchmark.serving}"
CACHE_DIR="${CACHE_DIR:-${ROOT_DIR}/sglang_nightly_cu13_cache}"

# Matrix and concurrency rules are defined in matrix.json by default.
MATRIX_FILE="${MATRIX_FILE:-${SCRIPT_DIR:-.}/matrix.json}"
MATRIX_MODE="${MATRIX_MODE:-Y}"

# Sampling density for concurrency.
# 0 = use the default heuristic in generate_scenarios.py (6-8 points).
# 2 = only test the low and high endpoints.
export CONCURRENCY_SAMPLES="${CONCURRENCY_SAMPLES:-2}"

# Per-scenario timeout to avoid hangs (seconds). Multi-node load is slower, so
# the default is raised.
SCENARIO_TIMEOUT_S="${SCENARIO_TIMEOUT_S:-2400}"

# GPU memory sampling interval (seconds).
GPU_MEM_SAMPLE_INTERVAL_S="${GPU_MEM_SAMPLE_INTERVAL_S:-1}"

# Dry-run mode: if 1, only log the server args and scenario plan without
# starting any server or sending requests.
DRY_RUN="${DRY_RUN:-0}"

# Per-config scenario limit for quick smoke tests. 0 = run all generated
# scenarios.
GRID_LIMIT="${GRID_LIMIT:-0}"

# PyTorch CUDA allocator setting for the SGLang server. expandable_segments
# reduces fragmentation from the GiB-scale indexer temporaries that OOM the
# SM120 torch fallback (fp8_paged_mqa_logits_torch_sm120) at ISL >= 4096.
PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:256}"

# Seconds to wait for node1 to come up and register before starting node0.
NODE1_READINESS_WAIT_S="${NODE1_READINESS_WAIT_S:-15}"