diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml index 753b8d45..59420739 100644 --- a/.github/workflows/pipeline.yml +++ b/.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] diff --git a/docs/installation.md b/docs/installation.md index 5eec393a..3380bcff 100644 --- a/docs/installation.md +++ b/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 diff --git a/executorlib/executor/flux.py b/executorlib/executor/flux.py index 33b45306..864548d6 100644 --- a/executorlib/executor/flux.py +++ b/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,6 +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). + 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 @@ -283,6 +284,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, @@ -317,6 +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). + 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 @@ -366,7 +369,7 @@ def __init__( cache_directory=cache_directory, resource_dict=resource_dict, flux_executor=None, - flux_executor_pmi_mode=None, + pmi_mode=pmi_mode, flux_executor_nesting=False, flux_log_files=False, pysqa_config_directory=pysqa_config_directory, @@ -384,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, @@ -405,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, @@ -434,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 @@ -467,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", []) @@ -476,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: diff --git a/executorlib/executor/single.py b/executorlib/executor/single.py index 9ad40c13..677782a6 100644 --- a/executorlib/executor/single.py +++ b/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, diff --git a/executorlib/executor/slurm.py b/executorlib/executor/slurm.py index c0353a97..3a4e202b 100644 --- a/executorlib/executor/slurm.py +++ b/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 diff --git a/executorlib/standalone/command.py b/executorlib/standalone/command.py index aa396caa..68af9abc 100644 --- a/executorlib/standalone/command.py +++ b/executorlib/standalone/command.py @@ -1,6 +1,7 @@ import importlib.util import os import sys +from typing import Optional def get_command_path(executable: str) -> str: @@ -16,24 +17,53 @@ def get_command_path(executable: str) -> str: return os.path.abspath(os.path.join(__file__, "..", "..", "backend", executable)) -def get_cache_execute_command(file_name: str, cores: int = 1) -> list: +def get_cache_execute_command( + file_name: str, + cores: int = 1, + backend: Optional[str] = None, + pmi_mode: Optional[str] = None, +) -> list: """ Get command to call backend as a list of two strings Args: 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"]. + 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 """ command_lst = [sys.executable] if cores > 1 and importlib.util.find_spec("mpi4py") is not None: - command_lst = ( - ["mpiexec", "-n", str(cores)] - + command_lst - + [get_command_path(executable="cache_parallel.py"), file_name] - ) + if backend is None: + command_lst = ( + ["mpiexec", "-n", str(cores)] + + command_lst + + [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 = ( + command_prepend + + command_lst + + [get_command_path(executable="cache_parallel.py"), file_name] + ) + elif backend == "flux": + 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] + ) + else: + raise ValueError(f"backend should be None, slurm or flux, not {backend}") elif cores > 1: raise ImportError( "mpi4py is required for parallel calculations. Please install mpi4py." diff --git a/executorlib/standalone/inputcheck.py b/executorlib/standalone/inputcheck.py index 56f39a5d..6f6ab763 100644 --- a/executorlib/standalone/inputcheck.py +++ b/executorlib/standalone/inputcheck.py @@ -146,10 +146,10 @@ def check_hostname_localhost(hostname_localhost: Optional[bool]) -> None: ) -def check_flux_executor_pmi_mode(flux_executor_pmi_mode: Optional[str]) -> None: - if flux_executor_pmi_mode is not None: +def check_pmi_mode(pmi_mode: Optional[str]) -> None: + if pmi_mode is not None: raise ValueError( - "The option to specify the flux pmi mode is not available with the pysqa based backend." + "The option to specify the pmi mode is not available on a local workstation, it requires SLURM or flux." ) diff --git a/executorlib/task_scheduler/file/shared.py b/executorlib/task_scheduler/file/shared.py index 0c5ac882..f2662e40 100644 --- a/executorlib/task_scheduler/file/shared.py +++ b/executorlib/task_scheduler/file/shared.py @@ -57,6 +57,7 @@ def execute_tasks_h5( pysqa_config_directory: Optional[str] = None, backend: Optional[str] = None, disable_dependencies: bool = False, + pmi_mode: Optional[str] = None, ) -> None: """ Execute tasks stored in a queue using HDF5 files. @@ -71,6 +72,7 @@ def execute_tasks_h5( pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend). backend (str, optional): name of the backend used to spawn tasks. disable_dependencies (boolean): Disable resolving future objects during the submission. + pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None (Flux only) Returns: None @@ -154,6 +156,8 @@ def execute_tasks_h5( command=get_cache_execute_command( file_name=file_name, cores=task_resource_dict["cores"], + backend=backend, + pmi_mode=pmi_mode, ), file_name=file_name, data_dict=data_dict, diff --git a/executorlib/task_scheduler/file/task_scheduler.py b/executorlib/task_scheduler/file/task_scheduler.py index 47bcda04..587b0e0a 100644 --- a/executorlib/task_scheduler/file/task_scheduler.py +++ b/executorlib/task_scheduler/file/task_scheduler.py @@ -3,11 +3,11 @@ from executorlib.standalone.inputcheck import ( check_executor, - check_flux_executor_pmi_mode, check_flux_log_files, check_hostname_localhost, check_max_workers_and_cores, check_nested_flux_executor, + check_pmi_mode, ) from executorlib.task_scheduler.base import TaskSchedulerBase from executorlib.task_scheduler.file.shared import execute_tasks_h5 @@ -34,6 +34,7 @@ def __init__( pysqa_config_directory: Optional[str] = None, backend: Optional[str] = None, disable_dependencies: bool = False, + pmi_mode: Optional[str] = None, ): """ Initialize the FileExecutor. @@ -48,6 +49,7 @@ def __init__( pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend). backend (str, optional): name of the backend used to spawn tasks. disable_dependencies (boolean): Disable resolving future objects during the submission. + pmi_mode (str): PMI interface to use (OpenMPI v5 requires pmix) default is None """ super().__init__(max_cores=None) default_resource_dict = { @@ -68,6 +70,7 @@ def __init__( "pysqa_config_directory": pysqa_config_directory, "backend": backend, "disable_dependencies": disable_dependencies, + "pmi_mode": pmi_mode, } self._set_process( Thread( @@ -83,8 +86,8 @@ def create_file_executor( backend: Optional[str] = None, max_cores: Optional[int] = None, cache_directory: Optional[str] = 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, pysqa_config_directory: Optional[str] = None, @@ -104,7 +107,8 @@ def create_file_executor( ) if cache_directory is not None: resource_dict["cache_directory"] = cache_directory - check_flux_executor_pmi_mode(flux_executor_pmi_mode=flux_executor_pmi_mode) + if backend is None: + check_pmi_mode(pmi_mode=pmi_mode) check_max_workers_and_cores(max_cores=max_cores, max_workers=max_workers) check_hostname_localhost(hostname_localhost=hostname_localhost) check_executor(executor=flux_executor) @@ -121,4 +125,5 @@ def create_file_executor( disable_dependencies=disable_dependencies, execute_function=execute_function, terminate_function=terminate_function, + pmi_mode=pmi_mode, ) diff --git a/executorlib/task_scheduler/interactive/fluxspawner.py b/executorlib/task_scheduler/interactive/fluxspawner.py index 9cb4ed55..848e7a8f 100644 --- a/executorlib/task_scheduler/interactive/fluxspawner.py +++ b/executorlib/task_scheduler/interactive/fluxspawner.py @@ -35,8 +35,8 @@ class FluxPythonSpawner(BaseSpawner): openmpi_oversubscribe (bool, optional): Whether to oversubscribe. Defaults to False. priority (int, optional): job urgency 0 (lowest) through 31 (highest) (default is 16). Priorities 0 through 15 are restricted to the instance owner. + pmi_mode (str, optional): The PMI option. Defaults to None. flux_executor (flux.job.FluxExecutor, optional): The FluxExecutor instance. Defaults to None. - flux_executor_pmi_mode (str, optional): The PMI option. Defaults to None. flux_executor_nesting (bool, optional): Whether to use nested FluxExecutor. Defaults to False. flux_log_files (bool, optional): Write flux stdout and stderr files. Defaults to False. """ @@ -51,8 +51,8 @@ def __init__( exclusive: bool = False, priority: Optional[int] = None, openmpi_oversubscribe: bool = False, + pmi_mode: Optional[str] = None, flux_executor: Optional[flux.job.FluxExecutor] = None, - flux_executor_pmi_mode: Optional[str] = None, flux_executor_nesting: bool = False, flux_log_files: bool = False, ): @@ -66,7 +66,7 @@ def __init__( self._num_nodes = num_nodes self._exclusive = exclusive self._flux_executor = flux_executor - self._flux_executor_pmi_mode = flux_executor_pmi_mode + self._pmi_mode = pmi_mode self._flux_executor_nesting = flux_executor_nesting self._flux_log_files = flux_log_files self._priority = priority @@ -109,8 +109,8 @@ def bootup( exclusive=self._exclusive, ) jobspec.environment = dict(os.environ) - if self._flux_executor_pmi_mode is not None: - jobspec.setattr_shell_option("pmi", self._flux_executor_pmi_mode) + if self._pmi_mode is not None: + jobspec.setattr_shell_option("pmi", self._pmi_mode) if self._cwd is not None: jobspec.cwd = self._cwd if self._flux_log_files and self._cwd is not None: diff --git a/executorlib/task_scheduler/interactive/slurmspawner.py b/executorlib/task_scheduler/interactive/slurmspawner.py index 8426012d..b6490657 100644 --- a/executorlib/task_scheduler/interactive/slurmspawner.py +++ b/executorlib/task_scheduler/interactive/slurmspawner.py @@ -7,17 +7,17 @@ def validate_max_workers(max_workers: int, cores: int, threads_per_core: int): - cores_total = int(os.environ["SLURM_NTASKS"]) * int( - os.environ["SLURM_CPUS_PER_TASK"] - ) - cores_requested = max_workers * cores * threads_per_core - if cores_total < cores_requested: - raise ValueError( - "The number of requested cores is larger than the available cores " - + str(cores_total) - + " < " - + str(cores_requested) - ) + env = os.environ + if "SLURM_NTASKS" in env and "SLURM_CPUS_PER_TASK" in env: + cores_total = int(env["SLURM_NTASKS"]) * int(env["SLURM_CPUS_PER_TASK"]) + cores_requested = max_workers * cores * threads_per_core + if cores_total < cores_requested: + raise ValueError( + "The number of requested cores is larger than the available cores " + + str(cores_total) + + " < " + + str(cores_requested) + ) class SrunSpawner(SubprocessSpawner): @@ -31,6 +31,7 @@ def __init__( exclusive: bool = False, openmpi_oversubscribe: bool = False, slurm_cmd_args: Optional[list[str]] = None, + pmi_mode: Optional[str] = None, ): """ Srun interface implementation. @@ -44,6 +45,7 @@ def __init__( 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 """ super().__init__( cwd=cwd, @@ -55,6 +57,7 @@ def __init__( self._slurm_cmd_args = slurm_cmd_args self._num_nodes = num_nodes self._exclusive = exclusive + self._pmi_mode = pmi_mode def generate_command(self, command_lst: list[str]) -> list[str]: """ @@ -75,6 +78,7 @@ def generate_command(self, command_lst: list[str]) -> list[str]: exclusive=self._exclusive, openmpi_oversubscribe=self._openmpi_oversubscribe, slurm_cmd_args=self._slurm_cmd_args, + pmi_mode=self._pmi_mode, ) return super().generate_command( command_lst=command_prepend_lst + command_lst, @@ -90,6 +94,7 @@ def generate_slurm_command( 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. @@ -103,6 +108,7 @@ def generate_slurm_command( 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. @@ -110,6 +116,8 @@ def generate_slurm_command( 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: diff --git a/notebooks/3-hpc-job.ipynb b/notebooks/3-hpc-job.ipynb index c0552c68..0c266722 100644 --- a/notebooks/3-hpc-job.ipynb +++ b/notebooks/3-hpc-job.ipynb @@ -1 +1,501 @@ -{"metadata":{"kernelspec":{"display_name":"Flux","language":"python","name":"flux"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.11"}},"nbformat_minor":5,"nbformat":4,"cells":[{"id":"87c3425d-5abe-4e0b-a948-e371808c322c","cell_type":"markdown","source":"# HPC Job Executor\nIn contrast to the [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html) which submits individual Python functions to HPC job schedulers, the HPC Job Executors take a given job allocation of the HPC job scheduler and executes Python functions with the resources available in this job allocation. In this regard it is similar to the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html) as it communicates with the individual Python processes using the [zero message queue](https://zeromq.org/), still it is more advanced as it can access the computational resources of all compute nodes of the given HPC job allocation and also provides the option to assign GPUs as accelerators for parallel execution.\n\nAvailable Functionality: \n* Submit Python functions with the [submit() function or the map() function](https://executorlib.readthedocs.io/en/latest/1-single-node.html#basic-functionality).\n* Support for parallel execution, either using the [message passing interface (MPI)](https://executorlib.readthedocs.io/en/latest/1-single-node.html#mpi-parallel-functions), [thread based parallelism](https://executorlib.readthedocs.io/en/latest/1-single-node.html#thread-parallel-functions) or by [assigning dedicated GPUs](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html#resource-assignment) to selected Python functions. All these resources assignments are handled via the [resource dictionary parameter resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary).\n* Performance optimization features, like [block allocation](https://executorlib.readthedocs.io/en/latest/1-single-node.html#block-allocation), [dependency resolution](https://executorlib.readthedocs.io/en/latest/1-single-node.html#dependencies) and [caching](https://executorlib.readthedocs.io/en/latest/1-single-node.html#cache).\n\nThe only parameter the user has to change is the `backend` parameter. ","metadata":{}},{"id":"8c788b9f-6b54-4ce0-a864-4526b7f6f170","cell_type":"markdown","source":"## SLURM\nWith the [Simple Linux Utility for Resource Management (SLURM)](https://slurm.schedmd.com/) currently being the most commonly used job scheduler, executorlib provides an interface to submit Python functions to SLURM. Internally, this is based on the [srun](https://slurm.schedmd.com/srun.html) command of the SLURM scheduler, which creates job steps in a given allocation. Given that all resource requests in SLURM are communicated via a central database a large number of submitted Python functions and resulting job steps can slow down the performance of SLURM. To address this limitation it is recommended to install the hierarchical job scheduler [flux](https://flux-framework.org/) in addition to SLURM, to use flux for distributing the resources within a given allocation. This configuration is discussed in more detail below in the section [SLURM with flux](https://executorlib.readthedocs.io/en/latest/3-hpc-job.html#slurm-with-flux).","metadata":{}},{"id":"133b751f-0925-4d11-99f0-3f8dd9360b54","cell_type":"code","source":"from executorlib import SlurmJobExecutor","metadata":{"trusted":true},"outputs":[],"execution_count":1},{"id":"9b74944e-2ccd-4cb0-860a-d876310ea870","cell_type":"markdown","source":"```python\nwith SlurmAllocationExecutor() as exe:\n future = exe.submit(sum, [1, 1])\n print(future.result())\n```","metadata":{}},{"id":"36e2d68a-f093-4082-933a-d95bfe7a60c6","cell_type":"markdown","source":"## SLURM with Flux \nAs discussed in the installation section it is important to select the [flux](https://flux-framework.org/) version compatible to the installation of a given HPC cluster. Which GPUs are available? Who manufactured these GPUs? Does the HPC use [mpich](https://www.mpich.org/) or [OpenMPI](https://www.open-mpi.org/) or one of their commercial counter parts like cray MPI or intel MPI? Depending on the configuration different installation options can be choosen, as explained in the [installation section](https://executorlib.readthedocs.io/en/latest/installation.html#hpc-job-executor).\n\nAfterwards flux can be started in an [sbatch](https://slurm.schedmd.com/sbatch.html) submission script using:\n```\nsrun flux start python \n```\nIn this Python script `` the `\"flux_allocation\"` backend can be used.","metadata":{}},{"id":"68be70c3-af18-4165-862d-7022d35bf9e4","cell_type":"markdown","source":"### Resource Assignment\nIndependent of the selected Executor [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html), [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html) or HPC job executor the assignment of the computational resources remains the same. They can either be specified in the `submit()` function by adding the resource dictionary parameter [resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary) or alternatively during the initialization of the `Executor` class by adding the resource dictionary parameter [resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary) there.\n\nThis functionality of executorlib is commonly used to rewrite individual Python functions to use MPI while the rest of the Python program remains serial.","metadata":{}},{"id":"8a2c08df-cfea-4783-ace6-68fcd8ebd330","cell_type":"code","source":"def calc_mpi(i):\n from mpi4py import MPI\n\n size = MPI.COMM_WORLD.Get_size()\n rank = MPI.COMM_WORLD.Get_rank()\n return i, size, rank","metadata":{"trusted":true},"outputs":[],"execution_count":2},{"id":"715e0c00-7b17-40bb-bd55-b0e097bfef07","cell_type":"markdown","source":"Depending on the choice of MPI version, it is recommended to specify the pmi standard which [flux](https://flux-framework.org/) should use internally for the resource assignment. For example for OpenMPI >=5 `\"pmix\"` is the recommended pmi standard.","metadata":{}},{"id":"5802c7d7-9560-4909-9d30-a915a91ac0a1","cell_type":"code","source":"from executorlib import FluxJobExecutor\n\nwith FluxJobExecutor(flux_executor_pmi_mode=\"pmix\") as exe:\n fs = exe.submit(calc_mpi, 3, resource_dict={\"cores\": 2})\n print(fs.result())","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"[(3, 2, 0), (3, 2, 1)]\n"}],"execution_count":3},{"id":"da862425-08b6-4ced-999f-89a74e85f410","cell_type":"markdown","source":"### Block Allocation\nThe block allocation for the HPC allocation mode follows the same implementation as the [block allocation for the Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#block-allocation). It starts by defining the initialization function `init_function()` which returns a dictionary which is internally used to look up input parameters for Python functions submitted to the `FluxJobExecutor` class. Commonly this functionality is used to store large data objects inside the Python process created for the block allocation, rather than reloading these Python objects for each submitted function.","metadata":{}},{"id":"cdc742c0-35f7-47ff-88c0-1b0dbeabe51b","cell_type":"code","source":"def init_function():\n return {\"j\": 4, \"k\": 3, \"l\": 2}","metadata":{"trusted":true},"outputs":[],"execution_count":4},{"id":"5ddf8343-ab2c-4469-ac9f-ee568823d4ad","cell_type":"code","source":"def calc_with_preload(i, j, k):\n return i + j + k","metadata":{"trusted":true},"outputs":[],"execution_count":5},{"id":"0da13efa-1941-416f-b9e6-bba15b5cdfa2","cell_type":"code","source":"with FluxJobExecutor(\n flux_executor_pmi_mode=\"pmix\",\n max_workers=2,\n init_function=init_function,\n block_allocation=True,\n) as exe:\n fs = exe.submit(calc_with_preload, 2, j=5)\n print(fs.result())","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"10\n"}],"execution_count":6},{"id":"82f3b947-e662-4a0d-b590-9475e0b4f7dd","cell_type":"markdown","source":"In this example the parameter `k` is used from the dataset created by the initialization function while the parameters `i` and `j` are specified by the call of the `submit()` function. \n\nWhen using the block allocation mode, it is recommended to set either the maxium number of workers using the `max_workers` parameter or the maximum number of CPU cores using the `max_cores` parameter to prevent oversubscribing the available resources. ","metadata":{}},{"id":"8ced8359-8ecb-480b-966b-b85d8446d85c","cell_type":"markdown","source":"### Dependencies\nPython functions with rather different computational resource requirements should not be merged into a single function. So to able to execute a series of Python functions which each depend on the output of the previous Python function executorlib internally handles the dependencies based on the [concurrent futures future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) objects from the Python standard library. This implementation is independent of the selected backend and works for HPC allocation mode just like explained in the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#dependencies) section.","metadata":{}},{"id":"bd26d97b-46fd-4786-9ad1-1e534b31bf36","cell_type":"code","source":"def add_funct(a, b):\n return a + b","metadata":{"trusted":true},"outputs":[],"execution_count":7},{"id":"1a2d440f-3cfc-4ff2-b74d-e21823c65f69","cell_type":"code","source":"with FluxJobExecutor(flux_executor_pmi_mode=\"pmix\") as exe:\n future = 0\n for i in range(1, 4):\n future = exe.submit(add_funct, i, future)\n print(future.result())","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"6\n"}],"execution_count":8},{"id":"f526c2bf-fdf5-463b-a955-020753138415","cell_type":"markdown","source":"### Caching\nFinally, also the caching is available for HPC allocation mode, in analogy to the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#cache). Again this functionality is not designed to identify function calls with the same parameters, but rather provides the option to reload previously cached results even after the Python processes which contained the executorlib `Executor` class is closed. As the cache is stored on the file system, this option can decrease the performance of executorlib. Consequently the caching option should primarily be used during the prototyping phase.","metadata":{}},{"id":"dcba63e0-72f5-49d1-ab04-2092fccc1c47","cell_type":"code","source":"with FluxJobExecutor(flux_executor_pmi_mode=\"pmix\", cache_directory=\"./file\") as exe:\n future_lst = [exe.submit(sum, [i, i]) for i in range(1, 4)]\n print([f.result() for f in future_lst])","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"[2, 4, 6]\n"}],"execution_count":9},{"id":"c3958a14-075b-4c10-9729-d1c559a9231c","cell_type":"code","source":"import os\nimport shutil\n\ncache_dir = \"./file\"\nif os.path.exists(cache_dir):\n print(os.listdir(cache_dir))\n try:\n shutil.rmtree(cache_dir)\n except OSError:\n pass","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"['sum89afbdf9da5eb1794f6976a3f01697c2_o.h5', 'sum0f7710227cda6456e5d07187702313f3_o.h5', 'sumf5ad27b855231a293ddd735a8554c9ea_o.h5']\n"}],"execution_count":10},{"id":"c24ca82d-60bd-4fb9-a082-bf9a81e838bf","cell_type":"markdown","source":"### Nested executors\nThe hierarchical nature of the [flux](https://flux-framework.org/) job scheduler allows the creation of additional executorlib Executors inside the functions submitted to the Executor. This hierarchy can be beneficial to separate the logic to saturate the available computational resources. ","metadata":{}},{"id":"06fb2d1f-65fc-4df6-9402-5e9837835484","cell_type":"code","source":"def calc_nested():\n from executorlib import FluxJobExecutor\n\n with FluxJobExecutor(flux_executor_pmi_mode=\"pmix\") as exe:\n fs = exe.submit(sum, [1, 1])\n return fs.result()","metadata":{"trusted":true},"outputs":[],"execution_count":11},{"id":"89b7d0fd-5978-4913-a79a-f26cc8047445","cell_type":"code","source":"with FluxJobExecutor(flux_executor_pmi_mode=\"pmix\", flux_executor_nesting=True) as exe:\n fs = exe.submit(calc_nested)\n print(fs.result())","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"2\n"}],"execution_count":12},{"id":"9f209925-1ce4-42e4-bbe5-becbb1f3cd79","cell_type":"markdown","source":"### Executor from Flux\nThe [flux framework](http://flux-framework.org/) provides its own [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) which can be used to submit shell scripts to the [flux framework](http://flux-framework.org/) for execution. The [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) returns its own representation of future objects which is incompatible with the [concurrent.futures.Future](https://docs.python.org/3/library/concurrent.futures.html) which is used by executorlib. Combining both provides the opportunity to link Python fucntions and external executables. For this purpose executorlib provides the option to use a [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) as an input for the `FluxJobExecutor`:","metadata":{}},{"id":"3df0357e-d936-4989-a271-d0b03c6d0b48","cell_type":"code","source":"from executorlib import FluxJobExecutor\nimport flux.job\n\nwith flux.job.FluxExecutor() as flux_executor:\n with FluxJobExecutor(flux_executor=flux_executor) as exe:\n future = exe.submit(sum, [1, 1])\n print(future.result())","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":"2\n"}],"execution_count":13},{"id":"34a8c690-ca5a-41d1-b38f-c67eff085750","cell_type":"markdown","source":"### Resource Monitoring\nFor debugging it is commonly helpful to keep track of the computational resources. [flux](https://flux-framework.org/) provides a number of features to analyse the resource utilization, so here only the two most commonly used ones are introduced. Starting with the option to list all the resources available in a given allocation with the `flux resource list` command:","metadata":{}},{"id":"7481eb0a-a41b-4d46-bb48-b4db299fcd86","cell_type":"code","source":"! flux resource list","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":" STATE NNODES NCORES NGPUS NODELIST\n free 1 24 0 jupyter-pyiron-executorlib-wx8wv67z\n allocated 0 0 0 \n down 0 0 0 \n"}],"execution_count":14},{"id":"08d98134-a0e0-4841-be82-e09e1af29e7f","cell_type":"markdown","source":"Followed by the list of jobs which were executed in a given flux session. This can be retrieved using the `flux jobs -a` command:","metadata":{}},{"id":"1ee6e147-f53a-4526-8ed0-fd036f2ee6bf","cell_type":"code","source":"! flux jobs -a","metadata":{"trusted":true},"outputs":[{"name":"stdout","output_type":"stream","text":" JOBID USER NAME ST NTASKS NNODES TIME INFO\n\u001b[01;32m ƒ66TjsQs jovyan python CD 1 1 0.149s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ4R3m4Sj jovyan flux CD 1 1 3.509s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ3N4Qc3y jovyan python CD 1 1 1.922s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ3DuUZ9y jovyan python CD 1 1 2.291s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ3DrWabH jovyan python CD 1 1 2.204s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ2z9sDYT jovyan python CD 1 1 0.271s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ2m9FX6w jovyan python CD 1 1 0.404s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ2dGdLJj jovyan python CD 1 1 0.346s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ29qrcvj jovyan python CD 1 1 0.848s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒ29tpbVR jovyan python CD 1 1 0.539s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m\u001b[01;32m ƒZsZ5QT jovyan python CD 2 1 0.966s jupyter-pyiron-executorlib-wx8wv67z\n\u001b[0;0m"}],"execution_count":15},{"id":"021f165b-27cc-4676-968b-cbcfd1f0210a","cell_type":"markdown","source":"## Flux\nWhile the number of HPC clusters which use [flux](https://flux-framework.org/) as primary job scheduler is currently still limited the setup and functionality provided by executorlib for running [SLURM with flux](https://executorlib.readthedocs.io/en/latest/3-hpc-job.html#slurm-with-flux) also applies to HPCs which use [flux](https://flux-framework.org/) as primary job scheduler.","metadata":{}},{"id":"04f03ebb-3f9e-4738-b9d2-5cb0db9b63c3","cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]} \ No newline at end of file +{ + "cells": [ + { + "cell_type": "markdown", + "id": "87c3425d-5abe-4e0b-a948-e371808c322c", + "metadata": {}, + "source": [ + "# HPC Job Executor\n", + "In contrast to the [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html) which submits individual Python functions to HPC job schedulers, the HPC Job Executors take a given job allocation of the HPC job scheduler and executes Python functions with the resources available in this job allocation. In this regard it is similar to the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html) as it communicates with the individual Python processes using the [zero message queue](https://zeromq.org/), still it is more advanced as it can access the computational resources of all compute nodes of the given HPC job allocation and also provides the option to assign GPUs as accelerators for parallel execution.\n", + "\n", + "Available Functionality: \n", + "* Submit Python functions with the [submit() function or the map() function](https://executorlib.readthedocs.io/en/latest/1-single-node.html#basic-functionality).\n", + "* Support for parallel execution, either using the [message passing interface (MPI)](https://executorlib.readthedocs.io/en/latest/1-single-node.html#mpi-parallel-functions), [thread based parallelism](https://executorlib.readthedocs.io/en/latest/1-single-node.html#thread-parallel-functions) or by [assigning dedicated GPUs](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html#resource-assignment) to selected Python functions. All these resources assignments are handled via the [resource dictionary parameter resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary).\n", + "* Performance optimization features, like [block allocation](https://executorlib.readthedocs.io/en/latest/1-single-node.html#block-allocation), [dependency resolution](https://executorlib.readthedocs.io/en/latest/1-single-node.html#dependencies) and [caching](https://executorlib.readthedocs.io/en/latest/1-single-node.html#cache).\n", + "\n", + "The only parameter the user has to change is the `backend` parameter. " + ] + }, + { + "cell_type": "markdown", + "id": "8c788b9f-6b54-4ce0-a864-4526b7f6f170", + "metadata": {}, + "source": [ + "## SLURM\n", + "With the [Simple Linux Utility for Resource Management (SLURM)](https://slurm.schedmd.com/) currently being the most commonly used job scheduler, executorlib provides an interface to submit Python functions to SLURM. Internally, this is based on the [srun](https://slurm.schedmd.com/srun.html) command of the SLURM scheduler, which creates job steps in a given allocation. Given that all resource requests in SLURM are communicated via a central database a large number of submitted Python functions and resulting job steps can slow down the performance of SLURM. To address this limitation it is recommended to install the hierarchical job scheduler [flux](https://flux-framework.org/) in addition to SLURM, to use flux for distributing the resources within a given allocation. This configuration is discussed in more detail below in the section [SLURM with flux](https://executorlib.readthedocs.io/en/latest/3-hpc-job.html#slurm-with-flux)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "133b751f-0925-4d11-99f0-3f8dd9360b54", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "from executorlib import SlurmJobExecutor" + ] + }, + { + "cell_type": "markdown", + "id": "9b74944e-2ccd-4cb0-860a-d876310ea870", + "metadata": {}, + "source": [ + "```python\n", + "with SlurmAllocationExecutor() as exe:\n", + " future = exe.submit(sum, [1, 1])\n", + " print(future.result())\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "36e2d68a-f093-4082-933a-d95bfe7a60c6", + "metadata": {}, + "source": [ + "## SLURM with Flux \n", + "As discussed in the installation section it is important to select the [flux](https://flux-framework.org/) version compatible to the installation of a given HPC cluster. Which GPUs are available? Who manufactured these GPUs? Does the HPC use [mpich](https://www.mpich.org/) or [OpenMPI](https://www.open-mpi.org/) or one of their commercial counter parts like cray MPI or intel MPI? Depending on the configuration different installation options can be choosen, as explained in the [installation section](https://executorlib.readthedocs.io/en/latest/installation.html#hpc-job-executor).\n", + "\n", + "Afterwards flux can be started in an [sbatch](https://slurm.schedmd.com/sbatch.html) submission script using:\n", + "```\n", + "srun flux start python \n", + "```\n", + "In this Python script `` the `\"flux_allocation\"` backend can be used." + ] + }, + { + "cell_type": "markdown", + "id": "68be70c3-af18-4165-862d-7022d35bf9e4", + "metadata": {}, + "source": [ + "### Resource Assignment\n", + "Independent of the selected Executor [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html), [HPC Cluster Executor](https://executorlib.readthedocs.io/en/latest/2-hpc-cluster.html) or HPC job executor the assignment of the computational resources remains the same. They can either be specified in the `submit()` function by adding the resource dictionary parameter [resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary) or alternatively during the initialization of the `Executor` class by adding the resource dictionary parameter [resource_dict](https://executorlib.readthedocs.io/en/latest/trouble_shooting.html#resource-dictionary) there.\n", + "\n", + "This functionality of executorlib is commonly used to rewrite individual Python functions to use MPI while the rest of the Python program remains serial." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "8a2c08df-cfea-4783-ace6-68fcd8ebd330", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "def calc_mpi(i):\n", + " from mpi4py import MPI\n", + "\n", + " size = MPI.COMM_WORLD.Get_size()\n", + " rank = MPI.COMM_WORLD.Get_rank()\n", + " return i, size, rank" + ] + }, + { + "cell_type": "markdown", + "id": "715e0c00-7b17-40bb-bd55-b0e097bfef07", + "metadata": {}, + "source": [ + "Depending on the choice of MPI version, it is recommended to specify the pmi standard which [flux](https://flux-framework.org/) should use internally for the resource assignment. For example for OpenMPI >=5 `\"pmix\"` is the recommended pmi standard." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5802c7d7-9560-4909-9d30-a915a91ac0a1", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[(3, 2, 0), (3, 2, 1)]\n" + ] + } + ], + "source": [ + "from executorlib import FluxJobExecutor\n", + "\n", + "with FluxJobExecutor(pmi_mode=\"pmix\") as exe:\n", + " fs = exe.submit(calc_mpi, 3, resource_dict={\"cores\": 2})\n", + " print(fs.result())" + ] + }, + { + "cell_type": "markdown", + "id": "da862425-08b6-4ced-999f-89a74e85f410", + "metadata": {}, + "source": [ + "### Block Allocation\n", + "The block allocation for the HPC allocation mode follows the same implementation as the [block allocation for the Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#block-allocation). It starts by defining the initialization function `init_function()` which returns a dictionary which is internally used to look up input parameters for Python functions submitted to the `FluxJobExecutor` class. Commonly this functionality is used to store large data objects inside the Python process created for the block allocation, rather than reloading these Python objects for each submitted function." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "cdc742c0-35f7-47ff-88c0-1b0dbeabe51b", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "def init_function():\n", + " return {\"j\": 4, \"k\": 3, \"l\": 2}" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "5ddf8343-ab2c-4469-ac9f-ee568823d4ad", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "def calc_with_preload(i, j, k):\n", + " return i + j + k" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0da13efa-1941-416f-b9e6-bba15b5cdfa2", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "10\n" + ] + } + ], + "source": [ + "with FluxJobExecutor(\n", + " pmi_mode=\"pmix\",\n", + " max_workers=2,\n", + " init_function=init_function,\n", + " block_allocation=True,\n", + ") as exe:\n", + " fs = exe.submit(calc_with_preload, 2, j=5)\n", + " print(fs.result())" + ] + }, + { + "cell_type": "markdown", + "id": "82f3b947-e662-4a0d-b590-9475e0b4f7dd", + "metadata": {}, + "source": [ + "In this example the parameter `k` is used from the dataset created by the initialization function while the parameters `i` and `j` are specified by the call of the `submit()` function. \n", + "\n", + "When using the block allocation mode, it is recommended to set either the maxium number of workers using the `max_workers` parameter or the maximum number of CPU cores using the `max_cores` parameter to prevent oversubscribing the available resources. " + ] + }, + { + "cell_type": "markdown", + "id": "8ced8359-8ecb-480b-966b-b85d8446d85c", + "metadata": {}, + "source": [ + "### Dependencies\n", + "Python functions with rather different computational resource requirements should not be merged into a single function. So to able to execute a series of Python functions which each depend on the output of the previous Python function executorlib internally handles the dependencies based on the [concurrent futures future](https://docs.python.org/3/library/concurrent.futures.html#future-objects) objects from the Python standard library. This implementation is independent of the selected backend and works for HPC allocation mode just like explained in the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#dependencies) section." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "bd26d97b-46fd-4786-9ad1-1e534b31bf36", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "def add_funct(a, b):\n", + " return a + b" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a2d440f-3cfc-4ff2-b74d-e21823c65f69", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "6\n" + ] + } + ], + "source": [ + "with FluxJobExecutor(pmi_mode=\"pmix\") as exe:\n", + " future = 0\n", + " for i in range(1, 4):\n", + " future = exe.submit(add_funct, i, future)\n", + " print(future.result())" + ] + }, + { + "cell_type": "markdown", + "id": "f526c2bf-fdf5-463b-a955-020753138415", + "metadata": {}, + "source": [ + "### Caching\n", + "Finally, also the caching is available for HPC allocation mode, in analogy to the [Single Node Executor](https://executorlib.readthedocs.io/en/latest/1-single-node.html#cache). Again this functionality is not designed to identify function calls with the same parameters, but rather provides the option to reload previously cached results even after the Python processes which contained the executorlib `Executor` class is closed. As the cache is stored on the file system, this option can decrease the performance of executorlib. Consequently the caching option should primarily be used during the prototyping phase." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dcba63e0-72f5-49d1-ab04-2092fccc1c47", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[2, 4, 6]\n" + ] + } + ], + "source": [ + "with FluxJobExecutor(pmi_mode=\"pmix\", cache_directory=\"./file\") as exe:\n", + " future_lst = [exe.submit(sum, [i, i]) for i in range(1, 4)]\n", + " print([f.result() for f in future_lst])" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "c3958a14-075b-4c10-9729-d1c559a9231c", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "['sum89afbdf9da5eb1794f6976a3f01697c2_o.h5', 'sum0f7710227cda6456e5d07187702313f3_o.h5', 'sumf5ad27b855231a293ddd735a8554c9ea_o.h5']\n" + ] + } + ], + "source": [ + "import os\n", + "import shutil\n", + "\n", + "cache_dir = \"./file\"\n", + "if os.path.exists(cache_dir):\n", + " print(os.listdir(cache_dir))\n", + " try:\n", + " shutil.rmtree(cache_dir)\n", + " except OSError:\n", + " pass" + ] + }, + { + "cell_type": "markdown", + "id": "c24ca82d-60bd-4fb9-a082-bf9a81e838bf", + "metadata": {}, + "source": [ + "### Nested executors\n", + "The hierarchical nature of the [flux](https://flux-framework.org/) job scheduler allows the creation of additional executorlib Executors inside the functions submitted to the Executor. This hierarchy can be beneficial to separate the logic to saturate the available computational resources. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06fb2d1f-65fc-4df6-9402-5e9837835484", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [ + "def calc_nested():\n", + " from executorlib import FluxJobExecutor\n", + "\n", + " with FluxJobExecutor(pmi_mode=\"pmix\") as exe:\n", + " fs = exe.submit(sum, [1, 1])\n", + " return fs.result()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89b7d0fd-5978-4913-a79a-f26cc8047445", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2\n" + ] + } + ], + "source": [ + "with FluxJobExecutor(pmi_mode=\"pmix\", flux_executor_nesting=True) as exe:\n", + " fs = exe.submit(calc_nested)\n", + " print(fs.result())" + ] + }, + { + "cell_type": "markdown", + "id": "9f209925-1ce4-42e4-bbe5-becbb1f3cd79", + "metadata": {}, + "source": [ + "### Executor from Flux\n", + "The [flux framework](http://flux-framework.org/) provides its own [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) which can be used to submit shell scripts to the [flux framework](http://flux-framework.org/) for execution. The [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) returns its own representation of future objects which is incompatible with the [concurrent.futures.Future](https://docs.python.org/3/library/concurrent.futures.html) which is used by executorlib. Combining both provides the opportunity to link Python fucntions and external executables. For this purpose executorlib provides the option to use a [FluxExecutor](https://flux-framework.readthedocs.io/projects/flux-core/en/latest/python/autogenerated/flux.job.executor.html#flux.job.executor.FluxExecutor) as an input for the `FluxJobExecutor`:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "3df0357e-d936-4989-a271-d0b03c6d0b48", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "2\n" + ] + } + ], + "source": [ + "from executorlib import FluxJobExecutor\n", + "import flux.job\n", + "\n", + "with flux.job.FluxExecutor() as flux_executor:\n", + " with FluxJobExecutor(flux_executor=flux_executor) as exe:\n", + " future = exe.submit(sum, [1, 1])\n", + " print(future.result())" + ] + }, + { + "cell_type": "markdown", + "id": "34a8c690-ca5a-41d1-b38f-c67eff085750", + "metadata": {}, + "source": [ + "### Resource Monitoring\n", + "For debugging it is commonly helpful to keep track of the computational resources. [flux](https://flux-framework.org/) provides a number of features to analyse the resource utilization, so here only the two most commonly used ones are introduced. Starting with the option to list all the resources available in a given allocation with the `flux resource list` command:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "7481eb0a-a41b-4d46-bb48-b4db299fcd86", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " STATE NNODES NCORES NGPUS NODELIST\n", + " free 1 24 0 jupyter-pyiron-executorlib-wx8wv67z\n", + " allocated 0 0 0 \n", + " down 0 0 0 \n" + ] + } + ], + "source": [ + "! flux resource list" + ] + }, + { + "cell_type": "markdown", + "id": "08d98134-a0e0-4841-be82-e09e1af29e7f", + "metadata": {}, + "source": [ + "Followed by the list of jobs which were executed in a given flux session. This can be retrieved using the `flux jobs -a` command:" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "1ee6e147-f53a-4526-8ed0-fd036f2ee6bf", + "metadata": { + "trusted": true + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " JOBID USER NAME ST NTASKS NNODES TIME INFO\n", + "\u001b[01;32m ƒ66TjsQs jovyan python CD 1 1 0.149s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ4R3m4Sj jovyan flux CD 1 1 3.509s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ3N4Qc3y jovyan python CD 1 1 1.922s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ3DuUZ9y jovyan python CD 1 1 2.291s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ3DrWabH jovyan python CD 1 1 2.204s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ2z9sDYT jovyan python CD 1 1 0.271s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ2m9FX6w jovyan python CD 1 1 0.404s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ2dGdLJj jovyan python CD 1 1 0.346s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ29qrcvj jovyan python CD 1 1 0.848s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒ29tpbVR jovyan python CD 1 1 0.539s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m\u001b[01;32m ƒZsZ5QT jovyan python CD 2 1 0.966s jupyter-pyiron-executorlib-wx8wv67z\n", + "\u001b[0;0m" + ] + } + ], + "source": [ + "! flux jobs -a" + ] + }, + { + "cell_type": "markdown", + "id": "021f165b-27cc-4676-968b-cbcfd1f0210a", + "metadata": {}, + "source": [ + "## Flux\n", + "While the number of HPC clusters which use [flux](https://flux-framework.org/) as primary job scheduler is currently still limited the setup and functionality provided by executorlib for running [SLURM with flux](https://executorlib.readthedocs.io/en/latest/3-hpc-job.html#slurm-with-flux) also applies to HPCs which use [flux](https://flux-framework.org/) as primary job scheduler." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "04f03ebb-3f9e-4738-b9d2-5cb0db9b63c3", + "metadata": { + "trusted": true + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Flux", + "language": "python", + "name": "flux" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/test_fluxclusterexecutor.py b/tests/test_fluxclusterexecutor.py index 27645d86..107d5add 100644 --- a/tests/test_fluxclusterexecutor.py +++ b/tests/test_fluxclusterexecutor.py @@ -41,6 +41,7 @@ def test_executor(self): resource_dict={"cores": 2, "cwd": "executorlib_cache"}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode=pmi, ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -54,6 +55,7 @@ def test_executor_no_cwd(self): resource_dict={"cores": 2}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode=pmi, ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -81,6 +83,7 @@ def test_executor_existing_files(self): resource_dict={"cores": 2, "cwd": "executorlib_cache"}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode=pmi, ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -99,6 +102,7 @@ def test_executor_existing_files(self): resource_dict={"cores": 2, "cwd": "executorlib_cache"}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode=pmi, ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) diff --git a/tests/test_fluxjobexecutor.py b/tests/test_fluxjobexecutor.py index 8cfa8e9a..03d56831 100644 --- a/tests/test_fluxjobexecutor.py +++ b/tests/test_fluxjobexecutor.py @@ -90,7 +90,7 @@ def test_flux_executor_parallel(self): resource_dict={"cores": 2}, flux_executor=self.executor, block_allocation=True, - flux_executor_pmi_mode=pmi, + pmi_mode=pmi, ) as exe: fs_1 = exe.submit(mpi_funct, 1) self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)]) @@ -102,7 +102,7 @@ def test_single_task(self): resource_dict={"cores": 2}, flux_executor=self.executor, block_allocation=True, - flux_executor_pmi_mode=pmi, + pmi_mode=pmi, ) as p: output = p.map(mpi_funct, [1, 2, 3]) self.assertEqual( diff --git a/tests/test_fluxpythonspawner.py b/tests/test_fluxpythonspawner.py index bf8eb939..01f1d160 100644 --- a/tests/test_fluxpythonspawner.py +++ b/tests/test_fluxpythonspawner.py @@ -82,7 +82,7 @@ def test_flux_executor_parallel(self): executor_kwargs={ "flux_executor": self.flux_executor, "cores": 2, - "flux_executor_pmi_mode": pmi, + "pmi_mode": pmi, }, spawner=FluxPythonSpawner, ) as exe: @@ -96,7 +96,7 @@ def test_single_task(self): executor_kwargs={ "flux_executor": self.flux_executor, "cores": 2, - "flux_executor_pmi_mode": pmi, + "pmi_mode": pmi, }, spawner=FluxPythonSpawner, ) as p: diff --git a/tests/test_slurmclusterexecutor.py b/tests/test_slurmclusterexecutor.py index 41d0f94b..4973037d 100644 --- a/tests/test_slurmclusterexecutor.py +++ b/tests/test_slurmclusterexecutor.py @@ -26,7 +26,7 @@ #SBATCH --job-name={{job_name}} #SBATCH --chdir={{working_directory}} #SBATCH --get-user-env=L -#SBATCH --cpus-per-task={{cores}} +#SBATCH --ntasks={{cores}} {{command}} """ @@ -50,6 +50,7 @@ def test_executor(self): resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode="pmi2", ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -63,6 +64,7 @@ def test_executor_no_cwd(self): resource_dict={"cores": 2, "submission_template": submission_template}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode="pmi2", ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -76,6 +78,7 @@ def test_executor_existing_files(self): resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode="pmi2", ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) @@ -94,6 +97,7 @@ def test_executor_existing_files(self): resource_dict={"cores": 2, "cwd": "executorlib_cache", "submission_template": submission_template}, block_allocation=False, cache_directory="executorlib_cache", + pmi_mode="pmi2", ) as exe: cloudpickle_register(ind=1) fs1 = exe.submit(mpi_funct, 1) diff --git a/tests/test_slurmjobexecutor.py b/tests/test_slurmjobexecutor.py index 5ef889d2..155b8b6d 100644 --- a/tests/test_slurmjobexecutor.py +++ b/tests/test_slurmjobexecutor.py @@ -14,15 +14,33 @@ def calc(i): return i +def mpi_funct(i): + from mpi4py import MPI + + size = MPI.COMM_WORLD.Get_size() + rank = MPI.COMM_WORLD.Get_rank() + return i, size, rank + + @unittest.skipIf( skip_slurm_test, "Slurm is not installed, so the Slurm tests are skipped." ) class TestSlurmBackend(unittest.TestCase): def test_slurm_executor_serial(self): - with SlurmJobExecutor() as exe: + with SlurmJobExecutor(resource_dict={"slurm_cmd_args": ["--mpi=pmi2"]}) as exe: fs_1 = exe.submit(calc, 1) fs_2 = exe.submit(calc, 2) self.assertEqual(fs_1.result(), 1) self.assertEqual(fs_2.result(), 2) self.assertTrue(fs_1.done()) self.assertTrue(fs_2.done()) + + def test_slurm_executor_parallel(self): + with SlurmJobExecutor( + max_cores=2, + resource_dict={"cores": 2, "slurm_cmd_args": ["--mpi=pmi2"]}, + block_allocation=True, + ) as exe: + fs_1 = exe.submit(mpi_funct, 1) + self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)]) + self.assertTrue(fs_1.done()) \ No newline at end of file diff --git a/tests/test_standalone_command.py b/tests/test_standalone_command.py new file mode 100644 index 00000000..d1bb55f1 --- /dev/null +++ b/tests/test_standalone_command.py @@ -0,0 +1,88 @@ +import os +import importlib.util +import sys +import unittest +from executorlib.standalone.command import get_cache_execute_command, get_interactive_execute_command + + +skip_mpi4py_test = importlib.util.find_spec("mpi4py") is None + + +class TestCommands(unittest.TestCase): + def test_get_interactive_execute_command_serial(self): + output = get_interactive_execute_command(cores=1) + self.assertEqual(output[0], sys.executable) + self.assertEqual(output[1].split(os.sep)[-1], "interactive_serial.py") + + @unittest.skipIf( + skip_mpi4py_test, + "mpi4py is not installed, so the mpi4py tests are skipped.", + ) + def test_get_interactive_execute_command_parallel(self): + output = get_interactive_execute_command(cores=2) + self.assertEqual(output[0], sys.executable) + self.assertEqual(output[1].split(os.sep)[-1], "interactive_parallel.py") + + def test_get_cache_execute_command_serial(self): + file_name = "test.txt" + output = get_cache_execute_command(cores=1, file_name=file_name) + self.assertEqual(output[0], sys.executable) + self.assertEqual(output[1].split(os.sep)[-1], "cache_serial.py") + self.assertEqual(output[2], file_name) + output = get_cache_execute_command(cores=1, file_name=file_name, backend="slurm") + self.assertEqual(output[0], sys.executable) + self.assertEqual(output[1].split(os.sep)[-1], "cache_serial.py") + self.assertEqual(output[2], file_name) + output = get_cache_execute_command(cores=1, file_name=file_name, backend="flux") + self.assertEqual(output[0], sys.executable) + self.assertEqual(output[1].split(os.sep)[-1], "cache_serial.py") + self.assertEqual(output[2], file_name) + + @unittest.skipIf( + skip_mpi4py_test, + "mpi4py is not installed, so the mpi4py tests are skipped.", + ) + def test_get_cache_execute_command_parallel(self): + file_name = "test.txt" + output = get_cache_execute_command(cores=2, file_name=file_name) + self.assertEqual(output[0], "mpiexec") + self.assertEqual(output[1], "-n") + self.assertEqual(output[2], str(2)) + self.assertEqual(output[3], sys.executable) + self.assertEqual(output[4].split(os.sep)[-1], "cache_parallel.py") + self.assertEqual(output[5], file_name) + output = get_cache_execute_command(cores=2, file_name=file_name, backend="slurm", pmi_mode="pmi2") + self.assertEqual(output[0], "srun") + self.assertEqual(output[1], "-n") + self.assertEqual(output[2], str(2)) + self.assertEqual(output[3], "--mpi=pmi2") + self.assertEqual(output[4], sys.executable) + self.assertEqual(output[5].split(os.sep)[-1], "cache_parallel.py") + self.assertEqual(output[6], file_name) + output = get_cache_execute_command(cores=2, file_name=file_name, backend="slurm") + self.assertEqual(output[0], "srun") + self.assertEqual(output[1], "-n") + self.assertEqual(output[2], str(2)) + self.assertEqual(output[3], sys.executable) + self.assertEqual(output[4].split(os.sep)[-1], "cache_parallel.py") + self.assertEqual(output[5], file_name) + output = get_cache_execute_command(cores=2, file_name=file_name, backend="flux") + self.assertEqual(output[0], "flux") + self.assertEqual(output[1], "run") + self.assertEqual(output[2], "-n") + self.assertEqual(output[3], str(2)) + self.assertEqual(output[4], sys.executable) + self.assertEqual(output[5].split(os.sep)[-1], "cache_parallel.py") + self.assertEqual(output[6], file_name) + output = get_cache_execute_command(cores=2, file_name=file_name, backend="flux", pmi_mode="pmix") + self.assertEqual(output[0], "flux") + self.assertEqual(output[1], "run") + self.assertEqual(output[2], "-o") + self.assertEqual(output[3], "pmi=pmix") + self.assertEqual(output[4], "-n") + self.assertEqual(output[5], str(2)) + self.assertEqual(output[6], sys.executable) + self.assertEqual(output[7].split(os.sep)[-1], "cache_parallel.py") + self.assertEqual(output[8], file_name) + with self.assertRaises(ValueError): + get_cache_execute_command(cores=2, file_name=file_name, backend="test") diff --git a/tests/test_standalone_inputcheck.py b/tests/test_standalone_inputcheck.py index d1d74df1..38fa896c 100644 --- a/tests/test_standalone_inputcheck.py +++ b/tests/test_standalone_inputcheck.py @@ -13,7 +13,7 @@ check_refresh_rate, check_resource_dict, check_resource_dict_is_empty, - check_flux_executor_pmi_mode, + check_pmi_mode, check_max_workers_and_cores, check_hostname_localhost, check_pysqa_config_directory, @@ -77,9 +77,9 @@ def test_check_plot_dependency_graph(self): with self.assertRaises(ValueError): check_plot_dependency_graph(plot_dependency_graph=True) - def test_check_flux_executor_pmi_mode(self): + def test_check_pmi_mode(self): with self.assertRaises(ValueError): - check_flux_executor_pmi_mode(flux_executor_pmi_mode="test") + check_pmi_mode(pmi_mode="test") def test_check_max_workers_and_cores(self): with self.assertRaises(ValueError): diff --git a/tests/test_standalone_interactive_backend.py b/tests/test_standalone_interactive_backend.py index cfa961af..c2306cae 100644 --- a/tests/test_standalone_interactive_backend.py +++ b/tests/test_standalone_interactive_backend.py @@ -84,6 +84,7 @@ def test_command_slurm_user_command(self): "2", "-D", os.path.abspath("."), + "--mpi=pmi2", "--gpus-per-task=1", "--oversubscribe", "--account=test", @@ -101,6 +102,7 @@ def test_command_slurm_user_command(self): gpus_per_core=1, openmpi_oversubscribe=True, slurm_cmd_args=["--account=test", "--job-name=executorlib"], + pmi_mode="pmi2", ) self.assertEqual( command_lst,