#!/bin/bash
# Run ON the Colab G4 VM (as root, which is what the Colab runtime gives you).
# Installs vLLM, the precompiled FlashInfer kernels, and the model.
#
# Order matters: install the kernels BEFORE the model is ever loaded.
# If FlashInfer has to JIT-compile fused_moe_120 on first use, nvcc/cicc run
# next to ~67 GB of pinned PLE host memory and the Linux OOM killer takes out
# VLLM::EngineCore. That is a host OOM, not a GPU OOM.
#
# Run detached and poll the log; the colab exec / ssh session may drop:
#   nohup bash /content/setup_vm.sh > /content/setup.log 2>&1 &
set -euo pipefail

MODEL="nvidia/Qwen3.8-Flash-Next-NVFP4"
REVISION="fc694b54fb0174e0913e6adf86691ef85a4ead47"   # the exact revision I ran
DEST="/content/models/Qwen3.8-Flash-Next-NVFP4"

echo "[1/4] python deps (observed: ~12-26 s)"
uv pip install --system vllm==0.32.0 hf_transfer

echo "[2/4] precompiled FlashInfer kernels, BEFORE any model load (observed: ~283 s)"
MAX_JOBS=1 vllm download-kernels

echo "[3/4] model download, pinned revision (observed: ~6.5 min for 124 GiB)"
mkdir -p /content/models
HF_HUB_ENABLE_HF_TRANSFER=1 hf download "$MODEL" --revision "$REVISION" --local-dir "$DEST"

echo "[4/4] drop page cache so the 124 GiB of just-written shards do not crowd the loader"
sync
echo 3 > /proc/sys/vm/drop_caches 2>/dev/null || true

du -sh /content/models
df -h / | tail -1
echo "SETUP_DONE"