#!/bin/bash
# Colab GPU VM setup for the decision-model run. Started in the background; logs to /content/logs.
#   1. hub / datasets packages, then the model download and the Jev baseline start in parallel
#   2. the Unsloth stack (same pins as Unsloth's Qwen3.8 (27B) Colab notebook)
#   3. an import check that FastDecisionModel is available
set -x
set -a
source /content/secrets.env
set +a
export PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_ROOT_USER_ACTION=ignore
MODEL_REPO="${MODEL_REPO:-Qwen/Qwen3.8-27B}"
MODEL_DIR="${MODEL_DIR:-/content/models/Qwen3.8-27B}"
mkdir -p /content/logs /content/out /content/models

nvidia-smi
df -h /content
free -g
python -c "import sys, torch; print('python', sys.version); print('torch', torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))"

pip install -q sentencepiece protobuf "datasets==4.3.0" "huggingface_hub>=0.34.0" hf_transfer

cd /content/dm
nohup python download.py "$MODEL_REPO" "$MODEL_DIR" > /content/logs/download.log 2>&1 &
nohup python jev_eval.py --out /content/out/jev_td_test.jsonl > /content/logs/jev.log 2>&1 &

pip install -q --no-deps unsloth_zoo bitsandbytes accelerate peft trl triton unsloth
unsloth install-kernels || echo "WARN: unsloth install-kernels failed"
pip install -q --no-deps --upgrade "torchao>=0.16.0"
pip install -q transformers==5.15.1
pip install -q --no-deps trl==0.22.2

python - <<'EOF'
import importlib
for name in ["torch", "unsloth", "unsloth_zoo", "transformers", "peft", "trl", "datasets",
             "huggingface_hub", "bitsandbytes", "accelerate", "fla", "causal_conv1d", "torchao", "triton"]:
    try:
        module = importlib.import_module(name)
        print("VERSION", name, getattr(module, "__version__", "?"))
    except Exception as exc:
        print("VERSION", name, "ERR", type(exc).__name__, str(exc)[:200])
from unsloth import DecisionTrainer, FastDecisionModel
print("FASTDECISION_OK")
EOF
echo SETUP_DONE