executorlib 0.0.8__py3-none-any.whl

executorlib/interactive/slurm.py

@@ -0,0 +1,109 @@
+import os
+from typing import Optional
+
+from executorlib.standalone.interactive.spawner import SubprocessSpawner
+
+SLURM_COMMAND = "srun"
+
+
+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)
+        )
+
+
+class SrunSpawner(SubprocessSpawner):
+    def __init__(
+        self,
+        cwd: Optional[str] = None,
+        cores: int = 1,
+        threads_per_core: int = 1,
+        gpus_per_core: int = 0,
+        openmpi_oversubscribe: bool = False,
+        slurm_cmd_args: list[str] = [],
+    ):
+        """
+        Srun interface implementation.
+
+        Args:
+            cwd (str, optional): The current working directory. Defaults to None.
+            cores (int, optional): The number of cores to use. Defaults to 1.
+            threads_per_core (int, optional): The number of threads per core. Defaults to 1.
+            gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
+            openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
+            slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
+        """
+        super().__init__(
+            cwd=cwd,
+            cores=cores,
+            openmpi_oversubscribe=openmpi_oversubscribe,
+            threads_per_core=threads_per_core,
+        )
+        self._gpus_per_core = gpus_per_core
+        self._slurm_cmd_args = slurm_cmd_args
+
+    def generate_command(self, command_lst: list[str]) -> list[str]:
+        """
+        Generate the command list for the Srun interface.
+
+        Args:
+            command_lst (list[str]): The command list.
+
+        Returns:
+            list[str]: The generated command list.
+        """
+        command_prepend_lst = generate_slurm_command(
+            cores=self._cores,
+            cwd=self._cwd,
+            threads_per_core=self._threads_per_core,
+            gpus_per_core=self._gpus_per_core,
+            openmpi_oversubscribe=self._openmpi_oversubscribe,
+            slurm_cmd_args=self._slurm_cmd_args,
+        )
+        return super().generate_command(
+            command_lst=command_prepend_lst + command_lst,
+        )
+
+
+def generate_slurm_command(
+    cores: int,
+    cwd: Optional[str],
+    threads_per_core: int = 1,
+    gpus_per_core: int = 0,
+    openmpi_oversubscribe: bool = False,
+    slurm_cmd_args: list[str] = [],
+) -> list[str]:
+    """
+    Generate the command list for the SLURM interface.
+
+    Args:
+        cores (int): The number of cores.
+        cwd (str): The current working directory.
+        threads_per_core (int, optional): The number of threads per core. Defaults to 1.
+        gpus_per_core (int, optional): The number of GPUs per core. Defaults to 0.
+        openmpi_oversubscribe (bool, optional): Whether to oversubscribe the cores. Defaults to False.
+        slurm_cmd_args (list[str], optional): Additional command line arguments. Defaults to [].
+
+    Returns:
+        list[str]: The generated command list.
+    """
+    command_prepend_lst = [SLURM_COMMAND, "-n", str(cores)]
+    if cwd is not None:
+        command_prepend_lst += ["-D", cwd]
+    if threads_per_core > 1:
+        command_prepend_lst += ["--cpus-per-task" + str(threads_per_core)]
+    if gpus_per_core > 0:
+        command_prepend_lst += ["--gpus-per-task=" + str(gpus_per_core)]
+    if openmpi_oversubscribe:
+        command_prepend_lst += ["--oversubscribe"]
+    if len(slurm_cmd_args) > 0:
+        command_prepend_lst += slurm_cmd_args
+    return command_prepend_lst

executorlib/standalone/__init__.py

@@ -0,0 +1,21 @@
+from executorlib.standalone.interactive.communication import (
+    SocketInterface,
+    interface_bootup,
+    interface_connect,
+    interface_receive,
+    interface_send,
+    interface_shutdown,
+)
+from executorlib.standalone.interactive.spawner import MpiExecSpawner
+from executorlib.standalone.thread import RaisingThread
+
+__all__ = [
+    "SocketInterface",
+    "interface_bootup",
+    "interface_connect",
+    "interface_send",
+    "interface_shutdown",
+    "interface_receive",
+    "RaisingThread",
+    "MpiExecSpawner",
+]

executorlib/standalone/command.py

@@ -0,0 +1,14 @@
+import os
+
+
+def get_command_path(executable: str) -> str:
+    """
+    Get path of the backend executable script
+
+    Args:
+        executable (str): Name of the backend executable script, either mpiexec.py or serial.py
+
+    Returns:
+        str: absolute path to the executable script
+    """
+    return os.path.abspath(os.path.join(__file__, "..", "..", "backend", executable))

