Move SLURM to separate module (#528)

* Move SLURM to separate module

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* clean up

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

executorlib/interactive/executor.py

@@ -7,6 +7,7 @@
     InteractiveStepExecutor,
     execute_tasks_with_dependencies,
 )
+from executorlib.interactive.slurm import SrunSpawner
 from executorlib.standalone.inputcheck import (
     check_command_line_argument_lst,
     check_executor,
@@ -17,10 +18,7 @@
     check_pmi,
     validate_number_of_cores,
 )
-from executorlib.standalone.interactive.spawner import (
-    MpiExecSpawner,
-    SrunSpawner,
-)
+from executorlib.standalone.interactive.spawner import MpiExecSpawner
 from executorlib.standalone.plot import (
     draw,
     generate_nodes_and_edges,

executorlib/interactive/slurm.py

@@ -0,0 +1,94 @@
+from typing import Optional
+
+from executorlib.standalone.interactive.spawner import SubprocessSpawner
+
+SLURM_COMMAND = "srun"
+
+
+class SrunSpawner(SubprocessSpawner):
+    def __init__(
+        self,
+        cwd: Optional[str] = None,
+        cores: int = 1,
+        threads_per_core: int = 1,
+        gpus_per_core: int = 0,
+        openmpi_oversubscribe: bool = False,
+        slurm_cmd_args: list[str] = [],
+    ):
+        """
+        Srun interface implementation.
+
+        Args:
+            cwd (str, optional): The current working directory. Defaults to None.
+            cores (int, optional): The number of cores to use. Defaults to 1.
+            threads_per_core (int, optional): The number of threads per core. Defaults to 1.
+            gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
+            openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
+            slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
+        """
+        super().__init__(
+            cwd=cwd,
+            cores=cores,
+            openmpi_oversubscribe=openmpi_oversubscribe,
+            threads_per_core=threads_per_core,
+        )
+        self._gpus_per_core = gpus_per_core
+        self._slurm_cmd_args = slurm_cmd_args
+
+    def generate_command(self, command_lst: list[str]) -> list[str]:
+        """
+        Generate the command list for the Srun interface.
+
+        Args:
+            command_lst (list[str]): The command list.
+
+        Returns:
+            list[str]: The generated command list.
+        """
+        command_prepend_lst = generate_slurm_command(
+            cores=self._cores,
+            cwd=self._cwd,
+            threads_per_core=self._threads_per_core,
+            gpus_per_core=self._gpus_per_core,
+            openmpi_oversubscribe=self._openmpi_oversubscribe,
+            slurm_cmd_args=self._slurm_cmd_args,
+        )
+        return super().generate_command(
+            command_lst=command_prepend_lst + command_lst,
+        )
+
+
+def generate_slurm_command(
+    cores: int,
+    cwd: str,
+    threads_per_core: int = 1,
+    gpus_per_core: int = 0,
+    openmpi_oversubscribe: bool = False,
+    slurm_cmd_args: list[str] = [],
+) -> list[str]:
+    """
+    Generate the command list for the SLURM interface.
+
+    Args:
+        cores (int): The number of cores.
+        cwd (str): The current working directory.
+        threads_per_core (int, optional): The number of threads per core. Defaults to 1.
+        gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
+        openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
+        slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
+
+    Returns:
+        list[str]: The generated command list.
+    """
+    command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
+    if cwd is not None:
+        command_prepend_lst += ["-D", cwd]
+    if threads_per_core > 1:
+        command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
+    if gpus_per_core > 0:
+        command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
+    if openmpi_oversubscribe:
+        command_prepend_lst += ["--oversubscribe"]
+    if len(slurm_cmd_args) > 0:
+        command_prepend_lst += slurm_cmd_args
+    return command_prepend_lst

executorlib/standalone/__init__.py

@@ -6,7 +6,7 @@
     interface_send,
     interface_shutdown,
 )
