Add function to shutdown tasks

Author: jan-janssen

executorlib/executor/flux.py

@@ -4,10 +4,12 @@
 from executorlib.standalone.inputcheck import (
     check_command_line_argument_lst,
     check_init_function,
+    check_log_obj_size,
     check_oversubscribe,
     check_plot_dependency_graph,
     check_pmi,
     check_refresh_rate,
+    check_terminate_tasks_on_shutdown,
     validate_number_of_cores,
 )
 from executorlib.task_scheduler.interactive.blockallocation import (
@@ -62,6 +64,7 @@ class FluxJobExecutor(BaseExecutor):
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
     Examples:
         ```
@@ -102,6 +105,7 @@ def __init__(
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        terminate_tasks_on_shutdown: bool = True,
     ):
         """
         The executorlib.FluxJobExecutor leverages either the message passing interface (MPI), the SLURM workload manager
@@ -147,6 +151,7 @@ def __init__(
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
         """
         default_resource_dict: dict = {
@@ -162,6 +167,9 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_terminate_tasks_on_shutdown(
+            terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
+        )
         if not disable_dependencies:
             super().__init__(
                 executor=DependencyTaskScheduler(
@@ -246,6 +254,8 @@ class FluxClusterExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
     Examples:
         ```
@@ -282,6 +292,8 @@ def __init__(
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        log_obj_size: bool = False,
+        terminate_tasks_on_shutdown: bool = True,
     ):
         """
         The executorlib.FluxClusterExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -323,6 +335,8 @@ def __init__(
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
         """
         default_resource_dict: dict = {
@@ -338,12 +352,20 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_log_obj_size(log_obj_size=log_obj_size)
         if not plot_dependency_graph:
             import pysqa  # noqa
 
             from executorlib.task_scheduler.file.task_scheduler import (
                 create_file_executor,
             )
+            if terminate_tasks_on_shutdown:
+                from executorlib.task_scheduler.file.queue_spawner import (
+                    terminate_with_pysqa,
+                )
+                terminate_function = terminate_with_pysqa
+            else:
+                terminate_function = None
 
             super().__init__(
                 executor=create_file_executor(
@@ -361,6 +383,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
+                    terminate_function=terminate_function,
                 )
             )
         else:

executorlib/executor/slurm.py

@@ -3,8 +3,10 @@
 from executorlib.executor.base import BaseExecutor
 from executorlib.standalone.inputcheck import (
     check_init_function,
+    check_log_obj_size,
     check_plot_dependency_graph,
     check_refresh_rate,
+    check_terminate_tasks_on_shutdown,
     validate_number_of_cores,
 )
 from executorlib.task_scheduler.interactive.blockallocation import (
@@ -58,6 +60,8 @@ class SlurmClusterExecutor(BaseExecutor):
         plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
     Examples:
         ```
@@ -94,6 +98,8 @@ def __init__(
         refresh_rate: float = 0.01,
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
+        log_obj_size: bool = False,
+        terminate_tasks_on_shutdown: bool = True,
     ):
         """
         The executorlib.SlurmClusterExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -135,6 +141,8 @@ def __init__(
             plot_dependency_graph (bool): Plot the dependencies of multiple future objects without executing them. For
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
+            log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
         """
         default_resource_dict: dict = {
@@ -150,12 +158,20 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_log_obj_size(log_obj_size=log_obj_size)
         if not plot_dependency_graph:
             import pysqa  # noqa
 
             from executorlib.task_scheduler.file.task_scheduler import (
                 create_file_executor,
             )
+            if terminate_tasks_on_shutdown:
+                from executorlib.task_scheduler.file.queue_spawner import (
+                    terminate_with_pysqa,
+                )
+                terminate_function = terminate_with_pysqa
+            else:
+                terminate_function = None
 
             super().__init__(
                 executor=create_file_executor(
@@ -173,6 +189,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
+                    terminate_function=terminate_function,
                 )
             )
         else:
@@ -239,6 +256,7 @@ class SlurmJobExecutor(BaseExecutor):
                                       debugging purposes and to get an overview of the specified dependencies.
         plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
         log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+        terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
     Examples:
         ```
@@ -275,6 +293,7 @@ def __init__(
         plot_dependency_graph: bool = False,
         plot_dependency_graph_filename: Optional[str] = None,
         log_obj_size: bool = False,
+        terminate_tasks_on_shutdown: bool = True,
     ):
         """
         The executorlib.SlurmJobExecutor leverages either the message passing interface (MPI), the SLURM workload
@@ -320,6 +339,7 @@ def __init__(
                                           debugging purposes and to get an overview of the specified dependencies.
             plot_dependency_graph_filename (str): Name of the file to store the plotted graph in.
             log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
+            terminate_tasks_on_shutdown (bool): Shutdown all tasks when the Executor is shutdown, this is the default.
 
         """
         default_resource_dict: dict = {
@@ -335,6 +355,9 @@ def __init__(
         resource_dict.update(
             {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
         )
+        check_terminate_tasks_on_shutdown(
+            terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
+        )
         if not disable_dependencies:
             super().__init__(
                 executor=DependencyTaskScheduler(

executorlib/standalone/inputcheck.py

@@ -194,7 +194,33 @@ def validate_number_of_cores(
 
 
 def check_file_exists(file_name: Optional[str]):
+    """
+    Check if file exists and raise a ValueError if it does not or file_name is None.
+    """
     if file_name is None:
         raise ValueError("file_name is not set.")
     if not os.path.exists(file_name):
         raise ValueError("file_name is not written to the file system.")
+
+
+def check_log_obj_size(log_obj_size: bool) -> None:
+    """
+    Check if log_obj_size is True and raise a ValueError if it is.
+    """
+    if log_obj_size:
+        raise ValueError(
+            "log_obj_size is not supported for the executorlib.SlurmClusterExecutor and executorlib.FluxClusterExecutor."
+            "Please use log_obj_size=False instead of log_obj_size=True."
+        )
+
+
+def check_terminate_tasks_on_shutdown(terminate_tasks_on_shutdown: bool) -> None:
+    """
+    Check if terminate_tasks_on_shutdown is False and raise a ValueError if it is.
+    """
+    if not terminate_tasks_on_shutdown:
+        raise ValueError(
+            "terminate_tasks_on_shutdown is not supported for the executorlib.SingleNodeExecutor, "
+            "executorlib.SlurmJobExecutor and executorlib.FluxJobExecutor."
+            "Please use terminate_tasks_on_shutdown=True instead of terminate_tasks_on_shutdown=False."
+        )

executorlib/task_scheduler/file/queue_spawner.py

@@ -81,6 +81,27 @@ def execute_with_pysqa(
     return queue_id
 
 
+def terminate_with_pysqa(
+    queue_id: int,
+    config_directory: Optional[str] = None,
+    backend: Optional[str] = None,
+):
+    """
+    Delete job from queuing system
+
+    Args:
+        queue_id (int): Queuing system ID of the job to delete.
+        config_directory (str, optional): path to the config directory.
+        backend (str, optional): name of the backend used to spawn tasks.
+    """
+    qa = QueueAdapter(
+        directory=config_directory,
+        queue_type=backend,
+        execute_command=_pysqa_execute_command,
+    )
+    qa.delete_job(process_id=queue_id)
+
+
 def _pysqa_execute_command(
     commands: str,
     working_directory: Optional[str] = None,

executorlib/task_scheduler/file/shared.py

@@ -10,6 +10,9 @@
 from executorlib.standalone.command import get_command_path
 from executorlib.standalone.serialize import serialize_funct_h5
 from executorlib.task_scheduler.file.hdf import dump, get_output
+from executorlib.task_scheduler.file.subprocess_spawner import (
+    terminate_subprocess
+)
 
 
 class FutureItem:
@@ -86,9 +89,16 @@ def execute_tasks_h5(
         with contextlib.suppress(queue.Empty):
             task_dict = future_queue.get_nowait()
         if task_dict is not None and "shutdown" in task_dict and task_dict["shutdown"]:
-            if terminate_function is not None:
+            if terminate_function is not None and terminate_function == terminate_subprocess:
                 for task in process_dict.values():
                     terminate_function(task=task)
+            elif terminate_function is not None:
+                for queue_id in process_dict.values():
+                    terminate_function(
+                        queue_id=queue_id,
+                        config_directory=pysqa_config_directory,
+                        backend=backend,
+                    )
             future_queue.task_done()
             future_queue.join()
             break

tests/test_standalone_inputcheck.py

@@ -18,6 +18,8 @@
     check_hostname_localhost,
     check_pysqa_config_directory,
     check_file_exists,
+    check_terminate_tasks_on_shutdown,
+    check_log_obj_size,
     validate_number_of_cores,
 )
 
@@ -119,3 +121,11 @@ def test_validate_number_of_cores(self):
             ),
             int,
         )
+
+    def test_check_log_obj_size(self):
+        with self.assertRaises(ValueError):
+            check_log_obj_size(log_obj_size=True)
+
+    def test_terminate_tasks_on_shutdown(self):
+        with self.assertRaises(ValueError):
+            check_terminate_tasks_on_shutdown(terminate_tasks_on_shutdown=False)