#!/usr/bin/env bash
# Start the SGLang server for Qwen3-235B-A22B on Kunlun P800 (TP=8, 8x XPU).
#
# Adapted from the proven qwen3-8b TP=1 launch on this image. Differences:
#   - new container exposing all 8 /dev/xpu* devices (the 8b container only had xpu0)
#   - --tp-size 8, XPU_VISIBLE_DEVICES=0..7
#   - context-length 8192 (8b used 4096)
#
# All XPU-specific env vars (XSGL_*, XMLIR_*) are taken verbatim from the 8b run.
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")"

# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
# Source config.env BEFORE platform.sh so our CONTAINER_NAME default wins
# over the platform's "sglang-dsv4-flash" default (avoids name collisions
# with the dsv4 experiments, which would docker rm -f this container).
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"

RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
LOG_DIR="${RESULT_ROOT}/logs"
mkdir -p "$LOG_DIR"

log_init "${LOG_DIR}/start_server.log"
log "starting server for ${EXPERIMENT_NAME}"
log "model: ${MODEL_PATH}"
log "tp: ${TP}  xpu: ${XPU_VISIBLE_DEVICES}"
log "port: ${PORT}"
log "container: ${CONTAINER_NAME}"
log "image: ${DOCKER_IMAGE}"

SERVER_LOG="${LOG_DIR}/server.log"

# --- 1. Ensure the container is running with all 8 XPU devices -------------
ensure_container() {
  if docker inspect "$CONTAINER_NAME" >/dev/null 2>&1; then
    log "container ${CONTAINER_NAME} exists; ensuring it is running"
    docker start "$CONTAINER_NAME" >/dev/null 2>&1 || true
  else
    log "creating container ${CONTAINER_NAME} from ${DOCKER_IMAGE}"
    local device_args=""
    local i
    for i in 0 1 2 3 4 5 6 7; do
      device_args="${device_args} --device /dev/xpu${i}:/dev/xpu${i}"
    done
    device_args="${device_args} --device /dev/xpuctrl:/dev/xpuctrl"

    docker run -d \
      --name "$CONTAINER_NAME" \
      --privileged \
      --network host \
      --ipc host \
      --security-opt label=disable \
      --shm-size 1g \
      ${device_args} \
      -v /data1:/data1 \
      -e XPU_VISIBLE_DEVICES="${XPU_VISIBLE_DEVICES}" \
      -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
      -e TZ=Asia/Shanghai \
      --entrypoint /bin/bash \
      "$DOCKER_IMAGE" \
      -lc "source /etc/profile >/dev/null 2>&1 || true; exec sleep infinity"
  fi
}

# --- 2. Launch sglang.launch_server inside the container -------------------
launch_sglang() {
  log "launching sglang launch_server (TP=${TP}) inside ${CONTAINER_NAME}"
  # Kill any prior sglang process inside the container.
  docker exec "$CONTAINER_NAME" bash -lc 'pkill -9 -f sglang.launch_server 2>/dev/null || true' || true

  docker exec -d "$CONTAINER_NAME" bash -lc "
    source /root/miniconda/etc/profile.d/conda.sh
    conda activate python310_torch25_cuda
    export XPU_VISIBLE_DEVICES=${XPU_VISIBLE_DEVICES}
    export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES}
    export XSGL_INTERTYPE_BFP16=1
    export XSGL_FAST_SWIGLU=1
    export XMLIR_FORCE_USE_XPU_GRAPH=1
    export XSGL_FUSE_SPLIT_NORM_ROPE_NEOX=1
    export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=3000000
    export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600
    exec python -m sglang.launch_server \
      --model-path ${MODEL_PATH} \
      --host 0.0.0.0 \
      --port ${PORT} \
      --tp-size ${TP} \
      --context-length ${CONTEXT_LENGTH} \
      --max-running-requests ${MAX_RUNNING_REQUESTS} \
      --attention-backend ${ATTENTION_BACKEND} \
      --trust-remote-code \
      --disable-radix-cache \
      --disable-custom-all-reduce \
      --chunked-prefill-size -1 \
      --page-size ${PAGE_SIZE} \
      --mem-fraction-static ${MEM_FRACTION_STATIC} \
      --max-prefill-tokens ${MAX_PREFILL_TOKENS} \
      --dtype ${DTYPE} \
      --cuda-graph-max-bs ${CUDA_GRAPH_MAX_BS} \
      --cuda-graph-bs ${CUDA_GRAPH_BS} \
      --watchdog-timeout 3000000 \
      > ${SERVER_LOG} 2>&1
  "
}

ensure_container

# --- 1b. Patch qwen3_moe.py bug in the image -------------------------------
# The Kunlun sglang image's qwen3_moe.py:1829 assigns to the read-only
# @property `routed_experts_weights_of_layer` instead of the private
# `_routed_experts_weights_of_layer`, raising
#   AttributeError: can't set attribute 'routed_experts_weights_of_layer'
# for any qwen3 MoE model (e.g. Qwen3-235B-A22B). Qwen3-8B is dense, so the
# proven 8b run never hit this path. Fix: assign to the private attr.
# Idempotent - safe to run on every start.
patch_qwen3_moe() {
  log "applying qwen3_moe.py patch in ${CONTAINER_NAME}"
  docker exec "$CONTAINER_NAME" bash -lc '
    F=/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/models/qwen3_moe.py
    if grep -q "self\.routed_experts_weights_of_layer = LazyValue(" "$F"; then
      cp "$F" "$F.bak.qwen3moe" 2>/dev/null || true
      sed -i "s/self\.routed_experts_weights_of_layer = LazyValue(/self._routed_experts_weights_of_layer = LazyValue(/" "$F"
      echo "PATCHED qwen3_moe.py (routed_experts_weights_of_layer -> _routed_experts_weights_of_layer)"
    else
      echo "qwen3_moe.py already patched"
    fi
  '
}
patch_qwen3_moe

launch_sglang

# 235B weights (~470GB) + cuda-graph capture for 16 batch sizes take a while.
HEALTH_TIMEOUT="${HEALTH_TIMEOUT:-1200}"
log "waiting for health on http://127.0.0.1:${PORT}/health (timeout=${HEALTH_TIMEOUT}s)"
if health_check 127.0.0.1 "$PORT" "$HEALTH_TIMEOUT"; then
  log "sglang server is READY at http://127.0.0.1:${PORT}"
  log "server log: ${SERVER_LOG}  (or: docker logs -f ${CONTAINER_NAME})"
  exit 0
else
  log "ERROR: server not healthy after ${HEALTH_TIMEOUT}s"
  log "----- last 200 lines of server log -----"
  tail -200 "$SERVER_LOG" 2>/dev/null || docker logs --tail 200 "$CONTAINER_NAME" 2>&1 | tail -200
  exit 1
fi