-from executorlib.standalone.interactive.spawner import MpiExecSpawner, SrunSpawner
+from executorlib.standalone.interactive.spawner import MpiExecSpawner
 from executorlib.standalone.thread import RaisingThread
 
 __all__ = [
@@ -18,5 +18,4 @@
     interface_receive,
     RaisingThread,
     MpiExecSpawner,
-    SrunSpawner,
 ]

executorlib/standalone/interactive/spawner.py

@@ -3,7 +3,6 @@
 from typing import Optional
 
 MPI_COMMAND = "mpiexec"
-SLURM_COMMAND = "srun"
 
 
 class BaseSpawner(ABC):
@@ -147,59 +146,6 @@ def generate_command(self, command_lst: list[str]) -> list[str]:
         )
 
 
-class SrunSpawner(SubprocessSpawner):
-    def __init__(
-        self,
-        cwd: Optional[str] = None,
-        cores: int = 1,
-        threads_per_core: int = 1,
-        gpus_per_core: int = 0,
-        openmpi_oversubscribe: bool = False,
-        slurm_cmd_args: list[str] = [],
-    ):
-        """
-        Srun interface implementation.
-
-        Args:
-            cwd (str, optional): The current working directory. Defaults to None.
-            cores (int, optional): The number of cores to use. Defaults to 1.
-            threads_per_core (int, optional): The number of threads per core. Defaults to 1.
-            gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
-            openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
-            slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
-        """
-        super().__init__(
-            cwd=cwd,
-            cores=cores,
-            openmpi_oversubscribe=openmpi_oversubscribe,
-            threads_per_core=threads_per_core,
-        )
-        self._gpus_per_core = gpus_per_core
-        self._slurm_cmd_args = slurm_cmd_args
-
-    def generate_command(self, command_lst: list[str]) -> list[str]:
-        """
-        Generate the command list for the Srun interface.
-
-        Args:
-            command_lst (list[str]): The command list.
-
-        Returns:
-            list[str]: The generated command list.
-        """
-        command_prepend_lst = generate_slurm_command(
-            cores=self._cores,
-            cwd=self._cwd,
-            threads_per_core=self._threads_per_core,
-            gpus_per_core=self._gpus_per_core,
-            openmpi_oversubscribe=self._openmpi_oversubscribe,
-            slurm_cmd_args=self._slurm_cmd_args,
-        )
-        return super().generate_command(
-            command_lst=command_prepend_lst + command_lst,
-        )
-
-
 def generate_mpiexec_command(
     cores: int, openmpi_oversubscribe: bool = False
 ) -> list[str]:
@@ -220,39 +166,3 @@ def generate_mpiexec_command(
         if openmpi_oversubscribe:
             command_prepend_lst += ["--oversubscribe"]
         return command_prepend_lst
-
-
-def generate_slurm_command(
-    cores: int,
-    cwd: str,
-    threads_per_core: int = 1,
-    gpus_per_core: int = 0,
-    openmpi_oversubscribe: bool = False,
-    slurm_cmd_args: list[str] = [],
-) -> list[str]:
-    """
-    Generate the command list for the SLURM interface.
-
-    Args:
-        cores (int): The number of cores.
-        cwd (str): The current working directory.
-        threads_per_core (int, optional): The number of threads per core. Defaults to 1.
-        gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
-        openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
-        slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
-
-    Returns:
-        list[str]: The generated command list.
-    """
-    command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
-    if cwd is not None:
-        command_prepend_lst += ["-D", cwd]
-    if threads_per_core > 1:
-        command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
-    if gpus_per_core > 0:
-        command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
-    if openmpi_oversubscribe:
-        command_prepend_lst += ["--oversubscribe"]
-    if len(slurm_cmd_args) > 0:
-        command_prepend_lst += slurm_cmd_args
-    return command_prepend_lst

tests/test_shared_backend.py

@@ -3,7 +3,8 @@
 import unittest
 
 from executorlib.standalone.interactive.backend import parse_arguments
-from executorlib.standalone.interactive.spawner import SrunSpawner, MpiExecSpawner
+from executorlib.standalone.interactive.spawner import MpiExecSpawner
+from executorlib.interactive.slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):