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 4adb073..59238f2 100644 --- a/tests/models/test_fairchem.py +++ b/tests/models/test_fairchem.py @@ -249,3 +249,70 @@ def test_calculation_required_tracks_the_cache(self): assert wrapper.calculation_required(atoms, ("energy",)) is False atoms.positions[0][0] += 0.1 assert wrapper.calculation_required(atoms, ("energy",)) is True + + +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)