Garden's Almanac of Matter Models | Checkpoint | Params | Polaris | Sophia | Perlmutter | Delta | Frontier | Della |
|---|---|---|---|---|---|---|---|
| orb-v2 | 25M | — | — | ○ | ○ | — | — |
| orb-d3-v2 | 25M | ● | ● | ○ | ● | ● | — |
| orb-mptraj-only-v2 | 25M | — | — | ○ | ○ | — | — |
| orb-v2:custom | — | — | — | — | — | — |
● verified, last 30 days ○ installed, not recently verified — not installed
No cluster has installed any checkpoint of this model yet.
Running this model
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4with RootstockCalculator(
5 cluster=YOUR_CLUSTER_ID, # eg, "polaris", "sophia"
6 checkpoint="orb-v2",
7 device="cuda",
8) as calc:
9 # model is now running in subprocess on compute node
10 atoms.calc = calc
11 atoms.get_potential_energy()
Environments
Rootstock runs each model family inside an isolated Python environment defined by a single file. This environment file includes the specific dependencies needed, plus a setup() function that loads the model and returns an ASE calculator. These files are usually almost identical for a given model family, but because of cluster-specific quirks (eg, an old CUDA driver) the dependencies and setup code can vary a bit.
orb_env.py
1# /// script
2# requires-python = ">=3.11"
3# dependencies = [
4# # >=0.5,<0.6: 0.5.5 is what the verified Delta env resolved — the v2
5# # loaders keep their single-return API through 0.5.x. 0.4.x is broken
6# # for us: it imports pynanoflann, which is git-only and undeclared, so
7# # a fresh build dies at import (Delta, 2026-07-31). 0.6 raises the
8# # Python floor to 3.12 — that line lives in orb_v3.py.
9# "orb-models>=0.5,<0.6",
10# "ase>=3.22",
11# "torch>=2.0",
12# # Not imported here — constrains orb-models' transitive dep. setup()'s
13# # no-lock serve path relies on cached_path returning local files without
14# # locking or writing, verified against exactly this version (#67).
15# "cached_path==1.8.10",
16# ]
17#
18# [tool.uv.sources]
19# torch = { index = "pytorch-cu128" }
20#
21# [[tool.uv.index]]
22# name = "pytorch-cu128"
23# url = "https://download.pytorch.org/whl/cu128"
24# explicit = true
25# ///
26"""Orb v2 env — kept only for the built-in-D3 dispersion variant.
27
28Orb v3 (orb_v3.py) is the primary orb env; it supersedes the v2 checkpoints
29except orb-d3-v2, which has no v3 equivalent (v3 ships no dispersion-corrected
30model). Catalog trimmed to that one id 2026-07-30. The two lines can't share
31an env: the v3 loaders need orb-models>=0.6, which raises the Python floor
32to 3.12 and torch to 2.8 (the v2 loaders here are fine through 0.5.x).
33"""
34
35import os
36import shutil
37import urllib.request
38from pathlib import Path
39
40CHECKPOINTS = {
41 "orb-d3-v2": "orb-d3-v2",
42 # Your own fine-tuned v2-architecture weights: pair with weights=
43 # (loaded via setup_from_path). v3 fine-tunes go to orb-v3:custom.
44 "orb-v2:custom": None,
45}
46
47
48def _default_weights_url(load_fn) -> str:
49 """The upstream URL baked into the loader's ``weights_path`` default."""
50 import inspect
51
52 default = inspect.signature(load_fn).parameters["weights_path"].default
53 if not isinstance(default, str) or not default.startswith(("http://", "https://")):
54 raise RuntimeError(
55 f"{load_fn.__name__} has no URL default for weights_path "
56 f"(got {default!r}); update this env file for the installed orb-models"
57 )
58 return default
59
60
61def _local_weights_path(url: str) -> Path:
62 """Where the checkpoint lives in the shared model cache."""
63 cache = Path(os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache")
64 return cache / "orb" / os.path.basename(url)
65
66
67def _fetch(url: str, dest: Path) -> None:
68 """Download ``url`` to ``dest`` atomically (tmp file + rename)."""
69 dest.parent.mkdir(parents=True, exist_ok=True)
70 tmp = dest.with_name(f"{dest.name}.tmp.{os.getpid()}")
71 try:
72 with urllib.request.urlopen(url) as resp, open(tmp, "wb") as out:
73 shutil.copyfileobj(resp, out)
74 os.replace(tmp, dest)
75 finally:
76 tmp.unlink(missing_ok=True)
77
78
79def setup(checkpoint: str, device: str = "cuda"):
80 import torch
81 from orb_models.forcefield import pretrained
82 from orb_models.forcefield.calculator import ORBCalculator
83
84 # orb-models exposes one function per checkpoint, e.g. pretrained.orb_v2().
85 fn_name = CHECKPOINTS[checkpoint].replace("-", "_")
86 load_fn = getattr(pretrained, fn_name)
87
88 # orb-models resolves its default weights URL through `cached_path`, which
89 # write-locks its cache dir even on warm hits — EACCES for anyone who can't
90 # write the shared install (Garden-AI/rootstock#67). Handed a *local* path
91 # instead, cached_path returns it without locking. So the weights are
92 # pre-fetched into the shared model cache at `rootstock add` time
93 # (maintainer, cache writable) and every later serve loads that file.
94 url = _default_weights_url(load_fn)
95 weights = _local_weights_path(url)
96 if not weights.exists():
97 _fetch(url, weights)
98
99 orbff = load_fn(weights_path=str(weights), device=torch.device(device))
100 return ORBCalculator(orbff, device=torch.device(device))
101
102
103def setup_from_path(path: str, device: str = "cuda", arch: str = "orb-v2"):
104 # Custom checkpoints (`:custom` ids with user weights). A weights file doesn't say
105 # which orb architecture produced it, so `arch` names the pretrained
106 # loader to instantiate — pass the right one at call time
107 # (setup_kwargs={"arch": ...} / --kwarg arch=...). Handing the loader a local path also means no
108 # network and no cached_path locking (see setup()).
109 import torch
110 from orb_models.forcefield import pretrained
111 from orb_models.forcefield.calculator import ORBCalculator
112
113 fn_name = arch.replace("-", "_")
114 try:
115 load_fn = getattr(pretrained, fn_name)
116 except AttributeError:
117 raise ValueError(
118 f"unknown orb architecture {arch!r}; expected a loader name from "
119 f"orb_models.forcefield.pretrained, e.g. orb-v2, orb-d3-v2"
120 ) from None
121
122 orbff = load_fn(weights_path=path, device=torch.device(device))
123 return ORBCalculator(orbff, device=torch.device(device))
124
Built on Polaris: 2026-07-31
Couldn't load the current environments from Rootstock.