executorlib/standalone/hdf.py

@@ -0,0 +1,116 @@
+import os
+from typing import Any, List, Optional, Tuple
+
+import cloudpickle
+import h5py
+import numpy as np
+
+group_dict = {
+    "fn": "function",
+    "args": "input_args",
+    "kwargs": "input_kwargs",
+    "output": "output",
+    "runtime": "runtime",
+    "queue_id": "queue_id",
+}
+
+
+def dump(file_name: Optional[str], data_dict: dict) -> None:
+    """
+    Dump data dictionary into HDF5 file
+
+    Args:
+        file_name (str): file name of the HDF5 file as absolute path
+        data_dict (dict): dictionary containing the python function to be executed {"fn": ..., "args": (), "kwargs": {}}
+    """
+    if file_name is not None:
+        with h5py.File(file_name, "a") as fname:
+            for data_key, data_value in data_dict.items():
+                if data_key in group_dict.keys():
+                    fname.create_dataset(
+                        name="/" + group_dict[data_key],
+                        data=np.void(cloudpickle.dumps(data_value)),
+                    )
+
+
+def load(file_name: str) -> dict:
+    """
+    Load data dictionary from HDF5 file
+
+    Args:
+        file_name (str): file name of the HDF5 file as absolute path
+
+    Returns:
+        dict: dictionary containing the python function to be executed {"fn": ..., "args": (), "kwargs": {}}
+    """
+    with h5py.File(file_name, "r") as hdf:
+        data_dict = {}
+        if "function" in hdf:
+            data_dict["fn"] = cloudpickle.loads(np.void(hdf["/function"]))
+        else:
+            raise TypeError("Function not found in HDF5 file.")
+        if "input_args" in hdf:
+            data_dict["args"] = cloudpickle.loads(np.void(hdf["/input_args"]))
+        else:
+            data_dict["args"] = ()
+        if "input_kwargs" in hdf:
+            data_dict["kwargs"] = cloudpickle.loads(np.void(hdf["/input_kwargs"]))
+        else:
+            data_dict["kwargs"] = {}
+        return data_dict
+
+
+def get_output(file_name: str) -> Tuple[bool, Any]:
+    """
+    Check if output is available in the HDF5 file
+
+    Args:
+        file_name (str): file name of the HDF5 file as absolute path
+
+    Returns:
+        Tuple[bool, object]: boolean flag indicating if output is available and the output object itself
+    """
+    with h5py.File(file_name, "r") as hdf:
+        if "output" in hdf:
+            return True, cloudpickle.loads(np.void(hdf["/output"]))
+        else:
+            return False, None
+
+
+def get_runtime(file_name: str) -> float:
+    """
+    Get run time from HDF5 file
+
+    Args:
+        file_name (str): file name of the HDF5 file as absolute path
+
+    Returns:
+        float: run time from the execution of the python function
+    """
+    with h5py.File(file_name, "r") as hdf:
+        if "runtime" in hdf:
+            return cloudpickle.loads(np.void(hdf["/runtime"]))
+        else:
+            return 0.0
+
+
+def get_queue_id(file_name: Optional[str]) -> Optional[int]:
+    if file_name is not None:
+        with h5py.File(file_name, "r") as hdf:
+            if "queue_id" in hdf:
+                return cloudpickle.loads(np.void(hdf["/queue_id"]))
+    return None
+
+
+def get_cache_data(cache_directory: str) -> List[dict]:
+    file_lst = []
+    for file_name in os.listdir(cache_directory):
+        with h5py.File(os.path.join(cache_directory, file_name), "r") as hdf:
+            file_content_dict = {
+                key: cloudpickle.loads(np.void(hdf["/" + key]))
+                for key in group_dict.values()
+                if key in hdf
+            }
+        file_content_dict["filename"] = file_name
+        file_lst.append(file_content_dict)
+    return file_lst

executorlib/standalone/inputcheck.py

