diff --git a/src/executorlib/executor/base.py b/src/executorlib/executor/base.py index 2925549f4..72e25ffd5 100644 --- a/src/executorlib/executor/base.py +++ b/src/executorlib/executor/base.py @@ -25,10 +25,22 @@ def __init__(self, executor: TaskSchedulerBase): @property def max_workers(self) -> Optional[int]: + """ + Return the number of parallel workers configured for this executor. + + Returns: + Optional[int]: The maximum number of parallel workers, or None if unconstrained. + """ return self._task_scheduler.max_workers @max_workers.setter def max_workers(self, max_workers: int): + """ + Set the number of parallel workers on the underlying task scheduler. + + Args: + max_workers (int): New maximum number of parallel workers. + """ self._task_scheduler.max_workers = max_workers @property diff --git a/src/executorlib/standalone/inputcheck.py b/src/executorlib/standalone/inputcheck.py index a7834b54c..283b45859 100644 --- a/src/executorlib/standalone/inputcheck.py +++ b/src/executorlib/standalone/inputcheck.py @@ -141,6 +141,16 @@ def check_init_function( def check_max_workers_and_cores( max_workers: Optional[int], max_cores: Optional[int] ) -> None: + """ + Check that neither max_workers nor max_cores is set when using a pysqa-based backend. + + Args: + max_workers (int, optional): Maximum number of workers. + max_cores (int, optional): Maximum number of cores. + + Raises: + ValueError: If max_workers or max_cores is not None. + """ if max_workers is not None: raise ValueError( "The number of workers cannot be controlled with the pysqa based backend." @@ -152,6 +162,15 @@ def check_max_workers_and_cores( def check_hostname_localhost(hostname_localhost: Optional[bool]) -> None: + """ + Check that hostname_localhost is not set when using a pysqa-based backend. + + Args: + hostname_localhost (bool, optional): Flag to use localhost for ZMQ connections. + + Raises: + ValueError: If hostname_localhost is not 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." @@ -159,6 +178,16 @@ def check_hostname_localhost(hostname_localhost: Optional[bool]) -> None: def check_restart_limit(restart_limit: int, block_allocation: bool = True) -> None: + """ + Check that restart_limit is only used together with block_allocation. + + Args: + restart_limit (int): Maximum number of times a worker process may be restarted. + block_allocation (bool): Whether block allocation is enabled. Defaults to True. + + Raises: + ValueError: If restart_limit is non-zero and block_allocation is False. + """ if not block_allocation and restart_limit != 0: raise ValueError( "The option to specify a restart limit for worker processes is only available with block_allocation=True." @@ -166,6 +195,15 @@ def check_restart_limit(restart_limit: int, block_allocation: bool = True) -> No def check_pmi_mode(pmi_mode: Optional[str]) -> None: + """ + Check that pmi_mode is not set on a local workstation without SLURM or flux. + + Args: + pmi_mode (str, optional): PMI interface name (e.g. "pmix", "pmi1", "pmi2"). + + Raises: + ValueError: If pmi_mode is not None. + """ if pmi_mode is not None: raise ValueError( "The option to specify the pmi mode is not available on a local workstation, it requires SLURM or flux." @@ -199,7 +237,22 @@ def validate_number_of_cores( set_local_cores: bool = False, ) -> int: """ - Validate the number of cores and return the appropriate value. + Validate the number of cores and return the number of workers to use. + + Precedence: max_cores / cores_per_worker > max_workers > CPU count (with warning). + + Args: + max_cores (int, optional): Total number of cores available. + max_workers (int, optional): Explicit number of parallel workers. + cores_per_worker (int, optional): Number of cores allocated to each worker. Defaults to 1. + set_local_cores (bool): When True, fall back to the local CPU count if neither + max_cores nor max_workers is given instead of raising an error. Defaults to False. + + Returns: + int: Number of parallel workers to use. + + Raises: + ValueError: If neither max_cores nor max_workers is set and set_local_cores is False. """ 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) diff --git a/src/executorlib/standalone/interactive/communication.py b/src/executorlib/standalone/interactive/communication.py index 68c6379a0..665d71f91 100644 --- a/src/executorlib/standalone/interactive/communication.py +++ b/src/executorlib/standalone/interactive/communication.py @@ -240,6 +240,9 @@ def interface_connect(host: str, port: str) -> tuple[zmq.Context, zmq.Socket]: Args: host (str): hostname of the host running the SocketInterface instance to connect to. port (str): port on the host the SocketInterface instance is running on. + + Returns: + Tuple[zmq.Context, zmq.Socket]: The ZMQ context and the connected PAIR socket. """ context = zmq.Context() socket = context.socket(zmq.PAIR) @@ -265,6 +268,9 @@ def interface_receive(socket: Optional[zmq.Socket]) -> dict: Args: socket (zmq.Socket): socket for the connection + + Returns: + dict: Deserialized dictionary received from the socket, or an empty dict if socket is None. """ if socket is not None: return cloudpickle.loads(socket.recv()) diff --git a/src/executorlib/standalone/select.py b/src/executorlib/standalone/select.py index f0edc082d..e3e3f5753 100644 --- a/src/executorlib/standalone/select.py +++ b/src/executorlib/standalone/select.py @@ -3,21 +3,51 @@ class FutureSelector(Future): + """ + A Future wrapper that returns a single element from a result that is a tuple, list, or dict. + + This is used by split_future() and get_item_from_future() to give callers an individual + Future-like object for each output of a function that returns a collection. + + Args: + future (Future): The underlying future whose result is a collection. + selector (int | str): Index (for sequences) or key (for mappings) used to extract + the desired element from the result. + """ + def __init__(self, future: Future, selector: int | str): + """ + Args: + future (Future): The underlying future whose result is a collection. + selector (int | str): Index or key used to extract an element from the result. + """ self._future = future self._selector = selector super().__init__() def __getattr__(self, attr: str) -> Any: + """Delegate attribute access to the wrapped future.""" return getattr(self._future, attr) def __setattr__(self, name: str, value: Any): + """Set attributes on the wrapped future, except for the two private instance variables.""" if name in ["_future", "_selector"]: super().__setattr__(name, value) else: setattr(self._future, name, value) def result(self, timeout: Optional[float] = None) -> Any: + """ + Return the selected element from the underlying future's result. + + Args: + timeout (float, optional): Maximum seconds to wait for the result. Defaults to None + (wait indefinitely). + + Returns: + Any: The element at position/key ``selector`` in the underlying result, or None if + the underlying result is None. + """ result = self._future.result(timeout=timeout) if result is not None: return result[self._selector] diff --git a/src/executorlib/standalone/serialize.py b/src/executorlib/standalone/serialize.py index 9311b7597..46c7b85a4 100644 --- a/src/executorlib/standalone/serialize.py +++ b/src/executorlib/standalone/serialize.py @@ -102,6 +102,18 @@ def _get_hash(binary: bytes) -> str: def _get_function_name(fn: Callable) -> str: + """ + Return a human-readable name for a callable. + + For regular functions and methods the ``__name__`` attribute is used. For callable objects + (classes implementing ``__call__``) the class name is extracted from ``__class__``. + + Args: + fn (Callable): The callable whose name should be determined. + + Returns: + str: A short name string suitable for use as part of a cache key prefix. + """ if hasattr(fn, "__name__"): return fn.__name__ else: diff --git a/src/executorlib/standalone/validate.py b/src/executorlib/standalone/validate.py index 9f4af396d..c77d7df1d 100644 --- a/src/executorlib/standalone/validate.py +++ b/src/executorlib/standalone/validate.py @@ -14,6 +14,25 @@ class ResourceDictValidation(BaseModel): + """ + Pydantic (or dataclass fallback) model for validating resource dictionaries passed to task submissions. + + Attributes: + cores (int, optional): Number of MPI cores to be used for each function call. + threads_per_core (int, optional): Number of OpenMP threads to be used for each function call. + gpus_per_core (int, optional): Number of GPUs per worker. + cwd (str, optional): Current working directory where the parallel python task is executed. + cache_key (str, optional): External cache key to identify tasks on the file system. + cache_directory (str, optional): The directory to store cache files. + num_nodes (int, optional): Number of compute nodes used for the evaluation. + exclusive (bool, optional): Reserve exclusive access to selected compute nodes. + error_log_file (str, optional): Path to the error log file. + run_time_max (int, optional): Maximum allowed execution time in seconds. + priority (int, optional): Queuing system priority for the task. + slurm_cmd_args (list[str], optional): Additional command line arguments for the srun call. + submission_template (str, optional): Template for queuing system job submission scripts. + """ + cores: Optional[int] = None threads_per_core: Optional[int] = None gpus_per_core: Optional[int] = None @@ -39,6 +58,18 @@ class Config: def _get_accepted_keys(class_type) -> list[str]: + """ + Return a list of accepted field names from a Pydantic model or dataclass. + + Args: + class_type: A Pydantic BaseModel subclass or a dataclass. + + Returns: + list[str]: Field names declared on the class. + + Raises: + TypeError: If the class type is neither a Pydantic model nor a dataclass. + """ if hasattr(class_type, "model_fields"): return list(class_type.model_fields.keys()) elif hasattr(class_type, "__dataclass_fields__"): @@ -47,10 +78,32 @@ def _get_accepted_keys(class_type) -> list[str]: def validate_resource_dict(resource_dict: dict) -> None: + """ + Validate a resource dictionary against the declared fields of ResourceDictValidation. + + Args: + resource_dict (dict): Dictionary of resource requirements to validate. Unknown keys raise + a validation error when pydantic is available. + + Raises: + ValidationError: If any key/value pair violates the ResourceDictValidation schema. + """ _ = ResourceDictValidation(**resource_dict) def validate_resource_dict_with_optional_keys(resource_dict: dict) -> None: + """ + Validate a resource dictionary, allowing unknown keys with a warning instead of an error. + + Known keys are validated against ResourceDictValidation. Unknown keys are collected and + emitted as a UserWarning so callers can detect unsupported options without hard failure. + + Args: + resource_dict (dict): Dictionary of resource requirements to validate. + + Raises: + ValidationError: If any known key/value pair violates the ResourceDictValidation schema. + """ accepted_keys = _get_accepted_keys(class_type=ResourceDictValidation) optional_lst = [key for key in resource_dict if key not in accepted_keys] validate_dict = { diff --git a/src/executorlib/task_scheduler/base.py b/src/executorlib/task_scheduler/base.py index 246e6e448..fd58e9586 100644 --- a/src/executorlib/task_scheduler/base.py +++ b/src/executorlib/task_scheduler/base.py @@ -15,6 +15,12 @@ def validate_resource_dict(resource_dict: dict): + """ + No-op resource dict validator used as the default when no validation is required. + + Args: + resource_dict (dict): Dictionary of resource requirements (ignored). + """ pass @@ -32,7 +38,13 @@ def __init__( validator: Callable = validate_resource_dict, ): """ - Initialize the ExecutorBase class. + Initialize the TaskSchedulerBase. + + Args: + max_cores (int, optional): Maximum number of cores available to the scheduler. + Tasks requesting more cores than this will be rejected. Defaults to None (unlimited). + validator (Callable): Function used to validate per-task resource dicts before + submission. Defaults to the no-op validate_resource_dict. """ cloudpickle_register(ind=3) self._process_kwargs: dict = {} @@ -43,10 +55,22 @@ def __init__( @property def max_workers(self) -> Optional[int]: + """ + Return the configured number of parallel workers, or None if unconstrained. + + Returns: + Optional[int]: The max_workers value stored in process kwargs, or None. + """ return self._process_kwargs.get("max_workers") @max_workers.setter def max_workers(self, max_workers: int): + """ + Setting max_workers after construction is not supported by the base scheduler. + + Raises: + NotImplementedError: Always. + """ raise NotImplementedError("The max_workers setter is not implemented.") @property diff --git a/src/executorlib/task_scheduler/file/shared.py b/src/executorlib/task_scheduler/file/shared.py index 91de4b326..e18a1bcca 100644 --- a/src/executorlib/task_scheduler/file/shared.py +++ b/src/executorlib/task_scheduler/file/shared.py @@ -354,6 +354,24 @@ def _cancel_processes( def _get_task_input( task_resource_dict: dict, executor_kwargs: dict ) -> tuple[dict, Optional[str], str, Optional[str]]: + """ + Merge per-task resource requirements with executor defaults and extract scheduling metadata. + + Executor-level kwargs fill in any keys not already present in the per-task resource dict. + The special keys ``cache_key``, ``cache_directory``, and ``error_log_file`` are popped from + the merged dict and returned separately so callers do not forward them to the backend. + + Args: + task_resource_dict (dict): Per-task resource dict from the submitted future, modified in place. + executor_kwargs (dict): Executor-level defaults (e.g. cores, cwd, cache_directory). + + Returns: + Tuple[dict, Optional[str], str, Optional[str]]: + - merged resource dict (without scheduling-only keys) + - cache_key (str or None) + - cache_directory (str, absolute path) + - error_log_file (str or None) + """ task_resource_dict.update( {k: v for k, v in executor_kwargs.items() if k not in task_resource_dict} ) @@ -364,6 +382,13 @@ def _get_task_input( def _cancel_futures(future_dict: dict): + """ + Cancel all pending futures in the dictionary. + + Args: + future_dict (dict): Mapping of task keys to Future objects. Already-done futures are + skipped; pending ones are cancelled. + """ for value in future_dict.values(): if not value.done(): value.cancel() @@ -380,6 +405,28 @@ def _shutdown_executor( backend: Optional[str] = None, refresh_rate: float = 0.01, ): + """ + Shut down the file-based executor, optionally waiting for or cancelling pending tasks. + + The four combinations of ``wait`` / ``cancel_futures`` mirror the semantics of + concurrent.futures.Executor.shutdown(): + + * wait=True, cancel_futures=False – poll until all tasks finish. + * wait=True, cancel_futures=True – cancel pending futures then wait for running ones. + * wait=False, cancel_futures=True – cancel everything immediately without blocking. + * wait=False, cancel_futures=False – detach tasks and mark futures cancelled. + + Args: + wait (bool): Whether to block until all outstanding tasks have resolved. + cancel_futures (bool): Whether to cancel futures that have not yet started. + memory_dict (dict): Mapping of task keys to their Future objects. + process_dict (dict): Mapping of task keys to process handles or queue IDs. + cache_dir_dict (dict): Mapping of task keys to the cache directory for each task. + terminate_function (Callable, optional): Function used to terminate running processes. + pysqa_config_directory (str, optional): Path to the pysqa config directory. + backend (str, optional): Name of the backend ("slurm", "flux", or None for subprocess). + refresh_rate (float): Polling interval in seconds when waiting for tasks. Defaults to 0.01. + """ if wait and not cancel_futures: while len(memory_dict) > 0: memory_dict = _refresh_memory_dict(