1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
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)
|
