#!/usr/bin/env bash
# Find the Total-TPS saturation concurrency for each SGLang TP/DP/ISL/OSL shape
# on Kunlun P800 (DeepSeek-V4-Flash-INT8).
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"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/adaptive_config.env"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/adaptive_bench_lib.sh"

ENGINE="sglang"
ENGINE_PORT="$SGLANG_PORT"
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/adaptive_results}"
ACTIVE_ENGINE_SERVER_LOG=""

PYTHON="${PYTHON:-$(command -v python3)}"

# The shared library's GPU monitor uses nvidia-smi; override it with an
# xpu-smi sampler that emits the same CSV shape.
adaptive_start_gpu_monitor() {
  local csv_path="$1"
  mkdir -p "$(dirname "$csv_path")"
  (
    echo "timestamp, index, memory.used [MiB], memory.total [MiB], utilization.gpu [%]"
    while true; do
      xpu-smi 2>/dev/null | awk -F'|' -v ts="$(date '+%Y/%m/%d %H:%M:%S')" '
        /^\|[[:space:]]*[0-9]+[[:space:]]/ {
          l=$0; sub(/^\|[[:space:]]*/, "", l); sub(/[[:space:]].*$/, "", l); idx=l
        }
        /MiB/ && $4 ~ /%/ {
          split($3, a, "/"); u=a[1]; t=a[2]
          gsub(/[^0-9.]/, "", u); gsub(/[^0-9.]/, "", t)
          r=$4; gsub(/[^0-9.]/, "", r)
          printf "%s, %s, %s MiB, %s MiB, %s %%\n", ts, idx, u, t, r
        }'
      sleep "$GPU_MEM_SAMPLE_INTERVAL_S"
    done
  ) >> "$csv_path" 2>/dev/null &
  echo $!
}

engine_is_healthy() {
  curl --fail --silent --show-error --max-time 5 \
    "http://127.0.0.1:${ENGINE_PORT}/health" >/dev/null 2>&1
}

engine_stop_server() {
  local tp="$1"
  local dp="$2"
  # The server only ever runs inside this experiment's dedicated container,
  # so removing the container is a complete and safe stop.
  log "stopping container ${CONTAINER_NAME} tp=${tp} dp=${dp}"
  docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
  ACTIVE_ENGINE_SERVER_LOG=""
  sleep 2
}

engine_build_server_args() {
  # Keep in sync with the launch args in start_sglang_docker.sh.
  local tp="$1"
  local dp="$2"
  local ep_args=""
  if (( dp == 1 )); then
    ep_args="--ep-size ${EP_SIZE:-$tp}"
  fi
  local args="--host 0.0.0.0 --port ${ENGINE_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} ${ep_args} --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} --chunked-prefill-size ${CHUNKED_PREFILL_SIZE} --max-prefill-tokens ${MAX_PREFILL_TOKENS} --speculative-algorithm EAGLE --speculative-num-steps ${SPEC_NUM_STEPS} --speculative-eagle-topk ${SPEC_EAGLE_TOPK} --speculative-num-draft-tokens ${SPEC_NUM_DRAFT_TOKENS}"
  if (( dp > 1 )); then
    args="${args} --dp-size ${dp}"
  fi
  printf '%s' "$args"
}

engine_start_server() {
  local tp="$1"
  local dp="$2"
  local outer_log="${ADAPTIVE_LOG_DIR}/sglang_tp${tp}_dp${dp}.server.outer.log"
  log "starting sglang server tp=${tp} dp=${dp}"
  bash "${SCRIPT_DIR}/start_sglang_docker.sh" "$tp" "$dp" >> "$outer_log" 2>&1
  if ! engine_is_healthy; then
    log "ERROR: sglang health check failed tp=${tp} dp=${dp}"
    return 1
  fi
  ACTIVE_ENGINE_SERVER_LOG=""
  log "sglang server healthy tp=${tp} dp=${dp} log=docker logs ${CONTAINER_NAME}"
}

engine_detect_oom() {
  local detail_log="$1"
  local tp="$2"
  local dp="$3"
  local pattern='CUDA out of memory|torch\.OutOfMemoryError|OutOfMemory|out of memory|OOM|RESOURCE_EXHAUSTED|Failed to allocate memory'
  if grep -Eiq "$pattern" "$detail_log" 2>/dev/null; then
    return 0
  fi
  docker logs --tail 2000 "$CONTAINER_NAME" 2>/dev/null | grep -Eiq "$pattern"
}

engine_run_bench() {
  local isl="$1"
  local osl="$2"
  local concurrency="$3"
  local num_prompts="$4"
  local output_file="$5"
  local container_output="/tmp/bench_outputs/adaptive_$(basename "$output_file")"
  if [[ "$output_file" == "/dev/null" ]]; then
    container_output="/dev/null"
  fi

  local -a bench_args=(
    --backend sglang
    --host 127.0.0.1
    --port "$ENGINE_PORT"
    --model "$MODEL_PATH"
    --dataset-name "$BENCH_DATASET_NAME"
    --random-input-len "$isl"
    --random-output-len "$osl"
    --random-range-ratio "$RANDOM_RANGE_RATIO"
    --num-prompts "$num_prompts"
    --max-concurrency "$concurrency"
    --request-rate 10000
    --output-file "$container_output"
    --output-details
  )
  if [[ "$BENCH_DATASET_NAME" == "random" ]]; then
    bench_args+=(--dataset-path "$CONTAINER_DATASET_PATH")
  else
    bench_args+=(--tokenize-prompt)
  fi

  docker exec "$CONTAINER_NAME" mkdir -p /tmp/bench_outputs 2>/dev/null || true
  docker exec "$CONTAINER_NAME" \
    env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
    "$CONTAINER_PYTHON" -m sglang.bench_serving "${bench_args[@]}" || return $?

  if [[ "$output_file" != "/dev/null" ]]; then
    docker cp "$CONTAINER_NAME:${container_output}" "$output_file" || return 1
  fi
}

export -f engine_run_bench
export ENGINE_PORT CONTAINER_NAME CONTAINER_PYTHON MODEL_PATH RESULT_BASE
export BENCH_DATASET_NAME CONTAINER_DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON

adaptive_main "$@"
