#!/bin/bash
# Comprehensive benchmark for sglang 8-card DeepSeek-V4-Flash
set -e

cd /data/user1/yy
mkdir -p bench_results logs

PYTHON="/data/user1/yy/envs/sglang/bin/python"
BENCH="-m sglang.bench_serving"
HOST="127.0.0.1"
PORT="30000"
MODEL="/data/models/DeepSeek-V4-Flash"
DATASET="/data/user1/yy/datasets/ShareGPT_V4.3_unfiltered_cleaned_split.json"
NUM_PROMPTS=2000
SEED=42

RESULT_DIR="bench_results/sglang_8card_systematic_$(date +%Y%m%d_%H%M%S)"
mkdir -p "${RESULT_DIR}"

SUMMARY_FILE="${RESULT_DIR}/summary.json"
LOG_FILE="${RESULT_DIR}/benchmark.log"

# Use a marker file for summary accumulation
SUMMARY_TMP="${RESULT_DIR}/.summary_tmp.json"
echo '[]' > "${SUMMARY_TMP}"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "${LOG_FILE}"
}

wait_for_server() {
    log "Waiting for sglang server at ${HOST}:${PORT} to become ready..."
    local retries=0
    local max_retries=120
    while ! curl -s "http://${HOST}:${PORT}/v1/models" > /dev/null 2>&1; do
        retries=$((retries + 1))
        if [ "${retries}" -ge "${max_retries}" ]; then
            log "ERROR: Server not ready after ${max_retries} retries"
            exit 1
        fi
        sleep 2
    done
    log "Server is ready."
}

run_bench() {
    local name="$1"
    shift
    local output_file="${RESULT_DIR}/${name}.json"
    local detail_log="${RESULT_DIR}/${name}.log"

    log "---------------------------------------------------"
    log "Running benchmark: ${name}"
    log "Args: $*"
    log "Output: ${output_file}"

    ${PYTHON} ${BENCH} \
        --backend sglang \
        --host "${HOST}" \
        --port "${PORT}" \
        --dataset-name sharegpt \
        --dataset-path "${DATASET}" \
        --num-prompts "${NUM_PROMPTS}" \
        --model "${MODEL}" \
        --output-file "${output_file}" \
        --output-details \
        --seed "${SEED}" \
        "$@" \
        > "${detail_log}" 2>&1

    local status=$?
    if [ ${status} -ne 0 ]; then
        log "ERROR: Benchmark ${name} failed with exit code ${status}"
        return ${status}
    fi

    log "Benchmark ${name} completed successfully."

    # Extract key metrics from the log
    local req_throughput=$(grep -oP 'Request throughput \(req/s\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local input_throughput=$(grep -oP 'Input token throughput \(tok/s\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local output_throughput=$(grep -oP 'Output token throughput \(tok/s\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local total_throughput=$(grep -oP 'Total token throughput \(tok/s\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local mean_ttft=$(grep -oP 'Mean TTFT \(ms\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local mean_tpot=$(grep -oP 'Mean TPOT \(ms\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local mean_e2e=$(grep -oP 'Mean E2E Latency \(ms\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local duration=$(grep -oP 'Benchmark duration \(s\):\s+\K[0-9.]+' "${detail_log}" || echo "null")
    local accept_length=$(grep -oP 'Accept length:\s+\K[0-9.]+' "${detail_log}" || echo "null")

    # Append to summary tmp
    ${PYTHON} - <<PYEOF
import json
entry = {
    "name": "${name}",
    "args": "$*",
    "duration_s": ${duration},
    "request_throughput": ${req_throughput},
    "input_token_throughput": ${input_throughput},
    "output_token_throughput": ${output_throughput},
    "total_token_throughput": ${total_throughput},
    "mean_ttft_ms": ${mean_ttft},
    "mean_tpot_ms": ${mean_tpot},
    "mean_e2e_latency_ms": ${mean_e2e},
    "accept_length": ${accept_length},
    "output_file": "${output_file}"
}
with open("${SUMMARY_TMP}", "r") as f:
    data = json.load(f)
data.append(entry)
with open("${SUMMARY_TMP}", "w") as f:
    json.dump(data, f, indent=2)
PYEOF
}

main() {
    log "Starting systematic benchmark for sglang 8-card DeepSeek-V4-Flash"
    log "Model: ${MODEL}"
    log "Dataset: ${DATASET}"
    log "Num prompts per test: ${NUM_PROMPTS}"
    log "Results will be saved to: ${RESULT_DIR}"

    wait_for_server

    # Concurrency sweep: simulate different real-world load levels
    log "=== Phase 1: Concurrency sweep ==="
    for c in 1 8 16 32 64 128; do
        run_bench "sharegpt_concurrency_${c}" --max-concurrency "${c}"
    done

    # Request-rate sweep: simulate Poisson-like traffic
    log "=== Phase 2: Request rate sweep ==="
    for r in 1 2 4 8 16; do
        run_bench "sharegpt_rps_${r}" --request-rate "${r}"
    done

    # Move final summary
    mv "${SUMMARY_TMP}" "${SUMMARY_FILE}"

    log "=== All benchmarks completed ==="
    log "Summary saved to: ${SUMMARY_FILE}"
    log "Detailed results in: ${RESULT_DIR}"

    # Print summary table
    ${PYTHON} - <<PYEOF
import json
with open("${SUMMARY_FILE}", "r") as f:
    data = json.load(f)

print("\n============ Benchmark Summary ============")
print(f"{'Name':<30} {'Req/s':>8} {'In tok/s':>10} {'Out tok/s':>10} {'Total tok/s':>12} {'TTFT':>8} {'TPOT':>8} {'E2E':>10}")
print("-" * 110)
for entry in data:
    print(f"{entry['name']:<30} {entry['request_throughput']:>8.2f} {entry['input_token_throughput']:>10.2f} {entry['output_token_throughput']:>10.2f} {entry['total_token_throughput']:>12.2f} {entry['mean_ttft_ms']:>8.1f} {entry['mean_tpot_ms']:>8.1f} {entry['mean_e2e_latency_ms']:>10.1f}")
PYEOF
}

main "$@"