[OwenPriceSkelly]

Summary

Re-adds nequip and allegro to the NVIDIA catalog (both were pruned on 2026-07-30 because they had no checkpoints to ship). They now host the pretrained foundation models published on nequip.net:

Env	Canonical ids
nequip	nequip-oam-s, nequip-oam-m, nequip-oam-l, nequip-oam-xl, nequip-mp-l, nequip:custom
allegro	allegro-oam-l, allegro-mp-l, allegro:custom

• Weights: fixed Zenodo uploads (the files nequip.net resolves mir-group/<model>:0.1 to), md5-checked and downloaded into the shared cache at rootstock add time. They're fetched directly because nequip's own nequip.net: loader queries the nequip.net API even when the file is already cached, which fails on offline workers.
• GPU: setup() compiles each model once with nequip's AOTInductor path (--target ase) for each combination of model, GPU architecture (sm), torch version and nequip version. The result is cached in {cache}/nequip/compiled/, or in 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 the startup timeout still leaves a usable cache entry. setup_kwargs={"compile": False} runs eager instead.
• CUDA toolkit: when none is loaded, CUDA headers come from the env's cuda-toolkit[cccl,nvcc] wheels, and -lcuda resolves through a temporary symlink to the driver library the process has loaded.
• CPU: runs eager, including rootstock add's device="cpu" download step. This uses nequip's private NequIPCalculator._from_saved_model, which is why nequip is pinned to <0.20.
• Custom weights: setup_from_path accepts .nequip.zip, .ckpt, or an already-compiled .nequip.pt2.
• modal_app.py: adds probe_nequip / probe_allegro, plus extra_index_url, cpu and timeout options on probe_image, so the probes build against the configs' cu128 torch.

Testing

• Local (macOS CPU): all 7 package md5s match. Every model loaded gave sensible energies, forces and stress on a rattled Cu supercell, and setup_from_path on the cached package matched setup(). The AOTInductor CPU compile also matched eager to within 2e-5 eV.
• Modal (A10G, 8 CPU, torch 2.10.0+cu128): probe_nequip passes for nequip-oam-s, -l and -xl, and probe_allegro passes for allegro-oam-l, including fresh compiles. Compile times: Allegro-OAM-L about 2 min, NequIP-OAM-L about 4.6 min, NequIP-OAM-XL about 6.4 min.
• pytest tests/sample_configs passes, including the setup_from_path signature contract. ruff check and ruff format --check are clean on the touched files.
• Not yet run on an HPC cluster. Staging for the ALCF install next.

Notes for review

• XL's first compile comes close to the 600 s default startup/verify timeout. The docstring tells users to raise timeout= / --verify-timeout.
• AOTInductor needs a C++17 host compiler on the compute node. Older system gccs may need a module load or CXX.

🤖 Generated with Claude Code