#!/usr/bin/env bash

# Real-only, pressure-matching probe for the next Fixed-PD workload. This
# intentionally profiles one anchor, then freezes the workload before any
# Frontier-vs-real selection comparison.

set -euo pipefail

OUT="${OUTPUT_ROOT:?OUTPUT_ROOT is required}"
RUNNER_DIR="${RUNNER_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
CLIENT="${CLIENT:-${RUNNER_DIR}/../frontier-phase-factorial-v0/qwen30_prefill_client.py}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
MODEL_ROOT="${MODEL_ROOT:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
FLASHINFER_WORKSPACE_BASE="${FLASHINFER_WORKSPACE_BASE:-${OUT}/flashinfer-workspace}"
GPU_IDS="${GPU_IDS:-0,1,2,3}"
TP="${TP:-4}"
MNS="${MNS:-64}"
REQUESTS="${REQUESTS:-257}"
GLOBAL_RATES="${GLOBAL_RATES:-4 8 12 16}"
SERVER_READY_ATTEMPTS="${SERVER_READY_ATTEMPTS:-180}"
PORT="${PORT:-8930}"
SERVED_MODEL="qwen3-30b-fixed-pd-pressure"
SERVER_PID=""

[[ "${TP}" == "4" ]] || { echo 'ERROR: this calibrated probe is TP4-only' >&2; exit 1; }
[[ "${MNS}" == "64" ]] || { echo 'ERROR: this calibrated probe is MNS64-only' >&2; exit 1; }
[[ "${REQUESTS}" =~ ^[1-9][0-9]*$ ]] || { echo 'ERROR: REQUESTS must be positive' >&2; exit 1; }
[[ -f "${CLIENT}" ]] || { echo "ERROR: client missing: ${CLIENT}" >&2; exit 1; }
[[ -f "${MODEL_ROOT}/config.json" ]] || { echo "ERROR: model missing: ${MODEL_ROOT}" >&2; exit 1; }

mkdir -p "${OUT}/provenance" "${OUT}/trials" "${FLASHINFER_WORKSPACE_BASE}"
exec > >(tee -a "${OUT}/controller.log") 2>&1

cleanup_server() {
  if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
    kill -TERM -- "-${SERVER_PID}" 2>/dev/null || true
    for _ in $(seq 1 30); do
      kill -0 "${SERVER_PID}" 2>/dev/null || break
      sleep 1
    done
    kill -KILL -- "-${SERVER_PID}" 2>/dev/null || true
  fi
  SERVER_PID=""
}
trap cleanup_server EXIT INT TERM

assert_idle() {
  nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader
  nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits \
    | awk '$1 > 16 {exit 1}'
}

wait_ready() {
  local target="$1"
  for _ in $(seq 1 "${SERVER_READY_ATTEMPTS}"); do
    if curl -fsS --max-time 2 "http://127.0.0.1:${PORT}/v1/models" > "${target}/models.json" 2>/dev/null; then
      return 0
    fi
    if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
      tail -200 "${target}/server.log" >&2 || true
      return 1
    fi
    sleep 3
  done
  echo "ERROR: vLLM did not become ready in $((SERVER_READY_ATTEMPTS * 3)) seconds" >&2
  return 1
}

start_server() {
  local target="$1"
  export TOKENIZERS_PARALLELISM=false
  export VLLM_USE_V1=1
  export TORCH_CUDA_ARCH_LIST=9.0
  export HF_HUB_OFFLINE=1
  export TRANSFORMERS_OFFLINE=1
  export FLASHINFER_WORKSPACE_BASE
  export HOME=/tmp/wjh
  export XDG_CACHE_HOME=/tmp/wjh/.cache
  export VLLM_CACHE_ROOT=/tmp/wjh/.cache/vllm
  export CUDA_VISIBLE_DEVICES="${GPU_IDS}"
  setsid "${VENV_ROOT}/bin/vllm" serve "${MODEL_ROOT}" \
    --host 127.0.0.1 --port "${PORT}" --served-model-name "${SERVED_MODEL}" \
    --tensor-parallel-size "${TP}" --gpu-memory-utilization 0.92 \
    --max-model-len 40960 --max-num-batched-tokens 8192 --max-num-seqs "${MNS}" \
    --no-enable-prefix-caching --enable-chunked-prefill --no-enable-log-requests \
    > "${target}/server.log" 2>&1 &
  SERVER_PID=$!
  wait_ready "${target}"
}

