nvidia_configs: add nequip and allegro envs hosting the nequip.net foundation models

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 `<family>: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

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)