env bootstrap (optional venv) + self-check in the ladder runner

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe

2 files changed, 32 insertions, 1 deletions

diff --git a/scripts/env_check.py b/scripts/env_check.py
new file mode 100644
index 0000000..92d18a1
--- /dev/null
+++ b/scripts/env_check.py
@@ -0,0 +1,23 @@
+"""Fast environment self-check with actionable messages. Exit 0 = ready to train."""
+import sys
+
+ok = True
+def item(name, fn, hint):
+    global ok
+    try:
+        msg = fn()
+        print(f"  [ok] {name}: {msg}")
+    except Exception as e:
+        ok = False
+        print(f"  [FAIL] {name}: {type(e).__name__}: {e}\n         -> {hint}")
+
+print("== zbp-scaling environment check")
+item("python >= 3.10", lambda: (sys.version.split()[0], 1/0 if sys.version_info < (3, 10) else "")[0], "use python3.10+")
+item("torch + CUDA", lambda: __import__("torch").__version__ + f", {__import__('torch').cuda.device_count()} GPU(s)"
+     + ("" if __import__("torch").cuda.is_available() else (_ for _ in ()).throw(RuntimeError("cuda not available"))),
+     "install a CUDA build of torch (pip install torch --index-url https://download.pytorch.org/whl/cu126) or load the cluster module")
+item("zbp_scaling package", lambda: __import__("zbp_scaling").__name__, "pip install -e .  (from the repo root)")
+item("numpy / yaml / datasets", lambda: ",".join(__import__(m).__name__ for m in ("numpy", "yaml", "datasets")), "pip install -e .")
+item("tiktoken GPT-2 vocab", lambda: f"{__import__('tiktoken').get_encoding('gpt2').n_vocab} tokens",
+     "first use needs network to fetch the BPE files; on air-gapped nodes pre-seed the tiktoken cache dir")
+print("== ready" if ok else "== NOT ready"); sys.exit(0 if ok else 1)
diff --git a/scripts/run_ladder.sh b/scripts/run_ladder.sh
index 65b43a9..17219bd 100755
--- a/scripts/run_ladder.sh
+++ b/scripts/run_ladder.sh
@@ -30,7 +30,15 @@ EXTRA="--set ${TOKENS:+tokens=$TOKENS} ${SET:-}"
 
 echo "== zbp-scaling ladder: sizes=[$SIZES] arms=[$ARMS] nproc=$NPROC data=$DATA out=$OUT ${TOKENS:+tokens=$TOKENS}"
 [ "$NPROC" -ge 1 ] || { echo "no GPUs detected; set NPROC"; exit 1; }
-python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q
+# --- environment bootstrap: current env first; VENV=1 creates ./.venv (fresh nodes) ---
+PY=${PY:-python3}
+if [ "${VENV:-0}" = 1 ]; then
+  [ -d .venv ] || "$PY" -m venv .venv
+  . .venv/bin/activate
+fi
+python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q || {
+  echo "pip install failed; on managed nodes run with VENV=1 (creates ./.venv) or install into your conda env"; exit 1; }
+python scripts/env_check.py || exit 1
 
 if [ ! -f "$DATA/train.bin" ]; then
   echo "== preparing data ($PREP_TOKENS tokens -> $DATA) — run this on the training node"