#!/usr/bin/env bash
# Kimi-K3 PD 分离部署编排脚本（MoonCake RDMA 传输）。
# 管理 mc-master（元数据服务）+ P 组（prefill）+ D 组（decode）+ router（MiniLB 入口）。
#
# 用法:
#   bash deploy_pd.sh start     # 按顺序启动: mc-master → P 组 → D 组 → router
#   bash deploy_pd.sh stop      # 逆序停止: router → D 组 → P 组 → mc-master
#   bash deploy_pd.sh status    # 查看各组件状态
#   bash deploy_pd.sh restart   # stop + start
#
# 依赖:
#   - 8 台节点镜像 lmsysorg/sglang:kimi-k3、patch、flashkda/mooncake wheel 已就位
#     （/data/flashkda_deploy/wheels/mooncake_transfer_engine_cuda13-0.3.12.post1-*.whl）
#   - 8 台节点可 ssh（NODE_SSH_USER），有 docker 权限
#   - 本脚本从任一可 ssh 全集群的节点执行
#
# 重要:
#   - 必须先启动 P 组再启动 D 组（prefill 需先注册 bootstrap）
#   - 不能设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments
#     （mooncake RDMA 注册 expandable GPU 段报 Bad address，见 GitHub #2511）
#   - P/D 组的容器启停复用 deploy 层（python -m sskj.deploy），profile 为单一事实来源
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"

MASTER_IP="174.1.60.1"
ROUTER_IP="174.1.60.5"
IMAGE="lmsysorg/sglang:kimi-k3"
MC_MASTER_CONTAINER="mc-master"
ROUTER_CONTAINER="router"
MODEL_PATH="/data/hf_models/Kimi-K3"

P_PROFILE="${ROOT_DIR}/deploy/profiles/pro6000/kimi3_pro6000_pd_prefill.env"
D_PROFILE="${ROOT_DIR}/deploy/profiles/pro6000/kimi3_pro6000_pd_decode.env"

P_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
D_HOSTS=(174.1.60.5 174.1.60.6 174.1.60.7 174.1.60.8)

log() { echo "[$(date +%H:%M:%S)] $*"; }

ssh_node() { ssh -o ConnectTimeout=15 -o StrictHostKeyChecking=no "$@"; }

deploy_util() {
  local profile="$1" container="$2" cmd="$3"
  PYTHONPATH="${ROOT_DIR}/src" python3 -m sskj.deploy "$cmd" \
    --profile "$profile" \
    --tp 32 --dp 1 --port 30000 --model-path "$MODEL_PATH" \
    --container-name "$container"
}

start_mc_master() {
  log "starting mc-master on ${MASTER_IP}"
  ssh_node "$MASTER_IP" "docker rm -f ${MC_MASTER_CONTAINER} >/dev/null 2>&1 || true; docker run -d --name ${MC_MASTER_CONTAINER} --network host ${IMAGE} mooncake_master" >/dev/null
  sleep 3
  if ssh_node "$MASTER_IP" "ss -tlnp 2>/dev/null | grep -q :50051"; then
    log "mc-master is up (50051)"
  else
    log "WARN: mc-master 50051 not listening"
  fi
}

start_p() {
  log "starting P 组 (prefill) via deploy layer"
  deploy_util "$P_PROFILE" "kimi3_pro6000_pd_prefill" start
}

start_d() {
  log "starting D 组 (decode) via deploy layer"
  deploy_util "$D_PROFILE" "kimi3_pro6000_pd_decode" start
}

start_router() {
  log "starting router (MiniLB) on ${ROUTER_IP}"
  ssh_node "$ROUTER_IP" "docker rm -f ${ROUTER_CONTAINER} >/dev/null 2>&1 || true; docker run -d --name ${ROUTER_CONTAINER} --network host ${IMAGE} python3 -m sglang_router.launch_router --pd-disaggregation --mini-lb --prefill http://${MASTER_IP}:30000 28800 --decode http://${ROUTER_IP}:30000 --host 0.0.0.0 --port 31000" >/dev/null
  sleep 3
  if ssh_node "$ROUTER_IP" "ss -tlnp 2>/dev/null | grep -q :31000"; then
    log "router is up (31000)"
  else
    log "WARN: router 31000 not listening"
  fi
}

stop_all() {
  log "stopping router (${ROUTER_IP})"
  ssh_node "$ROUTER_IP" "docker rm -f ${ROUTER_CONTAINER} >/dev/null 2>&1 || true" || true
  log "stopping D 组 via deploy layer"
  deploy_util "$D_PROFILE" "kimi3_pro6000_pd_decode" stop || true
  log "stopping P 组 via deploy layer"
  deploy_util "$P_PROFILE" "kimi3_pro6000_pd_prefill" stop || true
  log "stopping mc-master (${MASTER_IP})"
  ssh_node "$MASTER_IP" "docker rm -f ${MC_MASTER_CONTAINER} >/dev/null 2>&1 || true" || true
  log "all stopped"
}

status_all() {
  echo "=== mc-master ==="
  ssh_node "$MASTER_IP" "docker ps --format '{{.Names}} {{.Status}}' | grep ${MC_MASTER_CONTAINER} || echo stopped"
  echo "=== P 组 ==="
  deploy_util "$P_PROFILE" "kimi3_pro6000_pd_prefill" status || true
  echo "=== D 组 ==="
  deploy_util "$D_PROFILE" "kimi3_pro6000_pd_decode" status || true
  echo "=== router ==="
  ssh_node "$ROUTER_IP" "docker ps --format '{{.Names}} {{.Status}}' | grep ${ROUTER_CONTAINER} || echo stopped"
}

case "${1:-}" in
  start)   start_mc_master; start_p; sleep 30; start_d; sleep 30; start_router ;;
  stop)    stop_all ;;
  status)  status_all ;;
  restart) stop_all; sleep 5; start_mc_master; start_p; sleep 30; start_d; sleep 30; start_router ;;
  *) echo "用法: $0 {start|stop|status|restart}"; exit 1 ;;
esac
echo "done."