remove bad uma checkpoint and make task heads mandatory to specify

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
 

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 `<family>:custom` checkpoint id. Must be visible from the compute nodes; loaded through the env's `setup_from_path()` hook. No shipped weights are involved |
 

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)
 ```

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.

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,

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``.

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]}"
 

sample_model_configurations/_agent/probe.py

@@ -87,20 +87,34 @@ 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",
         choices=["molecule", "crystal", "slab_co"],
         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

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