#!/usr/bin/env bash
# Find the Total-TPS saturation concurrency for each vLLM TP/DP/ISL/OSL shape.
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="vllm"
ENGINE_PORT="$VLLM_PORT"
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/adaptive_results}"
ACTIVE_ENGINE_SERVER_LOG=""

if [[ -x "${VENV_CLIENT}/bin/python" ]]; then
  PYTHON="${VENV_CLIENT}/bin/python"
else
  PYTHON="$(command -v python3)"
fi

DOCKER_IMAGE="${DOCKER_IMAGE:-vllm-sm120-dsv4:0.25.1-fi0.6.14}"
DOCKER_CLIENT_IMAGE="${DOCKER_CLIENT_IMAGE:-sglang-sm120-dsv4:0.5.15.post1-fi0.6.14-sm120fix1}"

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"
  local pid_file="${RUNTIME_BASE}/${EXPERIMENT}_vllm_tp${tp}_dp${dp}.pid"

  if [[ -f "$pid_file" ]]; then
    local pid
    pid="$(cat "$pid_file")"
    if kill -0 "$pid" 2>/dev/null; then
      log "stopping vllm server pid=${pid} tp=${tp} dp=${dp}"
      kill "$pid" 2>/dev/null || true
      sleep 5
      kill -9 "$pid" 2>/dev/null || true
    fi
    rm -f "$pid_file"
  fi
  docker rm -f "${EXPERIMENT}_vllm_tp${tp}_dp${dp}" >/dev/null 2>&1 || true
  ACTIVE_ENGINE_SERVER_LOG=""
  sleep 2
}

engine_build_server_args() {
  local tp="$1"
  local dp="$2"
  local -a args=(
    vllm serve "$MODEL_PATH"
    --trust-remote-code
    --kv-cache-dtype "$KV_CACHE_DTYPE"
    --block-size "$BLOCK_SIZE"
    --tensor-parallel-size "$tp"
    --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION"
    --max-model-len "$MAX_MODEL_LEN"
    --max-num-seqs "$MAX_NUM_SEQS"
    --host 0.0.0.0
    --port "$ENGINE_PORT"
  )
  if (( dp > 1 )); then
    args+=(--data-parallel-size "$dp")
  fi
  printf '%q ' "${args[@]}"
}

engine_start_server() {
  local tp="$1"
  local dp="$2"
  local outer_log="${ADAPTIVE_LOG_DIR}/vllm_tp${tp}_dp${dp}.server.outer.log"
  log "starting vllm server tp=${tp} dp=${dp}"
  bash "${SCRIPT_DIR}/start_vllm_dp.sh" "$tp" "$dp" >> "$outer_log" 2>&1
  if ! engine_is_healthy; then
    log "ERROR: vllm health check failed tp=${tp} dp=${dp}"
    return 1
  fi
  ACTIVE_ENGINE_SERVER_LOG="$(
    find "${RUNTIME_BASE}/logs" -maxdepth 1 -type f \
      -name "${EXPERIMENT}_vllm*tp${tp}_dp${dp}_*.log" \
      -printf '%T@ %p\n' 2>/dev/null | sort -nr | head -n 1 | cut -d' ' -f2-
  )"
  log "vllm server healthy tp=${tp} dp=${dp} log=${ACTIVE_ENGINE_SERVER_LOG:-unknown}"
}

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'
  local outer_log="${ADAPTIVE_LOG_DIR}/vllm_tp${tp}_dp${dp}.server.outer.log"
  local -a logs=("$detail_log" "$outer_log")
  if [[ -n "$ACTIVE_ENGINE_SERVER_LOG" ]]; then
    logs+=("$ACTIVE_ENGINE_SERVER_LOG")
  fi
  grep -Eiq "$pattern" "${logs[@]}" 2>/dev/null
}

engine_run_bench() {
  local isl="$1"
  local osl="$2"
  local concurrency="$3"
  local num_prompts="$4"
  local output_file="$5"
  local warmup_requests
  warmup_requests="$(adaptive_warmup_request_count "$concurrency")"
  local -a bench_args=(
    --backend vllm
    --host 127.0.0.1
    --port "$ENGINE_PORT"
    --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
    --warmup-requests "$warmup_requests"
    --output-file "$output_file"
    --output-details
    --disable-tqdm
  )
  if [[ "$BENCH_DATASET_NAME" == "random" ]]; then
    bench_args+=(--dataset-path "$DATASET_PATH")
  elif [[ "$BENCH_DATASET_NAME" == "random-ids" ]]; then
    : # random-ids does not need --tokenize-prompt
  else
    bench_args+=(--tokenize-prompt)
  fi

  if [[ "$USE_DOCKER_CLIENT" == "1" ]]; then
    local -a volume_args=(-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_BASE}:${RESULT_BASE}")
    if [[ "$BENCH_DATASET_NAME" == "random" ]]; then
      volume_args+=(-v "${DATASET_PATH}:${DATASET_PATH}:ro")
    fi
    docker run --rm \
      --network host \
      "${volume_args[@]}" \
      -e PYTHONUNBUFFERED=1 \
      -e HF_HUB_OFFLINE=1 \
      -e TRANSFORMERS_OFFLINE=1 \
      --entrypoint python3 \
      "$DOCKER_CLIENT_IMAGE" \
      -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}"
  else
    "$PYTHON" -m "$SGLANG_BENCH_MODULE" "${bench_args[@]}"
  fi
}

export -f engine_run_bench
export ENGINE_PORT MODEL_PATH RESULT_BASE DOCKER_CLIENT_IMAGE USE_DOCKER_CLIENT
export BENCH_DATASET_NAME DATASET_PATH RANDOM_RANGE_RATIO BENCH_WARMUP_MAX_REQUESTS PYTHON SGLANG_BENCH_MODULE

export SEARCH_START_CONCURRENCY=16
export SEARCH_ADDEND=16
# If the initial concurrency violates the TTFT SLO, search downward. Stop at
# the first acceptable value (16 -> 8; only try 1 when 8 still violates it).
export SEARCH_INITIAL_BACKOFF_CONCURRENCIES="8 1"
# When concurrency 1 still has a severely excessive TTFT, stop the remaining
# shapes in this TP/DP group. Zero disables this rule.
export TTFT_GROUP_SKIP_MS="${TTFT_GROUP_SKIP_MS:-8000}"

adaptive_main "$@"
