#!/usr/bin/env sh
set -eu

CONDA_BIN=${CONDA_BIN:-}
if [ -z "$CONDA_BIN" ]; then
  if command -v conda >/dev/null 2>&1; then
    CONDA_BIN=$(command -v conda)
  elif [ -x /opt/miniconda/bin/conda ]; then
    CONDA_BIN=/opt/miniconda/bin/conda
  fi
fi

if [ -n "$CONDA_BIN" ]; then
  eval "$("$CONDA_BIN" shell.bash hook)"
  conda activate focallora4
fi

SCRIPT_DIR=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
MODEL_PATH=${MODEL_PATH:-/data/local/hujk/models/Llama-3.1-8B-Instruct}
OUT_DIR=${OUT_DIR:-$SCRIPT_DIR/test_outputs}
GPUS=${GPUS:-"0 1"}

mkdir -p "$OUT_DIR"

for gpu in $GPUS; do
  echo "[test_chunked_attention_llama] running on GPU $gpu"
  CUDA_VISIBLE_DEVICES="$gpu" python "$SCRIPT_DIR/test_chunked_attention.py" \
    --model-path "$MODEL_PATH" \
    --output "$OUT_DIR/chunked_attention_gpu${gpu}.json" \
    "$@"
done