@@ -0,0 +1,201 @@
+import inspect
+import multiprocessing
+import os.path
+from concurrent.futures import Executor
+from typing import Callable, List, Optional
+
+
+def check_oversubscribe(oversubscribe: bool) -> None:
+    """
+    Check if oversubscribe is True and raise a ValueError if it is.
+    """
+    if oversubscribe:
+        raise ValueError(
+            "Oversubscribing is not supported for the executorlib.flux.PyFLuxExecutor backend."
+            "Please use oversubscribe=False instead of oversubscribe=True."
+        )
+
+
+def check_command_line_argument_lst(command_line_argument_lst: List[str]) -> None:
+    """
+    Check if command_line_argument_lst is not empty and raise a ValueError if it is.
+    """
+    if len(command_line_argument_lst) > 0:
+        raise ValueError(
+            "The command_line_argument_lst parameter is not supported for the SLURM backend."
+        )
+
+
+def check_gpus_per_worker(gpus_per_worker: int) -> None:
+    """
+    Check if gpus_per_worker is not 0 and raise a TypeError if it is.
+    """
+    if gpus_per_worker != 0:
+        raise TypeError(
+            "GPU assignment is not supported for the executorlib.mpi.PyMPIExecutor backend."
+            "Please use gpus_per_worker=0 instead of gpus_per_worker="
+            + str(gpus_per_worker)
+            + "."
+        )
+
+
+def check_executor(executor: Executor) -> None:
+    """
+    Check if executor is not None and raise a ValueError if it is.
+    """
+    if executor is not None:
+        raise ValueError(
+            "The executor parameter is only supported for the flux framework backend."
+        )
+
+
+def check_nested_flux_executor(nested_flux_executor: bool) -> None:
+    """
+    Check if nested_flux_executor is True and raise a ValueError if it is.
+    """
+    if nested_flux_executor:
+        raise ValueError(
+            "The nested_flux_executor parameter is only supported for the flux framework backend."
+        )
+
+
+def check_resource_dict(function: Callable) -> None:
+    """
+    Check if the function has a parameter named 'resource_dict' and raise a ValueError if it does.
+    """
+    if "resource_dict" in inspect.signature(function).parameters.keys():
+        raise ValueError(
+            "The parameter resource_dict is used internally in executorlib, "
+            "so it cannot be used as a parameter in the submitted functions."
+        )
+
+
+def check_resource_dict_is_empty(resource_dict: dict) -> None:
+    """
+    Check if resource_dict is not empty and raise a ValueError if it is.
+    """
+    if len(resource_dict) > 0:
+        raise ValueError(
+            "When block_allocation is enabled, the resource requirements have to be defined on the executor level."
+        )
+
+
+def check_refresh_rate(refresh_rate: float) -> None:
+    """
+    Check if refresh_rate is not 0.01 and raise a ValueError if it is.
+    """
+    if refresh_rate != 0.01:
+        raise ValueError(
+            "The sleep_interval parameter is only used when disable_dependencies=False."
+        )
+
+
+def check_plot_dependency_graph(plot_dependency_graph: bool) -> None:
+    """
+    Check if plot_dependency_graph is True and raise a ValueError if it is.
+    """
+    if plot_dependency_graph:
+        raise ValueError(
+            "The plot_dependency_graph parameter is only used when disable_dependencies=False."
+        )
+
+
+def check_pmi(backend: Optional[str], pmi: Optional[str]) -> None:
+    """
+    Check if pmi is valid for the selected backend and raise a ValueError if it is not.
+    """
+    if backend is not None:
+        if backend != "flux_allocation" and pmi is not None:
+            raise ValueError(
+                "The pmi parameter is currently only implemented for flux."
+            )
+        elif backend == "flux_allocation" and pmi not in ["pmix", "pmi1", "pmi2", None]:
+            raise ValueError(
+                "The pmi parameter supports [pmix, pmi1, pmi2], but not: " + str(pmi)
+            )
+
+
+def check_init_function(
+    block_allocation: bool, init_function: Optional[Callable]
+) -> None:
+    """
+    Check if block_allocation is False and init_function is not None, and raise a ValueError if it is.
+    """
+    if not block_allocation and init_function is not None:
+        raise ValueError("")
+
+
+def check_max_workers_and_cores(
+    max_workers: Optional[int], max_cores: Optional[int]
+) -> None:
+    if max_workers is not None:
+        raise ValueError(
+            "The number of workers cannot be controlled with the pysqa based backend."
+        )
+    if max_cores is not None:
+        raise ValueError(
+            "The number of cores cannot be controlled with the pysqa based backend."
+        )
+
+
+def check_hostname_localhost(hostname_localhost: Optional[bool]) -> None:
+    if hostname_localhost is not None:
+        raise ValueError(
+            "The option to connect to hosts based on their hostname is not available with the pysqa based backend."
+        )
+
+
+def check_flux_executor_pmi_mode(flux_executor_pmi_mode: Optional[str]) -> None:
+    if flux_executor_pmi_mode is not None:
+        raise ValueError(
+            "The option to specify the flux pmi mode is not available with the pysqa based backend."
+        )
+
+
+def check_flux_log_files(flux_log_files: Optional[bool]) -> None:
+    """
+    Check if flux_log_files is True and raise a ValueError if it is.
+    """
+    if flux_log_files:
+        raise ValueError(
+            "The flux_log_files parameter is only supported for the flux framework backend."
+        )
+
+
+def check_pysqa_config_directory(pysqa_config_directory: Optional[str]) -> None:
+    """
+    Check if pysqa_config_directory is None and raise a ValueError if it is not.
+    """
+    if pysqa_config_directory is not None:
+        raise ValueError(
+            "pysqa_config_directory parameter is only supported for pysqa backend."
+        )
+
+
+def validate_number_of_cores(
+    max_cores: Optional[int] = None,
+    max_workers: Optional[int] = None,
+    cores_per_worker: Optional[int] = 1,
+    set_local_cores: bool = False,
+) -> int:
+    """
+    Validate the number of cores and return the appropriate value.
+    """
+    if max_cores is not None and max_workers is None and cores_per_worker is not None:
+        return int(max_cores / cores_per_worker)
+    elif max_workers is not None:
+        return int(max_workers)
+    else:
+        if max_cores is None and max_workers is None and not set_local_cores:
+            raise ValueError(
+                "Block allocation requires a fixed set of computational resources. Neither max_cores nor max_workers are defined."
+            )
+        else:
+            return multiprocessing.cpu_count()
+
+
+def check_file_exists(file_name: Optional[str]):
+    if file_name is None:
+        raise ValueError("file_name is not set.")
+    if not os.path.exists(file_name):
+        raise ValueError("file_name is not written to the file system.")

