#!/usr/bin/env bash
# Start P800 SGLang INT8 server in Docker for profiling.
set -Eeuo pipefail

TP="${1:-8}"
DP="${2:-1}"
if [[ -z "$TP" || -z "$DP" ]]; then
  echo "Usage: $0 <TP> <DP>"
  exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"
source "${SCRIPT_DIR}/config.env"

PORT="${SGLANG_PORT:-30015}"
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"
mkdir -p "${RUNTIME_BASE}/logs"
SERVER_LOG="${RUNTIME_BASE}/logs/${EXPERIMENT}_sglang_docker_tp${TP}_dp${DP}_$(date +%Y%m%d_%H%M%S).log"

mkdir -p "${PROFILE_OUTPUT_DIR}"
mkdir -p "${BENCH_OUTPUT_DIR}"

log "starting P800 SGLang INT8 server (tp=${TP}, dp=${DP})"
log "model: ${MODEL_PATH}"
log "port: ${PORT}"
log "container: ${CONTAINER_NAME}"

docker rm -f "$CONTAINER_NAME" 2>/dev/null || true

# Build device args
device_args=""
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"

# Environment variables
env_args=(
  -e XPU_VISIBLE_DEVICES="${XPU_VISIBLE_DEVICES}"
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"
  -e CUDA_DEVICE_ORDER=OAM_ID
  -e SGLANG_USE_TRANSFORMERS_V5_TOKENIZER=1
  -e XMLIR_FORCE_USE_XPU_GRAPH=1
  -e SGLANG_DSV4_MODE=2604
  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
  -e SGLANG_NSA_DUAL_STREAM=true
  -e SGLANG_NSA_QUANT_WQ_B_WK=false
  -e SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
  -e SGLANG_OPT_DEEPGEMM_HC_PRENORM=false
  -e SGLANG_OPT_USE_TILELANG_MHC_PRE=1
  -e SGLANG_OPT_USE_TILELANG_MHC_POST=1
  -e SGLANG_CLEAN_REQUEST_WHEN_RETRACT=1
  -e SGLANG_SET_CPU_AFFINITY=1
  -e SGLANG_OPT_USE_KLX_TOPK_KERNEL=1
  -e XSGL_INTERTYPE_BFP16=1
  -e ENABLE_FAST_BFP16_ATTN=1
  -e XSGL_USE_DEEP_GEMM_BMM=1
  -e XSGL_XDNN_QUANT=1
  -e XSGL_FUSE_RMS_NORM_QUANT=1
  -e XSGL_TRANSPOSE_MATMUL_WEIGHT=1
  -e XINFER_QUANT_SDNN=1
  -e XSGL_USE_MOE_SIGMOID_GROUP_TOPK_NORM=1
  -e XSGL_EARLY_FIRST_TOKEN=1
  -e XSGL_ENABLE_TGEMM_FP16=1
  -e SGLANG_ENABLE_SPEC_V2=True
  -e SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1
  -e PYTHONDONTWRITEBYTECODE=1
  -e XTORCH_OPS_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/xtorch_ops
  -e XPU_RUNTIME_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/xcudart/lib
  -e BKCL_TREE_THRESHOLD=1048576
  -e CUDA_ENABLE_P2P_NO_UVA=1
  -e NCCL_IB_GID_INDEX=3
  -e IS_DSV4=1
  -e MC_CUSTOM_TOPO_JSON=/workspace/nic_priority_matrix_test.json
  -e SGLANG_DSV4_FP4_EXPERTS=false
  -e SGLANG_APPLY_CONFIG_BACKUP=auto
  -e BKCL_ENABLE_XDR=1
  -e BKCL_RDMA_NICS=eth1,eth1,eth3,eth3,eth5,eth5,eth7,eth7
  -e BKCL_RDMA_VERBS=1
  -e XSGL_INT8_LM_HEAD=1
  -e SGLANG_P800_ALL_GATHER_FALLBACK=0
  -e XPU_ENABLE_PROFILER_TRACING=1
)

launch_args="--host 0.0.0.0 --port ${PORT} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --page-size 64 --mem-fraction-static ${MEM_FRACTION_STATIC} --tensor-parallel-size ${TP} --disable-shared-experts-fusion --quantization w8a8_int8 --kv-cache-dtype float16 --disable-piecewise-cuda-graph --cuda-graph-max-bs 32 --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging --context-length ${CONTEXT_LENGTH}"