run_client() {
  local target="$1" rate="$2"
  timeout --signal=TERM --kill-after=60s 1800 \
    "${VENV_ROOT}/bin/python" "${CLIENT}" \
      --port "${PORT}" --served-model "${SERVED_MODEL}" --model-path "${MODEL_ROOT}" \
      --rate "${rate}" --requests "${REQUESTS}" --input-tokens 4096 --output-tokens 256 \
      --timeout-seconds 1200 --output "${target}/result.json"
}

warmup_server() {
  local target="$1"
  timeout --signal=TERM --kill-after=60s 600 \
    "${VENV_ROOT}/bin/python" "${CLIENT}" \
      --port "${PORT}" --served-model "${SERVED_MODEL}" --model-path "${MODEL_ROOT}" \
      --rate 1 --requests 4 --input-tokens 512 --output-tokens 1 \
      --timeout-seconds 300 --output "${target}/result.json"
}

analyze() {
  "${VENV_ROOT}/bin/python" - "${OUT}" "${GLOBAL_RATES}" <<'PY'
import json
import math
import statistics
import sys
from pathlib import Path

root = Path(sys.argv[1])
rates = [float(value) for value in sys.argv[2].split()]
target = {"ttft_ms": 245.9527667526406, "tpot_ms": 13.178025610291787}

def p90(values):
    return sorted(values)[math.ceil(0.9 * len(values)) - 1]

rows = []
for rate in rates:
    label = f"r{rate:g}"
    trial_means = {"ttft_ms": [], "tpot_ms": [], "e2e_ms": []}
    pooled = {key: [] for key in trial_means}
    for trial in range(1, 4):
        path = root / "trials" / f"trial{trial}" / label / "result.json"
        payload = json.loads(path.read_text())
        workload = payload["workload"]
        if (float(workload["offered_request_rate"]) != rate or workload["request_count"] != 257
                or workload["input_tokens"] != 4096 or workload["output_tokens"] != 256
                or workload["prefix_caching"] is not False):
            raise ValueError(f"workload drift: {path}")
        requests = payload["requests"]
        if len(requests) != 257 or any(not request["success"] for request in requests):
            raise ValueError(f"incomplete client result: {path}")
        for key in pooled:
            values = [float(request[key]) for request in requests]
            pooled[key].extend(values)
            trial_means[key].append(statistics.mean(values))
    row = {
        "global_rate": rate,
        "per_gpu_rate": rate / 4.0,
        "requests_per_trial": 257,
        "trials": 3,
        "metrics": {
            key: {
                "pooled_mean_ms": statistics.mean(values),
                "pooled_p90_ms": p90(values),
                "trial_mean_stdev_ms": statistics.stdev(trial_means[key]),
            }
            for key, values in pooled.items()
        },
    }
    row["inflight_proxy"] = rate * row["metrics"]["e2e_ms"]["pooled_mean_ms"] / 1000.0
    row["relative_distance"] = math.sqrt(sum(
        ((row["metrics"][key]["pooled_mean_ms"] - target[key]) / target[key]) ** 2
        for key in target
    ))
    rows.append(row)

winner = min(rows, key=lambda row: (row["relative_distance"], row["global_rate"]))
payload = {
    "schema": "qwen30-fixed-pd-pressure-probe-v1",
    "target_trace_pd_tp4_mns64": target,
    "decision_rule": "minimum Euclidean distance of relative mean TTFT and TPOT errors",
    "rates": rows,
    "recommended_global_rate": winner["global_rate"],
    "recommended_per_gpu_rate": winner["per_gpu_rate"],
}
(root / "pressure-analysis.json").write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
lines = [
    "# Fixed-PD pressure probe",
    "",
    "| Global / per-GPU rps | TTFT mean / p90 (ms) | TPOT mean / p90 (ms) | E2E mean / p90 (ms) | In-flight proxy | Relative distance |",
    "|---|---:|---:|---:|---:|---:|",
]
for row in rows:
    metric = row["metrics"]
    lines.append(
        f"| {row['global_rate']:g} / {row['per_gpu_rate']:g} | "
        f"{metric['ttft_ms']['pooled_mean_ms']:.2f} / {metric['ttft_ms']['pooled_p90_ms']:.2f} | "
        f"{metric['tpot_ms']['pooled_mean_ms']:.2f} / {metric['tpot_ms']['pooled_p90_ms']:.2f} | "
        f"{metric['e2e_ms']['pooled_mean_ms']:.2f} / {metric['e2e_ms']['pooled_p90_ms']:.2f} | "
        f"{row['inflight_proxy']:.2f} | {row['relative_distance']:.3f} |"
    )
lines.extend([
    "",
    f"Recommended frozen rate: **{winner['global_rate']:g} global rps / {winner['per_gpu_rate']:g} rps per GPU**.",
    "Selection uses only pooled mean TTFT and TPOT; p90 and in-flight proxy are audit outputs.",
])
(root / "pressure-analysis.md").write_text("\n".join(lines) + "\n")
print(json.dumps(payload, sort_keys=True))
PY
}

