Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/executorlib/executor/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
55 changes: 54 additions & 1 deletion src/executorlib/standalone/inputcheck.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand All @@ -152,20 +162,48 @@ 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."
)


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."
)


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."
Expand Down Expand Up @@ -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)
Expand Down
6 changes: 6 additions & 0 deletions src/executorlib/standalone/interactive/communication.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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())
Expand Down
30 changes: 30 additions & 0 deletions src/executorlib/standalone/select.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
12 changes: 12 additions & 0 deletions src/executorlib/standalone/serialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
53 changes: 53 additions & 0 deletions src/executorlib/standalone/validate.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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__"):
Expand All @@ -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 = {
Expand Down
26 changes: 25 additions & 1 deletion src/executorlib/task_scheduler/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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 = {}
Expand All @@ -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
Expand Down
Loading
Loading