#!/usr/bin/env bash
# Start the SGLang server for Qwen3-235B-A22B on NVIDIA RTX 6000D (TP=8, 8 GPUs).
#
# NVIDIA counterpart of experiments/p800/qwen3_235b_p800_sglang_tp8/start_server.sh.
# Uses the standard sglang qwen3_moe path (no Kunlun patch, no XPU env).
# Launch pattern mirrors pro6000/dsv4_pro6000_sglang_tiny_1k_output/start_sglang_docker.sh.
set -Eeuo pipefail

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

# Source config before platform so our CONTAINER_NAME/PORT defaults win.
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
# 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}}"
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"
LOG_DIR="${RESULT_ROOT}/logs"
mkdir -p "$LOG_DIR" "${RUNTIME_BASE}/tmp" "$CACHE_DIR"

log_init "${LOG_DIR}/start_server.log"
log "starting server for ${EXPERIMENT_NAME} (NVIDIA, TP=${TP})"
log "model: ${MODEL_PATH}"
log "image: ${DOCKER_IMAGE}"
log "port: ${PORT}  container: ${CONTAINER_NAME}"

SERVER_LOG="${LOG_DIR}/server.log"
PID_FILE="${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid"
rm -f "$PID_FILE"

# Remove any stale container with the same name.
docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true

# Standard sglang launch for qwen3_moe on NVIDIA (bf16 auto, flashinfer attn).
# Mount /data so model + experiment results dir + dataset are all at the same
# path inside the container (bench client writes results there directly).
DTYPE_ARGS=()
if [[ "${DTYPE}" != "auto" ]]; then
  DTYPE_ARGS=(--dtype "${DTYPE}")
fi

nohup docker run --rm \
  --name "$CONTAINER_NAME" \
  --gpus all \
  --privileged \
  --ipc=host \
  --network host \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  --entrypoint python3 \
  -v /data:/data \
  -v "${CACHE_DIR}:/root/.cache" \
  -v "${RUNTIME_BASE}/tmp:/tmp" \
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
  -e PYTHONUNBUFFERED=1 \
  -e HF_HUB_OFFLINE=1 \
  -e TRANSFORMERS_OFFLINE=1 \
  -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
  "$DOCKER_IMAGE" \
  -m sglang.launch_server \
  --model-path "${MODEL_PATH}" \
  --trust-remote-code \
  --tp-size "${TP}" \
  --mem-fraction-static "${MEM_FRACTION_STATIC}" \
  --context-length "${CONTEXT_LENGTH}" \
  --max-running-requests "${MAX_RUNNING_REQUESTS}" \
  --host 0.0.0.0 \
  --port "${PORT}" \
  "${DTYPE_ARGS[@]}" \
  > "$SERVER_LOG" 2>&1 &

PID=$!
echo "$PID" > "$PID_FILE"
log "PID: ${PID}"
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}  (container: ${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