From 77ad65e949e4bc4bd2f9fdc4dcfe1863fad43b74 Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Mon, 5 Oct 2026 12:59:23 -0500 Subject: [PATCH 1/4] nvidia_configs: add nequip and allegro envs hosting the nequip.net foundation models nequip.py serves NequIP-OAM-{S,M,L,XL} and NequIP-MP-L; allegro.py serves Allegro-OAM-L and Allegro-MP-L; both declare a :custom entry with a setup_from_path hook (.nequip.zip, .ckpt, or a ready .nequip.pt2). Package files are pinned Zenodo artifacts, md5-checked and fetched directly into the shared cache: nequip's own nequip.net loader queries the nequip.net API even on a cache hit, which fails on offline workers. On GPU, setup() AOTInductor-compiles each model once per (model, sm arch, torch, nequip) into {cache}/nequip/compiled, or a per-user dir under TORCHINDUCTOR_CACHE_DIR when the shared cache is read-only. The compile runs in a child process that moves its own output into place, so a worker killed by a startup timeout still leaves a usable cache entry. When no CUDA toolkit is loaded, CUDA headers come from the env's cuda-toolkit[cccl,nvcc] wheels and -lcuda resolves through a temporary symlink to the loaded driver. CPU (including rootstock add's download step) runs eager. modal_app.py gains probe_nequip/probe_allegro plus extra_index_url, cpu and timeout options on probe_image, so the probes build against the configs' cu128 torch. Validated on Modal (A10G, 8 CPU): fresh compiles of Allegro-OAM-L ~2 min, NequIP-OAM-L ~4.6 min, NequIP-OAM-XL ~6.4 min. Co-Authored-By: Claude Opus 5.5 --- sample_model_configurations/README.md | 6 +- sample_model_configurations/modal_app.py | 51 +++- .../nvidia_configs/allegro.py | 282 +++++++++++++++++ .../nvidia_configs/nequip.py | 286 ++++++++++++++++++ 4 files changed, 621 insertions(+), 4 deletions(-) create mode 100644 sample_model_configurations/nvidia_configs/allegro.py create mode 100644 sample_model_configurations/nvidia_configs/nequip.py diff --git a/sample_model_configurations/README.md b/sample_model_configurations/README.md index 502d5d7..cdc2786 100644 --- a/sample_model_configurations/README.md +++ b/sample_model_configurations/README.md @@ -12,6 +12,7 @@ sample_model_configurations/ │ ├── failure_modes.md │ └── README.md └── nvidia_configs/ # configs targeting NVIDIA GPUs (the release catalog) + ├── allegro.py # Allegro-OAM-L / Allegro-MP-L (nequip.net, AOTInductor-compiled) ├── allscaip.py # AllScAIP (FAIRChem scalable-attention, OMol25) ├── ani.py # ANI-2x / ANI-1ccx / ANI-1x (organic molecules) ├── chgnet.py # CHGNet (inorganic, charge-informed) @@ -22,6 +23,7 @@ sample_model_configurations/ ├── mace.py # MACE foundations: MP-0, OFF23, MPA-0, MATPES, MH-1, OMOL ├── mace_polar.py # MACE-POLAR-1 (electrostatic MACE + graph-longrange) ├── mattersim.py # MatterSim-v1 (Microsoft universal) + ├── nequip.py # NequIP-OAM-{S,M,L,XL} / NequIP-MP-L (nequip.net, AOTInductor-compiled) ├── orb.py # Orb v2 (kept for orb-d3-v2 only; v3 is primary) ├── orb_v3.py # Orb v3 (Orbital Materials universal) ├── pet.py # UPET / PET-MAD successor (lab-cosmo) @@ -38,8 +40,8 @@ hardware-agnostic. Pruned 2026-07-30 (still on some clusters until their next cleanup): the legacy fairchem-core 1.x OC20 baselines beyond escn/equiformer (dimenet, gemnet, painn, schnet, scn), the BYO-weights-only envs with no shippable -checkpoints (allegro, nequip — the `:custom` mechanism is the -supported path for user weights now), and orb v2's superseded checkpoints +checkpoints (allegro, nequip — since re-added hosting the nequip.net +foundation models), and orb v2's superseded checkpoints (orb-v2, orb-mptraj-only-v2). > **Python floor caveat.** rootstock itself requires Python >=3.11, and the diff --git a/sample_model_configurations/modal_app.py b/sample_model_configurations/modal_app.py index 9d217b3..e807e16 100644 --- a/sample_model_configurations/modal_app.py +++ b/sample_model_configurations/modal_app.py @@ -63,9 +63,12 @@ def probe_image( *, python_version: str = "3.11", find_links: str | None = None, + extra_index_url: str | None = None, apt_packages: list[str] | None = None, no_deps: list[str] | None = None, gpu: str = "A10G", + cpu: float | None = None, + timeout: int = 900, ): """Decorator factory: builds a Modal image from deps and wires up app.function. @@ -75,7 +78,14 @@ def probe_image( img = modal.Image.debian_slim(python_version=python_version) if apt_packages: img = img.apt_install(*apt_packages) - img = img.uv_pip_install(*deps, **({"find_links": find_links} if find_links else {})) + index_opts = {} + if find_links: + index_opts["find_links"] = find_links + if extra_index_url: + # uv otherwise takes every package the extra index carries from it alone. + index_opts["extra_index_url"] = extra_index_url + index_opts["extra_options"] = "--index-strategy unsafe-best-match" + img = img.uv_pip_install(*deps, **index_opts) if no_deps: img = img.run_commands( f"/.uv/uv pip install --system --no-deps {' '.join(repr(p) for p in no_deps)}" @@ -88,7 +98,8 @@ def probe_image( gpu=gpu, volumes={CACHE_MOUNT: model_cache}, secrets=[HF_SECRET], - timeout=900, + timeout=timeout, + **({"cpu": cpu} if cpu else {}), ) @@ -382,3 +393,39 @@ def probe_uma_1p2p1(checkpoint: str = "uma-s-1p2p1", system: str = "crystal"): def probe_uma_medium(checkpoint: str = "uma-m-1p1", system: str = "crystal"): """Probe UMA Medium on OMAT-style bulk materials.""" return _run_probe_subprocess(checkpoint, system, kwargs=["task=omat"]) + + +# NequIP/Allegro compile models with AOTInductor on first GPU use, which needs +# a host C++ compiler (debian_slim ships none) and CUDA headers (from the +# nvidia-cuda-* wheels of the cu128 torch build, as in the configs). The +# compile is CPU-bound and takes many minutes for the larger models. +@probe_image( + "nequip.py", + ["nequip>=0.19,<0.20", "ase>=3.23", "torch>=2.6,<2.12", "cuda-toolkit[cccl,nvcc]>=12.8,<12.9"], + extra_index_url="https://download.pytorch.org/whl/cu128", + apt_packages=["build-essential"], + cpu=8, + timeout=3600, +) +def probe_nequip(checkpoint: str = "nequip-oam-l", system: str = "crystal"): + """Probe NequIP-OAM (AOTInductor-compiled on first load).""" + return _run_probe_subprocess(checkpoint, system) + + +@probe_image( + "allegro.py", + [ + "nequip>=0.19,<0.20", + "nequip-allegro>=0.8.3,<0.9", + "ase>=3.23", + "torch>=2.6,<2.12", + "cuda-toolkit[cccl,nvcc]>=12.8,<12.9", + ], + extra_index_url="https://download.pytorch.org/whl/cu128", + apt_packages=["build-essential"], + cpu=8, + timeout=3600, +) +def probe_allegro(checkpoint: str = "allegro-oam-l", system: str = "crystal"): + """Probe Allegro-OAM (AOTInductor-compiled on first load).""" + return _run_probe_subprocess(checkpoint, system) diff --git a/sample_model_configurations/nvidia_configs/allegro.py b/sample_model_configurations/nvidia_configs/allegro.py new file mode 100644 index 0000000..0414b29 --- /dev/null +++ b/sample_model_configurations/nvidia_configs/allegro.py @@ -0,0 +1,282 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# # setup() uses nequip's loader internals and keys its compile cache on the +# # nequip version, so stay within one minor release. +# "nequip>=0.19,<0.20", +# "nequip-allegro>=0.8.3,<0.9", +# "ase>=3.23", +# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit +# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA +# # newer than many cluster drivers support. +# "torch>=2.6,<2.12", +# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the +# # runtime headers come with torch); matches the cu128 torch build. +# "cuda-toolkit[cccl,nvcc]>=12.8,<12.9", +# ] +# +# [tool.uv.sources] +# torch = { index = "pytorch-cu128" } +# +# [[tool.uv.index]] +# name = "pytorch-cu128" +# url = "https://download.pytorch.org/whl/cu128" +# explicit = true +# /// +"""Allegro env — hosts the Allegro foundation models from nequip.net. + +Allegro is a strictly local E(3)-equivariant interatomic potential from the +Kozinsky group (Harvard), built on the NequIP framework; its locality is what +lets it scale across many GPUs in LAMMPS. Allegro-OAM-L is trained on OMat24 +and fine-tuned on sAlex + MPtrj; Allegro-MP-L is trained on MPtrj alone. All +are PBE-level and licensed CC-BY-4.0. Versions are pinned to fixed Zenodo +artifacts (md5-checked) so rebuilds serve the same weights; the package file +is downloaded into the shared model cache at `rootstock add` time and loaded +from there, so serving needs no network. + +Compilation: on GPU, models run through NequIP's AOTInductor path +(``nequip-compile --mode aotinductor --target ase``), the path NequIP +recommends for inference. A compiled model only runs on the hardware and +software it was compiled with, so setup() compiles once per (model, GPU +architecture, torch version, nequip version) and caches the result: in the +shared model cache when the caller can write there, otherwise in the caller's +TorchInductor cache directory. The first GPU load of a model on a new GPU type +therefore compiles first: on an 8-core node with an A10G, about 2 minutes for +Allegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs +into the worker startup timeout, raise it (``timeout=`` on +RootstockCalculator, ``--verify-timeout`` for verification); an interrupted +compile still finishes in the background and later loads reuse the compiled +file. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or +whatever ``CXX`` names) at compile time; CUDA headers come from the env's own +NVIDIA wheels when no CUDA toolkit is loaded. Pass +``setup_kwargs={"compile": False}`` to run eager instead. CPU loads always run +eager. + +Custom weights: ``allegro:custom`` with weights= accepts a packaged model +(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already +compiled for ASE on this hardware (``.nequip.pt2``). If your model's type +names are not chemical symbols, pass +``setup_kwargs={"chemical_species_to_atom_type_map": {"H": "", ...}}``. +""" + +import hashlib +import os +import shutil +import subprocess +import sys +import tempfile +import urllib.request +from pathlib import Path + +CHECKPOINTS = { + "allegro-oam-l": "Allegro-OAM-L-0.1", + "allegro-mp-l": "Allegro-MP-L-0.1", + # Your own Allegro model: pair with weights= (loaded via setup_from_path). + "allegro:custom": None, +} + +# Zenodo record and md5 of each pinned package file (the artifacts nequip.net +# resolves mir-group/:0.1 to). +_ARTIFACTS = { + "Allegro-OAM-L-0.1": ("16980200", "0db7f9b3c3a62e74d78b3fcf2973c462"), + "Allegro-MP-L-0.1": ("16980200", "c68d5345e925ee444543d6d4d968161f"), +} + + +def _shared_cache() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "nequip" + + +def _package_path(name: str) -> Path: + """The pinned package file in the shared cache, downloading it if missing.""" + path = _shared_cache() / "packages" / f"{name}.nequip.zip" + if path.is_file(): + return path + record, md5 = _ARTIFACTS[name] + url = f"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f".{path.name}.{os.getpid()}.part") + try: + digest = hashlib.md5() + with urllib.request.urlopen(url) as response, open(tmp, "wb") as f: + while chunk := response.read(1 << 20): + f.write(chunk) + digest.update(chunk) + if digest.hexdigest() != md5: + raise RuntimeError(f"{url}: md5 {digest.hexdigest()} does not match {md5}") + os.replace(tmp, path) + finally: + tmp.unlink(missing_ok=True) + return path + + +def _compiled_dirs() -> list[Path]: + """Where compiled models are looked up, in order; the first writable one + receives new compilations.""" + dirs = [_shared_cache() / "compiled"] + inductor = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + if inductor: + dirs.append(Path(inductor) / "nequip") + else: + dirs.append(Path.home() / ".cache" / "nequip" / "compiled") + return dirs + + +def _writable(directory: Path) -> bool: + probe = directory + while not probe.exists(): + probe = probe.parent + return os.access(probe, os.W_OK | os.X_OK) + + +def _hardware_tag(device: str) -> str: + import torch + + major, minor = torch.cuda.get_device_capability(torch.device(device)) + return f"sm{major}{minor}" + + +def _toolkit_env(env: dict, link_dir: Path) -> None: + """Point AOTInductor at CUDA headers and libcuda when no toolkit is loaded. + + The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links + ``-lcuda``. Compute nodes often expose a CUDA toolkit only through a + module, so when none is visible fall back to the headers in NVIDIA's CUDA + wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``), and + link against the driver library this process has loaded. + """ + if env.get("CUDA_HOME") or env.get("CUDA_PATH") or shutil.which("nvcc", path=env.get("PATH")): + return + import nvidia + + headers = next( + (h for base in nvidia.__path__ for h in Path(base).glob("*/include/cuda.h")), None + ) + if headers is None: + return + env["CUDA_HOME"] = str(headers.parent.parent) + extra = [ + str(h.parent.parent) + for base in nvidia.__path__ + for marker in ("crt/host_defines.h", "nv/target") + for h in Path(base).glob(f"*/include/{marker}") + ] + env["CPATH"] = os.pathsep.join(filter(None, [*extra, env.get("CPATH")])) + with open("/proc/self/maps") as maps: + libcuda = next((line.split()[-1] for line in maps if "/libcuda.so" in line), None) + if libcuda is not None: + (link_dir / "libcuda.so").symlink_to(libcuda) + env["LIBRARY_PATH"] = os.pathsep.join( + filter(None, [str(link_dir), env.get("LIBRARY_PATH")]) + ) + + +def _compiled_model(source: Path, key: str, device: str) -> Path: + """Path to ``source`` compiled for ASE on ``device``, compiling on a miss.""" + import nequip + import torch + + tag = f"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}" + name = f"{tag.replace('+', '_')}.nequip.pt2" + dirs = _compiled_dirs() + for directory in dirs: + if (directory / name).is_file(): + return directory / name + + target_dir = next((d for d in dirs if _writable(d)), None) + if target_dir is None: + raise RuntimeError(f"no writable directory to cache the compiled model in: {dirs}") + target_dir.mkdir(parents=True, exist_ok=True) + out = target_dir / name + tmp = target_dir / f".{tag}.{os.getpid()}.nequip.pt2" + + # Compile in a child process pinned to the same GPU, so the compiler's + # memory is released before the model is loaded. The child moves its + # output into place itself: if this worker is killed mid-compile (e.g. by + # a startup timeout), the compile still completes and lands in the cache. + env = dict(os.environ) + if ":" in device: + index = int(device.split(":", 1)[1]) + visible = env.get("CUDA_VISIBLE_DEVICES") + env["CUDA_VISIBLE_DEVICES"] = visible.split(",")[index] if visible else str(index) + compile_and_move = ( + "import os, sys; from nequip.scripts.compile import main; " + "main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])" + ) + cmd = [ + sys.executable, + "-c", + compile_and_move, + str(source), + str(tmp), + "--mode", + "aotinductor", + "--device", + "cuda", + "--target", + "ase", + str(out), + ] + print(f"nequip: compiling {source.name} for {tag} (one-time, several minutes)", file=sys.stderr) + try: + with tempfile.TemporaryDirectory() as link_dir: + _toolkit_env(env, Path(link_dir)) + result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr) + if result.returncode != 0: + raise RuntimeError( + f"nequip-compile failed (exit {result.returncode}); see the log above. " + 'Pass setup_kwargs={"compile": False} to run eager instead.' + ) + finally: + tmp.unlink(missing_ok=True) + return out + + +def _load(source: Path, key: str, device: str, compile: bool, kwargs: dict): + # torch's AOTInductor package loader reaches torch._inductor.codecache + # without importing it, which fails in a fresh process. + import torch._inductor.codecache # noqa: F401 + from nequip.integrations.ase import NequIPCalculator + + if source.name.endswith(".nequip.pt2"): + return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs) + if compile and device != "cpu": + compiled = _compiled_model(source, key, device) + return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs) + # Eager load of a package or checkpoint. nequip exposes this only as a + # private classmethod (its own tests use it); the public loader takes + # compiled models only. + return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs) + + +def setup(checkpoint: str, device: str = "cuda", compile: bool = True, **kwargs): + """ + Load an Allegro calculator for a pinned nequip.net foundation model. + + Args: + checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS. + device: PyTorch device string (e.g., "cuda", "cuda:0", "cpu"). + compile: Run the AOTInductor-compiled model on GPU (compiling and + caching it on first use). False runs eager PyTorch. + + Returns: + ASE-compatible calculator. + """ + name = CHECKPOINTS[checkpoint] + kwargs.setdefault("chemical_species_to_atom_type_map", True) + return _load(_package_path(name), name, device, compile, kwargs) + + +def setup_from_path(path: str, device: str = "cuda", compile: bool = True, **kwargs): + # Custom checkpoints (`:custom` ids with user weights): an Allegro .nequip.zip + # package or .ckpt checkpoint (compiled on first GPU use, cached under a + # key that includes the file's hash) or a ready .nequip.pt2. + source = Path(path) + key = source.name + if not key.endswith(".nequip.pt2") and compile and device != "cpu": + digest = hashlib.sha256() + with open(source, "rb") as f: + while chunk := f.read(1 << 20): + digest.update(chunk) + key = f"custom-{digest.hexdigest()[:16]}" + return _load(source, key, device, compile, kwargs) diff --git a/sample_model_configurations/nvidia_configs/nequip.py b/sample_model_configurations/nvidia_configs/nequip.py new file mode 100644 index 0000000..d0dbbf5 --- /dev/null +++ b/sample_model_configurations/nvidia_configs/nequip.py @@ -0,0 +1,286 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# # setup() uses nequip's loader internals and keys its compile cache on the +# # nequip version, so stay within one minor release. +# "nequip>=0.19,<0.20", +# "ase>=3.23", +# # AOTInductor compilation needs torch>=2.6. Listed here so the explicit +# # cu128 index below routes it: PyPI's recent torch wheels bundle a CUDA +# # newer than many cluster drivers support. +# "torch>=2.6,<2.12", +# # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the +# # runtime headers come with torch); matches the cu128 torch build. +# "cuda-toolkit[cccl,nvcc]>=12.8,<12.9", +# ] +# +# [tool.uv.sources] +# torch = { index = "pytorch-cu128" } +# +# [[tool.uv.index]] +# name = "pytorch-cu128" +# url = "https://download.pytorch.org/whl/cu128" +# explicit = true +# /// +"""NequIP env — hosts the NequIP-OAM foundation models from nequip.net. + +NequIP is an E(3)-equivariant message-passing interatomic potential from the +Kozinsky group (Harvard). The OAM models are trained on OMat24 and fine-tuned +on sAlex + MPtrj; NequIP-MP-L is trained on MPtrj alone. All are PBE-level and +licensed CC-BY-4.0. Versions are pinned to fixed Zenodo artifacts +(md5-checked) so rebuilds serve the same weights; the package file is +downloaded into the shared model cache at `rootstock add` time and loaded from +there, so serving needs no network. + +Compilation: on GPU, models run through nequip's AOTInductor path +(``nequip-compile --mode aotinductor --target ase``), the path NequIP +recommends for inference. A compiled model only runs on the hardware and +software it was compiled with, so setup() compiles once per (model, GPU +architecture, torch version, nequip version) and caches the result: in the +shared model cache when the caller can write there, otherwise in the caller's +TorchInductor cache directory. The first GPU load of a model on a new GPU type +therefore compiles first: on an 8-core node with an A10G, about 2 minutes for +Allegro-OAM-L, 5 for NequIP-OAM-L and 6-7 for NequIP-OAM-XL. If that load runs +into the worker startup timeout, raise it (``timeout=`` on +RootstockCalculator, ``--verify-timeout`` for verification); an interrupted +compile still finishes in the background and later loads reuse the compiled +file. AOTInductor needs a C++17 host compiler (``c++``/``gcc`` >= 8, or +whatever ``CXX`` names) at compile time; CUDA headers come from the env's own +NVIDIA wheels when no CUDA toolkit is loaded. Pass +``setup_kwargs={"compile": False}`` to run eager instead. CPU loads always run +eager. + +Custom weights: ``nequip:custom`` with weights= accepts a packaged model +(``.nequip.zip``), a training checkpoint (``.ckpt``), or a model already +compiled for ASE on this hardware (``.nequip.pt2``). If your model's type +names are not chemical symbols, pass +``setup_kwargs={"chemical_species_to_atom_type_map": {"H": "", ...}}``. +""" + +import hashlib +import os +import shutil +import subprocess +import sys +import tempfile +import urllib.request +from pathlib import Path + +CHECKPOINTS = { + "nequip-oam-s": "NequIP-OAM-S-0.1", + "nequip-oam-m": "NequIP-OAM-M-0.1", + "nequip-oam-l": "NequIP-OAM-L-0.1", + "nequip-oam-xl": "NequIP-OAM-XL-0.1", + "nequip-mp-l": "NequIP-MP-L-0.1", + # Your own NequIP model: pair with weights= (loaded via setup_from_path). + "nequip:custom": None, +} + +# Zenodo record and md5 of each pinned package file (the artifacts nequip.net +# resolves mir-group/:0.1 to). +_ARTIFACTS = { + "NequIP-OAM-S-0.1": ("18775904", "399a98bf36fc4550bc85d48f35451c6c"), + "NequIP-OAM-M-0.1": ("18775904", "e5b44583f4421be702b35d76503754cd"), + "NequIP-OAM-L-0.1": ("18775904", "67144367c710a70a53a8e21acf331980"), + "NequIP-OAM-XL-0.1": ("18775904", "3d2369c7238eb83a23141abdcb055a8f"), + "NequIP-MP-L-0.1": ("18775904", "624a01d4cbbae2d3bb75dc7074532d13"), +} + + +def _shared_cache() -> Path: + return Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "nequip" + + +def _package_path(name: str) -> Path: + """The pinned package file in the shared cache, downloading it if missing.""" + path = _shared_cache() / "packages" / f"{name}.nequip.zip" + if path.is_file(): + return path + record, md5 = _ARTIFACTS[name] + url = f"https://zenodo.org/api/records/{record}/files/{name}.nequip.zip/content" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f".{path.name}.{os.getpid()}.part") + try: + digest = hashlib.md5() + with urllib.request.urlopen(url) as response, open(tmp, "wb") as f: + while chunk := response.read(1 << 20): + f.write(chunk) + digest.update(chunk) + if digest.hexdigest() != md5: + raise RuntimeError(f"{url}: md5 {digest.hexdigest()} does not match {md5}") + os.replace(tmp, path) + finally: + tmp.unlink(missing_ok=True) + return path + + +def _compiled_dirs() -> list[Path]: + """Where compiled models are looked up, in order; the first writable one + receives new compilations.""" + dirs = [_shared_cache() / "compiled"] + inductor = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + if inductor: + dirs.append(Path(inductor) / "nequip") + else: + dirs.append(Path.home() / ".cache" / "nequip" / "compiled") + return dirs + + +def _writable(directory: Path) -> bool: + probe = directory + while not probe.exists(): + probe = probe.parent + return os.access(probe, os.W_OK | os.X_OK) + + +def _hardware_tag(device: str) -> str: + import torch + + major, minor = torch.cuda.get_device_capability(torch.device(device)) + return f"sm{major}{minor}" + + +def _toolkit_env(env: dict, link_dir: Path) -> None: + """Point AOTInductor at CUDA headers and libcuda when no toolkit is loaded. + + The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links + ``-lcuda``. Compute nodes often expose a CUDA toolkit only through a + module, so when none is visible fall back to the headers in NVIDIA's CUDA + wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``), and + link against the driver library this process has loaded. + """ + if env.get("CUDA_HOME") or env.get("CUDA_PATH") or shutil.which("nvcc", path=env.get("PATH")): + return + import nvidia + + headers = next( + (h for base in nvidia.__path__ for h in Path(base).glob("*/include/cuda.h")), None + ) + if headers is None: + return + env["CUDA_HOME"] = str(headers.parent.parent) + extra = [ + str(h.parent.parent) + for base in nvidia.__path__ + for marker in ("crt/host_defines.h", "nv/target") + for h in Path(base).glob(f"*/include/{marker}") + ] + env["CPATH"] = os.pathsep.join(filter(None, [*extra, env.get("CPATH")])) + with open("/proc/self/maps") as maps: + libcuda = next((line.split()[-1] for line in maps if "/libcuda.so" in line), None) + if libcuda is not None: + (link_dir / "libcuda.so").symlink_to(libcuda) + env["LIBRARY_PATH"] = os.pathsep.join( + filter(None, [str(link_dir), env.get("LIBRARY_PATH")]) + ) + + +def _compiled_model(source: Path, key: str, device: str) -> Path: + """Path to ``source`` compiled for ASE on ``device``, compiling on a miss.""" + import nequip + import torch + + tag = f"{key}-{_hardware_tag(device)}-torch{torch.__version__}-nequip{nequip.__version__}" + name = f"{tag.replace('+', '_')}.nequip.pt2" + dirs = _compiled_dirs() + for directory in dirs: + if (directory / name).is_file(): + return directory / name + + target_dir = next((d for d in dirs if _writable(d)), None) + if target_dir is None: + raise RuntimeError(f"no writable directory to cache the compiled model in: {dirs}") + target_dir.mkdir(parents=True, exist_ok=True) + out = target_dir / name + tmp = target_dir / f".{tag}.{os.getpid()}.nequip.pt2" + + # Compile in a child process pinned to the same GPU, so the compiler's + # memory is released before the model is loaded. The child moves its + # output into place itself: if this worker is killed mid-compile (e.g. by + # a startup timeout), the compile still completes and lands in the cache. + env = dict(os.environ) + if ":" in device: + index = int(device.split(":", 1)[1]) + visible = env.get("CUDA_VISIBLE_DEVICES") + env["CUDA_VISIBLE_DEVICES"] = visible.split(",")[index] if visible else str(index) + compile_and_move = ( + "import os, sys; from nequip.scripts.compile import main; " + "main(sys.argv[1:-1]); os.replace(sys.argv[2], sys.argv[-1])" + ) + cmd = [ + sys.executable, + "-c", + compile_and_move, + str(source), + str(tmp), + "--mode", + "aotinductor", + "--device", + "cuda", + "--target", + "ase", + str(out), + ] + print(f"nequip: compiling {source.name} for {tag} (one-time, several minutes)", file=sys.stderr) + try: + with tempfile.TemporaryDirectory() as link_dir: + _toolkit_env(env, Path(link_dir)) + result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr) + if result.returncode != 0: + raise RuntimeError( + f"nequip-compile failed (exit {result.returncode}); see the log above. " + 'Pass setup_kwargs={"compile": False} to run eager instead.' + ) + finally: + tmp.unlink(missing_ok=True) + return out + + +def _load(source: Path, key: str, device: str, compile: bool, kwargs: dict): + # torch's AOTInductor package loader reaches torch._inductor.codecache + # without importing it, which fails in a fresh process. + import torch._inductor.codecache # noqa: F401 + from nequip.integrations.ase import NequIPCalculator + + if source.name.endswith(".nequip.pt2"): + return NequIPCalculator.from_compiled_model(str(source), device=device, **kwargs) + if compile and device != "cpu": + compiled = _compiled_model(source, key, device) + return NequIPCalculator.from_compiled_model(str(compiled), device=device, **kwargs) + # Eager load of a package or checkpoint. nequip exposes this only as a + # private classmethod (its own tests use it); the public loader takes + # compiled models only. + return NequIPCalculator._from_saved_model(str(source), device=device, **kwargs) + + +def setup(checkpoint: str, device: str = "cuda", compile: bool = True, **kwargs): + """ + Load a NequIP calculator for a pinned nequip.net foundation model. + + Args: + checkpoint: Canonical checkpoint id, must be a key of CHECKPOINTS. + device: PyTorch device string (e.g., "cuda", "cuda:0", "cpu"). + compile: Run the AOTInductor-compiled model on GPU (compiling and + caching it on first use). False runs eager PyTorch. + + Returns: + ASE-compatible calculator. + """ + name = CHECKPOINTS[checkpoint] + kwargs.setdefault("chemical_species_to_atom_type_map", True) + return _load(_package_path(name), name, device, compile, kwargs) + + +def setup_from_path(path: str, device: str = "cuda", compile: bool = True, **kwargs): + # Custom checkpoints (`:custom` ids with user weights): a .nequip.zip + # package or .ckpt checkpoint (compiled on first GPU use, cached under a + # key that includes the file's hash) or a ready .nequip.pt2. + source = Path(path) + key = source.name + if not key.endswith(".nequip.pt2") and compile and device != "cpu": + digest = hashlib.sha256() + with open(source, "rb") as f: + while chunk := f.read(1 << 20): + digest.update(chunk) + key = f"custom-{digest.hexdigest()[:16]}" + return _load(source, key, device, compile, kwargs) From cfff34b8eddc4ed270e4d7a42a9f9ef4fd4b5e15 Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:27:20 -0500 Subject: [PATCH 2/4] nequip/allegro: add setuptools so AOTInductor's C++ builder can import torch.utils.cpp_extension imports setuptools at module level and uv venvs do not seed it, so the first GPU compile on the ALCF install died with ModuleNotFoundError before compiling anything. Modal's base image had it preinstalled, which is why the probes passed. Co-Authored-By: Claude Fable 5.1 --- sample_model_configurations/nvidia_configs/allegro.py | 3 +++ sample_model_configurations/nvidia_configs/nequip.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/sample_model_configurations/nvidia_configs/allegro.py b/sample_model_configurations/nvidia_configs/allegro.py index 0414b29..881800a 100644 --- a/sample_model_configurations/nvidia_configs/allegro.py +++ b/sample_model_configurations/nvidia_configs/allegro.py @@ -13,6 +13,9 @@ # # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the # # runtime headers come with torch); matches the cu128 torch build. # "cuda-toolkit[cccl,nvcc]>=12.8,<12.9", +# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which +# # imports setuptools at module level; uv venvs do not seed it. +# "setuptools>=70", # ] # # [tool.uv.sources] diff --git a/sample_model_configurations/nvidia_configs/nequip.py b/sample_model_configurations/nvidia_configs/nequip.py index d0dbbf5..fbbcd27 100644 --- a/sample_model_configurations/nvidia_configs/nequip.py +++ b/sample_model_configurations/nvidia_configs/nequip.py @@ -12,6 +12,9 @@ # # CUDA headers for compiling on nodes with no CUDA toolkit loaded (the # # runtime headers come with torch); matches the cu128 torch build. # "cuda-toolkit[cccl,nvcc]>=12.8,<12.9", +# # AOTInductor's C++ builder imports torch.utils.cpp_extension, which +# # imports setuptools at module level; uv venvs do not seed it. +# "setuptools>=70", # ] # # [tool.uv.sources] From 9de3de3a507571bdbb8c1836087fe26b47b63d3d Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:29:38 -0500 Subject: [PATCH 3/4] nequip/allegro: compile with a GNU-compatible host compiler AOTInductor builds its C++ wrapper with $CXX (default g++) using gcc-style flags. Polaris's default PrgEnv exports CXX=nvc++ and CC=nvc, so every compile on the ALCF install died with "nvc++-Error-Unknown switch". The compile child now keeps the caller's CXX only if it is gcc >= 9 or clang, otherwise takes the first qualifying g++-14..g++-9, g++ or clang++ on PATH, and sets CC to the matching C compiler for Triton's helper build. Co-Authored-By: Claude Opus 5.5 --- .../nvidia_configs/allegro.py | 48 +++++++++++++++++++ .../nvidia_configs/nequip.py | 48 +++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/sample_model_configurations/nvidia_configs/allegro.py b/sample_model_configurations/nvidia_configs/allegro.py index 881800a..d58aab8 100644 --- a/sample_model_configurations/nvidia_configs/allegro.py +++ b/sample_model_configurations/nvidia_configs/allegro.py @@ -174,6 +174,52 @@ def _toolkit_env(env: dict, link_dir: Path) -> None: ) +# Host compilers tried, in order, when CXX is unset or not GNU-compatible. +_CXX_CANDIDATES = ("g++-14", "g++-13", "g++-12", "g++-11", "g++-10", "g++-9", "g++", "clang++") + + +def _gnu_major(cxx: str) -> int | None: + """Major version of a GNU-compatible (gcc or clang) compiler, else None. + + AOTInductor passes gcc-style flags, which other compilers on HPC systems + (e.g. NVIDIA's nvc++) reject. + """ + try: + banner = subprocess.run([cxx, "--version"], capture_output=True, text=True).stdout + version = subprocess.run([cxx, "-dumpversion"], capture_output=True, text=True).stdout + except OSError: + return None + if "Free Software Foundation" not in banner and "clang version" not in banner: + return None + try: + return int(version.strip().split(".")[0]) + except ValueError: + return None + + +def _compiler_env(env: dict) -> None: + """Point CXX/CC at a GNU-compatible compiler new enough for C++17. + + Cluster compiler modules often export CXX/CC for a vendor compiler (e.g. + nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang; + otherwise use the first suitable candidate on PATH, with its matching C + compiler (Triton builds a small C helper with CC during compilation). + """ + path = os.pathsep.join(filter(None, [env.get("PATH"), "/usr/bin"])) + for name in filter(None, [env.get("CXX"), *_CXX_CANDIDATES]): + cxx = shutil.which(name, path=path) + major = _gnu_major(cxx) if cxx else None + if major is None or ("clang" not in Path(cxx).name and major < 9): + continue + env["CXX"] = cxx + cc = shutil.which( + Path(cxx).name.replace("clang++", "clang").replace("g++", "gcc"), path=path + ) + if cc: + env["CC"] = cc + return + + def _compiled_model(source: Path, key: str, device: str) -> Path: """Path to ``source`` compiled for ASE on ``device``, compiling on a miss.""" import nequip @@ -224,10 +270,12 @@ def _compiled_model(source: Path, key: str, device: str) -> Path: try: with tempfile.TemporaryDirectory() as link_dir: _toolkit_env(env, Path(link_dir)) + _compiler_env(env) result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr) if result.returncode != 0: raise RuntimeError( f"nequip-compile failed (exit {result.returncode}); see the log above. " + "It needs gcc >= 9 or clang (set CXX if neither is on PATH). " 'Pass setup_kwargs={"compile": False} to run eager instead.' ) finally: diff --git a/sample_model_configurations/nvidia_configs/nequip.py b/sample_model_configurations/nvidia_configs/nequip.py index fbbcd27..83e5d9b 100644 --- a/sample_model_configurations/nvidia_configs/nequip.py +++ b/sample_model_configurations/nvidia_configs/nequip.py @@ -178,6 +178,52 @@ def _toolkit_env(env: dict, link_dir: Path) -> None: ) +# Host compilers tried, in order, when CXX is unset or not GNU-compatible. +_CXX_CANDIDATES = ("g++-14", "g++-13", "g++-12", "g++-11", "g++-10", "g++-9", "g++", "clang++") + + +def _gnu_major(cxx: str) -> int | None: + """Major version of a GNU-compatible (gcc or clang) compiler, else None. + + AOTInductor passes gcc-style flags, which other compilers on HPC systems + (e.g. NVIDIA's nvc++) reject. + """ + try: + banner = subprocess.run([cxx, "--version"], capture_output=True, text=True).stdout + version = subprocess.run([cxx, "-dumpversion"], capture_output=True, text=True).stdout + except OSError: + return None + if "Free Software Foundation" not in banner and "clang version" not in banner: + return None + try: + return int(version.strip().split(".")[0]) + except ValueError: + return None + + +def _compiler_env(env: dict) -> None: + """Point CXX/CC at a GNU-compatible compiler new enough for C++17. + + Cluster compiler modules often export CXX/CC for a vendor compiler (e.g. + nvc++/nvc). Keep the caller's CXX when it is gcc >= 9 or clang; + otherwise use the first suitable candidate on PATH, with its matching C + compiler (Triton builds a small C helper with CC during compilation). + """ + path = os.pathsep.join(filter(None, [env.get("PATH"), "/usr/bin"])) + for name in filter(None, [env.get("CXX"), *_CXX_CANDIDATES]): + cxx = shutil.which(name, path=path) + major = _gnu_major(cxx) if cxx else None + if major is None or ("clang" not in Path(cxx).name and major < 9): + continue + env["CXX"] = cxx + cc = shutil.which( + Path(cxx).name.replace("clang++", "clang").replace("g++", "gcc"), path=path + ) + if cc: + env["CC"] = cc + return + + def _compiled_model(source: Path, key: str, device: str) -> Path: """Path to ``source`` compiled for ASE on ``device``, compiling on a miss.""" import nequip @@ -228,10 +274,12 @@ def _compiled_model(source: Path, key: str, device: str) -> Path: try: with tempfile.TemporaryDirectory() as link_dir: _toolkit_env(env, Path(link_dir)) + _compiler_env(env) result = subprocess.run(cmd, env=env, stdout=sys.stderr, stderr=sys.stderr) if result.returncode != 0: raise RuntimeError( f"nequip-compile failed (exit {result.returncode}); see the log above. " + "It needs gcc >= 9 or clang (set CXX if neither is on PATH). " 'Pass setup_kwargs={"compile": False} to run eager instead.' ) finally: From 2584fd8467da8e6d97a3143549a14565e1e00798 Mon Sep 17 00:00:00 2001 From: Owen Price Skelly <21372141+OwenPriceSkelly@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:11:13 -0500 Subject: [PATCH 4/4] nequip/allegro: only trust a CUDA_HOME that ships include/cuda.h With CUDA_HOME unset, torch derives it from nvcc on PATH. On Polaris that nvcc is the NVIDIA HPC SDK wrapper, so the AOTInductor build got -I .../hpc_sdk/.../compilers/include: the SDK's own immintrin.h (which g++ rejects) and no cuda.h. The compile child now uses CUDA_HOME, CUDA_PATH or nvcc's toolkit only when it has include/cuda.h, and otherwise sets CUDA_HOME to the CUDA wheel headers explicitly so torch does not fall back to nvcc. Co-Authored-By: Claude Opus 5.5 --- .../nvidia_configs/allegro.py | 22 ++++++++++++++----- .../nvidia_configs/nequip.py | 22 ++++++++++++++----- 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/sample_model_configurations/nvidia_configs/allegro.py b/sample_model_configurations/nvidia_configs/allegro.py index d58aab8..9b48288 100644 --- a/sample_model_configurations/nvidia_configs/allegro.py +++ b/sample_model_configurations/nvidia_configs/allegro.py @@ -140,16 +140,26 @@ def _hardware_tag(device: str) -> str: def _toolkit_env(env: dict, link_dir: Path) -> None: - """Point AOTInductor at CUDA headers and libcuda when no toolkit is loaded. + """Point AOTInductor at CUDA headers and libcuda. The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links - ``-lcuda``. Compute nodes often expose a CUDA toolkit only through a - module, so when none is visible fall back to the headers in NVIDIA's CUDA - wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``), and + ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which + goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC + SDK's, whose include dir has its own intrinsics headers and no cuda.h). + So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships + ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA + wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and link against the driver library this process has loaded. """ - if env.get("CUDA_HOME") or env.get("CUDA_PATH") or shutil.which("nvcc", path=env.get("PATH")): - return + candidates = [env.get("CUDA_HOME"), env.get("CUDA_PATH")] + nvcc = shutil.which("nvcc", path=env.get("PATH")) + if nvcc: + candidates.append(str(Path(nvcc).resolve().parent.parent)) + for home in filter(None, candidates): + if (Path(home) / "include" / "cuda.h").is_file(): + env["CUDA_HOME"] = home + return + env.pop("CUDA_PATH", None) import nvidia headers = next( diff --git a/sample_model_configurations/nvidia_configs/nequip.py b/sample_model_configurations/nvidia_configs/nequip.py index 83e5d9b..7c96ad0 100644 --- a/sample_model_configurations/nvidia_configs/nequip.py +++ b/sample_model_configurations/nvidia_configs/nequip.py @@ -144,16 +144,26 @@ def _hardware_tag(device: str) -> str: def _toolkit_env(env: dict, link_dir: Path) -> None: - """Point AOTInductor at CUDA headers and libcuda when no toolkit is loaded. + """Point AOTInductor at CUDA headers and libcuda. The C++ wrapper AOTInductor builds needs ``$CUDA_HOME/include`` and links - ``-lcuda``. Compute nodes often expose a CUDA toolkit only through a - module, so when none is visible fall back to the headers in NVIDIA's CUDA - wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``), and + ``-lcuda``. torch derives CUDA_HOME from ``nvcc`` on PATH when unset, which + goes wrong when that nvcc is a compiler-suite wrapper (e.g. the NVIDIA HPC + SDK's, whose include dir has its own intrinsics headers and no cuda.h). + So use CUDA_HOME, CUDA_PATH or nvcc's toolkit only if it really ships + ``include/cuda.h``; otherwise fall back to the headers in NVIDIA's CUDA + wheels (``nvidia/cuda_runtime/``, ``cuda_nvcc/`` and ``cuda_cccl/``) and link against the driver library this process has loaded. """ - if env.get("CUDA_HOME") or env.get("CUDA_PATH") or shutil.which("nvcc", path=env.get("PATH")): - return + candidates = [env.get("CUDA_HOME"), env.get("CUDA_PATH")] + nvcc = shutil.which("nvcc", path=env.get("PATH")) + if nvcc: + candidates.append(str(Path(nvcc).resolve().parent.parent)) + for home in filter(None, candidates): + if (Path(home) / "include" / "cuda.h").is_file(): + env["CUDA_HOME"] = home + return + env.pop("CUDA_PATH", None) import nvidia headers = next(