if [[ "$DP" -gt 1 ]]; then
  launch_args="${launch_args} --dp-size ${DP}"
fi

# Write the server entrypoint script to a file (avoids shell quoting issues)
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
ENTRYPOINT_FILE="${RUNTIME_BASE}/entrypoint_${TIMESTAMP}.sh"
cat > "$ENTRYPOINT_FILE" << 'ENTRYEOF'
#!/usr/bin/env bash
set -Eeuo pipefail
cd /workspace
find /root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
/root/miniconda/envs/python310_torch25_cuda/bin/pip install --upgrade safetensors -q 2>/dev/null
/root/miniconda/envs/python310_torch25_cuda/bin/python -c "import torch; torch.float8_e8m0fnu = torch.uint8; import runpy, sys; sys.argv[0] = 'sglang.launch_server'; runpy.run_module('sglang.launch_server', run_name='__main__')" LAUNCH_ARGS_PLACEHOLDER
ENTRYEOF

# Insert the launch args
sed -i "s|LAUNCH_ARGS_PLACEHOLDER|${launch_args}|g" "$ENTRYPOINT_FILE"
chmod +x "$ENTRYPOINT_FILE"

# Patch mounts
patch_mounts=(
  -v "${PATCH_ROOT}/nic_priority_matrix_test.json:/workspace/nic_priority_matrix_test.json:ro"
  -v "${PATCH_ROOT}/dummy_sharegpt.json:/workspace/dummy_sharegpt.json:ro"
  -v "${PATCH_ROOT}/sitecustomize_xpu.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sitecustomize.py:ro"
  -v "${PATCH_ROOT}/hf_transformers_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/utils/hf_transformers_utils.py:ro"
  -v "${PATCH_ROOT}/fp8_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/quantization/fp8_utils.py:ro"
  -v "${PATCH_ROOT}/config_backup_small_w8a8_int8.json:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/configs/config_backup_small.json:ro"
  -v "${PATCH_ROOT}/parallel_state.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/distributed/parallel_state.py:ro"
  -v "${PATCH_ROOT}/vocab_parallel_embedding.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/vocab_parallel_embedding.py:ro"
  -v "${ENTRYPOINT_FILE}:/workspace/entrypoint.sh:ro"
)
BENCH_SERVING_PATCH="${SCRIPT_DIR}/patches/bench_serving.py"
if [[ -f "${BENCH_SERVING_PATCH}" ]]; then
  patch_mounts+=(-v "${BENCH_SERVING_PATCH}:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/bench_serving.py:ro")
fi
if [[ -f "${DATASET_PATH}" ]]; then
  patch_mounts+=(-v "${DATASET_PATH}:${CONTAINER_DATASET_PATH}:ro")
fi
patch_mounts+=(-v "${PROFILE_OUTPUT_DIR}:/workspace/profile_output:rw")
patch_mounts+=(-v "${BENCH_OUTPUT_DIR}:/workspace/bench_output:rw")

echo "=== Starting P800 SGLang INT8 server in Docker (TP=${TP}, DP=${DP}) ==="
echo "Image: ${DOCKER_IMAGE}"
echo "Container name: ${CONTAINER_NAME}"
echo "Host port: ${PORT}"

docker run -d \
  --name "${CONTAINER_NAME}" \
  --privileged \
  --network host \
  --ipc host \
  ${device_args} \
  -v "${MODEL_PATH}:/models:ro" \
  -v "${MODEL_PATH}:${MODEL_PATH}:ro" \
  "${patch_mounts[@]}" \
  "${env_args[@]}" \
  "${DOCKER_IMAGE}" \
  bash /workspace/entrypoint.sh \
  >> "${SERVER_LOG}" 2>&1

log "container ${CONTAINER_NAME} started, waiting for health"
healthy=0
for ((i = 1; i <= 600; i++)); do
  if curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1; then
    healthy=1
    break
  fi
  if [[ "$(docker inspect -f '{{.State.Running}}' "${CONTAINER_NAME}" 2>/dev/null)" != "true" ]]; then
    log "ERROR: container exited during startup:"
    docker logs --tail 100 "${CONTAINER_NAME}" 2>&1 || true
    exit 1
  fi
  sleep 1
done
if [[ "$healthy" == "1" ]]; then
  log "container ${CONTAINER_NAME} is healthy"
  exit 0
else
  log "ERROR: container failed health check:"
  docker logs --tail 100 "${CONTAINER_NAME}" 2>&1 || true
  exit 1
fi