#!/usr/bin/env bash
# Kimi-K3 DP=2 部署编排脚本（方案 B：两个独立 TP32×EP32 实例 + router 负载均衡）。
#
# 架构:
#   实例 A: 174.1.60.5~8（TP32×EP32，现有 kimi3 profile）
#   实例 B: 174.1.60.1~4（TP32×EP32，instB profile）
#   router: sglang_router --worker-urls 指向两个实例的 head（.5:30000 与 .1:30000）
#   对外入口: .5:32000（router），实例间零通信。
#
# 用法:
#   bash deploy_dp2.sh start     # 起 实例 A → 实例 B → router
#   bash deploy_dp2.sh stop      # 停 router → 实例 B → 实例 A
#   bash deploy_dp2.sh status    # 查看状态
#   bash deploy_dp2.sh restart
#
# 依赖:
#   - 8 节点镜像 lmsysorg/sglang:kimi-k3、patch_k3_sm120.py 已就位
#   - .5 与 .1 上都有 sskj 仓库（deploy 层在各自 master 本地执行）
#   - 各节点 docker 权限、.5 可 ssh .1
set -Eeuo pipefail

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

IMAGE="lmsysorg/sglang:kimi-k3"
REPO_PATH="/data/yy/sskj"
MODEL_PATH="/data/hf_models/Kimi-K3"

A_HEAD="174.1.60.5"
B_HEAD="174.1.60.1"
ROUTER_IP="174.1.60.5"
ROUTER_CONTAINER="router_dp2"
ROUTER_PORT="32000"

PROFILE_A="pro6000/kimi3_pro6000_sglang_tp32ep32"
PROFILE_B="pro6000/kimi3_pro6000_sglang_tp32ep32_instB"

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

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

run_deploy_on() {
  local host="$1" profile="$2" container="$3" cmd="$4"
  ssh_node "$host" "cd ${REPO_PATH} && PYTHONPATH=src python3 -m sskj.deploy ${cmd} --profile ${profile} --tp 32 --dp 1 --port 30000 --model-path ${MODEL_PATH} --container-name ${container}"
}

start_instance_a() {
  log "starting 实例 A (prefill .5-.8) on ${A_HEAD}"
  run_deploy_on "$A_HEAD" "$PROFILE_A" "kimi3_pro6000_sglang_tp32ep32" start
}

start_instance_b() {
  log "starting 实例 B (.1-.4) on ${B_HEAD}"
  run_deploy_on "$B_HEAD" "$PROFILE_B" "kimi3_pro6000_sglang_tp32ep32_instB" start
}

start_router() {
  log "starting router (--worker-urls) on ${ROUTER_IP}:${ROUTER_PORT}"
  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 --worker-urls http://${B_HEAD}:30000 http://${A_HEAD}:30000 --host 0.0.0.0 --port ${ROUTER_PORT}" >/dev/null
  sleep 5
  if ssh_node "$ROUTER_IP" "ss -tlnp 2>/dev/null | grep -q :${ROUTER_PORT}"; then
    log "router is up (${ROUTER_PORT})"
  else
    log "WARN: router ${ROUTER_PORT} not listening"
  fi
}

stop_all() {
  log "stopping router"
  ssh_node "$ROUTER_IP" "docker rm -f ${ROUTER_CONTAINER} >/dev/null 2>&1 || true" || true
  log "stopping 实例 B"
  run_deploy_on "$B_HEAD" "$PROFILE_B" "kimi3_pro6000_sglang_tp32ep32_instB" stop || true
  log "stopping 实例 A"
  run_deploy_on "$A_HEAD" "$PROFILE_A" "kimi3_pro6000_sglang_tp32ep32" stop || true
  log "all stopped"
}

status_all() {
  echo "=== 实例 A (.5-.8) ==="
  run_deploy_on "$A_HEAD" "$PROFILE_A" "kimi3_pro6000_sglang_tp32ep32" status || true
  echo "=== 实例 B (.1-.4) ==="
  run_deploy_on "$B_HEAD" "$PROFILE_B" "kimi3_pro6000_sglang_tp32ep32_instB" status || true
  echo "=== router ==="
  ssh_node "$ROUTER_IP" "docker ps --format '{{.Names}} {{.Status}}' | grep ${ROUTER_CONTAINER} || echo stopped"
}

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