Skip to content

Reconcile the sample catalog with deployed truth; add a **kwargs escape hatch to every setup hook - #214

Merged
OwenPriceSkelly merged 5 commits into
mainfrom
upstream-env-fixes
Aug 12, 2026
Merged

Reconcile the sample catalog with deployed truth; add a **kwargs escape hatch to every setup hook#214
OwenPriceSkelly merged 5 commits into
mainfrom
upstream-env-fixes

Conversation

@OwenPriceSkelly

@OwenPriceSkelly OwenPriceSkelly commented Aug 12, 2026

Copy link
Copy Markdown
Member

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

  • Prune dead configs — 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 (allegro, nequip — <family>:custom is the supported path for user weights now). README tree and modal_app probes updated to match.
  • fairchem v2 indexed-device fix_fairchem_device() routes cuda:N through torch.cuda.set_device() because MLIPPredictUnit._setup_device asserts device 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 same torch.cuda API).
  • esen header — drop the fairchem-1.x pyg find-links dance (fairchem-core>=2.20 doesn't need it, and the find-links block universal lockfile resolution).
  • uma checkpoint table — document the uma-s-1p2 bug and the not-yet-released uma-s-1p2p1 replacement (re-add when the next fairchem-core ships).
  • orb v2/v3 — orb.py (both vendors) becomes the v2 remnant it is on every cluster: trimmed to orb-d3-v2 (no v3 dispersion model exists), orb-v2:custom, and orb-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.13 cap (dm-tree pin has no cp313 wheel), drops the omol ids (failing verify everywhere since 2026-05), and gains orb-v3:custom + setup_from_path(arch=, precision=).
  • New: OrbMol env (ROCm)amd_configs/orbmol.py on orb-models>=0.7,<0.8: orbmol-v1-conservative (upstream alias for orb-v3-conservative-omol), orbmol-v2, orbmol:custom (arch= kwarg). Wraps ORBCalculator to default atoms.info charge=0/spin=1 (the models raise without them), defaults edge_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.

**kwargs escape hatch

Every setup()/setup_from_path() now forwards **kwargs to the calculator constructor, so setup_kwargs= / --kwarg can override constructor knobs the config doesn't name. Config-chosen defaults (default_dtype etc.) move into kwargs.setdefault() so user overrides don't collide with named arguments. A new contract test in tests/sample_configs/test_setup_from_path_hooks.py pins the convention; the rootstock create skeleton and docs/environments.md document 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

OwenPriceSkelly and others added 5 commits August 12, 2026 14:00
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>
@OwenPriceSkelly
OwenPriceSkelly merged commit ccddc91 into main Aug 12, 2026
5 checks passed
@OwenPriceSkelly
OwenPriceSkelly deleted the upstream-env-fixes branch August 12, 2026 19:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant