feat: ban retired UMA checkpoints by md5

Refuse to load uma-s-1p1 and uma-s-1p2 inference checkpoints. The check
runs at the single chokepoint (MLIPPredictUnit.__init__) immediately
before torch.load, so every end-user inference path -- get_predict_unit,
FAIRChemCalculator.from_model_checkpoint, and direct construction -- is
covered. New bans can be added by appending an md5 to
_BANNED_CHECKPOINTS in predict.py.

Includes a parameterized integration test that loads the cached HF
artifacts and asserts BannedCheckpointError, plus a small utility
script for computing the md5 of a downloaded checkpoint.

Author: Michael Dzamba

scripts/download_uma_checkpoints.py

@@ -0,0 +1,41 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+
+from fairchem.core.calculate.pretrained_mlip import (
+    pretrained_checkpoint_path_from_name,
+)
+
+MODELS = ["uma-s-1p1", "uma-s-1p2"]
+
+
+def md5_of_file(path: str | Path, chunk_size: int = 1024 * 1024) -> str:
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def main() -> None:
+    for name in MODELS:
+        print(f"Downloading {name} ...")
+        path = pretrained_checkpoint_path_from_name(name)
+        size_mb = Path(path).stat().st_size / (1024 * 1024)
+        digest = md5_of_file(path)
+        print(f"  path: {path}")
+        print(f"  size: {size_mb:.1f} MB")
+        print(f"  md5:  {digest}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

src/fairchem/core/units/mlip_unit/predict.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import copy
+import hashlib
 import logging
 import math
 import os
@@ -56,6 +57,46 @@
     from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint
 
 
+class BannedCheckpointError(RuntimeError):
+    """
+    Raised when a checkpoint with a banned MD5 hash is loaded.
+    """
+
+
+# Map of banned checkpoint MD5 -> human-readable reason. Any inference
+# checkpoint whose md5 matches an entry here will be rejected before it is
+# deserialized by torch.load.
+_BANNED_CHECKPOINTS: dict[str, str] = {
+    # uma-s-1p1
+    "36a2f071350be0ee4c15e7ebdd16dde1": (
+        "uma-s-1p1 has been retired. Please upgrade to a newer UMA model "
+        "(see fairchem.core.calculate.pretrained_mlip.available_models)."
+    ),
+    # uma-s-1p2
+    "26ac47f57e7d68af9f031077cdc2cbe9": (
+        "uma-s-1p2 has been retired. Please upgrade to a newer UMA model "
+        "(see fairchem.core.calculate.pretrained_mlip.available_models)."
+    ),
+}
+
+
+def _md5_of_file(path: str, chunk_size: int = 1024 * 1024) -> str:
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _verify_checkpoint_not_banned(path: str) -> None:
+    digest = _md5_of_file(path)
+    if digest in _BANNED_CHECKPOINTS:
+        raise BannedCheckpointError(
+            f"Refusing to load banned checkpoint {path} (md5={digest}): "
+            f"{_BANNED_CHECKPOINTS[digest]}"
+        )
+
+
 def collate_predictions(predict_fn):
     @wraps(predict_fn)
     def collated_predict(
@@ -123,6 +164,7 @@ def __init__(
             )
 
         # Load checkpoint first to get model type
+        _verify_checkpoint_not_banned(inference_model_path)
         checkpoint = torch.load(
             inference_model_path, map_location="cpu", weights_only=False
         )

tests/core/units/mlip_unit/test_banned_checkpoints.py

@@ -0,0 +1,38 @@
+"""
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from __future__ import annotations
+
+import pytest
+from huggingface_hub import try_to_load_from_cache
+from huggingface_hub.file_download import _CACHED_NO_EXIST
+
+from fairchem.core._config import CACHE_DIR
+from fairchem.core.calculate.pretrained_mlip import _MODEL_CKPTS
+from fairchem.core.units.mlip_unit import MLIPPredictUnit
+from fairchem.core.units.mlip_unit.predict import BannedCheckpointError
+
+
+@pytest.mark.parametrize("model_name", ["uma-s-1p1", "uma-s-1p2"])
+def test_real_uma_checkpoint_is_banned(model_name):
+    """
+    The retired UMA checkpoints shipped on HuggingFace must trip the
+    BannedCheckpointError. Skipped when the file is not present in the
+    local HF cache so the test does not require network access in CI.
+    """
+    spec = _MODEL_CKPTS.checkpoints[model_name]
+    rel_path = f"{spec.subfolder}/{spec.filename}" if spec.subfolder else spec.filename
+    cached = try_to_load_from_cache(
+        repo_id=spec.repo_id,
+        filename=rel_path,
+        cache_dir=CACHE_DIR,
+    )
+    if cached is None or cached is _CACHED_NO_EXIST:
+        pytest.skip(f"{model_name} not in local HF cache at {CACHE_DIR}")
+
+    with pytest.raises(BannedCheckpointError):
+        MLIPPredictUnit(cached, device="cpu")