Merge branch 'main' into paper

Author: jan-janssen

.ci_support/check.py

@@ -0,0 +1,16 @@
+import tomlkit
+
+
+if __name__ == "__main__":
+    with open("pyproject.toml", "r") as f:
+        data = tomlkit.load(f)
+
+    lst = []
+    for sub_lst in data["project"]["optional-dependencies"].values():
+        for el in sub_lst:
+            lst.append(el)
+
+    data["project"]["dependencies"] += list(set(lst))
+
+    with open("pyproject.toml", "w") as f:
+        f.writelines(tomlkit.dumps(data))

.ci_support/environment-old.yml

@@ -2,7 +2,7 @@ channels:
 - conda-forge
 dependencies:
 - python
-- numpy
+- numpy =1.23.5
 - openmpi =4.1.4
 - cloudpickle =2.0.0
 - mpi4py =3.1.4

.github/workflows/pipeline.yml

@@ -95,7 +95,9 @@ jobs:
     - name: Setup
       shell: bash -l {0}
       run: |
-        pip install versioneer[toml]==0.29
+        pip install versioneer[toml]==0.29 tomlkit
+        python .ci_support/check.py
+        cat pyproject.toml
         pip install . --no-deps --no-build-isolation
         pip check
 

executorlib/base/executor.py

@@ -6,12 +6,12 @@
 from concurrent.futures import (
     Future,
 )
+from threading import Thread
 from typing import Callable, Optional, Union
 
 from executorlib.standalone.inputcheck import check_resource_dict
 from executorlib.standalone.queue import cancel_items_in_queue
 from executorlib.standalone.serialize import cloudpickle_register
-from executorlib.standalone.thread import RaisingThread
 
 
 class ExecutorBase(FutureExecutor):
@@ -29,7 +29,7 @@ def __init__(self, max_cores: Optional[int] = None):
         cloudpickle_register(ind=3)
         self._max_cores = max_cores
         self._future_queue: Optional[queue.Queue] = queue.Queue()
-        self._process: Optional[Union[RaisingThread, list[RaisingThread]]] = None
+        self._process: Optional[Union[Thread, list[Thread]]] = None
 
     @property
     def info(self) -> Optional[dict]:
@@ -40,13 +40,13 @@ def info(self) -> Optional[dict]:
             Optional[dict]: Information about the executor.
         """
         if self._process is not None and isinstance(self._process, list):
-            meta_data_dict = self._process[0].get_kwargs().copy()
+            meta_data_dict = self._process[0]._kwargs.copy()  # type: ignore
             if "future_queue" in meta_data_dict:
                 del meta_data_dict["future_queue"]
             meta_data_dict["max_workers"] = len(self._process)
             return meta_data_dict
         elif self._process is not None:
-            meta_data_dict = self._process.get_kwargs().copy()
+            meta_data_dict = self._process._kwargs.copy()  # type: ignore
             if "future_queue" in meta_data_dict:
                 del meta_data_dict["future_queue"]
             return meta_data_dict
@@ -138,13 +138,13 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
             cancel_items_in_queue(que=self._future_queue)
         if self._process is not None and self._future_queue is not None:
             self._future_queue.put({"shutdown": True, "wait": wait})
-            if wait and isinstance(self._process, RaisingThread):
+            if wait and isinstance(self._process, Thread):
                 self._process.join()
                 self._future_queue.join()
         self._process = None
         self._future_queue = None
 
-    def _set_process(self, process: RaisingThread):
+    def _set_process(self, process: Thread):
         """
         Set the process for the executor.
 

executorlib/cache/executor.py

@@ -1,4 +1,5 @@
 import os
+from threading import Thread
 from typing import Callable, Optional
 
 from executorlib.base.executor import ExecutorBase
@@ -15,7 +16,6 @@
     check_max_workers_and_cores,
     check_nested_flux_executor,
 )
-from executorlib.standalone.thread import RaisingThread
 
 try:
     from executorlib.cache.queue_spawner import execute_with_pysqa
@@ -64,7 +64,7 @@ def __init__(
         cache_directory_path = os.path.abspath(cache_directory)
         os.makedirs(cache_directory_path, exist_ok=True)
         self._set_process(
-            RaisingThread(
+            Thread(
                 target=execute_tasks_h5,
                 kwargs={
                     "future_queue": self._future_queue,

executorlib/cache/shared.py

@@ -115,8 +115,10 @@ def execute_tasks_h5(
                         ]
                     else:
                         if len(future_wait_key_lst) > 0:
-                            raise ValueError(
-                                "Future objects are not supported as input if disable_dependencies=True."
+                            task_dict["future"].set_exception(
+                                ValueError(
+                                    "Future objects are not supported as input if disable_dependencies=True."
+                                )
                             )
                         task_dependent_lst = []
                     process_dict[task_key] = execute_function(

executorlib/interactive/executor.py

@@ -1,4 +1,5 @@
 from concurrent.futures import Future
+from threading import Thread
 from typing import Any, Callable, Optional
 
 from executorlib.base.executor import ExecutorBase
@@ -8,7 +9,6 @@
     generate_nodes_and_edges,
     generate_task_hash,
 )
-from executorlib.standalone.thread import RaisingThread
 
 
 class ExecutorWithDependencies(ExecutorBase):
@@ -41,7 +41,7 @@ def __init__(
     ) -> None:
         super().__init__(max_cores=max_cores)
         self._set_process(
-            RaisingThread(
+            Thread(
                 target=execute_tasks_with_dependencies,
                 kwargs={
                     # Executor Arguments

executorlib/interactive/flux.py

@@ -29,6 +29,8 @@ class FluxPythonSpawner(BaseSpawner):
         cores (int, optional): The number of cores. Defaults to 1.
         threads_per_core (int, optional): The number of threads per base. Defaults to 1.
         gpus_per_core (int, optional): The number of GPUs per base. Defaults to 0.
+        num_nodes (int, optional): The number of compute nodes to use for executing the task. Defaults to None.
+        exclusive (bool): Whether to exclusively reserve the compute nodes, or allow sharing compute notes. Defaults to False.
         openmpi_oversubscribe (bool, optional): Whether to oversubscribe. Defaults to False.
         flux_executor (flux.job.FluxExecutor, optional): The FluxExecutor instance. Defaults to None.
         flux_executor_pmi_mode (str, optional): The PMI option. Defaults to None.
@@ -42,6 +44,8 @@ def __init__(
         cores: int = 1,
         threads_per_core: int = 1,
         gpus_per_core: int = 0,
+        num_nodes: Optional[int] = None,
+        exclusive: bool = False,
         openmpi_oversubscribe: bool = False,
         flux_executor: Optional[flux.job.FluxExecutor] = None,
         flux_executor_pmi_mode: Optional[str] = None,
@@ -55,6 +59,8 @@ def __init__(
         )
         self._threads_per_core = threads_per_core
         self._gpus_per_core = gpus_per_core
+        self._num_nodes = num_nodes
+        self._exclusive = exclusive
         self._flux_executor = flux_executor
         self._flux_executor_pmi_mode = flux_executor_pmi_mode
         self._flux_executor_nesting = flux_executor_nesting
@@ -85,17 +91,17 @@ def bootup(
                 num_tasks=self._cores,
                 cores_per_task=self._threads_per_core,
                 gpus_per_task=self._gpus_per_core,
-                num_nodes=None,
-                exclusive=False,
+                num_nodes=self._num_nodes,
+                exclusive=self._exclusive,
             )
         else:
             jobspec = flux.job.JobspecV1.from_nest_command(
                 command=command_lst,
                 num_slots=self._cores,
                 cores_per_slot=self._threads_per_core,
                 gpus_per_slot=self._gpus_per_core,
-                num_nodes=None,
-                exclusive=False,
+                num_nodes=self._num_nodes,
+                exclusive=self._exclusive,
             )
         jobspec.environment = dict(os.environ)
         if self._flux_executor_pmi_mode is not None:

executorlib/interactive/shared.py

@@ -3,7 +3,9 @@
 import queue
 import sys
 import time
-from concurrent.futures import Future
+from asyncio.exceptions import CancelledError
+from concurrent.futures import Future, TimeoutError
+from threading import Thread
 from time import sleep
 from typing import Any, Callable, Optional, Union
 
@@ -19,7 +21,6 @@
 )
 from executorlib.standalone.interactive.spawner import BaseSpawner, MpiExecSpawner
 from executorlib.standalone.serialize import serialize_funct_h5
-from executorlib.standalone.thread import RaisingThread
 
 
 class ExecutorBroker(ExecutorBase):
@@ -88,7 +89,7 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
         self._process = None
         self._future_queue = None
 
-    def _set_process(self, process: list[RaisingThread]):  # type: ignore
+    def _set_process(self, process: list[Thread]):  # type: ignore
         """
         Set the process for the executor.
 
@@ -148,7 +149,7 @@ def __init__(
         executor_kwargs["queue_join_on_shutdown"] = False
         self._set_process(
             process=[
-                RaisingThread(
+                Thread(
                     target=execute_parallel_tasks,
                     kwargs=executor_kwargs,
                 )
@@ -204,7 +205,7 @@ def __init__(
         executor_kwargs["max_cores"] = max_cores
         executor_kwargs["max_workers"] = max_workers
         self._set_process(
-            RaisingThread(
+            Thread(
                 target=execute_separate_tasks,
                 kwargs=executor_kwargs,
             )
@@ -361,15 +362,19 @@ def execute_tasks_with_dependencies(
             task_dict is not None and "fn" in task_dict and "future" in task_dict
         ):
             future_lst, ready_flag = _get_future_objects_from_input(task_dict=task_dict)
-            if len(future_lst) == 0 or ready_flag:
-                # No future objects are used in the input or all future objects are already done
-                task_dict["args"], task_dict["kwargs"] = _update_futures_in_input(
-                    args=task_dict["args"], kwargs=task_dict["kwargs"]
-                )
-                executor_queue.put(task_dict)
-            else:  # Otherwise add the function to the wait list
-                task_dict["future_lst"] = future_lst
-                wait_lst.append(task_dict)
+            exception_lst = _get_exception_lst(future_lst=future_lst)
+            if not _get_exception(future_obj=task_dict["future"]):
+                if len(exception_lst) > 0:
+                    task_dict["future"].set_exception(exception_lst[0])
+                elif len(future_lst) == 0 or ready_flag:
+                    # No future objects are used in the input or all future objects are already done
+                    task_dict["args"], task_dict["kwargs"] = _update_futures_in_input(
+                        args=task_dict["args"], kwargs=task_dict["kwargs"]
+                    )
+                    executor_queue.put(task_dict)
+                else:  # Otherwise add the function to the wait list
+                    task_dict["future_lst"] = future_lst
+                    wait_lst.append(task_dict)
             future_queue.task_done()
         elif len(wait_lst) > 0:
             number_waiting = len(wait_lst)
@@ -455,7 +460,10 @@ def _submit_waiting_task(wait_lst: list[dict], executor_queue: queue.Queue) -> l
     """
     wait_tmp_lst = []
     for task_wait_dict in wait_lst:
-        if all(future.done() for future in task_wait_dict["future_lst"]):
+        exception_lst = _get_exception_lst(future_lst=task_wait_dict["future_lst"])
+        if len(exception_lst) > 0:
+            task_wait_dict["future"].set_exception(exception_lst[0])
+        elif all(future.done() for future in task_wait_dict["future_lst"]):
             del task_wait_dict["future_lst"]
             task_wait_dict["args"], task_wait_dict["kwargs"] = _update_futures_in_input(
                 args=task_wait_dict["args"], kwargs=task_wait_dict["kwargs"]
@@ -483,6 +491,8 @@ def get_result(arg: Union[list[Future], Future]) -> Any:
             return arg.result()
         elif isinstance(arg, list):
             return [get_result(arg=el) for el in arg]
+        elif isinstance(arg, dict):
+            return {k: get_result(arg=v) for k, v in arg.items()}
         else:
             return arg
 
@@ -510,6 +520,8 @@ def find_future_in_list(lst):
                 future_lst.append(el)
             elif isinstance(el, list):
                 find_future_in_list(lst=el)
+            elif isinstance(el, dict):
+                find_future_in_list(lst=el.values())
 
     find_future_in_list(lst=task_dict["args"])
     find_future_in_list(lst=task_dict["kwargs"].values())
@@ -578,7 +590,7 @@ def _submit_function_to_separate_process(
             "init_function": None,
         }
     )
-    process = RaisingThread(
+    process = Thread(
         target=execute_parallel_tasks,
         kwargs=task_kwargs,
     )
@@ -599,14 +611,13 @@ def _execute_task(
         future_queue (Queue): Queue for receiving new tasks.
     """
     f = task_dict.pop("future")
-    if f.set_running_or_notify_cancel():
+    if not f.done() and f.set_running_or_notify_cancel():
         try:
             f.set_result(interface.send_and_receive_dict(input_dict=task_dict))
         except Exception as thread_exception:
             interface.shutdown(wait=True)
             future_queue.task_done()
             f.set_exception(exception=thread_exception)
-            raise thread_exception
         else:
             future_queue.task_done()
 
@@ -659,3 +670,15 @@ def _execute_task_with_cache(
         future = task_dict["future"]
         future.set_result(result)
         future_queue.task_done()
+
+
+def _get_exception_lst(future_lst: list[Future]) -> list:
+    return [f.exception() for f in future_lst if _get_exception(future_obj=f)]
+
+
+def _get_exception(future_obj: Future) -> bool:
+    try:
+        excp = future_obj.exception(timeout=10**-10)
+        return excp is not None and not isinstance(excp, CancelledError)
+    except TimeoutError:
+        return False