#!/bin/bash
set -e
export PATH=/usr/local/cuda/bin:$PATH
cd /root
pip install -q -U huggingface_hub hf_transfer 2>&1 | tail -1 || true
[ -d ingot ] || git clone --recurse-submodules --depth 1 https://github.com/DendriteLabs/ingot.git
mkdir -p /root/models
[ -d /root/models/teacher ] || hf download Qwen/Qwen3.8-27B --local-dir /root/models/teacher > /root/dl_teacher.log 2>&1 &
[ -d /root/models/king-ix ] || hf download dendriteholdings/Ingot-IX-27B-Q2 --local-dir /root/models/king-ix > /root/dl_king.log 2>&1 &
wait
echo DOWNLOADS_DONE
cd /root/ingot
cmake -S llama-prism -B llama-prism/build -DGGML_CUDA=ON -DLLAMA_CURL=OFF -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=90 > /root/cmake.log 2>&1
cmake --build llama-prism/build -j $(nproc) > /root/build_llama.log 2>&1
cmake -S src/eval_service/llama_states -B build/llama_states -DLLAMA_DIR=$PWD/llama-prism -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=90 >> /root/cmake.log 2>&1
cmake --build build/llama_states >> /root/build_llama.log 2>&1
echo LLAMA_DONE
uv venv --python 3.12 .venv
.venv/bin/pip install -q "vllm==0.31.0" flash-linear-attention ninja 2>&1 | tail -2
.venv/bin/pip install -q -e ".[eval]" 2>&1 | tail -2
SP=$(.venv/bin/python -c "import sysconfig; print(sysconfig.get_paths()['purelib'])")
grep -q _split_kv_cache_groups_to_single_layer "$SP/vllm/v1/core/kv_cache_utils.py" || patch --no-backup-if-mismatch -p1 -d "$SP" < scripts/vllm_pr56822.diff
CU=$SP/nvidia/cu13
[ -e "$CU/lib64" ] || ln -s lib "$CU/lib64"
(cd "$CU/lib" && for lib in libcudart libnvrtc libcublas libcublasLt; do ln -sfn "$lib.so.13" "$lib.so"; done; ln -sfn "$(ls libnvrtc-builtins.so.13.* | head -1)" libnvrtc-builtins.so)
.venv/bin/python -c "import torch, vllm; from vllm.v1.core import kv_cache_utils; assert hasattr(kv_cache_utils, '_split_kv_cache_groups_to_single_layer'); print('EVAL BOX OK', torch.cuda.get_device_name(0))"
echo SETUP_DONE