#!/usr/bin/env bash
# DeepSeek-V4.1-Flash UNCENSORED on vast.ai 4x RTX PRO 6000 (SM120) — whole onstart in one script.
# vast onstart:  curl -fsSL --retry 5 "$BASE/bootstrap.sh" | BASE="$BASE" bash
# Env from `vastai create --env`: API_KEY (optional), MODEL_REPO (optional), CTX/MEMFRAC/MAXREQ (optional).
set -uo pipefail
BASE=${BASE:?BASE url of the bundle is required}
W=/workspace
mkdir -p "$W/models"
exec > >(tee -a "$W/bootstrap.log") 2>&1
echo "[$(date -u +%T)] bootstrap start, bundle $BASE"

# lmsysorg/sglang dsv41 images ship /root with modes sshd StrictModes rejects; vast may write
# authorized_keys after onstart begins, so keep re-fixing for a couple of minutes.
fix_ssh() {
  chown root:root /root /root/.ssh /root/.ssh/authorized_keys 2>/dev/null
  chmod 700 /root /root/.ssh 2>/dev/null
  chmod 600 /root/.ssh/authorized_keys 2>/dev/null
}
fix_ssh
( for _ in $(seq 1 60); do fix_ssh; sleep 2; done ) &
env | grep _ >> /etc/environment

for f in serve.sh patch_sm120_extra_split.py SHA256SUMS; do
  curl -fsSL --retry 5 "$BASE/$f" -o "$W/$f" || { echo "fetch $f failed"; exit 1; }
done
(cd "$W" && sha256sum -c SHA256SUMS) || { echo "bundle checksum mismatch"; exit 1; }
chmod +x "$W/serve.sh"

# SM120 V4.1 prefill fix (extra KV cache 128 -> 64 page split); idempotent.
python3 "$W/patch_sm120_extra_split.py" || { echo "patch failed"; exit 1; }

MODEL_REPO=${MODEL_REPO:-dealignai/DeepSeek-V4.1-Flash-UNCENSORED-FP8}
M="$W/models/${MODEL_REPO##*/}"
pip install -q -U "huggingface_hub[hf_xet]" || true
export HF_XET_HIGH_PERFORMANCE=1
for try in 1 2 3; do
  echo "[$(date -u +%T)] download $MODEL_REPO (try $try)"
  hf download "$MODEL_REPO" --local-dir "$M" --max-workers 32 && break
  sleep 15
done
first=$(ls "$M"/model-*-of-*.safetensors 2>/dev/null | head -1)
[ -n "$first" ] || { echo "no shards downloaded"; exit 1; }
total=$(echo "$first" | sed -E 's/.*-of-0*([0-9]+)\.safetensors$/\1/')
n=$(ls "$M"/model-*-of-*.safetensors | wc -l)
[ "$n" = "$total" ] || { echo "shards $n/$total — download incomplete"; exit 1; }
echo "[$(date -u +%T)] model ok: $n/$total shards, $(du -sh "$M" | cut -f1)"

MODEL="$M" bash "$W/serve.sh"
for _ in $(seq 1 80); do
  grep -q "Uvicorn running" "$W/serve.log" 2>/dev/null && { echo "[$(date -u +%T)] SERVER READY"; exit 0; }
  pgrep -f "[s]glang.launch_server" >/dev/null || { echo "server died:"; tail -30 "$W/serve.log"; exit 1; }
  sleep 15
done
echo "server not ready after 20 min"; tail -20 "$W/serve.log"