#!/usr/bin/env bash
# One-time setup of a GPU machine for the labs. Everything lives under /nvme/wsp-lab (about 15 GB).
# From your laptop, in the project folder:
#   scp scripts/setup_vm.sh <host>:/tmp/ && ssh <host> bash /tmp/setup_vm.sh      # about 10 minutes
#   bash scripts/sync_vm.sh <host>                                                # copy the code
#   ssh <host> 'cd /nvme/wsp-lab/code && . scripts/env.sh &&
#     python scripts/prepare_data.py --exp8 external/lm_understand-logs/exp8_math --out /nvme/wsp-lab/data &&
#     python scripts/build_models.py && python scripts/build_models.py --variants configs/rebuilt.json'
# It needs no root: Python, packages, the model and every cache go into /nvme/wsp-lab.
set -euo pipefail
LAB=${LAB:-/nvme/wsp-lab}
UV=${UV:-$HOME/.local/bin/uv}
export UV_PYTHON_INSTALL_DIR=$LAB/python UV_CACHE_DIR=$LAB/uv-cache HF_HOME=$LAB/hf
mkdir -p "$LAB"/{models,data,runs,logs,code,nltk_data}
cd "$LAB"
[ -x "$UV" ] || { curl -LsSf https://astral.sh/uv/install.sh | sh; UV=$HOME/.local/bin/uv; }
[ -x venv/bin/python ] || "$UV" venv --python 3.12 venv
. venv/bin/activate
"$UV" pip install "vllm==0.33.0" "flash-linear-attention==0.6.0" "math-verify[antlr4_13_2]==0.9.0" \
  "lm-eval[ifeval]==0.4.13" datasets matplotlib scipy pandas
python -c "import nltk; nltk.download('punkt_tab', download_dir='$LAB/nltk_data')"
hf download Qwen/Qwen3.5-0.8B --revision 2fc06364715b967f1860aea9cf38778875588b17 \
  --local-dir "$LAB/models/Qwen3.5-0.8B"
python -c "import vllm, transformers, torch; print('vllm', vllm.__version__, 'transformers', transformers.__version__,
'torch', torch.__version__, 'cuda', torch.version.cuda)"
echo SETUP-DONE