#!/usr/bin/env bash
# Deploy Kimi-K3 standard PD on 601-608. Speculative decoding stays disabled.
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"

ACTION="${1:-status}"
RUN_ID="${RUN_ID:-kimi3-pd-standard-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
SUDO_PASSWORD_FILE="${SUDO_PASSWORD_FILE:-/data/hzy/.sudo_password}"

P_CONTAINER_PREFIX="${EXPERIMENT}_prefill"
D_CONTAINER_PREFIX="${EXPERIMENT}_decode"
MC_CONTAINER="${EXPERIMENT}_mc_master"
ROUTER_CONTAINER="${EXPERIMENT}_router"

mkdir -p "${RESULT_ROOT}"/{commands,logs,metadata,bench}

log() {
  printf '[%s] %s\n' "$(date '+%F %T')" "$*" | tee -a "${RESULT_ROOT}/orchestrator.log"
}

require_password() {
  if [[ -z "${SUDO_PASSWORD:-}" ]]; then
    [[ -r "${SUDO_PASSWORD_FILE}" ]] || {
      echo "ERROR: set SUDO_PASSWORD or create ${SUDO_PASSWORD_FILE}" >&2
      exit 2
    }
    IFS= read -r SUDO_PASSWORD <"${SUDO_PASSWORD_FILE}"
  fi
}

is_local() {
  [[ "$1" == "${P_HEAD}" ]]
}

remote() {
  local host="$1"
  shift
  if is_local "$host"; then
    "$@"
    return
  fi
  local command
  printf -v command '%q ' "$@"
  ssh "${SSH_OPTIONS[@]}" "${SSH_USER}@${host}" "${command}"
}

sudo_host() {
  local host="$1"
  shift
  require_password
  if is_local "$host"; then
    printf '%s\n' "${SUDO_PASSWORD}" | sudo -S -p '' -- "$@"
    return
  fi
  local command
  printf -v command '%q ' "$@"
  printf '%s\n' "${SUDO_PASSWORD}" | ssh "${SSH_OPTIONS[@]}" \
    "${SSH_USER}@${host}" "sudo -S -p '' -- ${command}"
}

node_suffix() {
  printf '%s' "${1##*.}"
}

container_name() {
  local group="$1" rank="$2"
  if [[ "$group" == "p" ]]; then
    printf '%s_node%s' "${P_CONTAINER_PREFIX}" "$rank"
  else
    printf '%s_node%s' "${D_CONTAINER_PREFIX}" "$rank"
  fi
}

build_image_on_host() {
  local host="$1" suffix build_dir
  suffix="$(node_suffix "$host")"
  build_dir="/tmp/${EXPERIMENT}_build"
  log "staging PD image build context on node ${suffix}"

  if is_local "$host"; then
    rm -rf "$build_dir"
    mkdir -p "$build_dir/patches"
    cp "${SCRIPT_DIR}/Dockerfile.pd" "$build_dir/Dockerfile"
    cp "${SCRIPT_DIR}/patches/31869.patch" "$build_dir/patches/31869.patch"
    cp "${SCRIPT_DIR}/patches/32797.patch" "$build_dir/patches/32797.patch"
    cp "${MOONCAKE_WHEEL}" "$build_dir/mooncake.whl"
  else
    ssh "${SSH_OPTIONS[@]}" "${SSH_USER}@${host}" \
      "rm -rf ${build_dir} && mkdir -p ${build_dir}/patches"
    scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/Dockerfile.pd" \
      "${SSH_USER}@${host}:${build_dir}/Dockerfile"
    scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/patches/31869.patch" \
      "${SSH_USER}@${host}:${build_dir}/patches/31869.patch"
    scp "${SSH_OPTIONS[@]}" "${SCRIPT_DIR}/patches/32797.patch" \
      "${SSH_USER}@${host}:${build_dir}/patches/32797.patch"
    scp "${SSH_OPTIONS[@]}" "${MOONCAKE_WHEEL}" \
      "${SSH_USER}@${host}:${build_dir}/mooncake.whl"
  fi

  sudo_host "$host" docker build \
    --build-arg "BASE_IMAGE=${BASE_IMAGE}" \
    -t "${PD_IMAGE}" "$build_dir" \
    >"${RESULT_ROOT}/logs/build_node${suffix}.log" 2>&1
  log "PD image built on node ${suffix}"
}

