more fixes - still need to remove executors from intital submission and handle dependencies

Author: jan-janssen

src/executorlib/backend/executor_nested.py

@@ -100,8 +100,8 @@ def submit(  # type: ignore
         self._future_dict[f] = id(f)
         self._tasks_dict[id(f)] = {
             "fn": fn,
-            "args": args,
-            "kwargs": kwargs,
+            "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]

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
 
@@ -102,6 +103,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 +249,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:
+            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,12 @@ def _execute_tasks_with_dependencies(
                 executor_queue=executor_queue,
                 refresh_rate=refresh_rate,
             )
+        elif task_return_dict is not None:
+            future_lookup_dict = {}  # we need this future dict later on to resolve the dependencies
+            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

@@ -1,6 +1,7 @@
 import queue
 from concurrent.futures import Future
 from threading import Thread
+from time import sleep
 from typing import Callable, Optional
 
 from executorlib.standalone.command import get_interactive_execute_command
@@ -61,6 +62,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 +79,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 +91,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 +120,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 +167,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 +217,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 +233,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 +270,7 @@ def _execute_task_in_thread(
         worker_id (int): Communicate the worker which ID was assigned to it for future reference and resource
                          distribution.
     """
-    status, _ = execute_task_dict(
+    status, nested_task_dict = execute_task_dict(
         task_dict=task_dict,
         future_obj=future_obj,
         interface=interface_bootup(
@@ -282,3 +290,5 @@ def _execute_task_in_thread(
         future_obj.set_exception(
             ExecutorlibSocketError("SocketInterface crashed during execution.")
         )
+    elif len(nested_task_dict) > 0:
+        return_queue.put(nested_task_dict)