Prepare FAIRChem Hydra entry points for Hydra 1.4

packages/fairchem-core/pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     "orjson",
     "tqdm",
     "submitit>=1.5.4",
-    "hydra-core",
+    "hydra-core>=1.3",
     "torchtnt",
     "pyyaml",
     "wandb",

src/fairchem/core/_cli.py

@@ -10,6 +10,7 @@
 import argparse
 import logging
 import os
+import tempfile
 from typing import TYPE_CHECKING
 
 import hydra
@@ -69,8 +70,22 @@ def get_hydra_config_from_yaml(
     os.environ["HYDRA_FULL_ERROR"] = "1"
     config_directory = os.path.dirname(os.path.abspath(config_yml))
     config_name = os.path.basename(config_yml)
-    hydra.initialize_config_dir(config_directory, version_base="1.1")
-    cfg = hydra.compose(config_name=config_name, overrides=overrides_args)
+
+    if config_name.endswith(".yml"):
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_config_name = f"{os.path.splitext(config_name)[0]}.yaml"
+            temp_config_path = os.path.join(temp_dir, temp_config_name)
+            with open(config_yml) as source, open(temp_config_path, "w") as target:
+                target.write(source.read())
+
+            with hydra.initialize_config_dir(temp_dir, version_base="1.3"):
+                cfg = hydra.compose(
+                    config_name=temp_config_name,
+                    overrides=overrides_args,
+                )
+    else:
+        with hydra.initialize_config_dir(config_directory, version_base="1.3"):
+            cfg = hydra.compose(config_name=config_name, overrides=overrides_args)
     # merge default structured config with initialized job object
     cfg = OmegaConf.merge({"job": OmegaConf.structured(JobConfig)}, cfg)
     # canonicalize config (remove top level keys that just used replacing variables)

tests/core/scripts/test_sweep_inference_benchmark.py

@@ -0,0 +1,40 @@
+"""
+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 textwrap
+
+import hydra
+
+from fairchem.core.scripts.sweep_inference_benchmark import load_config
+
+
+def test_load_config_hydra_composition_and_defaults(tmp_path):
+    """
+    Test that benchmark config composition works and initializes JobConfig metadata.
+    """
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    config_path = tmp_path / "benchmark_config.yaml"
+    config_path.write_text(
+        textwrap.dedent(
+            """
+            benchmark:
+              base_natoms: 32
+            """
+        )
+    )
+
+    cfg = load_config(
+        str(config_path),
+        overrides=["+job.scheduler.ranks_per_node=2"],
+    )
+
+    assert cfg.benchmark.base_natoms == 32
+    assert cfg.job.scheduler.ranks_per_node == 2
+    assert cfg.job.metadata is not None
+    assert cfg.job.run_name is not None

tests/lammps/test_lammps_fc.py

@@ -8,12 +8,18 @@
 from __future__ import annotations
 
 import os
+import sys
 import tempfile
 
+import hydra
 import numpy as np
 import pytest
 from ase import Atoms
-from fairchem.lammps.lammps_fc import restricted_cell_from_lammps_box
+
+pytest.importorskip("lammps")
+
+from fairchem.lammps import lammps_fc  # noqa: E402
+from fairchem.lammps.lammps_fc import restricted_cell_from_lammps_box  # noqa: E402
 
 
 def create_lammps_data_file(filepath, positions, cell, atom_types, masses):
@@ -230,3 +236,63 @@ def test_cell_conversion_preserves_volume(box_name, boxlo, boxhi, xy, yz, xz):
         f"Volume mismatch for {box_name}: boxlo={boxlo}, boxhi={boxhi}, xy={xy}, yz={yz}, xz={xz}.\n"
         f"Expected: {expected_volume}, Actual: {actual_volume}"
     )
+
+
+def test_lammps_hydra_entrypoint_startup(monkeypatch, tmp_path):
+    """
+    Test that the Hydra-decorated LAMMPS entry point composes config and starts.
+    """
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    lammps_state = {}
+    run_args = {}
+
+    def fake_instantiate(_cfg):
+        return object()
+
+    def fake_run_lammps_with_fairchem(
+        predictor,
+        lammps_input_path,
+        task_name,
+        charge=0,
+        spin=0,
+    ):
+        run_args.update(
+            {
+                "predictor": predictor,
+                "lammps_input_path": lammps_input_path,
+                "task_name": task_name,
+                "charge": charge,
+                "spin": spin,
+            }
+        )
+
+        class DummyLammps:
+            pass
+
+        lmp = DummyLammps()
+        lmp._predictor = predictor
+        lammps_state["lmp"] = lmp
+        return lmp
+
+    monkeypatch.setattr(lammps_fc.hydra.utils, "instantiate", fake_instantiate)
+    monkeypatch.setattr(
+        lammps_fc, "run_lammps_with_fairchem", fake_run_lammps_with_fairchem
+    )
+
+    old_argv = sys.argv[:]
+    try:
+        sys.argv = [
+            "lammps_fc.py",
+            f"hydra.run.dir={tmp_path}",
+            "hydra.output_subdir=null",
+        ]
+        lammps_fc.main()
+    finally:
+        sys.argv = old_argv
+        hydra.core.global_hydra.GlobalHydra.instance().clear()
+
+    assert run_args["lammps_input_path"] == "lammps_in_example.file"
+    assert run_args["task_name"] == "omol"
+    assert run_args["charge"] == 0
+    assert run_args["spin"] == 0
+    assert not hasattr(lammps_state["lmp"], "_predictor")