# Source this on a GPU machine before running anything:  . /nvme/wsp-lab/code/scripts/env.sh
# It activates the lab's Python and keeps every cache inside /nvme/wsp-lab.
export LAB=/nvme/wsp-lab
. "$LAB/venv/bin/activate"
export HF_HOME=$LAB/hf NLTK_DATA=$LAB/nltk_data VLLM_CACHE_ROOT=$LAB/vllm-cache \
  FLASHINFER_WORKSPACE_BASE=$LAB TRITON_CACHE_DIR=$LAB/triton-cache TORCHINDUCTOR_CACHE_DIR=$LAB/inductor-cache \
  XDG_CACHE_HOME=$LAB/xdg-cache UV_CACHE_DIR=$LAB/uv-cache MPLCONFIGDIR=$LAB/mpl TOKENIZERS_PARALLELISM=false
# FlashInfer compiles its sampling kernel on first use and needs the cuRAND headers. Some machines' CUDA lacks
# them; the venv's nvidia-curand wheel has a copy, which we expose through a folder that holds only those files.
if [ ! -e "${CUDA_HOME:-/usr/local/cuda}/include/curand.h" ]; then
  _inc=$LAB/cuda-extra-include
  if [ ! -e "$_inc/curand.h" ]; then
    mkdir -p "$_inc" && ln -sf "$LAB"/venv/lib/python3.12/site-packages/nvidia/cu13/include/curand*.h "$_inc"/
  fi
  export FLASHINFER_EXTRA_CUDAFLAGS="-isystem $_inc"
fi