build_image_all() {
  local host
  local -a build_nodes
  read -r -a build_nodes <<<"${BUILD_NODES_OVERRIDE:-${ALL_NODES[*]}}"
  for host in "${build_nodes[@]}"; do
    build_image_on_host "$host"
  done
}

preflight() {
  require_password
  local host suffix
  : >"${RESULT_ROOT}/metadata/preflight.tsv"
  printf 'node\timage\tmodel\trdma\tgpu_processes\tports\n' \
    >>"${RESULT_ROOT}/metadata/preflight.tsv"
  for host in "${ALL_NODES[@]}"; do
    suffix="$(node_suffix "$host")"
    log "preflight node ${suffix}"
    sudo_host "$host" docker image inspect "$PD_IMAGE" >/dev/null
    remote "$host" test -d "$MODEL_PATH"
    remote "$host" test -e /dev/infiniband/uverbs0
    local gpu_pids ports
    gpu_pids="$(remote "$host" nvidia-smi --query-compute-apps=pid --format=csv,noheader | xargs || true)"
    ports="$(remote "$host" sh -c "ss -lnt | grep -E ':(20000|28800|30000|31000|50051)[[:space:]]' || true")"
    if [[ -n "$gpu_pids" ]]; then
      echo "ERROR: node ${suffix} still has GPU processes: ${gpu_pids}" >&2
      return 1
    fi
    if [[ -n "$ports" ]]; then
      echo "ERROR: node ${suffix} has occupied PD ports: ${ports}" >&2
      return 1
    fi
    printf '%s\tok\tok\tok\t0\tfree\n' "$suffix" \
      >>"${RESULT_ROOT}/metadata/preflight.tsv"
  done

  local hash_cmd
  hash_cmd="sha256sum /sgl-workspace/sglang/python/sglang/srt/disaggregation/prefill.py /sgl-workspace/sglang/python/sglang/srt/disaggregation/decode.py /sgl-workspace/sglang/python/sglang/srt/disaggregation/utils.py /sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py /sgl-workspace/sglang/python/sglang/srt/managers/scheduler.py; python3 -m pip list | grep -E 'sglang|flashinfer|mooncake-transfer-engine'"
  sudo_host "$P_HEAD" docker run --rm --entrypoint bash "$PD_IMAGE" -lc "$hash_cmd" \
    >"${RESULT_ROOT}/metadata/pd_image_versions.txt"
  grep -q 'mooncake-transfer-engine-cuda13.*0.3.12.post1' \
    "${RESULT_ROOT}/metadata/pd_image_versions.txt"
  log "preflight passed on all 8 nodes"
}

common_docker_args() {
  local host="$1"
  local suffix
  suffix="$(node_suffix "$host")"
  COMMON_DOCKER_ARGS=(
    --gpus all --network host --ipc=host --ulimit memlock=-1
    --device /dev/infiniband --shm-size 32g --entrypoint bash
    -v "${MODEL_PATH}:${MODEL_PATH}:ro"
    -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
    -e "SGLANG_HOST_IP=174.1.60.${suffix}"
    -e GLOO_SOCKET_IFNAME=bond0
    -e NCCL_SOCKET_IFNAME=bond1
    -e "NCCL_IB_HCA=${NCCL_IB_HCAS}"
    -e NCCL_IB_GID_INDEX=3
    -e NCCL_IB_TIMEOUT=22
    -e NCCL_IB_RETRY_CNT=7
    -e NCCL_CUMEM_ENABLE=1
    -e "MOONCAKE_MASTER=${MC_MASTER}"
    -e MOONCAKE_PROTOCOL=rdma
    -e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0
    -e SGLANG_MOE_FUSED_GATE_RADIX=1
    -e FLASHINFER_DISABLE_JIT=1
    -e FLASHINFER_DISABLE_VERSION_CHECK=1
  )
}

