#!/usr/bin/env bash
# Start P800 SGLang INT8 server for the long-context matrix.
# Usage: start_sglang.sh <max_context_len> <max_running>
set -Eeuo pipefail

MAX_CONTEXT_LEN="${1:-}"
MAX_RUNNING="${2:-}"
if [[ -z "$MAX_CONTEXT_LEN" || -z "$MAX_RUNNING" ]]; then
  echo "Usage: $0 <max_context_len> <max_running>"
  exit 1
fi

if [[ "$MAX_RUNNING" -lt 2 ]]; then
  echo "ERROR: max_running must be >= 2 on P800 INT8 (1 causes req_to_token_pool allocation failures)"
  exit 1
fi

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

RESULT_ROOT="${RESULT_ROOT:-/tmp/${EXPERIMENT}}"
SERVER_LOG="${RESULT_ROOT}/logs/server_${MAX_CONTEXT_LEN}.outer.log"
mkdir -p "$(dirname "$SERVER_LOG")"

log "starting P800 SGLang INT8 server (context-length=${MAX_CONTEXT_LEN}, max-running=${MAX_RUNNING})"
log "model: ${MODEL_PATH}"
log "port: ${PORT}"
log "server log: ${SERVER_LOG}"

# Stop any existing container with the same 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 required by the P800 SGLang INT8 image.
env_args=(
  -e XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
  -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
  # INT8 specific
  -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
)

# Launch args. Keep the proven P800 INT8 command from dsv4_p800_max_context_length,
# but make --context-length and --max-running-requests dynamic per group.
# Note: chunked prefill is required for stable concurrent long-context prefill on P800;
# speculative decoding is disabled as requested.
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 --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.8 --max-prefill-tokens ${MAX_CONTEXT_LEN} --max-running-requests ${MAX_RUNNING} --tensor-parallel-size 8 --ep-size 8 --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 ${MAX_CONTEXT_LEN}"

# Base64-encode the bootstrap command to avoid host-shell quoting issues.
server_cmd=$(cat <<EOF
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
/root/miniconda/envs/python310_torch25_cuda/bin/pip install https://files.pythonhosted.org/packages/14/8b/2a1333a6455c6fad401c2285dee6f58016c55b1cb44cae3a31f8a9cc7d83/apache_tvm_ffi-0.1.0b2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -q
/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}
EOF
)
server_cmd_b64=$(printf '%s' "$server_cmd" | base64 -w0)

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"
)

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 -c "echo '${server_cmd_b64}' | base64 -d | bash" \
  >> "${SERVER_LOG}" 2>&1

log "container ${CONTAINER_NAME} started, waiting for health"
if health_check 127.0.0.1 "$PORT" 600; then
  log "container ${CONTAINER_NAME} is healthy"
else
  log "ERROR: container ${CONTAINER_NAME} failed health check"
  exit 1
fi