Refactoring: Definition of task spawner (#797)

* Refactoring: Definition of task spawner

* Format black

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

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

---------

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

Author: 3 people

executorlib/__init__.py

@@ -55,7 +55,7 @@ def terminate_tasks_in_cache(
         config_directory (str, optional): path to the config directory.
         backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"].
     """
-    from executorlib.task_scheduler.file.queue_spawner import terminate_tasks_in_cache
+    from executorlib.task_scheduler.file.spawner_pysqa import terminate_tasks_in_cache
 
     return terminate_tasks_in_cache(
         cache_directory=cache_directory,

executorlib/executor/flux.py

@@ -458,7 +458,7 @@ def create_flux_executor(
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
     """
-    from executorlib.task_scheduler.interactive.fluxspawner import (
+    from executorlib.task_scheduler.interactive.spawner_flux import (
         FluxPythonSpawner,
         validate_max_workers,
     )

executorlib/executor/single.py

@@ -314,7 +314,7 @@ def __init__(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
         if not plot_dependency_graph:
-            from executorlib.task_scheduler.file.subprocess_spawner import (
+            from executorlib.task_scheduler.file.spawner_subprocess import (
                 execute_in_subprocess,
             )
             from executorlib.task_scheduler.file.task_scheduler import (

executorlib/executor/slurm.py

@@ -13,7 +13,7 @@
 )
 from executorlib.task_scheduler.interactive.dependency import DependencyTaskScheduler
 from executorlib.task_scheduler.interactive.onetoone import OneProcessTaskScheduler
-from executorlib.task_scheduler.interactive.slurmspawner import (
+from executorlib.task_scheduler.interactive.spawner_slurm import (
     SrunSpawner,
     validate_max_workers,
 )

executorlib/standalone/command.py

@@ -3,6 +3,8 @@
 import sys
 from typing import Optional
 
+SLURM_COMMAND = "srun"
+
 
 def get_command_path(executable: str) -> str:
     """
@@ -112,3 +114,51 @@ def get_interactive_execute_command(
     else:
         command_lst += [get_command_path(executable="interactive_serial.py")]
     return command_lst
+
+
+def generate_slurm_command(
+    cores: int,
+    cwd: Optional[str],
+    threads_per_core: int = 1,
+    gpus_per_core: int = 0,
+    num_nodes: Optional[int] = None,
+    exclusive: bool = False,
+    openmpi_oversubscribe: bool = False,
+    slurm_cmd_args: Optional[list[str]] = None,
+    pmi_mode: Optional[str] = None,
+) -> 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.
+        num_nodes (int, optional): The number of compute nodes to use for executing the task. Defaults to None.
+        exclusive (bool): Whether to exclusively reserve the compute nodes, or allow sharing compute notes. Defaults to False.
+        openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
+        slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
+
+    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 pmi_mode is not None:
+        command_prepend_lst += ["--mpi=" + pmi_mode]
+    if num_nodes is not None:
+        command_prepend_lst += ["-N", str(num_nodes)]
+    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 exclusive:
+        command_prepend_lst += ["--exact"]
+    if openmpi_oversubscribe:
+        command_prepend_lst += ["--oversubscribe"]
+    if slurm_cmd_args is not None and len(slurm_cmd_args) > 0:
+        command_prepend_lst += slurm_cmd_args
+    return command_prepend_lst

executorlib/standalone/slurm_command.py

@@ -7,7 +7,7 @@
 from executorlib.standalone.command import get_cache_execute_command
 from executorlib.standalone.hdf import get_cache_files, get_output
 from executorlib.standalone.serialize import serialize_funct
-from executorlib.task_scheduler.file.subprocess_spawner import terminate_subprocess
+from executorlib.task_scheduler.file.spawner_subprocess import terminate_subprocess
 
 
 class FutureItem:

executorlib/task_scheduler/file/shared.py

@@ -11,14 +11,14 @@
 )
 from executorlib.task_scheduler.base import TaskSchedulerBase
 from executorlib.task_scheduler.file.shared import execute_tasks_h5
-from executorlib.task_scheduler.file.subprocess_spawner import (
+from executorlib.task_scheduler.file.spawner_subprocess import (
     execute_in_subprocess,
     terminate_subprocess,
 )
 
 try:
     from executorlib.standalone.scheduler import terminate_with_pysqa
-    from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+    from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
 except ImportError:
     # If pysqa is not available fall back to executing tasks in a subprocess
     execute_with_pysqa = execute_in_subprocess  # type: ignore