Rename flux_executor_pmi_mode to pmi_mode (#762)

* Debug slurm pmi options

* Update pipeline.yml

* disable MPI parallel test

* sinfo

* update output

* add mpi parallel test

* core validation is only possible on compute node not on login node

* fix test

* enforce pmix for testing

* Update slurmspawner.py

* downgrade to openmpi

* remove pmi setting

* use slurm environment

* Try pmix again

* Update slurmspawner.py

* Update slurmspawner.py

* Update pipeline.yml

* Update pipeline.yml

* Update slurmspawner.py

* Update pipeline.yml

* use slrum args

* extend test

* fix tests

* Add executor_pmi_mode option

* check all files

* fixes

* extend tests

* rename to pmi_mode

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

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

* fixes

---------

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

Author: jan-janssen

.github/workflows/pipeline.yml

@@ -280,6 +280,8 @@ jobs:
     - uses: actions/checkout@v4
     - uses: koesterlab/setup-slurm-action@v1
       timeout-minutes: 5
+    - name: ubnuntu install
+      run: sudo apt install -y mpich
     - name: Conda config
       shell: bash -l {0}
       run: echo -e "channels:\n  - conda-forge\n" > .condarc
@@ -295,8 +297,10 @@ jobs:
       run: |
         pip install . --no-deps --no-build-isolation
         cd tests
-        python -m unittest test_slurmclusterexecutor.py
+        sinfo -o "%n %e %m %a %c %C"
+        srun --mpi=list
         python -m unittest test_slurmjobexecutor.py
+        python -m unittest test_slurmclusterexecutor.py
 
   unittest_mpich:
     needs: [black]

docs/installation.md

@@ -120,7 +120,7 @@ For the version 5 of openmpi the backend changed to `pmix`, this requires the ad
 ```
 conda install -c conda-forge flux-core flux-sched flux-pmix openmpi>=5 executorlib
 ```
-In addition, the `flux_executor_pmi_mode="pmix"` parameter has to be set for the `FluxJobExecutor` or the 
+In addition, the `pmi_mode="pmix"` parameter has to be set for the `FluxJobExecutor` or the 
 `FluxClusterExecutor` to switch to `pmix` as backend.
 
 ### Test Flux Framework

executorlib/executor/flux.py

@@ -43,8 +43,8 @@ class FluxJobExecutor(BaseExecutor):
                                                   compute notes. Defaults to False.
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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.
         flux_log_files (bool, optional): Write flux stdout and stderr files. Defaults to False.
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
@@ -93,8 +93,8 @@ def __init__(
         cache_directory: Optional[str] = None,
         max_cores: Optional[int] = None,
         resource_dict: Optional[dict] = None,
+        pmi_mode: Optional[str] = None,
         flux_executor=None,
-        flux_executor_pmi_mode: Optional[str] = None,
         flux_executor_nesting: bool = False,
         flux_log_files: bool = False,
         hostname_localhost: Optional[bool] = None,
@@ -130,8 +130,8 @@ def __init__(
                                                       compute notes. Defaults to False.
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
+            pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
             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.
             flux_log_files (bool, optional): Write flux stdout and stderr files. Defaults to False.
             hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
@@ -175,8 +175,8 @@ def __init__(
                         cache_directory=cache_directory,
                         max_cores=max_cores,
                         resource_dict=resource_dict,
+                        pmi_mode=pmi_mode,
                         flux_executor=flux_executor,
-                        flux_executor_pmi_mode=flux_executor_pmi_mode,
                         flux_executor_nesting=flux_executor_nesting,
                         flux_log_files=flux_log_files,
                         hostname_localhost=hostname_localhost,
@@ -199,8 +199,8 @@ def __init__(
                     cache_directory=cache_directory,
                     max_cores=max_cores,
                     resource_dict=resource_dict,
+                    pmi_mode=pmi_mode,
                     flux_executor=flux_executor,
-                    flux_executor_pmi_mode=flux_executor_pmi_mode,
                     flux_executor_nesting=flux_executor_nesting,
                     flux_log_files=flux_log_files,
                     hostname_localhost=hostname_localhost,
@@ -236,7 +236,7 @@ class FluxClusterExecutor(BaseExecutor):
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
         pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
-        flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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
@@ -284,7 +284,7 @@ def __init__(
         max_cores: Optional[int] = None,
         resource_dict: Optional[dict] = None,
         pysqa_config_directory: Optional[str] = None,
-        flux_executor_pmi_mode: Optional[str] = None,
+        pmi_mode: Optional[str] = None,
         hostname_localhost: Optional[bool] = None,
         block_allocation: bool = False,
         init_function: Optional[Callable] = None,
@@ -319,7 +319,7 @@ def __init__(
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
             pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
-            flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
+            pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
             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
@@ -369,7 +369,7 @@ def __init__(
                     cache_directory=cache_directory,
                     resource_dict=resource_dict,
                     flux_executor=None,
-                    flux_executor_pmi_mode=flux_executor_pmi_mode,
+                    pmi_mode=pmi_mode,
                     flux_executor_nesting=False,
                     flux_log_files=False,
                     pysqa_config_directory=pysqa_config_directory,
@@ -387,8 +387,8 @@ def __init__(
                         cache_directory=cache_directory,
                         max_cores=max_cores,
                         resource_dict=resource_dict,
+                        pmi_mode=None,
                         flux_executor=None,
-                        flux_executor_pmi_mode=None,
                         flux_executor_nesting=False,
                         flux_log_files=False,
                         hostname_localhost=hostname_localhost,
@@ -408,8 +408,8 @@ def create_flux_executor(
     max_cores: Optional[int] = None,
     cache_directory: Optional[str] = None,
     resource_dict: Optional[dict] = None,
+    pmi_mode: Optional[str] = None,
     flux_executor=None,
-    flux_executor_pmi_mode: Optional[str] = None,
     flux_executor_nesting: bool = False,
     flux_log_files: bool = False,
     hostname_localhost: Optional[bool] = None,
@@ -437,8 +437,8 @@ def create_flux_executor(
                                                   compute notes. Defaults to False.
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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.
         flux_log_files (bool, optional): Write flux stdout and stderr files. Defaults to False.
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
@@ -470,7 +470,7 @@ def create_flux_executor(
     resource_dict["hostname_localhost"] = hostname_localhost
     resource_dict["log_obj_size"] = log_obj_size
     check_init_function(block_allocation=block_allocation, init_function=init_function)
-    check_pmi(backend="flux_allocation", pmi=flux_executor_pmi_mode)
+    check_pmi(backend="flux_allocation", pmi=pmi_mode)
     check_oversubscribe(oversubscribe=resource_dict.get("openmpi_oversubscribe", False))
     check_command_line_argument_lst(
         command_line_argument_lst=resource_dict.get("slurm_cmd_args", [])
@@ -479,8 +479,8 @@ def create_flux_executor(
         del resource_dict["openmpi_oversubscribe"]
     if "slurm_cmd_args" in resource_dict:
         del resource_dict["slurm_cmd_args"]
+    resource_dict["pmi_mode"] = pmi_mode
     resource_dict["flux_executor"] = flux_executor
-    resource_dict["flux_executor_pmi_mode"] = flux_executor_pmi_mode
     resource_dict["flux_executor_nesting"] = flux_executor_nesting
     resource_dict["flux_log_files"] = flux_log_files
     if block_allocation:

executorlib/executor/single.py

@@ -329,7 +329,7 @@ def __init__(
                     cache_directory=cache_directory,
                     resource_dict=resource_dict,
                     flux_executor=None,
-                    flux_executor_pmi_mode=None,
+                    pmi_mode=None,
                     flux_executor_nesting=False,
                     flux_log_files=False,
                     pysqa_config_directory=None,

executorlib/executor/slurm.py

@@ -44,6 +44,7 @@ class SlurmClusterExecutor(BaseExecutor):
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
         pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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
@@ -91,6 +92,7 @@ def __init__(
         max_cores: Optional[int] = None,
         resource_dict: Optional[dict] = None,
         pysqa_config_directory: Optional[str] = None,
+        pmi_mode: Optional[str] = None,
         hostname_localhost: Optional[bool] = None,
         block_allocation: bool = False,
         init_function: Optional[Callable] = None,
@@ -125,6 +127,7 @@ def __init__(
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
             pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
+            pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
             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
@@ -173,8 +176,8 @@ def __init__(
                     max_cores=max_cores,
                     cache_directory=cache_directory,
                     resource_dict=resource_dict,
+                    pmi_mode=pmi_mode,
                     flux_executor=None,
-                    flux_executor_pmi_mode=None,
                     flux_executor_nesting=False,
                     flux_log_files=False,
                     pysqa_config_directory=pysqa_config_directory,
@@ -232,6 +235,7 @@ class SlurmJobExecutor(BaseExecutor):
                                                   compute notes. Defaults to False.
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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
@@ -278,6 +282,7 @@ def __init__(
         cache_directory: Optional[str] = None,
         max_cores: Optional[int] = None,
         resource_dict: Optional[dict] = None,
+        pmi_mode: Optional[str] = None,
         hostname_localhost: Optional[bool] = None,
         block_allocation: bool = False,
         init_function: Optional[Callable] = None,
@@ -315,6 +320,7 @@ def __init__(
                                                       compute notes. Defaults to False.
                                   - error_log_file (str): Name of the error log file to use for storing exceptions
                                                           raised by the Python functions submitted to the Executor.
+            pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
             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
@@ -356,6 +362,7 @@ def __init__(
                         cache_directory=cache_directory,
                         max_cores=max_cores,
                         resource_dict=resource_dict,
+                        pmi_mode=pmi_mode,
                         hostname_localhost=hostname_localhost,
                         block_allocation=block_allocation,
                         init_function=init_function,
@@ -376,6 +383,7 @@ def __init__(
                     cache_directory=cache_directory,
                     max_cores=max_cores,
                     resource_dict=resource_dict,
+                    pmi_mode=pmi_mode,
                     hostname_localhost=hostname_localhost,
                     block_allocation=block_allocation,
                     init_function=init_function,
@@ -389,6 +397,7 @@ def create_slurm_executor(
     max_cores: Optional[int] = None,
     cache_directory: Optional[str] = None,
     resource_dict: Optional[dict] = None,
+    pmi_mode: Optional[str] = None,
     hostname_localhost: Optional[bool] = None,
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
@@ -418,6 +427,7 @@ def create_slurm_executor(
                                                   compute notes. Defaults to False.
                               - error_log_file (str): Name of the error log file to use for storing exceptions raised
                                                       by the Python functions submitted to the Executor.
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None
         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
@@ -441,6 +451,7 @@ def create_slurm_executor(
     resource_dict["cache_directory"] = cache_directory
     resource_dict["hostname_localhost"] = hostname_localhost
     resource_dict["log_obj_size"] = log_obj_size
+    resource_dict["pmi_mode"] = pmi_mode
     check_init_function(block_allocation=block_allocation, init_function=init_function)
     if block_allocation:
         resource_dict["init_function"] = init_function

executorlib/standalone/command.py

@@ -21,7 +21,7 @@ def get_cache_execute_command(
     file_name: str,
     cores: int = 1,
     backend: Optional[str] = None,
-    flux_executor_pmi_mode: Optional[str] = None,
+    pmi_mode: Optional[str] = None,
 ) -> list:
     """
     Get command to call backend as a list of two strings
@@ -30,7 +30,7 @@ def get_cache_execute_command(
         file_name (str): The name of the file.
         cores (int, optional): Number of cores used to execute the task. Defaults to 1.
         backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"].
-        flux_executor_pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
+        pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only)
 
     Returns:
         list[str]: List of strings containing the python executable path and the backend script to execute
@@ -44,25 +44,21 @@ def get_cache_execute_command(
                 + [get_command_path(executable="cache_parallel.py"), file_name]
             )
         elif backend == "slurm":
+            command_prepend = ["srun", "-n", str(cores)]
+            if pmi_mode is not None:
+                command_prepend += ["--mpi=" + pmi_mode]
             command_lst = (
-                ["srun", "-n", str(cores)]
+                command_prepend
                 + command_lst
                 + [get_command_path(executable="cache_parallel.py"), file_name]
             )
         elif backend == "flux":
-            if flux_executor_pmi_mode is not None:
-                flux_command = [
-                    "flux",
-                    "run",
-                    "-o",
-                    "pmi=" + flux_executor_pmi_mode,
-                    "-n",
-                    str(cores),
-                ]
-            else:
-                flux_command = ["flux", "run", "-n", str(cores)]
+            flux_command = ["flux", "run"]
+            if pmi_mode is not None:
+                flux_command += ["-o", "pmi=" + pmi_mode]
             command_lst = (
                 flux_command
+                + ["-n", str(cores)]
                 + command_lst
                 + [get_command_path(executable="cache_parallel.py"), file_name]
             )