From 52d36f966e43fbb70c72639c369f0d0f0c117692 Mon Sep 17 00:00:00 2001 From: Ard van Noordenne Date: Wed, 26 Aug 2026 17:04:13 +0200 Subject: [PATCH] models: accept a local Fairchem checkpoint in batch mode _load_fairchem_torchsim called _verify_model_name_and_cache_dir before branching on model_path. That helper raises when model_name is empty, so a config supplying only a local checkpoint was rejected: ValueError: Model name must be specified in the configuration _load_fairchem_calculator checks model_path first and has no such problem, so the same config optimized fine in sequential mode and failed in batch mode. Resolve the name only on the branch that needs it, matching the ASE loader. model_name is still required when no checkpoint is given -- deferring that to FairChemModel would report it from somewhere much less obvious. --- src/gpuma/models/fairchem.py | 7 +++- tests/models/test_fairchem.py | 67 +++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/src/gpuma/models/fairchem.py b/src/gpuma/models/fairchem.py index 767cb2e..b001d64 100644 --- a/src/gpuma/models/fairchem.py +++ b/src/gpuma/models/fairchem.py @@ -67,7 +67,6 @@ def _load_fairchem_torchsim(config: Config) -> Any: _load_hf_token_to_env(config) model_path = _verify_model_path(config) - model_name, model_cache_dir = _verify_model_name_and_cache_dir(config) # Fairchem internally only accepts "cuda" or "cpu"; _setup_fairchem_device # calls torch.cuda.set_device(N) when a specific GPU is requested so that # Fairchem's internal device resolution picks the correct GPU. @@ -75,8 +74,14 @@ def _load_fairchem_torchsim(config: Config) -> Any: torch_device = torch.device(backend_device) if model_path: + # model_name is deliberately not resolved on this branch. + # _verify_model_name_and_cache_dir raises when model_name is empty, so + # calling it unconditionally rejected a config that supplies only a + # local checkpoint -- which _load_fairchem_calculator accepts. The same + # config then worked in sequential mode and failed in batch mode. uma_model = FairChemModel(model=model_path, task_name="omol", device=torch_device) else: + model_name, model_cache_dir = _verify_model_name_and_cache_dir(config) uma_model = FairChemModel( model=model_name, model_cache_dir=model_cache_dir, diff --git a/tests/models/test_fairchem.py b/tests/models/test_fairchem.py index 25757e5..59b3236 100644 --- a/tests/models/test_fairchem.py +++ b/tests/models/test_fairchem.py @@ -155,3 +155,70 @@ def test_fairchem_model_names_exist(self): """AVAILABLE_FAIRCHEM_MODELS contains all expected UMA models.""" assert len(AVAILABLE_FAIRCHEM_MODELS) >= 3 assert "uma-s-1p2" in AVAILABLE_FAIRCHEM_MODELS + + +class TestFairchemLocalCheckpoint: + """A config supplying only ``model_path`` must work on both paths. + + These stub out ``FairChemModel``/``pretrained_mlip`` rather than downloading + a checkpoint: the behaviour under test is purely which config fields each + loader insists on *before* it constructs anything, so the real model would + add a token requirement without adding coverage. + """ + + @staticmethod + def _checkpoint_only_config(tmp_path): + ckpt = tmp_path / "local.pt" + ckpt.write_bytes(b"") + return Config({ + "model": {"model_type": "fairchem", "model_name": "", "model_path": str(ckpt)}, + "technical": {"device": "cpu"}, + }) + + def test_torchsim_accepts_checkpoint_without_model_name(self, tmp_path, monkeypatch): + """Batch mode used to demand model_name even with a checkpoint given. + + _verify_model_name_and_cache_dir raises on an empty model_name, and it + was called unconditionally, so this config failed in batch mode while + working in sequential mode. + """ + import torch_sim.models.fairchem as ts_fairchem + + sentinel = object() + seen = {} + + def fake_model(**kwargs): + seen.update(kwargs) + return sentinel + + monkeypatch.setattr(ts_fairchem, "FairChemModel", fake_model) + + from gpuma.models.fairchem import _load_fairchem_torchsim + + assert _load_fairchem_torchsim(self._checkpoint_only_config(tmp_path)) is sentinel + assert seen["model"].name == "local.pt" + + def test_calculator_accepts_checkpoint_without_model_name(self, tmp_path, monkeypatch): + """The sequential path already worked; pin it so the two stay in step.""" + import fairchem.core as fc + + sentinel = object() + monkeypatch.setattr( + fc.pretrained_mlip, "load_predict_unit", lambda **kw: sentinel + ) + monkeypatch.setattr(fc, "FAIRChemCalculator", lambda **kw: sentinel) + + from gpuma.models.fairchem import _load_fairchem_calculator + + assert _load_fairchem_calculator(self._checkpoint_only_config(tmp_path)) is sentinel + + def test_model_name_still_required_without_a_checkpoint(self): + """Dropping the check entirely would defer the error to a confusing place.""" + from gpuma.models.fairchem import _load_fairchem_torchsim + + config = Config({ + "model": {"model_type": "fairchem", "model_name": ""}, + "technical": {"device": "cpu"}, + }) + with pytest.raises(ValueError, match="Model name must be specified"): + _load_fairchem_torchsim(config)