Reconcile the sample catalog with deployed truth; add a **kwargs escape hatch to every setup hook - #214
Merged
Merged
Conversation
Remove the seven configs dropped from every deployed set on 2026-07-30: the fairchem-core 1.x OC20 baselines beyond escn/equiformer (dimenet, gemnet, painn, schnet, scn) and the BYO-weights-only envs with no shippable checkpoints (allegro, nequip — '<family>:custom' is the supported path for user weights now). Update the catalog README tree and modal_app probes to match. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fairchem v2's MLIPPredictUnit._setup_device asserts device is 'cpu' or 'cuda' and resolves the actual GPU via torch.cuda.current_device(), so an indexed 'cuda:N' has to travel through torch's current-device state rather than the argument. Multi-checkpoint verification on a multi-GPU node hands each worker an indexed device — that killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync (4x A100, VERIFY_JOBS=4). Deployed as a hotfix on the ALCF install since 2026-08-06; applied to the ROCm configs too (same torch.cuda API on AMD). Also carries the rest of the deployed esen/uma deltas: esen drops the fairchem-1.x pyg find-links dance (fairchem-core>=2.20 doesn't need it and the find-links block universal lockfile resolution), and uma's checkpoint table documents the uma-s-1p2 bug and the not-yet-released uma-s-1p2p1 replacement. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
orb.py (nvidia + amd) is now the v2 remnant it is in every deployed set: trimmed to orb-d3-v2 (the one v2 checkpoint with no v3 equivalent — v3 ships no dispersion-corrected model), ':custom' renamed orb-v2:custom to make room for orb-v3:custom, and orb-models pinned >=0.5,<0.6 (0.4.x imports git-only pynanoflann and dies at import on a fresh build; 0.6 raises the Python floor to 3.12 and belongs to orb_v3.py). orb_v3.py picks up the deployed extras: a <3.13 cap (orb-models pins dm-tree==0.1.8, no cp313 wheel), the omol ids dropped from the catalog 2026-07-30 (failing verify on every cluster since 2026-05), an orb-v3:custom entry, and setup_from_path(arch=, precision=) so custom v3 weights can name the architecture that produced them. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Every config's setup() and setup_from_path() now takes **kwargs and forwards it to the calculator constructor, giving users an escape hatch (setup_kwargs= / --kwarg) for constructor knobs the config doesn't name explicitly. Where a config sets its own default (default_dtype/dtype), it moves into kwargs.setdefault() so a user override doesn't collide with a named argument — this also subsumes the default_dtype passthrough deployed on ALCF's mace.py. A new contract test pins the convention for every current and future config, and the docs/environments.md contract and 'rootstock create' skeleton document it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
New amd_configs env on orb-models>=0.7,<0.8 (0.7.0 first ships these models): orbmol-v1-conservative (upstream alias for orb-v3-conservative-omol), orbmol-v2, and orbmol:custom (arch= kwarg). Wraps ORBCalculator to default atoms.info charge=0/spin=1 — the models raise without them. edge_method defaults to knn_scipy because 0.7's knn_alchemi default launches CUDA Warp kernels and ROCm torch self-reports as cuda; orbmol-v2's periodic-system electrostatics (PME) are expected to fail on ROCm, non-periodic is pure torch. torch pinned to the rocm6.4 index. Follows the setup/**kwargs convention from this branch. Not yet deployed anywhere — ships to Frontier next deploy. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upstreams everything the deployed clusters are running that the repo catalog didn't have yet (sources: the
staging/snapshots regenerated from the admin-dashboard manifests, 2026-08-04/06), plus one convention change.Catalog reconciliation
<family>:customis the supported path for user weights now). README tree and modal_app probes updated to match._fairchem_device()routescuda:Nthroughtorch.cuda.set_device()becauseMLIPPredictUnit._setup_deviceassertsdevice in ["cpu", "cuda"]and resolves the GPU from torch state. Indexed devices are exactly what multi-checkpoint verification hands workers on a multi-GPU node; this killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync (4x A100, VERIFY_JOBS=4). Hotfixed on the ALCF install since 2026-08-06; applied here to nvidia + amd allscaip/esen/uma (ROCm torch exposes the sametorch.cudaAPI).fairchem-core>=2.20doesn't need it, and the find-links block universal lockfile resolution).orb-v2:custom, andorb-models>=0.5,<0.6(0.4.x dies at import on fresh builds — git-only undeclared pynanoflann; 0.6 needs Python 3.12 and belongs to orb_v3). orb_v3.py picks up the deployed<3.13cap (dm-tree pin has no cp313 wheel), drops the omol ids (failing verify everywhere since 2026-05), and gainsorb-v3:custom+setup_from_path(arch=, precision=).amd_configs/orbmol.pyonorb-models>=0.7,<0.8:orbmol-v1-conservative(upstream alias for orb-v3-conservative-omol),orbmol-v2,orbmol:custom(arch=kwarg). Wraps ORBCalculator to defaultatoms.infocharge=0/spin=1 (the models raise without them), defaultsedge_method="knn_scipy"on ROCm (0.7's knn_alchemi default launches CUDA Warp kernels and ROCm torch self-reports as cuda), torch pinned to the rocm6.4 index. Not yet deployed — ships to Frontier on the next deploy.**kwargsescape hatchEvery
setup()/setup_from_path()now forwards**kwargsto the calculator constructor, sosetup_kwargs=/--kwargcan override constructor knobs the config doesn't name. Config-chosen defaults (default_dtypeetc.) move intokwargs.setdefault()so user overrides don't collide with named arguments. A new contract test intests/sample_configs/test_setup_from_path_hooks.pypins the convention; therootstock createskeleton anddocs/environments.mddocument it.Deliberately not included: Delta's tace cu128 torch-index pin (cluster-specific — its driver caps at CUDA 12.8) and the polaris-only ani/pet variants (cluster-specific by construction).
Tests:
ruff check,ruff format --check,ty check, and the full pytest suite (875 passed) are green.🤖 Generated with Claude Code