[feature] Add option to specify number of nodes (#565)

* [feature] Add option to specify number of nodes

* extend test

* Add DocStrings

Author: jan-janssen

executorlib/interactive/flux.py

@@ -29,6 +29,8 @@ class FluxPythonSpawner(BaseSpawner):
         cores (int, optional): The number of cores. Defaults to 1.
         threads_per_core (int, optional): The number of threads per base. Defaults to 1.
         gpus_per_core (int, optional): The number of GPUs per base. 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. Defaults to False.
         flux_executor (flux.job.FluxExecutor, optional): The FluxExecutor instance. Defaults to None.
         flux_executor_pmi_mode (str, optional): The PMI option. Defaults to None.
@@ -42,6 +44,8 @@ def __init__(
         cores: int = 1,
         threads_per_core: int = 1,
         gpus_per_core: int = 0,
+        num_nodes: Optional[int] = None,
+        exclusive: bool = False,
         openmpi_oversubscribe: bool = False,
         flux_executor: Optional[flux.job.FluxExecutor] = None,
         flux_executor_pmi_mode: Optional[str] = None,
@@ -55,6 +59,8 @@ def __init__(
         )
         self._threads_per_core = threads_per_core
         self._gpus_per_core = gpus_per_core
+        self._num_nodes = num_nodes
+        self._exclusive = exclusive
         self._flux_executor = flux_executor
         self._flux_executor_pmi_mode = flux_executor_pmi_mode
         self._flux_executor_nesting = flux_executor_nesting
@@ -85,17 +91,17 @@ def bootup(
                 num_tasks=self._cores,
                 cores_per_task=self._threads_per_core,
                 gpus_per_task=self._gpus_per_core,
-                num_nodes=None,
-                exclusive=False,
+                num_nodes=self._num_nodes,
+                exclusive=self._exclusive,
             )
         else:
             jobspec = flux.job.JobspecV1.from_nest_command(
                 command=command_lst,
                 num_slots=self._cores,
                 cores_per_slot=self._threads_per_core,
                 gpus_per_slot=self._gpus_per_core,
-                num_nodes=None,
-                exclusive=False,
+                num_nodes=self._num_nodes,
+                exclusive=self._exclusive,
             )
         jobspec.environment = dict(os.environ)
         if self._flux_executor_pmi_mode is not None:

executorlib/interactive/slurm.py

@@ -27,6 +27,8 @@ def __init__(
         cores: int = 1,
         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,
     ):
@@ -38,6 +40,8 @@ def __init__(
             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.
+            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 [].
         """
@@ -49,6 +53,8 @@ def __init__(
         )
         self._gpus_per_core = gpus_per_core
         self._slurm_cmd_args = slurm_cmd_args
+        self._num_nodes = num_nodes
+        self._exclusive = exclusive
 
     def generate_command(self, command_lst: list[str]) -> list[str]:
         """
@@ -65,6 +71,8 @@ def generate_command(self, command_lst: list[str]) -> list[str]:
             cwd=self._cwd,
             threads_per_core=self._threads_per_core,
             gpus_per_core=self._gpus_per_core,
+            num_nodes=self._num_nodes,
+            exclusive=self._exclusive,
             openmpi_oversubscribe=self._openmpi_oversubscribe,
             slurm_cmd_args=self._slurm_cmd_args,
         )
@@ -78,6 +86,8 @@ def generate_slurm_command(
     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,
 ) -> list[str]:
@@ -89,6 +99,8 @@ def generate_slurm_command(
         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 [].
 
@@ -98,10 +110,14 @@ def generate_slurm_command(
     command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
     if cwd is not None:
         command_prepend_lst += ["-D", cwd]
+    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:

executorlib/interfaces/flux.py

@@ -43,9 +43,10 @@ class FluxJobExecutor:
                               - threads_per_core (int): number of OpenMP threads to be used for each function call
                               - gpus_per_core (int): number of GPUs per worker - defaults to 0
                               - cwd (str/None): current working directory where the parallel python task is executed
-                              - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI and
-                                                              SLURM only) - default False
-                              - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM only)
+                              - 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.
         flux_executor (flux.job.FluxExecutor): Flux Python interface to submit the workers to flux
         flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
         flux_executor_nesting (bool): Provide hierarchically nested Flux job scheduler inside the submitted function.
@@ -147,10 +148,10 @@ def __new__(
                                   - threads_per_core (int): number of OpenMP threads to be used for each function call
                                   - gpus_per_core (int): number of GPUs per worker - defaults to 0
                                   - cwd (str/None): current working directory where the parallel python task is executed
-                                  - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI
-                                                                  and SLURM only) - default False
-                                  - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM
-                                                           only)
+                                  - 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.
             flux_executor (flux.job.FluxExecutor): Flux Python interface to submit the workers to flux
             flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
             flux_executor_nesting (bool): Provide hierarchically nested Flux job scheduler inside the submitted function.
@@ -444,10 +445,10 @@ def create_flux_executor(
                               - threads_per_core (int): number of OpenMP threads to be used for each function call
                               - gpus_per_core (int): number of GPUs per worker - defaults to 0
                               - cwd (str/None): current working directory where the parallel python task is executed
-                              - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI
-                                                              and SLURM only) - default False
-                              - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM
-                                                       only)
+                              - 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.
         flux_executor (flux.job.FluxExecutor): Flux Python interface to submit the workers to flux
         flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
         flux_executor_nesting (bool): Provide hierarchically nested Flux job scheduler inside the submitted function.

executorlib/interfaces/slurm.py

@@ -227,6 +227,10 @@ class SlurmJobExecutor:
                               - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI and
                                                               SLURM only) - default False
                               - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM only)
+                              - 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.
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
                                       context of an HPC cluster this essential to be able to communicate to an
                                       Executor running on a different compute node within the same allocation. And
@@ -320,6 +324,10 @@ def __new__(
                                                                   and SLURM only) - default False
                                   - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM
                                                            only)
+                                  - 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.
             hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
                                       context of an HPC cluster this essential to be able to communicate to an
                                       Executor running on a different compute node within the same allocation. And
@@ -409,6 +417,10 @@ def create_slurm_executor(
                                                               and SLURM only) - default False
                               - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM
                                                        only)
+                              - 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.
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
                                   context of an HPC cluster this essential to be able to communicate to an
                                   Executor running on a different compute node within the same allocation. And

tests/test_pysqa_subprocess.py

@@ -51,9 +51,11 @@ def test_generate_slurm_command(self):
             cwd="/tmp/test",
             threads_per_core=2,
             gpus_per_core=1,
+            num_nodes=1,
+            exclusive=True,
             openmpi_oversubscribe=True,
             slurm_cmd_args=["--help"],
         )
-        self.assertEqual(len(command_lst), 9)
-        reply_lst = ['srun', '-n', '1', '-D', '/tmp/test', '--cpus-per-task=2', '--gpus-per-task=1', '--oversubscribe', '--help']
+        self.assertEqual(len(command_lst), 12)
+        reply_lst = ['srun', '-n', '1', '-D', '/tmp/test', '-N', '1', '--cpus-per-task=2', '--gpus-per-task=1', '--exact', '--oversubscribe', '--help']
         self.assertEqual(command_lst, reply_lst)