#!/usr/bin/env bash
set -Eeuo pipefail

# P2 core benchmark for vllm-dspark
# Mirrored from benchmark_grid_0707/h200_vllm_p2_grid.sh

BACKEND="vllm"
PORT="${PORT:-30004}"
MODEL="/data/models/DeepSeek-V4-Flash-DSpark"
BENCH_PY="/data/user1/yy/envs/sglang/bin/python"
DATASET="/data/user1/yy/datasets/ShareGPT_filtered_chat.json"
RESULT_ROOT="${RESULT_ROOT:-/data/user1/yy/bench_results/dspark_grid}"
WARMUP_REQUESTS="${WARMUP_REQUESTS:-100}"
PHASE="p2_core"
RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"

num_prompts() {
  local c="$1"
  if (( c <= 4 )); then
    local n=$((c * 8))
    (( n < 32 )) && n=32
    echo "${n}"
  elif (( c <= 32 )); then
    echo $((c * 16))
  else
    echo $((c * 8))
  fi
}

run_case() {
  local scenario="$1"
  local input_len="$2"
  local output_len="$3"
  local concurrencies="$4"
  for concurrency in ${concurrencies}; do
    local prompts
    prompts="$(num_prompts "${concurrency}")"
    local out_dir="${RESULT_ROOT}/${PHASE}/${scenario}/${RUN_ID}"
    mkdir -p "${out_dir}"
    echo "phase=${PHASE} scenario=${scenario} input=${input_len} output=${output_len} concurrency=${concurrency} prompts=${prompts}"
    "${BENCH_PY}" -m sglang.bench_serving \
      --backend "${BACKEND}" \
      --host 127.0.0.1 \
      --port "${PORT}" \
      --model "${MODEL}" \
      --dataset-name random \
      --dataset-path "${DATASET}" \
      --random-input-len "${input_len}" \
      --random-output-len "${output_len}" \
      --num-prompts "${prompts}" \
      --max-concurrency "${concurrency}" \
      --warmup-requests "${WARMUP_REQUESTS}" \
      2>&1 | tee "${out_dir}/c${concurrency}.log"
  done
}

curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null
run_case "chat_short" 512 256 "1 8 16 32 64"
run_case "chat_standard" 1000 256 "64"
run_case "generation_standard" 1000 1000 "64"
run_case "rag_medium" 4000 512 "1 4 8 16 32"
run_case "summarization" 8000 1000 "32"
run_case "long_context_probe" 16000 512 "1 2 4 8 16"