start_node() {
  local group="$1" rank="$2" host name bootstrap
  local -a nodes launch_args cmd
  if [[ "$group" == "p" ]]; then
    nodes=("${P_NODES[@]}")
    host="${nodes[$rank]}"
    name="$(container_name p "$rank")"
    launch_args=(
      --model-path "$MODEL_PATH" --served-model-name "$SERVED_MODEL_NAME"
      --tp-size "$P_TP" --pp-size "$P_PP" --ep-size "$P_EP"
      --nnodes 4 --node-rank "$rank" --dist-init-addr "${P_HEAD}:${P_DIST_PORT}"
      --trust-remote-code --moe-runner-backend flashinfer_mxfp4
      --chunked-prefill-size "$P_CHUNKED_PREFILL_SIZE"
      --page-size "$P_PAGE_SIZE"
      --mem-fraction-static "$P_MEM_FRACTION_STATIC"
      --cuda-graph-max-bs-decode 16
      --mamba-radix-cache-strategy extra_buffer_lazy
      --disable-radix-cache --dist-timeout 3600
      --mamba-full-memory-ratio "$P_MAMBA_FULL_MEMORY_RATIO"
      --disaggregation-transfer-backend mooncake
      --disaggregation-bootstrap-port "$P_BOOTSTRAP_PORT"
      --disaggregation-ib-device "$DISAGG_IB_DEVICES"
      --disaggregation-mode prefill
      --host 0.0.0.0 --port "$P_PORT"
    )
  else
    nodes=("${D_NODES[@]}")
    host="${nodes[$rank]}"
    name="$(container_name d "$rank")"
    launch_args=(
      --model-path "$MODEL_PATH" --served-model-name "$SERVED_MODEL_NAME"
      --tp-size "$D_TP" --pp-size "$D_PP" --ep-size "$D_EP"
      --nnodes 4 --node-rank "$rank" --dist-init-addr "${D_HEAD}:${D_DIST_PORT}"
      --trust-remote-code --moe-runner-backend marlin
      --mem-fraction-static "$D_MEM_FRACTION_STATIC"
      --cuda-graph-max-bs-decode 16
      --mamba-radix-cache-strategy extra_buffer_lazy
      --disable-radix-cache --dist-timeout 3600
      --mamba-full-memory-ratio "$D_MAMBA_FULL_MEMORY_RATIO"
      --page-size "$D_PAGE_SIZE"
      --disaggregation-transfer-backend mooncake
      --disaggregation-bootstrap-port "$P_BOOTSTRAP_PORT"
      --disaggregation-ib-device "$DISAGG_IB_DEVICES"
      --disaggregation-mode decode
      --host 0.0.0.0 --port "$D_PORT"
    )
  fi

  printf -v bootstrap '%q ' python3 -m sglang.launch_server "${launch_args[@]}"
  bootstrap="exec ${bootstrap}"
  common_docker_args "$host"
  cmd=(docker run -d --name "$name" "${COMMON_DOCKER_ARGS[@]}" "$PD_IMAGE" -lc "$bootstrap")
  printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/commands/${group}_node${rank}.cmd.txt"
  printf '\n' >>"${RESULT_ROOT}/commands/${group}_node${rank}.cmd.txt"
  sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
  sudo_host "$host" "${cmd[@]}" >/dev/null
  log "started ${group} rank ${rank} on node $(node_suffix "$host")"
}

wait_health() {
  local label="$1" host="$2" port="$3" elapsed
  for ((elapsed = 1; elapsed <= HEALTH_WAIT_S; elapsed++)); do
    if curl --fail --silent --max-time 5 "http://${host}:${port}/health" >/dev/null 2>&1; then
      log "${label} healthy after ${elapsed}s"
      return 0
    fi
    if (( elapsed % 30 == 0 )); then
      log "waiting for ${label}, elapsed=${elapsed}s"
      collect_group_logs "${label}_starting"
      if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|AssertionError|RuntimeError' \
          "${RESULT_ROOT}/logs/${label}_starting_"*.log 2>/dev/null; then
        log "ERROR: fatal pattern found while starting ${label}"
        return 1
      fi
    fi
    sleep 1
  done
  return 1
}

start_group() {
  local group="$1" rank
  for rank in 1 2 3; do
    start_node "$group" "$rank"
  done
  sleep 5
  start_node "$group" 0
  if [[ "$group" == "p" ]]; then
    wait_health p "$P_HEAD" "$P_PORT"
  else
    wait_health d "$D_HEAD" "$D_PORT"
  fi
  collect_group_logs "${group}_healthy"
}