{
  echo "FIXED_PD_PRESSURE_PROBE_LAUNCH_ECHO host=$(hostname) model=${MODEL_ROOT} engine=vLLM-0.20.0+cu129 dtype=BF16 config=TP${TP}_MNS${MNS}_MBT8192 gpus=${GPU_IDS} prefix=false shape=4096_to_256 requests_per_rate=${REQUESTS} global_rates={${GLOBAL_RATES}} rate_contract=global_rate_divided_by_TP trials=3 fresh_server=true metric_target=TracePD_TP4_MNS64_meanTTFT245.95ms_meanTPOT13.18ms expected_wall=12-20m expected_cost=0.8-1.4_H20-GPUh output=${OUT}"
  date -u +START_UTC=%Y-%m-%dT%H:%M:%SZ
  assert_idle
  sha256sum "${BASH_SOURCE[0]}" "${CLIENT}" "${MODEL_ROOT}/config.json" > "${OUT}/provenance/input.sha256"
  "${VENV_ROOT}/bin/vllm" --version > "${OUT}/provenance/vllm.version"
  "${VENV_ROOT}/bin/python" -c 'import torch, transformers, vllm; print(f"torch={torch.__version__}"); print(f"transformers={transformers.__version__}"); print(f"vllm={vllm.__version__}")' > "${OUT}/provenance/runtime.versions"
  nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total --format=csv,noheader > "${OUT}/provenance/gpus.before.csv"

  declare -a ORDERS=("4 8 12 16" "16 12 8 4" "8 16 4 12")
  for trial in 1 2 3; do
    trial_root="${OUT}/trials/trial${trial}"
    mkdir -p "${trial_root}"
    echo "TRIAL_START trial=${trial} order=${ORDERS[$((trial - 1))]}"
    start_server "${trial_root}"
    warmup_server "${trial_root}/warmup"
    for rate in ${ORDERS[$((trial - 1))]}; do
      rate_root="${trial_root}/r${rate}"
      mkdir -p "${rate_root}"
      echo "RATE_START trial=${trial} global_rate=${rate} per_gpu_rate=$(awk -v value="${rate}" 'BEGIN {printf "%.3f", value / 4}')"
      run_client "${rate_root}" "${rate}"
      echo "RATE_COMPLETE trial=${trial} global_rate=${rate}"
    done
    cleanup_server
    assert_idle
    echo "TRIAL_COMPLETE trial=${trial}"
  done
  analyze
  find "${OUT}" -type f ! -path '*/provenance/artifacts.sha256' -print0 | sort -z | xargs -0 sha256sum > "${OUT}/provenance/artifacts.sha256"
  nvidia-smi --query-gpu=index,name,uuid,driver_version,memory.total --format=csv,noheader > "${OUT}/provenance/gpus.after.csv"
  date -u +END_UTC=%Y-%m-%dT%H:%M:%SZ
  echo 'FIXED_PD_PRESSURE_PROBE_COMPLETE'
} >> "${OUT}/controller.log" 2>&1