Rename ParallelMLIPPredictUnit (#1582)

* Update README.md

* rename ParallelMLIPPredictUnit

* update

Author: rayg1234

README.md

@@ -187,7 +187,7 @@ If you have multiple gpus (or multiple nodes), we handle all the parallelism for
 pip install fairchem-core[extras]
 ```
 
-```
+```python
 from ase import units
 from ase.md.langevin import Langevin
 from fairchem.core import pretrained_mlip, FAIRChemCalculator

docs/core/common_tasks/ase_calculator.md

@@ -106,7 +106,7 @@ predictor = pretrained_mlip.get_predict_unit(
 
 UMA supports Graph Parallel inference natively. The graph is chunked into each rank and both the forward and backwards communication is handled by the built-in graph parallel algorithm with torch distributed. Because Multi-GPU inference requires special setup of communication protocols within a node and across nodes, we leverage [ray](https://www.ray.io/) to launch Ray Actors for each GPU-rank under the hood. This allows us to seemlessly scale to any infrastructure that can run Ray.
 
-To make things simple for the user that wants to run multi-gpu inference locally, we provide a drop-in replacement for MLIPPredictUnit, called [ParallelMLIPPredictUnitRay](https://github.com/facebookresearch/fairchem/blob/85bd83535fedbc1d99eee4c12e175603ccc44ef7/src/fairchem/core/units/mlip_unit/predict.py#L415)
+To make things simple for the user that wants to run multi-gpu inference locally, we provide a drop-in replacement for MLIPPredictUnit, called [ParallelMLIPPredictUnit](https://github.com/facebookresearch/fairchem/blob/85bd83535fedbc1d99eee4c12e175603ccc44ef7/src/fairchem/core/units/mlip_unit/predict.py#L415)
 
 To enable this you need to install Ray manually or through the fairchem extra dependencies option
 

docs/core/common_tasks/lammps.md

@@ -47,7 +47,7 @@ lmp_fc lmp_in="lammps_in_example.file"  task_name="omol"
 ```
 
 ## Multi-GPU parallelism
-Our LAMMPs integration is fully compatible out the box with our Multi-GPU inference API. The only change required is to pass it the ParallelMLIPPredictUnitRay [here](https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/lammps/lammps_fc_config.yaml#L20) instead of the regular predict unit when initializing the lammps fairchem script. No need to install anything new such as Kokkos or add communication code.
+Our LAMMPs integration is fully compatible out the box with our Multi-GPU inference API. The only change required is to pass it the ParallelMLIPPredictUnit [here](https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/lammps/lammps_fc_config.yaml#L20) instead of the regular predict unit when initializing the lammps fairchem script. No need to install anything new such as Kokkos or add communication code.
 
 For example:
 ```

src/fairchem/core/calculate/pretrained_mlip.py

@@ -85,7 +85,7 @@ def get_predict_unit(
         device: Optional torch device to load the model onto. If None, uses the default device.
         cache_dir: Path to folder where model files will be stored. Default is "~/.cache/fairchem"
         workers: Number of parallel workers for prediction unit. Default is 1. If greater than 1,
-            we will instantiate a ParallelMLIPPredictUnitRay instead of the normal predict unit.
+            we will instantiate a ParallelMLIPPredictUnit instead of the normal predict unit.
 
     Returns:
         An initialized MLIPPredictUnit ready for making predictions.

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

@@ -41,7 +41,7 @@ def load_predict_unit(
         device: Optional torch device to load the model onto.
         atom_refs: Optional dictionary of isolated atom reference energies.
         workers: Number of parallel workers for prediction unit. Default is 1. If greater than 1,
-            we will instantiate a ParallelMLIPPredictUnitRay instead of the normal predict unit.
+            we will instantiate a ParallelMLIPPredictUnit instead of the normal predict unit.
 
     Returns:
         A MLIPPredictUnit instance ready for inference
@@ -54,9 +54,9 @@ def load_predict_unit(
     inference_settings = guess_inference_settings(inference_settings)
     overrides = overrides or {"backbone": {"always_use_pbc": False}}
     if workers > 1:
-        from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnitRay
+        from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnit
 
-        return ParallelMLIPPredictUnitRay(
+        return ParallelMLIPPredictUnit(
             path,
             device=device,
             inference_settings=inference_settings,

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

@@ -412,7 +412,7 @@ def predict(self, data: AtomicData) -> dict[str, torch.tensor] | None:
 
 
 @requires(ray_installed, message="Requires `ray` to be installed")
-class ParallelMLIPPredictUnitRay(MLIPPredictUnitProtocol):
+class ParallelMLIPPredictUnit(MLIPPredictUnitProtocol):
     def __init__(
         self,
         inference_model_path: str,

src/fairchem/lammps/lammps_fc_config.yaml

@@ -18,7 +18,7 @@ local_predict_unit:
 
 # Use parallel predict unit for multi-gpu
 parallel_predict_unit:
-  _target_: fairchem.core.units.mlip_unit.predict.ParallelMLIPPredictUnitRay
+  _target_: fairchem.core.units.mlip_unit.predict.ParallelMLIPPredictUnit
   inference_model_path:
     _target_: fairchem.core.calculate.pretrained_mlip.pretrained_checkpoint_path_from_name
     model_name: "uma-s-1p1"

tests/core/units/mlip_unit/test_predict.py

@@ -10,7 +10,7 @@
 from fairchem.core.calculate.pretrained_mlip import pretrained_checkpoint_path_from_name
 from fairchem.core.datasets.atomic_data import AtomicData, atomicdata_list_to_batch
 from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
-from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnitRay
+from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnit
 from tests.conftest import seed_everywhere
 
 ATOL = 1e-5
@@ -144,7 +144,7 @@ def test_parallel_predict_unit(workers, device):
     atomic_data = AtomicData.from_ase(atoms, task_name=["omat"])
 
     seed_everywhere(seed)
-    ppunit = ParallelMLIPPredictUnitRay(
+    ppunit = ParallelMLIPPredictUnit(
         inference_model_path=model_path,
         device=device,
         inference_settings=ifsets,