diff --git a/CLAUDE.md b/CLAUDE.md index 41372d8..4fb1e92 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -141,7 +141,9 @@ with RootstockCalculator( energy = atoms.get_potential_energy() # Forward extra kwargs to the env's setup() function. Cannot contain -# "checkpoint" or "device" — those are passed at the top level. +# "checkpoint" or "device" — those are passed at the top level. Multi-head +# models (UMA `task`, MACE-MH-1 `head`) REQUIRE a head selection here — no +# default head; the env's VERIFY_KWARGS covers smoke-test/add instead. with RootstockCalculator( cluster="delta", checkpoint="uma-s-1p1", @@ -177,7 +179,7 @@ with RootstockCalculator( | `esen` | `esen-md-direct-all-omol`, `esen-sm-conserving-all-omol`, `esen-sm-direct-all-omol` | | `orb` | `orb-v2` | | `tensornet` | `tensornet-matpes-pbe-2025-2` | -| `uma` | `uma-s-1p1` | +| `uma` | `uma-s-1p1`, `uma-s-1p2p1`, `uma-m-1p1` | ## Build Process diff --git a/docs/api.md b/docs/api.md index ce69a92..4d58964 100644 --- a/docs/api.md +++ b/docs/api.md @@ -32,7 +32,7 @@ with RootstockCalculator( | `root` | `str` | Yes* | Custom install-root path instead of a known cluster | | `cache_root` | `str` | No | Override path for the model-weight cache and redirected `HOME`. When omitted, the install's own declaration (`{root}/layout.json`) decides, falling back to the cluster registry for legacy roots, then to `root` | | `device` | `str` | No | `"cuda"` (default) or `"cpu"` | -| `setup_kwargs` | `dict` | No | Extra keyword arguments forwarded to the env's `setup()` function (e.g., `{"task": "omol"}`). Cannot contain `checkpoint` or `device`. For a `:custom` checkpoint they go to `setup_from_path()` instead and cannot contain `path` | +| `setup_kwargs` | `dict` | No | Extra keyword arguments forwarded to the env's `setup()` function (e.g., `{"task": "omol"}`). Multi-head models (UMA's `task`, MACE-MH-1's `head`) require a head selection here and error without one — there is no default head. Cannot contain `checkpoint` or `device`. For a `:custom` checkpoint they go to `setup_from_path()` instead and cannot contain `path` | | `timeout` | `float` | No | Socket timeout in seconds for worker operations (default 600, matching checkpoint verification — so the first real force call, which may pay for `torch.compile` or large neighbor lists, runs under the envelope verification exercised) | | `weights` | `str \| Path` | With `:custom` | Path to your own weights file (e.g. a fine-tune of one of the family's shipped checkpoints). Required with — and only valid with — a `:custom` checkpoint id. Must be visible from the compute nodes; loaded through the env's `setup_from_path()` hook. No shipped weights are involved | diff --git a/docs/environments.md b/docs/environments.md index fc12655..010f6ae 100644 --- a/docs/environments.md +++ b/docs/environments.md @@ -149,6 +149,12 @@ Signature: `setup(checkpoint: str, device: str = "cuda", **extra)`. - `device`: PyTorch device. - Optional extra kwargs are forwarded from `RootstockCalculator(setup_kwargs=...)` and `rootstock add --kwarg KEY=VAL`. +A kwarg that selects among a model's task heads (UMA's `task`, MACE-MH-1's +`head`) should be **required, not defaulted**: give it a `None` default and +raise a `ValueError` naming the valid choices when it's missing. A silent +default head means users unknowingly compute with the wrong physics. +Declare `VERIFY_KWARGS` (below) so verification still runs without user input. + Return: an ASE-compatible calculator. ### `setup_from_path()` function (optional — enables custom checkpoints) @@ -170,15 +176,16 @@ Signature: `setup_from_path(path: str, device: str = "cuda", **extra)`. usually a *different* upstream call than loading a registry name — e.g. FAIRChem's `load_predict_unit(path)` vs `get_predict_unit(name)` — which is why this is a separate function rather than a path-shaped `checkpoint`. -- `device`, extra kwargs: as for `setup()`. Give extras defaults where - possible; a fine-tune needing a required kwarg still works (users pass it - via `setup_kwargs=` / `--kwarg`), but defaults make the common case - friendlier. +- `device`, extra kwargs: as for `setup()`. Give extras defaults where a + default is *correct* for any weights file; a head-selection kwarg should + default to `None` and be forwarded — the upstream library errors when the + fine-tune actually needs one, and users pass it via `setup_kwargs=` / + `--kwarg` (it forwards to `setup_from_path()` for `:custom` checkpoints). Return: an ASE-compatible calculator. ```python -def setup_from_path(path: str, device: str = "cuda", task: str = "omat"): +def setup_from_path(path: str, device: str = "cuda", task: str | None = None): from fairchem.core import FAIRChemCalculator from fairchem.core.units.mlip_unit import load_predict_unit @@ -222,6 +229,31 @@ Absent `CLUSTERS` (the normal case) means the env serves every cluster its install does. Like `CHECKPOINTS`, the list is AST-parsed — string literals only, and an empty list is an authoring error. +### `VERIFY_KWARGS` dict (optional — verification kwargs for required-selection envs) + +When `setup()` *requires* a kwarg (a task-head selection with no default), +verification needs a way to pick one: `rootstock smoke-test` and a bare +`rootstock add` call `setup()` with no extra kwargs, and would otherwise fail +on exactly the error the requirement exists to raise. Declare a module-level + +```python +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, # the weights= smoke-test leg needs one too +} +``` + +keyed by canonical checkpoint id (`:custom` entries included — the nightly +weights= leg re-loads a shipped checkpoint's weights through +`setup_from_path()` and compares against that checkpoint's baseline, so give +both the same selection). Values are literal dicts of setup kwargs, AST-parsed +like `CHECKPOINTS` — no names or calls. + +Verification-only: explicit kwargs (`rootstock add --kwarg ...`) always win, +and `RootstockCalculator` never reads it — users still select explicitly. +Checkpoints without an entry verify with no extra kwargs, so most envs never +declare this. + ## Examples ### MACE (MP-0 and OFF23 in one env) @@ -256,17 +288,31 @@ def setup(checkpoint: str, device: str = "cuda"): ### UMA (FAIRChem) -`setup()` accepts an extra `task` kwarg. Users pass `setup_kwargs={"task": "omol"}` to `RootstockCalculator`, or `--kwarg task=omol` to `rootstock add`. +UMA is multi-task: `setup()` requires an explicit `task` and errors without +one. Users pass `setup_kwargs={"task": "omol"}` to `RootstockCalculator`, or +`--kwarg task=omol` to `rootstock add`; `VERIFY_KWARGS` picks the head for +verification. ```python CHECKPOINTS = { "uma-s-1p1": "uma-s-1p1", } +UMA_TASKS = ("omat", "omol", "oc20", "odac", "omc") + +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, +} + -def setup(checkpoint: str, device: str = "cuda", task: str = "omat"): +def setup(checkpoint: str, device: str = "cuda", task: str | None = None): from fairchem.core import FAIRChemCalculator, pretrained_mlip + if task is None: + raise ValueError( + f"{checkpoint} is multi-task and has no default head - select one " + f'with setup_kwargs={{"task": ...}}: one of {", ".join(UMA_TASKS)}' + ) predictor = pretrained_mlip.get_predict_unit(CHECKPOINTS[checkpoint], device=device) return FAIRChemCalculator(predictor, task_name=task) ``` diff --git a/docs/nightly-smoke-testing.md b/docs/nightly-smoke-testing.md index 31e0954..f37af1e 100644 --- a/docs/nightly-smoke-testing.md +++ b/docs/nightly-smoke-testing.md @@ -124,8 +124,10 @@ Override via env at submit time (SLURM: inline `VAR=val sbatch …`; PBS: > `date -d` in the recipes is GNU date (present on Linux HPC login/compute > nodes). The recipes are not meant to run on macOS. -> **Smoke-test always uses default kwargs.** `smoke-test` calls each env's -> `setup()` with no extra kwargs, so a checkpoint that only works with -> non-default kwargs (e.g. a UMA checkpoint needing `task=omol`) will appear -> failing even though `add` succeeded. The remedy is to make the preferred -> kwargs the env's default in the env file. +> **Smoke-test passes no kwargs of its own.** `smoke-test` calls each env's +> `setup()` with no extra kwargs; when a checkpoint needs one (e.g. UMA's +> required `task`), verification falls back to the env's `VERIFY_KWARGS` +> entry for that checkpoint (see `docs/environments.md`). A checkpoint whose +> env requires a kwarg but declares no `VERIFY_KWARGS` entry will appear +> failing even though `add --kwarg ...` succeeded — the remedy is to declare +> the entry, never to give `setup()` a silent default. diff --git a/rootstock/commands/smoke_test.py b/rootstock/commands/smoke_test.py index 5bd0975..76f06ee 100644 --- a/rootstock/commands/smoke_test.py +++ b/rootstock/commands/smoke_test.py @@ -413,7 +413,7 @@ def cmd_smoke_test(args) -> int: env_name=env_name, checkpoint=ckpt_name, device=device, - setup_kwargs={}, # smoke-test always uses env defaults; see design §7.2 + setup_kwargs={}, # empty → verify falls back to the env's VERIFY_KWARGS cache_root=cache_root, weights_capture_path=str(capture_path), results=run_results, @@ -493,7 +493,7 @@ def cmd_smoke_test(args) -> int: env_name=leg.env_name, checkpoint=leg.custom_id, device=device, - setup_kwargs={}, + setup_kwargs={}, # empty → VERIFY_KWARGS fallback (keyed by the :custom id) cache_root=cache_root, checkpoint_path=str(weights_path), results=run_results, diff --git a/rootstock/environment.py b/rootstock/environment.py index 64eb5d6..cff9c1a 100644 --- a/rootstock/environment.py +++ b/rootstock/environment.py @@ -465,6 +465,69 @@ def parse_clusters_list(env_source_path: Path) -> list[str] | None: return None +def parse_verify_kwargs(env_source_path: Path) -> dict[str, dict]: + """AST-extract the optional module-level ``VERIFY_KWARGS`` dict literal. + + ``VERIFY_KWARGS = {"uma-s-1p1": {"task": "omat"}}`` names the setup + kwargs verification uses for a checkpoint when the caller supplies none — + the escape hatch for envs whose ``setup()`` *requires* a kwarg (e.g. a + multi-head model that errors without an explicit head selection). + Verification-only: ``RootstockCalculator`` never reads it, so users still + have to select explicitly. Keys are canonical checkpoint ids (including + ``:custom`` entries, for the weights= smoke-test leg); values are literal + dicts of JSON-safe scalars, same shape as ``setup_kwargs``. + + Raises ValueError on a non-literal or wrongly-shaped declaration, + matching ``parse_checkpoints_dict``. Absent means ``{}``. + """ + tree = ast.parse(env_source_path.read_text(), filename=str(env_source_path)) + for node in tree.body: + if isinstance(node, ast.Assign): + targets, value = node.targets, node.value + elif isinstance(node, ast.AnnAssign) and node.value is not None: + targets, value = [node.target], node.value + else: + continue + if not ( + len(targets) == 1 + and isinstance(targets[0], ast.Name) + and targets[0].id == "VERIFY_KWARGS" + ): + continue + try: + parsed = ast.literal_eval(value) + except ValueError: + raise ValueError( + f"{env_source_path}: VERIFY_KWARGS must be a dict literal " + f"(no names, calls, or comprehensions)." + ) from None + if not isinstance(parsed, dict): + raise ValueError(f"{env_source_path}: VERIFY_KWARGS must be a dict literal.") + for ckpt_id, kwargs in parsed.items(): + if not isinstance(ckpt_id, str): + raise ValueError( + f"{env_source_path}: VERIFY_KWARGS keys must be checkpoint-id strings." + ) + if not isinstance(kwargs, dict) or not all(isinstance(k, str) for k in kwargs): + raise ValueError( + f"{env_source_path}: VERIFY_KWARGS['{ckpt_id}'] must be a " + f"dict of setup kwargs (string keys)." + ) + return parsed + return {} + + +def verify_kwargs_for(root: Path | str, env_name: str, checkpoint: str) -> dict: + """The built env's ``VERIFY_KWARGS`` entry for one checkpoint, ``{}`` when + the source is missing or declares none. Malformed declarations raise — + an authoring error should surface at verification, not degrade into a + cryptic model-load failure.""" + env_source = Path(root) / "envs" / env_name / "env_source.py" + if not env_source.exists(): + return {} + return parse_verify_kwargs(env_source).get(checkpoint, {}) + + def declares_setup_from_path(env_source_path: Path) -> bool: """ Return True when the env source declares a module-level ``setup_from_path``. diff --git a/rootstock/verify.py b/rootstock/verify.py index 0964a86..b7155b1 100644 --- a/rootstock/verify.py +++ b/rootstock/verify.py @@ -85,7 +85,9 @@ def verify_checkpoint( env_name: Name of pre-built environment (e.g., "mace"). checkpoint: Canonical checkpoint id passed to setup(). device: PyTorch device (e.g., "cuda", "cpu"). - setup_kwargs: Extra keyword arguments forwarded to setup(). + setup_kwargs: Extra keyword arguments forwarded to setup(). When + empty, the env's declared ``VERIFY_KWARGS`` for this + checkpoint (if any) are used instead. cache_root: Optional separate root for the model-weight cache and redirected HOME. Defaults to ``root``. checkpoint_path: Path to a local (user-registered) weights file; the @@ -103,9 +105,15 @@ def verify_checkpoint( On failure, success is False and error_message is a short string describing what went wrong. """ + from .environment import verify_kwargs_for from .server import RootstockServer - setup_kwargs = setup_kwargs or {} + # No explicit kwargs → fall back to the env's declared VERIFY_KWARGS for + # this checkpoint. This is how an env whose setup() *requires* a kwarg + # (e.g. UMA's task, MACE-MH-1's head — no default, explicit selection + # only) still verifies under smoke-test and a bare `rootstock add`. + # Explicit kwargs always win, and the calculator path never reads this. + setup_kwargs = setup_kwargs or verify_kwargs_for(root, env_name, checkpoint) atoms = _smoke_test_atoms() socket_name = f"rootstock_verify_{uuid.uuid4().hex[:8]}" diff --git a/sample_model_configurations/_agent/probe.py b/sample_model_configurations/_agent/probe.py index 909cf2a..83f77ef 100644 --- a/sample_model_configurations/_agent/probe.py +++ b/sample_model_configurations/_agent/probe.py @@ -87,6 +87,14 @@ def main() -> int: help="Checkpoint/model arg passed to setup(). Empty = setup default.", ) parser.add_argument("--device", default="cuda") + parser.add_argument( + "--kwarg", + action="append", + default=[], + metavar="KEY=VAL", + help="Extra setup() kwarg (repeatable), e.g. --kwarg task=omat for " + "models that require an explicit head selection.", + ) parser.add_argument( "--system", default="molecule", @@ -94,13 +102,19 @@ def main() -> int: help="Probe system to compute one forward pass on.", ) args = parser.parse_args() + setup_kwargs = {} + for spec in args.kwarg: + key, sep, value = spec.partition("=") + if not sep or not key: + parser.error(f"--kwarg expects key=value, got {spec!r}") + setup_kwargs[key] = value config_path = Path(args.config).resolve() overall_t0 = time.time() t0 = overall_t0 print( f"PROBE: config={config_path} checkpoint={args.checkpoint!r} " - f"device={args.device} system={args.system}", + f"device={args.device} system={args.system} kwargs={setup_kwargs}", flush=True, ) @@ -111,7 +125,11 @@ def main() -> int: atoms = build_system(args.system) t0 = stage(f"build_system:{args.system}:{len(atoms)}atoms", t0) - calc = setup(args.checkpoint, args.device) if args.checkpoint else setup(device=args.device) + calc = ( + setup(args.checkpoint, args.device, **setup_kwargs) + if args.checkpoint + else setup(device=args.device, **setup_kwargs) + ) t0 = stage("setup_calculator", t0) atoms.calc = calc diff --git a/sample_model_configurations/amd_configs/mace.py b/sample_model_configurations/amd_configs/mace.py index 67760f4..6fecbbe 100644 --- a/sample_model_configurations/amd_configs/mace.py +++ b/sample_model_configurations/amd_configs/mace.py @@ -34,7 +34,8 @@ Multi-head checkpoints select a head via the `head` kwarg on setup() (setup_kwargs={"head": ...} / --kwarg head=...), named by upstream's training -corpus — see MH1_HEADS; omat_pbe is the default. +corpus — see MH1_HEADS. Selection is required: the multi-head checkpoint +has no default head, and setup() errors when none is given. The OMOL checkpoint expects `charge` and `spin` in `atoms.info`. """ @@ -72,11 +73,23 @@ "matpes_r2scan", ) +# Verification-only head selection for the multi-head checkpoint: +# smoke-test and a bare `rootstock add` verify with this +# (setup() itself has no default head). +VERIFY_KWARGS = { + "mace-mh-1": {"head": "omat_pbe"}, +} + def setup(checkpoint: str, device: str = "cuda", head: str | None = None, **kwargs): arg = CHECKPOINTS[checkpoint] if arg.startswith("mh:"): - head = head or "omat_pbe" + if head is None: + raise ValueError( + f"{checkpoint} is multi-head and has no default — select one " + f"with setup_kwargs={{'head': ...}} (or --kwarg head=...): " + f"one of {', '.join(MH1_HEADS)}" + ) if head not in MH1_HEADS: raise ValueError(f"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}") from mace.calculators import mace_mp diff --git a/sample_model_configurations/amd_configs/uma.py b/sample_model_configurations/amd_configs/uma.py index c61da15..0ac71ed 100644 --- a/sample_model_configurations/amd_configs/uma.py +++ b/sample_model_configurations/amd_configs/uma.py @@ -1,7 +1,8 @@ # /// script # requires-python = ">=3.11" # dependencies = [ -# "fairchem-core>=2.20", +# # 2.22 is the first release with uma-s-1p2p1 in the registry. +# "fairchem-core>=2.22", # "ase>=3.22", # "torch>=2.4.0", # # torch's ROCm wheels depend on this; it lives only on the ROCm @@ -23,16 +24,45 @@ fairchem-core v2 is a plain PyPI install (no torch-geometric/pyg-find-links), so the only ROCm change is the torch wheel index. Requires HF_TOKEN for the gated facebook/UMA checkpoints. + +UMA is multi-task: setup() requires an explicit `task` +(setup_kwargs={"task": ...} / --kwarg task=...) and errors without one. +Verification picks its own head via VERIFY_KWARGS below. """ CHECKPOINTS = { "uma-s-1p1": "uma-s-1p1", - "uma-s-1p2": "uma-s-1p2", + # uma-s-1p2 had a known major bug; uma-s-1p2p1 is the fixed, + # upstream-recommended small model and replaces it here. + "uma-s-1p2p1": "uma-s-1p2p1", "uma-m-1p1": "uma-m-1p1", # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path). "uma:custom": None, } +UMA_TASKS = ("omat", "omol", "oc20", "odac", "omc") + +# Verification-only head selection: smoke-test and a bare `rootstock add` +# verify with these (setup() itself has no default task). +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, + "uma-s-1p2p1": {"task": "omat"}, + "uma-m-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, +} + + +def _require_task(task, checkpoint): + if task is None: + raise ValueError( + f"{checkpoint} is multi-task and has no default head - select one " + f'with setup_kwargs={{"task": ...}} (or --kwarg task=...): ' + f"one of {', '.join(UMA_TASKS)}" + ) + if task not in UMA_TASKS: + raise ValueError(f"unknown task {task!r}; expected one of {', '.join(UMA_TASKS)}") + return task + def _fairchem_device(device: str) -> str: """Translate an indexed device ("cuda:2") into what fairchem v2 accepts. @@ -55,7 +85,8 @@ def _fairchem_device(device: str) -> str: return device -def setup(checkpoint: str, device: str = "cuda", task: str = "omat", **kwargs): +def setup(checkpoint: str, device: str = "cuda", task: str | None = None, **kwargs): + task = _require_task(task, checkpoint) from fairchem.core import FAIRChemCalculator, pretrained_mlip predictor = pretrained_mlip.get_predict_unit( @@ -64,7 +95,7 @@ def setup(checkpoint: str, device: str = "cuda", task: str = "omat", **kwargs): return FAIRChemCalculator(predictor, task_name=task, **kwargs) -def setup_from_path(path: str, device: str = "cuda", task: str = "omat", **kwargs): +def setup_from_path(path: str, device: str = "cuda", task: str | None = None, **kwargs): # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through # load_predict_unit, not the registry-name lookup setup() uses. from fairchem.core import FAIRChemCalculator diff --git a/sample_model_configurations/aurora_configs/mace.py b/sample_model_configurations/aurora_configs/mace.py index ad3781d..1cde5b3 100644 --- a/sample_model_configurations/aurora_configs/mace.py +++ b/sample_model_configurations/aurora_configs/mace.py @@ -36,7 +36,8 @@ Multi-head checkpoints select a head via the `head` kwarg on setup() (setup_kwargs={"head": ...} / --kwarg head=...), named by upstream's training -corpus - see MH1_HEADS; omat_pbe is the default. +corpus - see MH1_HEADS. Selection is required: the multi-head checkpoint +has no default head, and setup() errors when none is given. The OMOL checkpoint expects `charge` and `spin` in `atoms.info`. """ @@ -74,11 +75,23 @@ "matpes_r2scan", ) +# Verification-only head selection for the multi-head checkpoint: +# smoke-test and a bare `rootstock add` verify with this +# (setup() itself has no default head). +VERIFY_KWARGS = { + "mace-mh-1": {"head": "omat_pbe"}, +} + def setup(checkpoint: str, device: str = "xpu", head: str | None = None, **kwargs): arg = CHECKPOINTS[checkpoint] if arg.startswith("mh:"): - head = head or "omat_pbe" + if head is None: + raise ValueError( + f"{checkpoint} is multi-head and has no default - select one " + f"with setup_kwargs={{'head': ...}} (or --kwarg head=...): " + f"one of {', '.join(MH1_HEADS)}" + ) if head not in MH1_HEADS: raise ValueError(f"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}") from mace.calculators import mace_mp diff --git a/sample_model_configurations/aurora_configs/uma.py b/sample_model_configurations/aurora_configs/uma.py index 1f0fa49..8575ea6 100644 --- a/sample_model_configurations/aurora_configs/uma.py +++ b/sample_model_configurations/aurora_configs/uma.py @@ -1,7 +1,8 @@ # /// script # requires-python = ">=3.11" # dependencies = [ -# "fairchem-core>=2.20", +# # 2.22 is the first release with uma-s-1p2p1 in the registry. +# "fairchem-core>=2.22", # "ase>=3.22", # # Intel XPU (Aurora PVC) torch build. >=2.13: older XPU wheels (e.g. 2.8) # # have far slower FP64 kernels -- UMA's first forward took >40 min on @@ -35,16 +36,45 @@ Pin one PVC tile with ZE_AFFINITY_MASK in the job (the worker inherits it). Requires HF_TOKEN for the gated facebook/UMA checkpoints (download on a login node; `rootstock add ... --no-verify`, then verify on a compute node). + +UMA is multi-task: setup() requires an explicit `task` +(setup_kwargs={"task": ...} / --kwarg task=...) and errors without one. +Verification picks its own head via VERIFY_KWARGS below. """ CHECKPOINTS = { "uma-s-1p1": "uma-s-1p1", - "uma-s-1p2": "uma-s-1p2", + # uma-s-1p2 had a known major bug; uma-s-1p2p1 is the fixed, + # upstream-recommended small model and replaces it here. + "uma-s-1p2p1": "uma-s-1p2p1", "uma-m-1p1": "uma-m-1p1", # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path). "uma:custom": None, } +UMA_TASKS = ("omat", "omol", "oc20", "odac", "omc") + +# Verification-only head selection: smoke-test and a bare `rootstock add` +# verify with these (setup() itself has no default task). +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, + "uma-s-1p2p1": {"task": "omat"}, + "uma-m-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, +} + + +def _require_task(task, checkpoint): + if task is None: + raise ValueError( + f"{checkpoint} is multi-task and has no default head - select one " + f'with setup_kwargs={{"task": ...}} (or --kwarg task=...): ' + f"one of {', '.join(UMA_TASKS)}" + ) + if task not in UMA_TASKS: + raise ValueError(f"unknown task {task!r}; expected one of {', '.join(UMA_TASKS)}") + return task + def _enable_xpu() -> None: """Teach FairChem's predict unit to accept device="xpu". @@ -76,7 +106,8 @@ def _fp64_settings(): return InferenceSettings(base_precision_dtype=torch.float64, tf32=False) -def setup(checkpoint: str, device: str = "xpu", task: str = "omat", **kwargs): +def setup(checkpoint: str, device: str = "xpu", task: str | None = None, **kwargs): + task = _require_task(task, checkpoint) _enable_xpu() from fairchem.core import FAIRChemCalculator, pretrained_mlip @@ -86,7 +117,7 @@ def setup(checkpoint: str, device: str = "xpu", task: str = "omat", **kwargs): return FAIRChemCalculator(predictor, task_name=task, **kwargs) -def setup_from_path(path: str, device: str = "xpu", task: str = "omat", **kwargs): +def setup_from_path(path: str, device: str = "xpu", task: str | None = None, **kwargs): # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads # through load_predict_unit, not the registry-name lookup setup() uses. _enable_xpu() diff --git a/sample_model_configurations/modal_app.py b/sample_model_configurations/modal_app.py index c02243a..9d217b3 100644 --- a/sample_model_configurations/modal_app.py +++ b/sample_model_configurations/modal_app.py @@ -92,7 +92,9 @@ def probe_image( ) -def _run_probe_subprocess(checkpoint: str, system: str, device: str = "cuda") -> int: +def _run_probe_subprocess( + checkpoint: str, system: str, device: str = "cuda", kwargs: list[str] | None = None +) -> int: """ Run probe.py as a subprocess so its STAGE markers stream to stdout in realtime. Called from inside an @app.function whose image has the right @@ -123,6 +125,8 @@ def _run_probe_subprocess(checkpoint: str, system: str, device: str = "cuda") -> ] if checkpoint: cmd += ["--checkpoint", checkpoint] + for spec in kwargs or []: + cmd += ["--kwarg", spec] print(f"PROBE_CMD: {' '.join(cmd)}", flush=True) result = subprocess.run(cmd, env=sub_env) if result.returncode != 0: @@ -352,29 +356,29 @@ def probe_escn(checkpoint: str = "escn-l6-m2-lay12-s2ef-oc20-all-md", system: st @probe_image( "uma.py", - ["fairchem-core>=2.20", "ase>=3.22", "torch>=2.4.0"], + ["fairchem-core>=2.22", "ase>=3.22", "torch>=2.4.0"], python_version="3.11", ) def probe_uma_small(checkpoint: str = "uma-s-1p1", system: str = "crystal"): """Probe UMA Small on OMAT-style bulk materials.""" - return _run_probe_subprocess(checkpoint, system) + return _run_probe_subprocess(checkpoint, system, kwargs=["task=omat"]) @probe_image( "uma.py", - ["fairchem-core>=2.20", "ase>=3.22", "torch>=2.4.0"], + ["fairchem-core>=2.22", "ase>=3.22", "torch>=2.4.0"], python_version="3.11", ) -def probe_uma_1p2(checkpoint: str = "uma-s-1p2", system: str = "crystal"): - """Probe UMA Small v1.2 (latest, fixes the uma-s-1 extensivity bug).""" - return _run_probe_subprocess(checkpoint, system) +def probe_uma_1p2p1(checkpoint: str = "uma-s-1p2p1", system: str = "crystal"): + """Probe UMA Small v1.2.1 (latest; fixes uma-s-1p2's known bug).""" + return _run_probe_subprocess(checkpoint, system, kwargs=["task=omat"]) @probe_image( "uma.py", - ["fairchem-core>=2.20", "ase>=3.22", "torch>=2.4.0"], + ["fairchem-core>=2.22", "ase>=3.22", "torch>=2.4.0"], python_version="3.11", ) 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) + return _run_probe_subprocess(checkpoint, system, kwargs=["task=omat"]) diff --git a/sample_model_configurations/nvidia_configs/mace.py b/sample_model_configurations/nvidia_configs/mace.py index bda1b1a..42fb630 100644 --- a/sample_model_configurations/nvidia_configs/mace.py +++ b/sample_model_configurations/nvidia_configs/mace.py @@ -18,7 +18,8 @@ Multi-head checkpoints select a head via the `head` kwarg on setup() (setup_kwargs={"head": ...} / --kwarg head=...), named by upstream's training -corpus — see MH1_HEADS; omat_pbe is the default. +corpus — see MH1_HEADS. Selection is required: the multi-head checkpoint +has no default head, and setup() errors when none is given. The OMOL checkpoint expects `charge` and `spin` in `atoms.info`. """ @@ -56,11 +57,23 @@ "matpes_r2scan", ) +# Verification-only head selection for the multi-head checkpoint: +# smoke-test and a bare `rootstock add` verify with this +# (setup() itself has no default head). +VERIFY_KWARGS = { + "mace-mh-1": {"head": "omat_pbe"}, +} + def setup(checkpoint: str, device: str = "cuda", head: str | None = None, **kwargs): arg = CHECKPOINTS[checkpoint] if arg.startswith("mh:"): - head = head or "omat_pbe" + if head is None: + raise ValueError( + f"{checkpoint} is multi-head and has no default — select one " + f"with setup_kwargs={{'head': ...}} (or --kwarg head=...): " + f"one of {', '.join(MH1_HEADS)}" + ) if head not in MH1_HEADS: raise ValueError(f"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}") from mace.calculators import mace_mp diff --git a/sample_model_configurations/nvidia_configs/uma.py b/sample_model_configurations/nvidia_configs/uma.py index 8c2aca2..2e279ac 100644 --- a/sample_model_configurations/nvidia_configs/uma.py +++ b/sample_model_configurations/nvidia_configs/uma.py @@ -1,7 +1,8 @@ # /// script # requires-python = ">=3.11" # dependencies = [ -# "fairchem-core>=2.20", +# # 2.22 is the first release with uma-s-1p2p1 in the registry. +# "fairchem-core>=2.22", # "ase>=3.22", # "torch>=2.4.0", # ] @@ -11,24 +12,34 @@ fairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so this env is a plain PyPI install. The original uma-s-1 had an extensivity bug and was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1. + +UMA is multi-task: every calculation runs one of the model's task heads, and +there is no sensible task-agnostic default — setup() requires an explicit +`task` (setup_kwargs={"task": ...} / --kwarg task=...) and errors without one. +Verification picks its own head via VERIFY_KWARGS below. """ CHECKPOINTS = { "uma-s-1p1": "uma-s-1p1", - # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the - # upstream-recommended small model. 1p2 stays listed for reproducibility - # of existing runs. - "uma-s-1p2": "uma-s-1p2", - # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any - # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the - # 2026-07-30 sync failed on exactly this). Re-add when the next - # fairchem-core ships, and bump the dependency floor to that version. - # "uma-s-1p2p1": "uma-s-1p2p1", + # uma-s-1p2 had a known major bug; uma-s-1p2p1 is the fixed, + # upstream-recommended small model and replaces it here. + "uma-s-1p2p1": "uma-s-1p2p1", "uma-m-1p1": "uma-m-1p1", # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path). "uma:custom": None, } +UMA_TASKS = ("omat", "omol", "oc20", "odac", "omc") + +# Verification-only head selection: smoke-test and a bare `rootstock add` +# verify with these (setup() itself has no default task). +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, + "uma-s-1p2p1": {"task": "omat"}, + "uma-m-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, +} + def _fairchem_device(device: str) -> str: """Translate an indexed device ("cuda:2") into what fairchem v2 accepts. @@ -49,7 +60,20 @@ def _fairchem_device(device: str) -> str: return device -def setup(checkpoint: str, device: str = "cuda", task: str = "omat", **kwargs): +def _require_task(task: str | None, checkpoint: str) -> str: + if task is None: + raise ValueError( + f"{checkpoint} is multi-task and has no default head — select one " + f'with setup_kwargs={{"task": ...}} (or --kwarg task=...): ' + f"one of {', '.join(UMA_TASKS)}" + ) + if task not in UMA_TASKS: + raise ValueError(f"unknown task {task!r}; expected one of {', '.join(UMA_TASKS)}") + return task + + +def setup(checkpoint: str, device: str = "cuda", task: str | None = None, **kwargs): + task = _require_task(task, checkpoint) from fairchem.core import FAIRChemCalculator, pretrained_mlip predictor = pretrained_mlip.get_predict_unit( @@ -58,9 +82,11 @@ def setup(checkpoint: str, device: str = "cuda", task: str = "omat", **kwargs): return FAIRChemCalculator(predictor, task_name=task, **kwargs) -def setup_from_path(path: str, device: str = "cuda", task: str = "omat", **kwargs): +def setup_from_path(path: str, device: str = "cuda", task: str | None = None, **kwargs): # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through - # load_predict_unit, not the registry-name lookup setup() uses. + # load_predict_unit, not the registry-name lookup setup() uses. No task + # validation here — a fine-tune may carry its own task names, or a single + # task (fairchem itself errors when a multi-task model gets task_name=None). from fairchem.core import FAIRChemCalculator from fairchem.core.units.mlip_unit import load_predict_unit diff --git a/tests/verify/test_verify_kwargs.py b/tests/verify/test_verify_kwargs.py new file mode 100644 index 0000000..0afcb53 --- /dev/null +++ b/tests/verify/test_verify_kwargs.py @@ -0,0 +1,160 @@ +"""VERIFY_KWARGS: parsing, lookup, and the verify_checkpoint fallback. + +Envs whose ``setup()`` requires a kwarg (a task-head selection with no +default — UMA's ``task``, MACE-MH-1's ``head``) declare a module-level +``VERIFY_KWARGS`` so smoke-test and a bare ``rootstock add`` can still +verify. Explicit kwargs always win; the calculator path never reads it. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from rootstock import verify +from rootstock.environment import parse_verify_kwargs, verify_kwargs_for + +DECLARING = """CHECKPOINTS = {"uma-s-1p1": "uma-s-1p1", "uma:custom": None} + +VERIFY_KWARGS = { + "uma-s-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, +} + +def setup(checkpoint, device="cuda", task=None): + return None + +def setup_from_path(path, device="cuda", task=None): + return None +""" + +PLAIN = """CHECKPOINTS = {"mace-mp-0-medium": "medium"} + +def setup(checkpoint, device="cuda"): + return None +""" + + +def _write(tmp_path: Path, source: str) -> Path: + path = tmp_path / "env.py" + path.write_text(source) + return path + + +def _install(root: Path, name: str, source: str) -> None: + env_dir = root / "envs" / name + env_dir.mkdir(parents=True) + (env_dir / "env_source.py").write_text(source) + + +# --- parse_verify_kwargs ----------------------------------------------------- + + +def test_parse_absent_means_empty(tmp_path): + assert parse_verify_kwargs(_write(tmp_path, PLAIN)) == {} + + +def test_parse_declaration(tmp_path): + parsed = parse_verify_kwargs(_write(tmp_path, DECLARING)) + assert parsed == { + "uma-s-1p1": {"task": "omat"}, + "uma:custom": {"task": "omat"}, + } + + +def test_parse_annotated_assignment(tmp_path): + source = 'VERIFY_KWARGS: dict = {"a": {"task": "omat"}}\n' + assert parse_verify_kwargs(_write(tmp_path, source)) == {"a": {"task": "omat"}} + + +def test_parse_rejects_non_dict(tmp_path): + with pytest.raises(ValueError, match="dict literal"): + parse_verify_kwargs(_write(tmp_path, "VERIFY_KWARGS = ['nope']\n")) + + +def test_parse_rejects_non_literal(tmp_path): + source = 'X = {"task": "omat"}\nVERIFY_KWARGS = {"a": X}\n' + with pytest.raises(ValueError, match="dict literal"): + parse_verify_kwargs(_write(tmp_path, source)) + + +def test_parse_rejects_non_dict_values(tmp_path): + with pytest.raises(ValueError, match="setup kwargs"): + parse_verify_kwargs(_write(tmp_path, 'VERIFY_KWARGS = {"a": "omat"}\n')) + + +# --- verify_kwargs_for ------------------------------------------------------- + + +def test_lookup_declared_entry(tmp_path): + _install(tmp_path, "uma", DECLARING) + assert verify_kwargs_for(tmp_path, "uma", "uma-s-1p1") == {"task": "omat"} + assert verify_kwargs_for(tmp_path, "uma", "uma:custom") == {"task": "omat"} + + +def test_lookup_undeclared_checkpoint_is_empty(tmp_path): + _install(tmp_path, "uma", DECLARING) + assert verify_kwargs_for(tmp_path, "uma", "uma-m-1p1") == {} + + +def test_lookup_missing_source_is_empty(tmp_path): + assert verify_kwargs_for(tmp_path, "ghost", "uma-s-1p1") == {} + + +# --- verify_checkpoint fallback ---------------------------------------------- + + +@pytest.fixture +def capturing_server(monkeypatch): + """Stub RootstockServer that records ctor kwargs and verifies cleanly.""" + captured: list[dict] = [] + + class _Server: + def __init__(self, **kwargs): + captured.append(kwargs) + + def start(self): + pass + + def calculate(self, **_): + forces = np.array([[0.1, -0.2, 0.0], [-0.05, 0.1, 0.0], [-0.05, 0.1, 0.0]]) + return -14.0, forces, np.zeros((3, 3)) + + def stop(self): + pass + + monkeypatch.setattr("rootstock.server.RootstockServer", _Server) + return captured + + +def test_empty_kwargs_fall_back_to_declaration(tmp_path, capturing_server): + _install(tmp_path, "uma", DECLARING) + ok, err = verify.verify_checkpoint( + root=tmp_path, env_name="uma", checkpoint="uma-s-1p1", device="cpu", setup_kwargs={} + ) + assert ok, err + assert capturing_server[0]["setup_kwargs"] == {"task": "omat"} + + +def test_explicit_kwargs_win(tmp_path, capturing_server): + _install(tmp_path, "uma", DECLARING) + ok, err = verify.verify_checkpoint( + root=tmp_path, + env_name="uma", + checkpoint="uma-s-1p1", + device="cpu", + setup_kwargs={"task": "omol"}, + ) + assert ok, err + assert capturing_server[0]["setup_kwargs"] == {"task": "omol"} + + +def test_no_declaration_stays_empty(tmp_path, capturing_server): + _install(tmp_path, "mace", PLAIN) + ok, err = verify.verify_checkpoint( + root=tmp_path, env_name="mace", checkpoint="mace-mp-0-medium", device="cpu" + ) + assert ok, err + assert capturing_server[0]["setup_kwargs"] == {}