From f68f6b00f047d9412113d7c2627bd7f6e96f2d62 Mon Sep 17 00:00:00 2001 From: yurenh Date: Mon, 31 Aug 2026 18:33:06 -0500 Subject: env bootstrap (optional venv) + self-check in the ladder runner Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GkgLsACEF6CCP7EUfA5fZe --- README.md | 3 ++- scripts/env_check.py | 23 +++++++++++++++++++++++ scripts/run_ladder.sh | 10 +++++++++- 3 files changed, 34 insertions(+), 2 deletions(-) create mode 100644 scripts/env_check.py diff --git a/README.md b/README.md index e0a2446..f7f803d 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,8 @@ pip install -e . # or: pip install -e .[dev] && pytest tests/ ## Run (one click, on the H200 node) ``` -./scripts/run_ladder.sh # all sizes x arms, sequential, resume-safe +./scripts/run_ladder.sh # env check+install -> data prep -> all sizes x arms (resume-safe) +VENV=1 ./scripts/run_ladder.sh # fresh node: bootstrap ./.venv first SIZES="s60" ARMS="bp zbp_n16" TOKENS=3e8 ./scripts/run_ladder.sh # any subset / budget DRY=1 ./scripts/run_ladder.sh # print the plan ``` 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" -- cgit v1.2.3