start_mc_master() {
  sudo_host "$P_HEAD" docker rm -f "$MC_CONTAINER" >/dev/null 2>&1 || true
  sudo_host "$P_HEAD" docker run -d --name "$MC_CONTAINER" --network host --gpus all \
    --entrypoint /usr/local/bin/mooncake_master "$PD_IMAGE" >/dev/null
  for _ in $(seq 1 30); do
    remote "$P_HEAD" sh -c "ss -lnt | grep -q ':50051 '" && {
      log "Mooncake master listening on 50051"
      return 0
    }
    sleep 1
  done
  return 1
}

start_router() {
  local -a cmd=(
    docker run -d --name "$ROUTER_CONTAINER" --network host
    --entrypoint python3 "$PD_IMAGE" -m sglang_router.launch_router
    --pd-disaggregation --mini-lb
    --prefill "http://${P_HEAD}:${P_PORT}" "$P_BOOTSTRAP_PORT"
    --decode "http://${D_HEAD}:${D_PORT}"
    --host 0.0.0.0 --port "$ROUTER_PORT"
  )
  printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/commands/router.cmd.txt"
  printf '\n' >>"${RESULT_ROOT}/commands/router.cmd.txt"
  sudo_host "$ROUTER_HOST" docker rm -f "$ROUTER_CONTAINER" >/dev/null 2>&1 || true
  sudo_host "$ROUTER_HOST" "${cmd[@]}" >/dev/null
  wait_health router "$ROUTER_HOST" "$ROUTER_PORT"
}

collect_group_logs() {
  local label="$1" group rank host name
  for group in p d; do
    for rank in 0 1 2 3; do
      if [[ "$group" == "p" ]]; then
        host="${P_NODES[$rank]}"
      else
        host="${D_NODES[$rank]}"
      fi
      name="$(container_name "$group" "$rank")"
      sudo_host "$host" docker logs "$name" \
        >"${RESULT_ROOT}/logs/${label}_${group}_node${rank}.log" 2>&1 || true
    done
  done
  sudo_host "$P_HEAD" docker logs "$MC_CONTAINER" \
    >"${RESULT_ROOT}/logs/${label}_mc_master.log" 2>&1 || true
  sudo_host "$ROUTER_HOST" docker logs "$ROUTER_CONTAINER" \
    >"${RESULT_ROOT}/logs/${label}_router.log" 2>&1 || true
}

status_all() {
  local host
  for host in "${ALL_NODES[@]}"; do
    echo "===== node $(node_suffix "$host") ====="
    sudo_host "$host" docker ps --format '{{.Names}}|{{.Image}}|{{.Status}}' \
      --filter "name=${EXPERIMENT}"
    remote "$host" nvidia-smi --query-gpu=index,memory.used,utilization.gpu \
      --format=csv,noheader
  done
  for endpoint in \
    "p=http://${P_HEAD}:${P_PORT}/health" \
    "d=http://${D_HEAD}:${D_PORT}/health" \
    "router=http://${ROUTER_HOST}:${ROUTER_PORT}/health"; do
    local name="${endpoint%%=*}" url="${endpoint#*=}"
    if curl --fail --silent --max-time 5 "$url" >/dev/null 2>&1; then
      echo "${name}=healthy"
    else
      echo "${name}=down"
    fi
  done
}

stop_all() {
  local rank host name
  sudo_host "$ROUTER_HOST" docker rm -f "$ROUTER_CONTAINER" >/dev/null 2>&1 || true
  for rank in 0 1 2 3; do
    host="${D_NODES[$rank]}"
    name="$(container_name d "$rank")"
    sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
  done
  for rank in 0 1 2 3; do
    host="${P_NODES[$rank]}"
    name="$(container_name p "$rank")"
    sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
  done
  sudo_host "$P_HEAD" docker rm -f "$MC_CONTAINER" >/dev/null 2>&1 || true
  log "standard PD service stopped"
}

start_all() {
  preflight
  start_mc_master
  start_group p
  start_group d
  start_router
  collect_group_logs ready
  status_all | tee "${RESULT_ROOT}/metadata/status_ready.txt"
  log "standard PD service ready at http://${ROUTER_HOST}:${ROUTER_PORT}"
}

case "$ACTION" in
  build-image) build_image_all ;;
  preflight) preflight ;;
  start) start_all ;;
  stop) stop_all ;;
  restart) stop_all; start_all ;;
  status) status_all ;;
  logs) collect_group_logs manual ;;
  *)
    echo "Usage: $0 {build-image|preflight|start|stop|restart|status|logs}" >&2
    exit 2
    ;;
esac