Fairchem ray serve (#3228)

zulissimeta · pre-commit-ci[bot] · vineetbansal · web-flow · commit 7cd3800dd92f · 2026-06-08T23:06:15.000-04:00
## Summary of Changes Adds an option to use a batching and distributed inference system for all fairchem MLIP inference when 1. Running the jobs/flows within a ray server and 2. When FAIRCHEM_RAY_SERVE_BATCHING quacc setting is true Relies on fairchem>=2.21 ### Requirements - [x] My PR is focused on a [single feature addition or bugfix](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/getting-started/best-practices-for-pull-requests#write-small-prs). - [x] My PR has relevant, comprehensive [unit tests](https://quantum-accelerators.github.io/quacc/dev/contributing.html#unit-tests). - [x] My PR is on a [custom branch](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-and-deleting-branches-within-your-repository) (i.e. is _not_ named `main`). Note: If you are an external contributor, you will see a comment from [@buildbot-princeton](https://github.com/buildbot-princeton). This is solely for the maintainers. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Vineet Bansal <vineetbansal@protonmail.com> Co-authored-by: Andrew S. Rosen <asrosen93@gmail.com>
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -403,7 +403,7 @@ jobs:
       - name: Install pip packages
         run: |
           pip install uv
-          uv pip install --system -r tests/requirements.txt -r tests/requirements-fairchem.txt "quacc[dev] @ ."
+          uv pip install --system -r tests/requirements.txt -r tests/requirements-fairchem.txt -r tests/requirements-ray.txt "quacc[dev] @ ."
 
       - name: Set coverage core
         run: |
diff --git a/pyproject.toml b/pyproject.toml
@@ -43,7 +43,7 @@ dependencies = [
 [project.optional-dependencies]
 dask = ["dask[distributed]>=2023.12.1", "dask-jobqueue>=0.8.2"]
 defects = ["pymatgen-analysis-defects>=2024.10.22", "shakenbreak>=3.2.0"]
-fairchem = ["fairchem-data-omat>=0.2", "fairchem-data-oc>=1.0.2", "fairchem-core>=2.2.0"]
+fairchem = ["fairchem-data-omat>=0.2", "fairchem-data-oc>=1.0.2", "fairchem-core>=2.21.0"]
 jobflow = ["jobflow>=0.3.0", "jobflow-remote>=1.0.0"]
 orb = ["orb-models>=0.4.1"]
 mace = ["mace-torch>=0.3.3", "mace-models>=0.1.6"]
diff --git a/src/quacc/recipes/mlp/_base.py b/src/quacc/recipes/mlp/_base.py
@@ -82,11 +82,16 @@ def pick_calculator(
     """
     import torch
 
-    if not torch.cuda.is_available():
-        LOGGER.warning("CUDA is not available to PyTorch. Calculations will be slow.")
+    from quacc import get_settings
 
+    settings = get_settings()
     method = method.lower()
 
+    # Skip CUDA warning for rayserve batching mode (inference happens on remote GPU)
+    use_ray_serve = method == "fairchem" and settings.FAIRCHEM_RAY_SERVE_BATCHING
+    if not use_ray_serve and not torch.cuda.is_available():
+        LOGGER.warning("CUDA is not available to PyTorch. Calculations will be slow.")
+
     if "m3gnet" in method or "chgnet" in method or "tensornet" in method:
         import matgl
         from matgl import __version__
@@ -132,7 +137,67 @@ def pick_calculator(
     elif method.lower() == "fairchem":
         from fairchem.core import FAIRChemCalculator, __version__
 
-        calc = FAIRChemCalculator.from_model_checkpoint(**calc_kwargs)
+        # Check if Ray Serve batching is enabled AND Ray is actually available
+        if use_ray_serve:
+            try:
+                import ray
+
+                if not ray.is_initialized():
+                    LOGGER.warning(
+                        "FAIRCHEM_RAY_SERVE_BATCHING is enabled but Ray is not initialized. "
+                        "Falling back to local inference. To use Ray Serve batching, ensure "
+                        "your flow is running with a Ray cluster (e.g., SlurmRayTaskRunner with "
+                        "start_inference_server=True)."
+                    )
+                    use_ray_serve = False
+            except ImportError:
+                LOGGER.warning(
+                    "FAIRCHEM_RAY_SERVE_BATCHING is enabled but Ray is not installed. "
+                    "Falling back to local inference."
+                )
+                use_ray_serve = False
+
+        if use_ray_serve:
+            # Use Ray Serve multiplexed deployment for inference. The deployment
+            # is expected to already be running on the cluster (typically started
+            # by get_local_inference_raycluster / get_slurm_inference_raycluster with
+            # setup_multiplexed_batch_predict_server). We connect by deployment
+            # name and route requests to the appropriate model via the
+            # multiplexed_model_id, which has the form
+            # "<checkpoint_name_or_path>:<inference_settings>".
+            from fairchem.core.units.mlip_unit.predict import BatchServerPredictUnit
+
+            calc_kwargs = calc_kwargs.copy()  # Don't modify the original kwargs
+
+            # Determine model identifier: prefer name_or_path (local checkpoint)
+            # over model_id/checkpoint
+            name_or_path = calc_kwargs.pop("name_or_path", None)
+            if name_or_path is not None:
+                checkpoint_id = str(name_or_path)
+            else:
+                checkpoint_id = calc_kwargs.pop("model_id", None) or calc_kwargs.pop(
+                    "checkpoint", "uma-s-1p1"
+                )
+
+            inference_settings = calc_kwargs.pop("inference_settings", "default")
+            task_name = calc_kwargs.pop("task_name")
+
+            # Drop kwargs only meaningful when loading the checkpoint locally
+            calc_kwargs.pop("device", None)
+            calc_kwargs.pop("overrides", None)
+            calc_kwargs.pop("seed", None)
+
+            multiplexed_model_id = f"{checkpoint_id}:{inference_settings}"
+
+            mlip_unit = BatchServerPredictUnit.from_deployment_connection_info(
+                deployment_name="multiplexed-predict-server",
+                multiplexed_model_id=multiplexed_model_id,
+            )
+
+            calc = FAIRChemCalculator(predict_unit=mlip_unit, task_name=task_name)
+        else:
+            # Use local inference
+            calc = FAIRChemCalculator.from_model_checkpoint(**calc_kwargs)
 
     else:
         raise ValueError(f"Unrecognized {method=}.")
diff --git a/src/quacc/settings.py b/src/quacc/settings.py
@@ -126,6 +126,23 @@ class QuaccSettings(BaseSettings):
         description="Whether to resolve all futures in flow results to data and fail if not possible",
     )
 
+    # ---------------------------
+    # FAIRChem Settings
+    # ---------------------------
+    FAIRCHEM_RAY_SERVE_BATCHING: bool = Field(
+        False,
+        description=(
+            """
+            Whether to use Ray Serve for batched FAIRChem inference. When True,
+            the 'fairchem' method in pick_calculator will use RayServeMLIPUnit
+            to send inference requests to a Ray Serve deployment instead of
+            running inference locally. This enables efficient batched inference
+            across multiple concurrent calculations. Requires a Ray Serve inference
+            server to be running (e.g., via SlurmRayTaskRunner with start_inference_server=True).
+            """
+        ),
+    )
+
     # ---------------------------
     # ORCA Settings
     # ---------------------------
diff --git a/tests/core/recipes/mlp_recipes/test_fairchem_ray_serve.py b/tests/core/recipes/mlp_recipes/test_fairchem_ray_serve.py
@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+from importlib.util import find_spec
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("ray")
+pytest.importorskip("fairchem.core")
+
+if find_spec("fairchem"):
+    from huggingface_hub.utils._auth import get_token
+
+    if not get_token():
+        pytest.skip(
+            "HF_TOKEN required for fairchem ray-serve tests", allow_module_level=True
+        )
+
+import numpy as np
+from ase.build import bulk
+
+from quacc import get_settings
+from quacc.recipes.mlp._base import pick_calculator
+from quacc.recipes.mlp.core import relax_job, static_job
+from quacc.recipes.mlp.elastic import elastic_tensor_flow
+
+
+@pytest.fixture(scope="module")
+def ray_serve_cluster():
+    """Spin up a local Ray cluster + multiplexed fairchem serve deployment.
+
+    The deployment is named ``predict-server`` to match what
+    ``quacc.recipes.mlp._base`` connects to when
+    ``FAIRCHEM_RAY_SERVE_BATCHING`` is enabled.
+    """
+    import ray
+    from fairchem.core.components.batch_server import (
+        setup_multiplexed_batch_predict_server,
+    )
+    from ray import serve
+
+    started_ray = False
+    if not ray.is_initialized():
+        ray.init(
+            include_dashboard=False,
+            ignore_reinit_error=True,
+            log_to_driver=False,
+            configure_logging=False,
+        )
+        started_ray = True
+
+    handle = setup_multiplexed_batch_predict_server(
+        deployment_name="predict-server",
+        route_prefix="/predict-server",
+        num_replicas=1,
+        ray_actor_options={"num_cpus": 1},
+    )
+    try:
+        yield handle
+    finally:
+        try:
+            serve.shutdown()
+        finally:
+            if started_ray and ray.is_initialized():
+                ray.shutdown()
+
+
+@pytest.fixture
+def _batching(monkeypatch, request, ray_serve_cluster):
+    settings = get_settings()
+    monkeypatch.setattr(
+        settings, "FAIRCHEM_RAY_SERVE_BATCHING", request.param, raising=False
+    )
+    # pick_calculator is wrapped by freezeargs(lru_cache(...)); clear the
+    # underlying lru_cache so each parametrize value rebuilds the calc
+    # under the current FAIRCHEM_RAY_SERVE_BATCHING setting.
+    pick_calculator.__wrapped__.cache_clear()
+    yield request.param
+    pick_calculator.__wrapped__.cache_clear()
+
+
+def _assert_calc_path(batching):
+    """Verify the calculator built during the test actually used the
+    expected code path (Ray Serve batching vs. local inference)."""
+    from fairchem.core.units.mlip_unit.predict import BatchServerPredictUnit
+
+    calc = pick_calculator(
+        method="fairchem", name_or_path="uma-s-1p1", task_name="omat"
+    )
+    used_batching = isinstance(getattr(calc, "predictor", None), BatchServerPredictUnit)
+    assert used_batching is batching, (
+        f"Expected FAIRCHEM_RAY_SERVE_BATCHING={batching} path, "
+        f"but BatchServerPredictUnit usage was {used_batching}."
+    )
+
+
+@pytest.mark.parametrize("_batching", [False, True], indirect=True)
+def test_static_job(tmp_path, monkeypatch, _batching):
+    monkeypatch.chdir(tmp_path)
+    torch.set_default_dtype(torch.float32)
+
+    atoms = bulk("Cu")
+    output = static_job(
+        atoms, method="fairchem", name_or_path="uma-s-1p1", task_name="omat"
+    )
+    assert output["results"]["energy"] == pytest.approx(-3.7501682869643735, rel=1e-3)
+    assert np.shape(output["results"]["forces"]) == (1, 3)
+    assert output["atoms"] == atoms
+    _assert_calc_path(_batching)
+
+
+@pytest.mark.parametrize("_batching", [False, True], indirect=True)
+def test_relax_job(tmp_path, monkeypatch, _batching):
+    monkeypatch.chdir(tmp_path)
+    torch.set_default_dtype(torch.float32)
+
+    atoms = bulk("Cu") * (2, 2, 2)
+    atoms[0].position += 0.1
+    output = relax_job(
+        atoms, method="fairchem", name_or_path="uma-s-1p1", task_name="omat"
+    )
+    assert output["results"]["energy"] == pytest.approx(-30.001143639922756, rel=1e-3)
+    assert np.shape(output["results"]["forces"]) == (8, 3)
+    assert output["atoms"] != atoms
+    assert output["atoms"].get_volume() == pytest.approx(atoms.get_volume())
+    _assert_calc_path(_batching)
+
+
+@pytest.mark.parametrize("_batching", [False, True], indirect=True)
+def test_relax_cell_job(tmp_path, monkeypatch, _batching):
+    monkeypatch.chdir(tmp_path)
+    torch.set_default_dtype(torch.float32)
+
+    atoms = bulk("Cu") * (2, 2, 2)
+    atoms[0].position += 0.1
+    output = relax_job(
+        atoms,
+        method="fairchem",
+        relax_cell=True,
+        name_or_path="uma-s-1p1",
+        task_name="omat",
+    )
+    assert output["results"]["energy"] == pytest.approx(-30.005004590392726, rel=1e-3)
+    assert np.shape(output["results"]["forces"]) == (8, 3)
+    assert output["atoms"] != atoms
+    assert output["atoms"].get_volume() != pytest.approx(atoms.get_volume())
+    _assert_calc_path(_batching)
+
+
+@pytest.mark.parametrize("_batching", [False, True], indirect=True)
+def test_elastic_jobs(tmp_path, monkeypatch, _batching):
+    monkeypatch.chdir(tmp_path)
+    torch.set_default_dtype(torch.float32)
+
+    atoms = bulk("Cu")
+    outputs = elastic_tensor_flow(
+        atoms,
+        run_static=False,
+        pre_relax=True,
+        job_params={
+            "all": {
+                "method": "fairchem",
+                "name_or_path": "uma-s-1p1",
+                "task_name": "omat",
+            },
+            "relax_job": {"opt_params": {"fmax": 0.01}},
+        },
+    )
+    assert outputs["deformed_results"][0]["atoms"].get_volume() != pytest.approx(
+        atoms.get_volume()
+    )
+    assert outputs["undeformed_result"]["results"]["stress"] == pytest.approx(
+        0, abs=1e-2
+    )
+    assert outputs["elasticity_doc"].bulk_modulus.voigt == pytest.approx(151.367, abs=2)
+    _assert_calc_path(_batching)
diff --git a/tests/core/recipes/mlp_recipes/test_fairchem_ray_serve_units.py b/tests/core/recipes/mlp_recipes/test_fairchem_ray_serve_units.py
diff --git a/tests/requirements-ray.txt b/tests/requirements-ray.txt
