#!/usr/bin/env bash
# Builds koboldcpp's CUDA backend from a release tag with machine code for the GPUs storiesbyme rents, and uploads
# the result to a Hugging Face repo. The pod image downloads that tarball instead of compiling (gpu-image/Dockerfile).
# Runs as a Hugging Face Job inside nvidia/cuda:<version>-devel-ubuntu22.04; see README.md for the command.
# Env: KCPP_VERSION (e.g. 1.121), HF_REPO (e.g. user/repo), HF_TOKEN (write access; pass as a Job secret).
set -euo pipefail
: "${KCPP_VERSION:?}" "${HF_REPO:?}" "${HF_TOKEN:?}"
NAME="koboldcpp-${KCPP_VERSION}-cu12-sm86-sm89"
JOBS="${CPU_CORES:-}"; JOBS="${JOBS%%[!0-9]*}"; JOBS="${JOBS:-$(nproc)}"   # nproc reports the host's cores
log() { echo "[$(date -u +%H:%M:%S)] $*"; }

export DEBIAN_FRONTEND=noninteractive
apt-get update -qq && apt-get install -y -qq --no-install-recommends git ca-certificates python3-pip binutils >/dev/null
nvcc --version | tail -2; gcc --version | head -1
git clone -q --depth 1 --branch "v${KCPP_VERSION}" https://github.com/LostRuins/koboldcpp /src
cd /src
log "source: $(git log --oneline -1)"

# The release compiles its GPU code with `code=compute_NN`, which emits only PTX (portable intermediate code), so the
# NVIDIA driver has to compile it on a pod's first launch: about 55 s per model. `code=sm_NN` emits machine code.
# Targets: sm_86 (A40, RTX A6000), sm_89 (L40, L40S, RTX 6000 Ada), and compute_80 PTX as the fallback for other cards.
CU12='/^else ifdef LLAMA_ARCHES_CU12/,/KCPP_LIMIT_CUDA_MAX_ARCH=800/'
sed -i "${CU12}{/compute_50\|compute_61\|compute_75/d; s|code=compute_80 |code=compute_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_89,code=sm_89 |}" Makefile
sed -n "${CU12}p" Makefile | tee /tmp/BUILD_FLAGS
[ "$(grep -c gencode /tmp/BUILD_FLAGS)" = 1 ] && grep -q 'code=sm_86 .*code=sm_89' /tmp/BUILD_FLAGS \
  || { echo "Makefile no longer has the CUDA 12 flag block this script edits"; exit 1; }

TOTAL=$(make -n LLAMA_CUBLAS=1 LLAMA_PORTABLE=1 LLAMA_ARCHES_CU12=1 koboldcpp_cublas 2>/dev/null | grep -cE ' -c ' || true)
log "building with -j$JOBS, about $TOTAL files to compile"
( while sleep 60; do log "PROGRESS $(find . -name '*.o' | wc -l)/$TOTAL files"; done ) &
PROGRESS_PID=$!
make -j"$JOBS" LLAMA_CUBLAS=1 LLAMA_PORTABLE=1 LLAMA_ARCHES_CU12=1 koboldcpp_cublas > /tmp/make.log 2>&1 \
  || { kill $PROGRESS_PID; tail -60 /tmp/make.log; exit 1; }
kill $PROGRESS_PID
log "build done: $(ls -la koboldcpp_cublas.so)"
log "PTX targets in the library: $(strings -n 6 koboldcpp_cublas.so | grep -oE '^\.target sm_[0-9]+' | sort -u | tr '\n' ' ')"

# The same files the release's PyInstaller bundle carries, minus the non-CUDA backends.
mkdir -p /out/koboldcpp
cp -r koboldcpp.py json_to_gbnf.py kcpp_adapters embd_res koboldcpp_cublas.so \
      LICENSE.md MIT_LICENSE_GGML_SDCPP_LLAMACPP_ONLY.md /out/koboldcpp/
git rev-parse HEAD > /out/koboldcpp/SOURCE_COMMIT
cp /tmp/BUILD_FLAGS /out/koboldcpp/BUILD_FLAGS
tar -C /out -czf "/out/$NAME.tar.gz" koboldcpp
( cd /out && sha256sum "$NAME.tar.gz" | tee "$NAME.tar.gz.sha256" )
ls -la /out

pip3 install -q huggingface_hub
hf upload "$HF_REPO" "/out/$NAME.tar.gz" "$NAME.tar.gz" --commit-message "koboldcpp $KCPP_VERSION, CUDA 12, sm_86 + sm_89"
hf upload "$HF_REPO" "/out/$NAME.tar.gz.sha256" "$NAME.tar.gz.sha256" --commit-message "checksum for $NAME"
log "UPLOAD_DONE $NAME.tar.gz"