Merge pull request #120 from al-rigazzi/auto_launcher

Automatic Launcher Detection (#111)

Add automatic detection of available launcher.
This allows the creation of totally portable base
RunSettings.

[ committed by @al-rigazzi ]
[ reviewed by @Spartee ]

smartsim/experiment.py

@@ -39,6 +39,7 @@
 from .settings import settings
 from .utils import get_logger
 from .utils.helpers import colorize, init_default
+from .wlm import detect_launcher
 
 logger = get_logger(__name__)
 
@@ -65,7 +66,9 @@ def __init__(self, name, exp_path=None, launcher="local"):
         :param exp_path: path to location of ``Experiment`` directory if generated
         :type exp_path: str, optional
         :param launcher: type of launcher being used, options are "slurm", "pbs",
-                         "cobalt", "lsf", or "local". Defaults to "local"
+                         "cobalt", "lsf", or "local". If set to "auto",
+                         an attempt will be made to find an available launcher on the system.
+                         Defaults to "local"
         :type launcher: str, optional
         """
         self.name = name
@@ -76,6 +79,10 @@ def __init__(self, name, exp_path=None, launcher="local"):
                 raise NotADirectoryError("Experiment path provided does not exist")
             exp_path = osp.abspath(exp_path)
         self.exp_path = init_default(osp.join(getcwd(), name), exp_path, str)
+
+        if launcher == "auto":
+            launcher = detect_launcher()
+
         self._control = Controller(launcher=launcher)
         self._launcher = launcher.lower()
 

smartsim/launcher/util/shell.py

@@ -1,5 +1,5 @@
 import time
-from subprocess import PIPE, Popen, TimeoutExpired
+from subprocess import PIPE, TimeoutExpired
 
 import psutil
 

smartsim/settings/settings.py

@@ -26,6 +26,7 @@
 
 from ..error import SmartSimError
 from ..utils.helpers import is_valid_cmd
+from ..wlm import detect_launcher
 from . import *
 
 
@@ -36,7 +37,8 @@ def create_batch_settings(
 
     See Experiment.create_batch_settings for details
 
-    :param launcher: launcher for this experiment
+    :param launcher: launcher for this experiment, if set to 'auto', 
+                     an attempt will be made to find an available launcher on the system
     :type launcher: str
     :param nodes: number of nodes for batch job, defaults to 1
     :type nodes: int, optional
@@ -60,6 +62,9 @@ def create_batch_settings(
         "lsf": BsubBatchSettings,
     }
 
+    if launcher == "auto":
+        launcher = detect_launcher()
+
     if launcher == "local":
         raise SmartSimError("Local launcher does not support batch workloads")
 
@@ -96,7 +101,8 @@ def create_run_settings(
 
     See Experiment.create_run_settings docstring for more details
 
-    :param launcher: launcher to create settings for
+    :param launcher: launcher to create settings for, if set to 'auto', 
+                     an attempt will be made to find an available launcher on the system
     :type launcher: str
     :param run_command: command to run the executable
     :type run_command: str
@@ -129,6 +135,9 @@ def create_run_settings(
         "lsf": ["jsrun", "mpirun"],
     }
 
+    if launcher == "auto":
+        launcher = detect_launcher()
+
     def _detect_command(launcher):
         if launcher in by_launcher:
             for cmd in by_launcher[launcher]:

smartsim/utils/helpers.py

@@ -73,7 +73,7 @@ def init_default(default, init_value, expected_type=None):
 def expand_exe_path(exe):
     """Takes an executable and returns the full path to that executable
 
-    :param exe: exectable or file
+    :param exe: executable or file
     :type exe: str
     """
 

smartsim/wlm.py

@@ -0,0 +1,34 @@
+from shutil import which
+from subprocess import run
+
+
+def detect_launcher():
+    """Detect available launcher."""
+    # Precedence: PBS, Cobalt, LSF, Slurm, local
+    if which("qsub") and which("qstat") and which("qdel"):
+        qsub_version = run(
+            ["qsub", "--version"], shell=False, capture_output=True, encoding="utf-8"
+        )
+        if "pbs" in (qsub_version.stdout).lower():
+            return "pbs"
+        if "cobalt" in (qsub_version.stdout).lower():
+            return "cobalt"
+    if (
+        which("bsub")
+        and which("jsrun")
+        and which("jslist")
+        and which("bjobs")
+        and which("bkill")
+    ):
+        return "lsf"
+    if (
+        which("sacct")
+        and which("srun")
+        and which("salloc")
+        and which("sbatch")
+        and which("scancel")
+        and which("sstat")
+        and which("sinfo")
+    ):
+        return "slurm"
+    return "local"

tests/test_configs/smartredis/consumer.py

@@ -5,6 +5,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+
 from smartredis import Client
 
 if __name__ == "__main__":

tests/test_configs/smartredis/producer.py

@@ -5,6 +5,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+
 from smartredis import Client
 
 

tests/test_experiment.py

@@ -123,3 +123,16 @@ def test_summary(fileutils):
     assert m.type == row["Entity-Type"]
     assert 0 == int(row["RunID"])
     assert 0 == int(row["Returncode"])
+
+
+def test_launcher_detection(wlmutils):
+    exp = Experiment("test-launcher-detection", launcher="auto")
+
+    # We check whether the right launcher is found. But if
+    # the test launcher was set to local, we tolerate finding
+    # another one (this cannot be avoided)
+    if (
+        exp._launcher != wlmutils.get_test_launcher()
+        and wlmutils.get_test_launcher() != "local"
+    ):
+        assert False

tests/test_smartredis.py

@@ -22,8 +22,9 @@
 
 shouldrun = True
 try:
-    import smartredis
     import torch
+
+    import smartredis
 except ImportError:
     shouldrun = False