File size: 3,141 Bytes
7985e7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | #!/usr/bin/env python3
"""Check that `direct-d3il` can actually run this folder on this machine.
python tools/setup_avoiding_gpu.py # check, and fix what it can
python tools/setup_avoiding_gpu.py --check # report only
Unlike the maze2d application there is nothing to download: the 96 expert
episodes are vendored under DiRecT/d3il/environments/dataset/data/avoiding/ and
the planner is in checkpoints/. This only verifies the environment.
"""
import argparse, os, subprocess, sys
HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ENV = os.environ.get("AVOIDING_ENV") or os.path.dirname(os.path.dirname(sys.executable))
PY_, PIP = f"{ENV}/bin/python", f"{ENV}/bin/pip"
NEEDED = ("torch", "numpy", "gym", "mujoco", "casadi", "einops", "matplotlib", "pyyaml")
CUDA_INDEX = "https://download.pytorch.org/whl/cu128"
def gpu_ok():
"""True when the installed torch can drive THIS machine's GPU.
Compares the device's own compute capability against torch's compiled arch
list rather than hardcoding one: cu124 runs an H100 (sm_90) but not a 5090
(sm_120), and the reverse check would force a needless reinstall.
"""
probe = ("import torch\n"
"if not torch.cuda.is_available():\n"
" print('no-cuda', torch.__version__); raise SystemExit\n"
"a = 'sm_%d%d' % torch.cuda.get_device_capability(0)\n"
"print(a, a in torch.cuda.get_arch_list(), torch.__version__)")
out = subprocess.run([PY_, "-c", probe], capture_output=True, text=True).stdout.strip()
return (" True " in f" {out} "), out or "torch unusable"
def have(pkg):
return subprocess.run([PIP, "show", pkg], capture_output=True).returncode == 0
def main():
ap = argparse.ArgumentParser(); ap.add_argument("--check", action="store_true")
a = ap.parse_args()
ok, detail = gpu_ok()
missing = [p for p in NEEDED if not have(p)]
ckpt = os.path.join(HERE, "checkpoints", "model_ema_10.pth")
data = os.path.join(HERE, "DiRecT", "d3il", "environments", "dataset", "data", "avoiding")
print(f"env : {ENV}")
print(f"torch on this GPU : {'OK' if ok else 'NO'} ({detail})")
print(f"missing packages : {missing or 'none'}")
print(f"planner present : {'OK' if os.path.exists(ckpt) else 'NO'} ({ckpt})")
print(f"dataset present : {'OK' if os.path.isdir(data) else 'NO'} ({data})")
if a.check:
return
if not ok:
print(" installed torch cannot drive this GPU; switching to cu128")
subprocess.run([PIP, "install", "--index-url", CUDA_INDEX, "--force-reinstall",
"--no-deps", "torch==2.11.0+cu128"], check=True)
subprocess.run([PIP, "install", "--index-url", CUDA_INDEX, "torch==2.11.0+cu128"], check=True)
if missing:
subprocess.run([PIP, "install", *missing], check=True)
ok, detail = gpu_ok()
print(f"\nfinal: torch {'OK' if ok else 'NO'} ({detail})")
if not ok:
sys.exit("[FAIL] torch still cannot see the GPU")
print("[done] ready to fit dynamics and evaluate")
if __name__ == "__main__":
main()
|