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

source /usr/local/Ascend/ascend-toolkit/set_env.sh 2>/dev/null || source /usr/local/Ascend/cann/set_env.sh 2>/dev/null || true
source /data/vllm-ascend-smoke/venv/bin/activate

export VLLM_TARGET_DEVICE=ascend
export ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-0}
export PYTORCH_NPU_ALLOC_CONF=${PYTORCH_NPU_ALLOC_CONF:-expandable_segments:True}
export VLLM_LOGGING_LEVEL=${VLLM_LOGGING_LEVEL:-INFO}

MODEL=${MODEL:-/data/vllm-ascend-smoke/models/Qwen/Qwen2.5-0.5B-Instruct}
ITERATIONS=${ITERATIONS:-10}
LOGDIR=${LOGDIR:-/data/vllm-ascend-smoke/logs/ttft-$(date +%Y%m%d-%H%M%S)}
mkdir -p "$LOGDIR"
echo "$LOGDIR" > /data/vllm-ascend-smoke/logs/latest-ttft-dir.txt

python /data/vllm-ascend-smoke/ttft_compare.py \
  --model "$MODEL" \
  --iterations "$ITERATIONS" \
  --case baseline_0_1:0:1:noop \
  --case prefetch_4_1:4:1:prefetch \
  --case sync_4_1:4:1:sync \
  2>&1 | tee "$LOGDIR/ttft-compare.log"

grep -H "TTFT_RESULT" "$LOGDIR"/ttft-compare.log
