spb/wpai0001_uqo

1#!/usr/bin/env python32"""Environment audit: hardware, Python stack, backends, cached open-weight checkpoints, cluster nodes.34Writes docs/ENVIRONMENT_AUDIT.json (and prints a short report). Cluster probing is optional5(``--cluster``) and uses the user's existing SSH aliases; nodes that do not answer are reported as6unreachable rather than assumed online. No secrets are read or written.7"""8from __future__ import annotations910import argparse11import json12import os13import subprocess14import sys15from pathlib import Path1617ROOT = Path(__file__).resolve().parents[1]18sys.path.insert(0, str(ROOT / "src"))19from wp_ai_0001.instrumentation.provenance import hardware, library_versions, utc_now  # noqa: E4022021NODES = ["M3U96a", "M3U96b", "M2U64", "M4M64a", "M4M64b", "M4M64c", "M5M36", "M4M36", "M4BP48", "m4mm", "m4mn"]22REMOTE_CMD = (23    'echo "$(sysctl -n machdep.cpu.brand_string)|$(( $(sysctl -n hw.memsize)/1073741824 ))|'24    '$(df -k ~ | tail -1 | awk \'{print int($4/1048576)}\')|$(python3 --version 2>&1 | cut -d" " -f2)|'25    '$(python3 -c "import mlx_lm;print(mlx_lm.__version__)" 2>/dev/null || echo none)|'26    '$(python3 -c "import torch;print(torch.__version__)" 2>/dev/null || echo none)|'27    '$(ls ~/.cache/huggingface/hub 2>/dev/null | grep -c models--)"'28)293031def hf_cache() -> list[dict]:32    hub = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) / "hub"33    out = []34    if hub.exists():35        for d in sorted(hub.glob("models--*")):36            size = sum(f.stat().st_size for f in d.rglob("*") if f.is_file())37            out.append({"repo": d.name[len("models--"):].replace("--", "/"), "size_gb": round(size / 2**30, 2)})38    return out394041def backends() -> dict:42    b: dict = {}43    try:44        import mlx.core as mx4546        b["mlx"] = {"version": mx.__version__, "metal": bool(mx.metal.is_available())}47    except Exception as e:  # noqa: BLE00148        b["mlx"] = {"error": repr(e)}49    try:50        import torch5152        b["torch"] = {"version": torch.__version__, "mps": bool(torch.backends.mps.is_available()), "cuda": bool(torch.cuda.is_available())}53    except Exception as e:  # noqa: BLE00154        b["torch"] = {"error": repr(e)}55    return b565758def probe_nodes() -> dict:59    res = {}60    for n in NODES:61        try:62            out = subprocess.run(["ssh", "-o", "ConnectTimeout=6", "-o", "BatchMode=yes", n, REMOTE_CMD], capture_output=True, text=True, timeout=30)63            line = out.stdout.strip().splitlines()[-1] if out.stdout.strip() else ""64            if out.returncode != 0 or "|" not in line:65                res[n] = {"reachable": False, "error": (out.stderr or out.stdout).strip()[:120]}66                continue67            chip, mem, free, py, mlxlm, torch_v, hf = line.split("|")68            res[n] = {"reachable": True, "chip": chip, "memory_gb": int(mem), "free_disk_gb": int(free), "python": py, "mlx_lm": mlxlm, "torch": torch_v, "hf_models_cached": int(hf)}69        except Exception as e:  # noqa: BLE00170            res[n] = {"reachable": False, "error": repr(e)[:120]}71    return res727374def main() -> None:75    ap = argparse.ArgumentParser()76    ap.add_argument("--cluster", action="store_true")77    args = ap.parse_args()78    report = {"utc": utc_now(), "hardware": hardware(), "libraries": library_versions(), "backends": backends(), "hf_cache": hf_cache()}79    if args.cluster:80        report["cluster"] = probe_nodes()81    (ROOT / "docs" / "ENVIRONMENT_AUDIT.json").write_text(json.dumps(report, indent=2, default=str))82    print(json.dumps({k: report[k] for k in ("hardware", "backends")}, indent=2, default=str))83    print("cached checkpoints:", [m["repo"] for m in report["hf_cache"]])84    if args.cluster:85        for n, r in report["cluster"].items():86            print(f"  {n}: {'OK ' + r['chip'] + ' ' + str(r['memory_gb']) + 'GB mlx_lm=' + r['mlx_lm'] if r['reachable'] else 'UNREACHABLE'}")87    print("written docs/ENVIRONMENT_AUDIT.json")888990if __name__ == "__main__":91    main()92