fix tests

src/fairchem/core/common/test_utils.py

@@ -44,32 +44,29 @@ class PGConfig:
     use_gp: bool = True
 
 
-def _init_env_rank_and_launch_test(
+def init_env_rank_and_launch_test(
     rank: int,
     pg_setup_params: PGConfig,
     mp_output_dict: dict[int, object],
     test_method: callable,
     args: list[object],
     kwargs: dict[str, object],
-    init_process_group: bool = False,
 ) -> None:
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = pg_setup_params.port
     os.environ["WORLD_SIZE"] = str(pg_setup_params.world_size)
     os.environ["LOCAL_RANK"] = str(rank)
     os.environ["RANK"] = str(rank)
-
     mp_output_dict[rank] = test_method(*args, **kwargs)  # pyre-fixme
 
 
-def _init_pg_and_rank_and_launch_test(
+def init_pg_and_rank_and_launch_test(
     rank: int,
     pg_setup_params: PGConfig,
     mp_output_dict: dict[int, object],
     test_method: callable,
     args: list[object],
     kwargs: dict[str, object],
-    init_process_group: bool = False,
 ) -> None:
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = pg_setup_params.port

tests/core/e2e/test_s2ef.py

@@ -11,7 +11,7 @@
 import yaml
 from fairchem.core.common.test_utils import (
     PGConfig,
-    _init_env_rank_and_launch_test,
+    init_env_rank_and_launch_test,
     spawn_multi_process,
 )
 from fairchem.core.scripts.make_lmdb_sizes import get_lmdb_sizes_parser, make_lmdb_sizes
@@ -121,12 +121,12 @@ def _run_main(
 
     if world_size > 0:
         pg_config = PGConfig(
-            backend="gloo", world_size=2, gp_group_size=1, use_gp=False
+            backend="gloo", world_size=world_size, gp_group_size=1, use_gp=False
         )
         spawn_multi_process(
             pg_config,
             Runner(distributed=True),
-            _init_env_rank_and_launch_test,
+            init_env_rank_and_launch_test,
             config,
         )
     else:

tests/core/models/test_equiformer_v2.py

@@ -18,7 +18,11 @@
 from torch.nn.parallel.distributed import DistributedDataParallel
 
 from fairchem.core.common.registry import registry
-from fairchem.core.common.test_utils import PGConfig, spawn_multi_process
+from fairchem.core.common.test_utils import (
+    PGConfig,
+    init_pg_and_rank_and_launch_test,
+    spawn_multi_process,
+)
 from fairchem.core.common.utils import load_state_dict, setup_imports
 from fairchem.core.datasets import data_list_collater
 from fairchem.core.models.equiformer_v2.so3 import (
@@ -140,7 +144,9 @@ def test_energy_force_shape(self, snapshot):
     def test_ddp(self, snapshot):
         data_dist = self.data.clone().detach()
         config = PGConfig(backend="gloo", world_size=1, gp_group_size=1, use_gp=False)
-        output = spawn_multi_process(config, _runner, data_dist)
+        output = spawn_multi_process(
+            config, _runner, init_pg_and_rank_and_launch_test, data_dist
+        )
         assert len(output) == 1
         energy, forces = output[0]["energy"], output[0]["forces"]
         assert snapshot == energy.shape
@@ -151,7 +157,9 @@ def test_ddp(self, snapshot):
     def test_gp(self, snapshot):
         data_dist = self.data.clone().detach()
         config = PGConfig(backend="gloo", world_size=2, gp_group_size=2, use_gp=True)
-        output = spawn_multi_process(config, _runner, data_dist)
+        output = spawn_multi_process(
+            config, _runner, init_pg_and_rank_and_launch_test, data_dist
+        )
         assert len(output) == 2
         energy, forces = output[0]["energy"], output[0]["forces"]
         assert snapshot == energy.shape