#!/usr/bin/env bash

set -euo pipefail

OUTPUT_ROOT="${OUTPUT_ROOT:-$(pwd)/artifacts/decode-attention-profile-20260716}"
FRONTIER_ROOT="${FRONTIER_ROOT:-/home/admin/cpfs/wjh/frontier-community-qwen235-smoke-20260715/Frontier-d9cfeb6-best-effort-v6-batched-lanes-r3}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh-frontier-vllm0102-smoke/.venv}"
PROFILE_ROOT="${OUTPUT_ROOT}/profiles"
LOG_DIR="${OUTPUT_ROOT}/logs"
PROVENANCE_DIR="${OUTPUT_ROOT}/provenance"
MODEL="Qwen3-235B-A22B-FP8"

mkdir -p "${PROFILE_ROOT}" "${LOG_DIR}" "${PROVENANCE_DIR}"
exec > >(tee -a "${LOG_DIR}/profile.log") 2>&1

if [[ -z "${CUDA_VISIBLE_DEVICES:-}" ]]; then
  echo "ERROR: CUDA_VISIBLE_DEVICES must contain the fleet-allocated GPU" >&2
  exit 1
fi
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES}"
if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
  echo "ERROR: expected exactly one GPU, got ${CUDA_VISIBLE_DEVICES}" >&2
  exit 1
fi

echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} operator=FlashInfer_attention phase=decode TP_workers=4 batch_sizes=1,2 kv_sizes=2048,2176 block=16 measurement=CUDA_EVENT output=${OUTPUT_ROOT} expected_wall=5-10m hard_wall=900s hard_gpu_cap=0.25_H20h"
date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader

test -x "${VENV_ROOT}/bin/python"
test -f "${FRONTIER_ROOT}/pyproject.toml"
test -f "${FRONTIER_ROOT}/data/config/models/${MODEL}.json"
sha256sum run_decode_attention_profile.sh > "${PROVENANCE_DIR}/source.sha256"

export PYTHONPATH="${FRONTIER_ROOT}"
export TOKENIZERS_PARALLELISM=false
export TORCH_CUDA_ARCH_LIST=9.0

cd "${FRONTIER_ROOT}"
timeout --signal=TERM --kill-after=30s 600 \
  "${VENV_ROOT}/bin/python" -m frontier.profiling.attention.main \
  --disable_ray \
  --models "${MODEL}" \
  --num_gpus 1 \
  --max_model_len 40960 \
  --max_seq_len 2176 \
  --min_batch_size 1 \
  --max_batch_size 2 \
  --batch_size_list 1 2 \
  --decode_kv_cache_size_list 2048 2176 \
  --num_tensor_parallel_workers 4 \
  --max_pipeline_parallel_size 1 \
  --attention_backend FLASHINFER \
  --block_size 16 \
  --profile_only_decode \
  --device h20 \
  --profile_method cuda_event \
  --output_dir "${PROFILE_ROOT}" \
  --yes

ATTENTION_CSV="${PROFILE_ROOT}/compute/h20/${MODEL}/attention.csv"
test -s "${ATTENTION_CSV}"
"${VENV_ROOT}/bin/python" - "${ATTENTION_CSV}" \
  > "${PROVENANCE_DIR}/coverage.json" <<'PY'
import json
import sys

import pandas as pd

path = sys.argv[1]
frame = pd.read_csv(path)
decode = frame[frame["is_prefill"] == False]  # noqa: E712
payload = {
    "path": path,
    "row_count": len(frame),
    "decode_row_count": len(decode),
    "batch_sizes": sorted(int(value) for value in decode["batch_size"].unique()),
    "kv_cache_sizes": sorted(int(value) for value in decode["kv_cache_size"].unique()),
    "attn_decode_median_non_null": int(
        decode["time_stats.attn_decode.median"].notna().sum()
    ),
}
print(json.dumps(payload, indent=2, sort_keys=True))
if payload["decode_row_count"] < 4 or payload["attn_decode_median_non_null"] < 4:
    raise SystemExit(1)
PY

sha256sum \
  "${ATTENTION_CSV}" \
  "${PROVENANCE_DIR}/coverage.json" \
  "${PROVENANCE_DIR}/source.sha256" \
  > "${PROVENANCE_DIR}/artifacts.sha256"
nvidia-smi --query-gpu=index,name,memory.used,utilization.gpu --format=csv,noheader
date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
echo "DECODE_ATTENTION_PROFILE_COMPLETE"