[Feature] Set wait=False during initialization of Executor (#893)

* Only terminate tasks on queue when cancel_futures is set to true

* Add cancel_futures_on_shutdown to FileTaskScheduler

* Add cancel_futures_on_shutdown in __init__()

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update flux.py

* Update slurm.py

* Update test_standalone_inputcheck.py

* fix

* fixes

* try different tests

* rename from cancel_futures_on_shutdown to wait

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* continue renaming

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

src/executorlib/executor/flux.py

@@ -9,6 +9,7 @@
     check_plot_dependency_graph,
     check_pmi,
     check_refresh_rate,
+    check_wait_on_shutdown,
     validate_number_of_cores,
 )
 from executorlib.task_scheduler.interactive.blockallocation import (
@@ -67,6 +68,7 @@ class FluxJobExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -108,6 +110,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.FluxJobExecutor leverages either the message passing interface (MPI), the SLURM workload manager
@@ -156,6 +159,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -187,6 +191,7 @@ def __init__(
                         block_allocation=block_allocation,
                         init_function=init_function,
                         log_obj_size=log_obj_size,
+                        wait=wait,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -212,6 +217,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     log_obj_size=log_obj_size,
+                    wait=wait,
                 )
             )
 
@@ -261,6 +267,7 @@ class FluxClusterExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -300,6 +307,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.FluxClusterExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -346,6 +354,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -405,6 +414,7 @@ def __init__(
                         block_allocation=block_allocation,
                         init_function=init_function,
                         disable_dependencies=disable_dependencies,
+                        wait=wait,
                     )
                 )
         else:
@@ -445,6 +455,7 @@ def create_flux_executor(
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
+    wait: bool = True,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a flux executor
@@ -483,6 +494,7 @@ def create_flux_executor(
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -504,6 +516,7 @@ def create_flux_executor(
     check_command_line_argument_lst(
         command_line_argument_lst=resource_dict.get("slurm_cmd_args", [])
     )
+    check_wait_on_shutdown(wait_on_shutdown=wait)
     if "openmpi_oversubscribe" in resource_dict:
         del resource_dict["openmpi_oversubscribe"]
     if "slurm_cmd_args" in resource_dict:

src/executorlib/executor/single.py

@@ -7,6 +7,7 @@
     check_init_function,
     check_plot_dependency_graph,
     check_refresh_rate,
+    check_wait_on_shutdown,
     validate_number_of_cores,
 )
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
@@ -60,6 +61,7 @@ class SingleNodeExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -97,6 +99,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.SingleNodeExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -142,6 +145,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -169,6 +173,7 @@ def __init__(
                         block_allocation=block_allocation,
                         init_function=init_function,
                         log_obj_size=log_obj_size,
+                        wait=wait,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -190,6 +195,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     log_obj_size=log_obj_size,
+                    wait=wait,
                 )
             )
 
@@ -232,6 +238,7 @@ class TestClusterExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -269,6 +276,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.api.TestClusterExecutor is designed to test the file based communication used in the
@@ -307,6 +315,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -346,6 +355,7 @@ def __init__(
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
                     execute_function=execute_in_subprocess,
+                    wait=wait,
                 )
             )
         else:
@@ -379,6 +389,7 @@ def create_single_node_executor(
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
+    wait: bool = True,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a single node executor
@@ -413,6 +424,7 @@ def create_single_node_executor(
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -429,6 +441,7 @@ def create_single_node_executor(
     check_command_line_argument_lst(
         command_line_argument_lst=resource_dict.get("slurm_cmd_args", [])
     )
+    check_wait_on_shutdown(wait_on_shutdown=wait)
     if "threads_per_core" in resource_dict:
         del resource_dict["threads_per_core"]
     if "gpus_per_core" in resource_dict:

src/executorlib/executor/slurm.py

@@ -6,6 +6,7 @@
     check_log_obj_size,
     check_plot_dependency_graph,
     check_refresh_rate,
+    check_wait_on_shutdown,
     validate_number_of_cores,
 )
 from executorlib.task_scheduler.interactive.blockallocation import (
@@ -65,6 +66,7 @@ class SlurmClusterExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -104,6 +106,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.SlurmClusterExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -150,6 +153,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -210,6 +214,7 @@ def __init__(
                         block_allocation=block_allocation,
                         init_function=init_function,
                         disable_dependencies=disable_dependencies,
+                        wait=wait,
                     )
                 )
         else:
@@ -281,6 +286,7 @@ class SlurmJobExecutor(BaseExecutor):
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         export_workflow_filename (str): Name of the file to store the exported workflow graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Examples:
         ```
@@ -319,6 +325,7 @@ def __init__(
         plot_dependency_graph_filename: Optional[str] = None,
         export_workflow_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        wait: bool = True,
     ):
         """
         The executorlib.SlurmJobExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -368,6 +375,7 @@ def __init__(
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             export_workflow_filename (str): Name of the file to store the exported workflow graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
         """
         default_resource_dict: dict = {
@@ -396,6 +404,7 @@ def __init__(
                         block_allocation=block_allocation,
                         init_function=init_function,
                         log_obj_size=log_obj_size,
+                        wait=wait,
                     ),
                     max_cores=max_cores,
                     refresh_rate=refresh_rate,
@@ -418,6 +427,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     log_obj_size=log_obj_size,
+                    wait=wait,
                 )
             )
 
@@ -432,6 +442,7 @@ def create_slurm_executor(
     block_allocation: bool = False,
     init_function: Optional[Callable] = None,
     log_obj_size: bool = False,
+    wait: bool = True,
 ) -> Union[OneProcessTaskScheduler, BlockAllocationTaskScheduler]:
     """
     Create a SLURM executor
@@ -471,6 +482,7 @@ def create_slurm_executor(
                                     of the individual function.
         init_function (None): optional function to preset arguments for functions which are submitted later
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        wait (bool): Whether to wait for the completion of all tasks before shutting down the executor.
 
     Returns:
         InteractiveStepExecutor/ InteractiveExecutor
@@ -483,6 +495,7 @@ def create_slurm_executor(
     resource_dict["log_obj_size"] = log_obj_size
     resource_dict["pmi_mode"] = pmi_mode
     check_init_function(block_allocation=block_allocation, init_function=init_function)
+    check_wait_on_shutdown(wait_on_shutdown=wait)
     if block_allocation:
         resource_dict["init_function"] = init_function
         max_workers = validate_number_of_cores(

src/executorlib/standalone/inputcheck.py

@@ -17,6 +17,18 @@ def check_oversubscribe(oversubscribe: bool) -> None:
         )
 
 
+def check_wait_on_shutdown(
+    wait_on_shutdown: bool,
+) -> None:
+    """
+    Check if wait_on_shutdown is False and raise a ValueError if it is.
+    """
+    if not wait_on_shutdown:
+        raise ValueError(
+            "The wait_on_shutdown parameter is only supported for the executorlib.FluxClusterExecutor and executorlib.SlurmClusterExecutor."
+        )
+
+
 def check_command_line_argument_lst(command_line_argument_lst: list[str]) -> None:
     """
     Check if command_line_argument_lst is not empty and raise a ValueError if it is.

src/executorlib/task_scheduler/base.py

@@ -198,7 +198,9 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
         if cancel_futures and self._future_queue is not None:
             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})
+            self._future_queue.put(
+                {"shutdown": True, "wait": wait, "cancel_futures": cancel_futures}
+            )
             if wait and isinstance(self._process, Thread):
                 self._process.join()
                 self._future_queue.join()