executorlib/standalone/interactive/backend.py

@@ -0,0 +1,98 @@
+import inspect
+from typing import Any, Callable, Optional
+
+
+def call_funct(
+    input_dict: dict, funct: Optional[Callable] = None, memory: Optional[dict] = None
+) -> Any:
+    """
+    Call function from dictionary
+
+    Args:
+        input_dict (dict): dictionary containing the function 'fn', its arguments 'args' and keyword arguments 'kwargs'
+        funct (Callable, optional): function to be evaluated if it is not included in the input dictionary
+        memory (dict, optional): variables stored in memory which can be used as keyword arguments
+
+    Returns:
+        Any: Result of the function
+    """
+    if funct is None:
+
+        def funct(*args, **kwargs):
+            return args[0].__call__(*args[1:], **kwargs)
+
+    funct_args = inspect.getfullargspec(input_dict["fn"]).args
+    if memory is not None:
+        input_dict["kwargs"].update(
+            _update_dict_delta(
+                dict_input=memory,
+                dict_output=input_dict["kwargs"],
+                keys_possible_lst=funct_args,
+            )
+        )
+    return funct(input_dict["fn"], *input_dict["args"], **input_dict["kwargs"])
+
+
+def parse_arguments(argument_lst: list[str]) -> dict:
+    """
+    Simple function to parse command line arguments
+
+    Args:
+        argument_lst (list): list of arguments as strings
+
+    Returns:
+        dict: dictionary with the parsed arguments and their corresponding values
+    """
+    return update_default_dict_from_arguments(
+        argument_lst=argument_lst,
+        argument_dict={
+            "zmqport": "--zmqport",
+            "host": "--host",
+        },
+        default_dict={"host": "localhost"},
+    )
+
+
+def update_default_dict_from_arguments(
+    argument_lst: list[str], argument_dict: dict, default_dict: dict
+) -> dict:
+    """
+    Update default dictionary with values from command line arguments
+
+    Args:
+        argument_lst (list[str]): List of arguments as strings
+        argument_dict (dict): Dictionary mapping argument names to their corresponding command line flags
+        default_dict (dict): Default dictionary to be updated
+
+    Returns:
+        dict: Updated default dictionary
+    """
+    default_dict.update(
+        {
+            k: argument_lst[argument_lst.index(v) + 1]
+            for k, v in argument_dict.items()
+            if v in argument_lst
+        }
+    )
+    return default_dict
+
+
+def _update_dict_delta(
+    dict_input: dict, dict_output: dict, keys_possible_lst: list[str]
+) -> dict:
+    """
+    Update dictionary with values from another dictionary, only if the keys are present in a given list
+
+    Args:
+        dict_input (dict): Input dictionary
+        dict_output (dict): Output dictionary to be updated
+        keys_possible_lst (list[str]): List of possible keys to be updated
+
+    Returns:
+        dict: Updated dictionary
+    """
+    return {
+        k: v
+        for k, v in dict_input.items()
+        if k in keys_possible_lst and k not in dict_output.keys()
+    }