#!/bin/bash
# vast.ai node bootstrap for jspace (see vast-ai-notes.md for pitfalls)
set -x
source /venv/main/bin/activate
uv pip install -q "transformers==5.13.*" datasets accelerate hf_transfer
export HF_HOME=/dev/shm/hf; mkdir -p /dev/shm/hf
cd /dev/shm && git clone https://git.draic.info/nils/jspace.git 2>/dev/null || (cd jspace && git pull origin main)
cd /dev/shm/jspace
rclone copy jspace:jspace/results-loop /dev/shm/jspace/results-loop/ --include "*.json" --include "adapter_code_e399.pt"
export HF_HUB_ENABLE_HF_TRANSFER=1
# downloads can hang mid-transfer on some vast.ai networks (xet stall):
# retry loop, resume is free
# PIN revisions: upstream updated gemma-4-12B-it's chat template (added a
# thought channel) mid-project; unpinned downloads silently change behavior.
for spec in "google/gemma-4-12B-it 0e2b1058541244490925fbacf8972041435691ac" \
            "google/gemma-4-E2B-it 9dbdf8a839e4e9e0eb56ed80cc8886661d3817cf"; do
  set -- $spec
  for i in 1 2 3 4 5; do
    timeout 900 hf download "$1" --revision "$2" && break
    echo "RETRY $i: $1"; sleep 5
  done
  echo "$2" > "$HF_HOME/hub/models--${1//\//--}/refs/main"
done
# probe gate: one verified generation before anything batch
export HF_HUB_OFFLINE=1 JLENS_MODEL=google/gemma-4-12B-it
python - <<'PY'
import torch, sys
sys.path.insert(0, "/dev/shm/jspace/scripts"); sys.path.insert(0, "/dev/shm/jspace")
from jlens.core import load_model
m, tok = load_model(dtype=torch.bfloat16)
ids = tok.apply_chat_template([{"role":"user","content":"What is 2+2? Just the number."}],
      return_tensors="pt", add_generation_prompt=True, return_dict=True)["input_ids"].cuda()
out = m.generate(ids, max_new_tokens=8, do_sample=False)
txt = tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
print("PROBE:", repr(txt)); assert "4" in txt, "PROBE FAILED"
PY
echo NODE-READY