Refactoring: Definition of task spawner

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):

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/task_scheduler/file/shared.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/task_scheduler.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

executorlib/task_scheduler/interactive/slurmspawner.py → executorlib/task_scheduler/interactive/spawner_slurm.py

@@ -1,8 +1,8 @@
 import os
 from typing import Optional
 
+from executorlib.standalone.command import generate_slurm_command
 from executorlib.standalone.interactive.spawner import SubprocessSpawner
-from executorlib.standalone.slurm_command import generate_slurm_command
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):

tests/test_cache_fileexecutor_mpi.py

@@ -5,7 +5,7 @@
 
 try:
  from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
- from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
+ from executorlib.task_scheduler.file.spawner_subprocess import execute_in_subprocess
 
  skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -7,7 +7,7 @@
 from time import sleep
 
 try:
- from executorlib.task_scheduler.file.subprocess_spawner import (
+ from executorlib.task_scheduler.file.spawner_subprocess import (
  execute_in_subprocess,
  terminate_subprocess,
  )

tests/test_fluxclusterexecutor.py

@@ -12,7 +12,7 @@
  import flux.job
  from executorlib import terminate_tasks_in_cache
  from executorlib.standalone.hdf import dump
- from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+ from executorlib.task_scheduler.file.spawner_pysqa import execute_with_pysqa
  from executorlib.standalone.scheduler import terminate_with_pysqa
 
  skip_flux_test = "FLUX_URI" not in os.environ

tests/test_fluxjobexecutor.py

@@ -8,7 +8,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_fluxjobexecutor_plot.py

@@ -10,7 +10,7 @@
 try:
  import pygraphviz
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_graphviz_flux_test = "FLUX_URI" not in os.environ
 except ImportError:

tests/test_fluxpythonspawner.py

@@ -12,7 +12,7 @@
 
 try:
  import flux.job
- from executorlib.task_scheduler.interactive.fluxspawner import FluxPythonSpawner
+ from executorlib.task_scheduler.interactive.spawner_flux import FluxPythonSpawner
 
  skip_flux_test = "FLUX_URI" not in os.environ
  pmi = os.environ.get("EXECUTORLIB_PMIX", None)

tests/test_interactive_slurmspawner.py

@@ -1,5 +1,5 @@
 import unittest
-from executorlib.standalone.slurm_command import generate_slurm_command
+from executorlib.standalone.command import generate_slurm_command
 
 try:
  from executorlib.standalone.scheduler import pysqa_execute_command

tests/test_standalone_interactive_backend.py

@@ -4,7 +4,7 @@
 
 from executorlib.standalone.interactive.backend import parse_arguments
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
-from executorlib.task_scheduler.interactive.slurmspawner import SrunSpawner
+from executorlib.task_scheduler.interactive.spawner_slurm import SrunSpawner
 
 
 class TestParser(unittest.TestCase):