fix bug

Author: rayg1234

src/fairchem/core/_cli.py

@@ -124,7 +124,7 @@ def main(
             # the hands all responsibility the user, ie they must initialize ray
             runner: Runner = hydra.utils.instantiate(cfg.runner, _recursive_=False)
             runner.run()
-        elif scheduler_cfg.ranks_per_node > 1:
+        else:
             from fairchem.core.launchers.slurm_launch import local_launch
 
             # else launch locally using torch elastic or local mode

src/fairchem/core/launchers/slurm_launch.py

@@ -269,25 +269,22 @@ def local_launch(cfg: DictConfig, log_dir: str):
     Launch locally with torch elastic (for >1 workers) or just single process
     """
     scheduler_cfg = cfg.job.scheduler
-    from torch.distributed.launcher.api import LaunchConfig, elastic_launch
-
-    from fairchem.core.launchers import slurm_launch
-
-    launch_config = LaunchConfig(
-        min_nodes=1,
-        max_nodes=1,
-        nproc_per_node=scheduler_cfg.ranks_per_node,
-        rdzv_backend="c10d",
-        max_restarts=0,
-    )
-    elastic_launch(launch_config, slurm_launch.runner_wrapper)(cfg)
-    if "reducer" in cfg:
-        elastic_launch(launch_config, slurm_launch.runner_wrapper)(cfg, RunType.REDUCE)
+    if scheduler_cfg.ranks_per_node > 1:
+        from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+
+        launch_config = LaunchConfig(
+            min_nodes=1,
+            max_nodes=1,
+            nproc_per_node=scheduler_cfg.ranks_per_node,
+            rdzv_backend="c10d",
+            max_restarts=0,
+        )
+        elastic_launch(launch_config, runner_wrapper)(cfg)
+        if "reducer" in cfg:
+            elastic_launch(launch_config, runner_wrapper)(cfg, RunType.REDUCE)
     else:
         logging.info("Running in local mode without elastic launch")
-        from fairchem.core.launchers import slurm_launch
-
         distutils.setup_env_local()
-        slurm_launch.runner_wrapper(cfg)
+        runner_wrapper(cfg)
         if "reducer" in cfg:
-            slurm_launch.runner_wrapper(cfg, RunType.REDUCE)
+            runner_wrapper(cfg, RunType.REDUCE)

tests/core/conftest.py

@@ -6,21 +6,10 @@
 """
 
 from __future__ import annotations
-from fairchem.core.units.mlip_unit.mlip_unit import (
-    UNIT_INFERENCE_CHECKPOINT,
-    UNIT_RESUME_CONFIG,
-)
-
 
-from tests.core.units.mlip_unit.create_fake_dataset import (
-    create_fake_uma_dataset,
-)
 import os
 import tempfile
-
-from tests.core.testing_utils import launch_main
 from itertools import product
-import logging
 from random import choice
 from typing import TYPE_CHECKING
 
@@ -33,6 +22,14 @@
 from syrupy.extensions.amber import AmberSnapshotExtension
 
 from fairchem.core.datasets import AseDBDataset
+from fairchem.core.units.mlip_unit.mlip_unit import (
+    UNIT_INFERENCE_CHECKPOINT,
+    UNIT_RESUME_CONFIG,
+)
+from tests.core.testing_utils import launch_main
+from tests.core.units.mlip_unit.create_fake_dataset import (
+    create_fake_uma_dataset,
+)
 
 if TYPE_CHECKING:
     from syrupy.types import SerializableData
@@ -222,7 +219,7 @@ def dummy_binary_dataset(dummy_binary_dataset_path):
 def run_around_tests():
     # If debugging GPU memory issues, uncomment this print statement
     # to get full GPU memory allocations before each test runs
-    #print(torch.cuda.memory_summary())
+    # print(torch.cuda.memory_summary())
     yield
     torch.cuda.empty_cache()
 
@@ -343,7 +340,6 @@ def conserving_mole_checkpoint(fake_uma_dataset):
     return inference_checkpoint_pt, checkpoint_state_yaml
 
 
-
 @pytest.fixture(scope="session")
 def fake_uma_dataset():
     with tempfile.TemporaryDirectory() as tempdirname: