[Feature] Nested executor

src/executorlib/backend/executor_nested.py

@@ -0,0 +1,113 @@
+from concurrent.futures import Future
+from typing import Callable, Optional
+
+from executorlib.standalone.interactive.arguments import (
+    get_future_objects_from_input,
+)
+
+
+class BackendExecutor:
+    def __init__(self):
+        self._tasks_dict = {}
+        self._future_dict = {}
+
+    @property
+    def tasks(self):
+        return self._tasks_dict
+
+    def batched(
+        self,
+        iterable: list[Future],
+        n: int,
+    ) -> list[Future]:
+        """
+        Batch futures from the iterable into tuples of length n. The last batch may be shorter than n.
+
+        Args:
+            iterable (list): list of future objects to batch based on which future objects finish first
+            n (int): badge size
+
+        Returns:
+            list[Future]: list of future objects one for each batch
+        """
+        raise NotImplementedError("The batched method is not implemented.")
+
+    def map(
+        self,
+        fn: Callable,
+        *iterables,
+        timeout: Optional[float] = None,
+        chunksize: int = 1,
+    ):
+        """Returns an iterator equivalent to map(fn, iter).
+
+        Args:
+            fn: A callable that will take as many arguments as there are
+                passed iterables.
+            timeout: The maximum number of seconds to wait. If None, then there
+                is no limit on the wait time.
+            chunksize: The size of the chunks the iterable will be broken into
+                before being passed to a child process. This argument is only
+                used by ProcessPoolExecutor; it is ignored by
+                ThreadPoolExecutor.
+
+        Returns:
+            An iterator equivalent to: map(func, *iterables) but the calls may
+            be evaluated out-of-order.
+
+        Raises:
+            TimeoutError: If the entire result iterator could not be generated
+                before the given timeout.
+            Exception: If fn(*args) raises for any values.
+        """
+        raise NotImplementedError("The map method is not implemented.")
+
+    def submit(  # type: ignore
+        self,
+        fn: Callable,
+        /,
+        *args,
+        resource_dict: Optional[dict] = None,
+        **kwargs,
+    ) -> Future:
+        """
+        Submits a callable to be executed with the given arguments.
+
+        Schedules the callable to be executed as fn(*args, **kwargs) and returns
+        a Future instance representing the execution of the callable.
+
+        Args:
+            fn (callable): function to submit for execution
+            args: arguments for the submitted function
+            kwargs: keyword arguments for the submitted function
+            resource_dict (dict): A dictionary of resources required by the task. With the following keys:
+                              - cores (int): number of MPI cores to be used for each function call
+                              - threads_per_core (int): number of OpenMP threads to be used for each function call
+                              - gpus_per_core (int): number of GPUs per worker - defaults to 0
+                              - cwd (str/None): current working directory where the parallel python task is executed
+                              - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI and
+                                                              SLURM only) - default False
+                              - slurm_cmd_args (list): Additional command line arguments for the srun call (SLURM only)
+                              - error_log_file (str): Name of the error log file to use for storing exceptions raised
+                                                      by the Python functions submitted to the Executor.
+
+        Returns:
+            Future: A Future representing the given call.
+        """
+        if resource_dict is None:
+            resource_dict = {}
+        f: Future = Future()
+        self._future_dict[f] = id(f)
+        self._tasks_dict[id(f)] = {
+            "fn": fn,
+            "args": args,  # at the momemt the future objects are not removed from the args
+            "kwargs": kwargs,  # at the momemt the future objects are not removed from the kwargs
+            "resource_dict": resource_dict,
+            "dependencies": [
+                self._future_dict[future_dependency]
+                for future_dependency in get_future_objects_from_input(
+                    args=args, kwargs=kwargs
+                )
+            ],
+        }
+        return f

src/executorlib/backend/interactive_parallel.py

@@ -6,6 +6,7 @@
 import cloudpickle
 import zmq
 
+from executorlib.backend.executor_nested import BackendExecutor
 from executorlib.standalone.error import backend_write_error_file
 from executorlib.standalone.interactive.backend import call_funct, parse_arguments
 from executorlib.standalone.interactive.communication import (
@@ -43,7 +44,11 @@ def main() -> None:
             host=argument_dict["host"], port=argument_dict["zmqport"]
         )
 
-    memory = {"executorlib_worker_id": int(argument_dict["worker_id"])}
+    nested_executor = BackendExecutor()
+    memory = {
+        "executorlib_worker_id": int(argument_dict["worker_id"]),
+        "executorlib_executor": nested_executor,
+    }
 
     # required for flux interface - otherwise the current path is not included in the python path
     cwd = abspath(".")
@@ -90,7 +95,13 @@ def main() -> None:
             else:
                 # Send output
                 if mpi_rank_zero:
-                    interface_send(socket=socket, result_dict={"result": output_reply})
+                    interface_send(
+                        socket=socket,
+                        result_dict={
+                            "result": output_reply,
+                            "tasks_nested": nested_executor.tasks,
+                        },
+                    )
         elif (
             "init" in input_dict
             and input_dict["init"]

src/executorlib/backend/interactive_serial.py

@@ -2,6 +2,7 @@
 from os.path import abspath
 from typing import Optional
 
+from executorlib.backend.executor_nested import BackendExecutor
 from executorlib.standalone.error import backend_write_error_file
 from executorlib.standalone.interactive.backend import call_funct, parse_arguments
 from executorlib.standalone.interactive.communication import (
@@ -29,7 +30,11 @@ def evaluate_cmd(argument_lst: Optional[list[str]] = None):
         host=argument_dict["host"], port=argument_dict["zmqport"]
     )
 
-    memory = {"executorlib_worker_id": int(argument_dict["worker_id"])}
+    nested_executor = BackendExecutor()
+    memory = {
+        "executorlib_worker_id": int(argument_dict["worker_id"]),
+        "executorlib_executor": nested_executor,
+    }
 
     # required for flux interface - otherwise the current path is not included in the python path
     cwd = abspath(".")
@@ -65,7 +70,13 @@ def evaluate_cmd(argument_lst: Optional[list[str]] = None):
                 )
             else:
                 # Send output
-                interface_send(socket=socket, result_dict={"result": output})
+                interface_send(
+                    socket=socket,
+                    result_dict={
+                        "result": output,
+                        "tasks_nested": nested_executor.tasks,
+                    },
+                )
         elif (
             "init" in input_dict
             and input_dict["init"]

src/executorlib/task_scheduler/base.py

@@ -50,6 +50,7 @@ def __init__(
         self._process_kwargs: dict = {}
         self._max_cores = max_cores
         self._future_queue: Optional[queue.Queue] = queue.Queue()
+        self._return_queue: Optional[queue.Queue] = queue.Queue()
         self._process: Optional[Union[Thread, list[Thread]]] = None
         self._validator = validator
 
@@ -84,6 +85,8 @@ def info(self) -> Optional[dict]:
         meta_data_dict = self._process_kwargs.copy()
         if "future_queue" in meta_data_dict:
             del meta_data_dict["future_queue"]
+        if "return_queue" in meta_data_dict:
+            del meta_data_dict["return_queue"]
         if self._process is not None and isinstance(self._process, list):
             meta_data_dict["max_workers"] = len(self._process)
             return meta_data_dict
@@ -102,6 +105,16 @@ def future_queue(self) -> Optional[queue.Queue]:
         """
         return self._future_queue
 
+    @property
+    def return_queue(self) -> Optional[queue.Queue]:
+        """
+        Get the return queue.
+
+        Returns:
+            queue.Queue: The return queue.
+        """
+        return self._return_queue
+
     def batched(
         self,
         iterable: list[Future],
@@ -238,8 +251,10 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
             if isinstance(self._process, Thread):
                 self._process.join()
                 self._future_queue.join()
+                self._return_queue.join()
         self._process = None
         self._future_queue = None
+        self._return_queue = None
 
     def _set_process(self, process: Thread):
         """

src/executorlib/task_scheduler/interactive/blockallocation.py

@@ -75,6 +75,7 @@ def __init__(
             max_cores=executor_kwargs.get("max_cores"), validator=validator
         )
         executor_kwargs["future_queue"] = self._future_queue
+        executor_kwargs["return_queue"] = self._return_queue
         executor_kwargs["spawner"] = spawner
         executor_kwargs["queue_join_on_shutdown"] = False
         executor_kwargs["restart_limit"] = restart_limit
@@ -217,6 +218,7 @@ def _set_process(self, process: list[Thread]):  # type: ignore
 
 def _execute_multiple_tasks(
     future_queue: queue.Queue,
+    return_queue: queue.Queue,
     cores: int = 1,
     spawner: type[BaseSpawner] = MpiExecSpawner,
     hostname_localhost: Optional[bool] = None,
@@ -240,6 +242,7 @@ def _execute_multiple_tasks(
 
     Args:
         future_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process
+        return_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process
         cores (int): defines the total number of MPI ranks to use
         spawner (BaseSpawner): Spawner to start process on selected compute resources
         hostname_localhost (boolean): use localhost instead of the hostname to establish the zmq connection. In the
@@ -317,7 +320,7 @@ def _execute_multiple_tasks(
                     f.set_exception(exception=interface_initialization_exception)
                 else:
                     # The interface failed during the execution
-                    interface.status = execute_task_dict(
+                    interface.status, nested_task_dict = execute_task_dict(
                         task_dict=task_dict,
                         future_obj=f,
                         interface=interface,
@@ -329,6 +332,8 @@ def _execute_multiple_tasks(
                         reset_task_dict(
                             future_obj=f, future_queue=future_queue, task_dict=task_dict
                         )
+                    elif len(nested_task_dict) > 0:
+                        return_queue.put(nested_task_dict)
                 task_done(future_queue=future_queue)
 
 

src/executorlib/task_scheduler/interactive/dependency.py

@@ -257,6 +257,10 @@ def _execute_tasks_with_dependencies(
             task_dict = future_queue.get_nowait()
         except queue.Empty:
             task_dict = None
+        try:
+            task_return_dict = executor.return_queue.get_nowait()
+        except (queue.Empty, AttributeError):
+            task_return_dict = None
         if (  # shutdown the executor
             task_dict is not None and "shutdown" in task_dict and task_dict["shutdown"]
         ):
@@ -324,6 +328,16 @@ def _execute_tasks_with_dependencies(
                 executor_queue=executor_queue,
                 refresh_rate=refresh_rate,
             )
+        elif task_return_dict is not None:
+            # we need this future dict later on to resolve the dependencies
+            future_lookup_dict = {}
+            for k, v in task_return_dict.items():
+                future_lookup_dict[k] = executor.submit(
+                    fn=v["fn"],
+                    *v["args"],
+                    **v["kwargs"],
+                    resource_dict=v["resource_dict"],
+                )
         else:
             # If there is nothing else to do, sleep for a moment
             sleep(refresh_rate)

src/executorlib/task_scheduler/interactive/onetoone.py

@@ -61,6 +61,7 @@ def __init__(
         executor_kwargs.update(
             {
                 "future_queue": self._future_queue,
+                "return_queue": self._return_queue,
                 "spawner": spawner,
                 "max_cores": max_cores,
                 "max_workers": max_workers,
@@ -77,6 +78,7 @@ def __init__(
 
 def _execute_single_task(
     future_queue: queue.Queue,
+    return_queue: queue.Queue,
     spawner: type[BaseSpawner] = MpiExecSpawner,
     max_cores: Optional[int] = None,
     max_workers: Optional[int] = None,
@@ -88,6 +90,7 @@ def _execute_single_task(
 
     Args:
         future_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process
+        return_queue (queue.Queue): task queue of dictionary objects which are submitted to the parallel process
         spawner (BaseSpawner): Interface to start process on selected compute resources
         max_cores (int): defines the number cores which can be used in parallel
         max_workers (int): for backwards compatibility with the standard library, max_workers also defines the number of
@@ -116,6 +119,7 @@ def _execute_single_task(
         elif "fn" in task_dict and "future" in task_dict:
             process, active_task_dict = _wrap_execute_task_in_separate_process(
                 task_dict=task_dict,
+                return_queue=return_queue,
                 active_task_dict=active_task_dict,
                 spawner=spawner,
                 executor_kwargs=kwargs,
@@ -162,6 +166,7 @@ def _wait_for_free_slots(
 
 def _wrap_execute_task_in_separate_process(
     task_dict: dict,
+    return_queue: queue.Queue,
     active_task_dict: dict,
     spawner: type[BaseSpawner],
     executor_kwargs: dict,
@@ -211,6 +216,7 @@ def _wrap_execute_task_in_separate_process(
     task_kwargs.update(
         {
             "task_dict": task_dict,
+            "return_queue": return_queue,
             "spawner": spawner,
             "hostname_localhost": hostname_localhost,
             "future_obj": f,
@@ -226,6 +232,7 @@ def _wrap_execute_task_in_separate_process(
 
 def _execute_task_in_thread(
     task_dict: dict,
+    return_queue: queue.Queue,
     future_obj: Future,
     cores: int = 1,
     spawner: type[BaseSpawner] = MpiExecSpawner,
@@ -262,7 +269,7 @@ def _execute_task_in_thread(
         worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource
                          distribution.
     """
-    if not execute_task_dict(
+    status, nested_task_dict = execute_task_dict(
         task_dict=task_dict,
         future_obj=future_obj,
         interface=interface_bootup(
@@ -277,7 +284,10 @@ def _execute_task_in_thread(
         cache_directory=cache_directory,
         cache_key=cache_key,
         error_log_file=error_log_file,
-    ):
+    )
+    if status is False:
         future_obj.set_exception(
             ExecutorlibSocketError("SocketInterface crashed during execution.")
         )
+    elif len(nested_task_dict) > 0:
+        return_queue.put(nested_task_dict)