#!/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 && rclone copy jspace:jspace/git/jspace.bundle /dev/shm/
git clone /dev/shm/jspace.bundle jspace 2>/dev/null || (cd jspace && git pull /dev/shm/jspace.bundle 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
hf download google/gemma-4-12B-it
hf download google/gemma-4-E2B-it
# 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
