Move cleanup of previous jobs into execute function

executorlib/executor/flux.py

@@ -9,7 +9,6 @@
     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 (
@@ -64,7 +63,6 @@ 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:
         ```
@@ -105,7 +103,6 @@ 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
@@ -151,7 +148,6 @@ 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 = {
@@ -167,9 +163,6 @@ 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(
@@ -255,7 +248,6 @@ class FluxClusterExecutor(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:
         ```
@@ -336,7 +328,6 @@ 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 = {
@@ -376,7 +367,6 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
-                    terminate_tasks_on_shutdown=terminate_tasks_on_shutdown,
                 )
             )
         else:

executorlib/executor/slurm.py

@@ -6,7 +6,6 @@
     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 (
@@ -61,7 +60,6 @@ class SlurmClusterExecutor(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:
         ```
@@ -142,7 +140,6 @@ 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 = {
@@ -182,7 +179,6 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
-                    terminate_tasks_on_shutdown=terminate_tasks_on_shutdown,
                 )
             )
         else:
@@ -249,7 +245,6 @@ 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:
         ```
@@ -332,7 +327,6 @@ 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 = {
@@ -348,9 +342,6 @@ 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

@@ -212,15 +212,3 @@ def check_log_obj_size(log_obj_size: bool) -> None:
             "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/hdf.py

@@ -101,7 +101,7 @@ def get_queue_id(file_name: Optional[str]) -> Optional[int]:
     Returns:
         int: queuing system id from the execution of the python function
     """
-    if file_name is not None:
+    if file_name is not None and os.path.exists(file_name):
         with h5py.File(file_name, "r") as hdf:
             if "queue_id" in hdf:
                 return cloudpickle.loads(np.void(hdf["/queue_id"]))

executorlib/task_scheduler/file/queue_spawner.py

@@ -10,9 +10,10 @@
 
 def execute_with_pysqa(
     command: list,
+    file_name: str,
+    data_dict: dict,
     cache_directory: str,
     task_dependent_lst: Optional[list[int]] = None,
-    file_name: Optional[str] = None,
     resource_dict: Optional[dict] = None,
     config_directory: Optional[str] = None,
     backend: Optional[str] = None,
@@ -22,9 +23,10 @@ def execute_with_pysqa(
 
     Args:
         command (list): The command to be executed.
+        file_name (str): Name of the HDF5 file which contains the Python function
+        data_dict (dict): dictionary containing the python function to be executed {"fn": ..., "args": (), "kwargs": {}}
         cache_directory (str): The directory to store the HDF5 files.
         task_dependent_lst (list): A list of subprocesses that the current subprocess depends on. Defaults to [].
-        file_name (str): Name of the HDF5 file which contains the Python function
         resource_dict (dict): resource dictionary, which defines the resources used for the execution of the function.
                               Example resource dictionary: {
                                   cwd: None,
@@ -37,13 +39,24 @@ def execute_with_pysqa(
     """
     if task_dependent_lst is None:
         task_dependent_lst = []
-    check_file_exists(file_name=file_name)
-    queue_id = get_queue_id(file_name=file_name)
     qa = QueueAdapter(
         directory=config_directory,
         queue_type=backend,
         execute_command=_pysqa_execute_command,
     )
+    queue_id = get_queue_id(file_name=file_name)
+    if (
+        os.path.exists(file_name)
+        and queue_id is None
+        or os.path.exists(file_name)
+        and queue_id is not None
+        and qa.get_status_of_job(process_id=queue_id) is None
+    ):
+        os.remove(file_name)
+        dump(file_name=file_name, data_dict=data_dict)
+    elif not os.path.exists(file_name):
+        dump(file_name=file_name, data_dict=data_dict)
+    check_file_exists(file_name=file_name)
     if queue_id is None or qa.get_status_of_job(process_id=queue_id) is None:
         if resource_dict is None:
             resource_dict = {}

executorlib/task_scheduler/file/shared.py

@@ -9,7 +9,7 @@
 from executorlib.standalone.cache import get_cache_files
 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.hdf import get_output
 from executorlib.task_scheduler.file.subprocess_spawner import terminate_subprocess
 
 
@@ -87,16 +87,17 @@ 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"]:
-            while len(memory_dict) > 0:
-                memory_dict = {
-                    key: _check_task_output(
-                        task_key=key,
-                        future_obj=value,
-                        cache_directory=cache_dir_dict[key],
-                    )
-                    for key, value in memory_dict.items()
-                    if not value.done()
-                }
+            if task_dict["wait"]:
+                while len(memory_dict) > 0:
+                    memory_dict = {
+                        key: _check_task_output(
+                            task_key=key,
+                            future_obj=value,
+                            cache_directory=cache_dir_dict[key],
+                        )
+                        for key, value in memory_dict.items()
+                        if not value.done()
+                    }
             if (
                 terminate_function is not None
                 and terminate_function == terminate_subprocess
@@ -137,9 +138,6 @@ def execute_tasks_h5(
                     cache_directory, task_key + "_o.h5"
                 ) not in get_cache_files(cache_directory=cache_directory):
                     file_name = os.path.join(cache_directory, task_key + "_i.h5")
-                    if os.path.exists(file_name):
-                        os.remove(file_name)
-                    dump(file_name=file_name, data_dict=data_dict)
                     if not disable_dependencies:
                         task_dependent_lst = [
                             process_dict[k] for k in future_wait_key_lst
@@ -158,6 +156,7 @@ def execute_tasks_h5(
                             cores=task_resource_dict["cores"],
                         ),
                         file_name=file_name,
+                        data_dict=data_dict,
                         task_dependent_lst=task_dependent_lst,
                         resource_dict=task_resource_dict,
                         config_directory=pysqa_config_directory,

executorlib/task_scheduler/file/subprocess_spawner.py

@@ -1,40 +1,47 @@
+import os
 import subprocess
 import time
 from typing import Optional
 
 from executorlib.standalone.inputcheck import check_file_exists
+from executorlib.task_scheduler.file.hdf import dump
 
 
 def execute_in_subprocess(
     command: list,
+    file_name: str,
+    data_dict: dict,
+    cache_directory: Optional[str] = None,
     task_dependent_lst: Optional[list] = None,
-    file_name: Optional[str] = None,
     resource_dict: Optional[dict] = None,
     config_directory: Optional[str] = None,
     backend: Optional[str] = None,
-    cache_directory: Optional[str] = None,
 ) -> subprocess.Popen:
     """
     Execute a command in a subprocess.
 
     Args:
         command (list): The command to be executed.
-        task_dependent_lst (list): A list of subprocesses that the current subprocess depends on. Defaults to [].
         file_name (str): Name of the HDF5 file which contains the Python function
+        data_dict (dict): dictionary containing the python function to be executed {"fn": ..., "args": (), "kwargs": {}}
+        cache_directory (str): The directory to store the HDF5 files.
+        task_dependent_lst (list): A list of subprocesses that the current subprocess depends on. Defaults to [].
         resource_dict (dict): resource dictionary, which defines the resources used for the execution of the function.
                               Example resource dictionary: {
                                   cwd: None,
                               }
         config_directory (str, optional): path to the config directory.
         backend (str, optional): name of the backend used to spawn tasks.
-        cache_directory (str): The directory to store the HDF5 files.
 
     Returns:
         subprocess.Popen: The subprocess object.
 
     """
     if task_dependent_lst is None:
         task_dependent_lst = []
+    if os.path.exists(file_name):
+        os.remove(file_name)
+    dump(file_name=file_name, data_dict=data_dict)
     check_file_exists(file_name=file_name)
     while len(task_dependent_lst) > 0:
         task_dependent_lst = [

executorlib/task_scheduler/file/task_scheduler.py

@@ -95,7 +95,6 @@ def create_file_executor(
     init_function: Optional[Callable] = None,
     disable_dependencies: bool = False,
     execute_function: Callable = execute_with_pysqa,
-    terminate_tasks_on_shutdown: bool = True,
 ):
     if block_allocation:
         raise ValueError(
@@ -113,12 +112,10 @@ def create_file_executor(
     check_executor(executor=flux_executor)
     check_nested_flux_executor(nested_flux_executor=flux_executor_nesting)
     check_flux_log_files(flux_log_files=flux_log_files)
-    if terminate_tasks_on_shutdown and execute_function != execute_in_subprocess:
+    if execute_function != execute_in_subprocess:
         terminate_function = terminate_with_pysqa  # type: ignore
-    elif terminate_tasks_on_shutdown and execute_function == execute_in_subprocess:
-        terminate_function = terminate_subprocess  # type: ignore
     else:
-        terminate_function = None  # type: ignore
+        terminate_function = terminate_subprocess  # type: ignore
     return FileTaskScheduler(
         resource_dict=resource_dict,
         pysqa_config_directory=pysqa_config_directory,

tests/test_cache_fileexecutor_mpi.py

@@ -1,13 +1,11 @@
 import importlib.util
-import os
 import shutil
 import unittest
 
-from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
-
 
 try:
     from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler
+    from executorlib.task_scheduler.file.subprocess_spawner import execute_in_subprocess
 
     skip_h5py_test = False
 except ImportError:

tests/test_cache_fileexecutor_serial.py

@@ -5,12 +5,11 @@
 import unittest
 from threading import Thread
 
-from executorlib.task_scheduler.file.subprocess_spawner import (
-    execute_in_subprocess,
-    terminate_subprocess,
-)
-
 try:
+    from executorlib.task_scheduler.file.subprocess_spawner import (
+        execute_in_subprocess,
+        terminate_subprocess,
+    )
     from executorlib.task_scheduler.file.task_scheduler import FileTaskScheduler, create_file_executor
     from executorlib.task_scheduler.file.shared import execute_tasks_h5
 
@@ -203,10 +202,18 @@ def test_executor_function_dependence_args(self):
     def test_execute_in_subprocess_errors(self):
         with self.assertRaises(ValueError):
             execute_in_subprocess(
-                file_name=__file__, command=[], config_directory="test"
+                file_name=os.path.join(__file__, "..", "executorlib_cache", "test.h5"),
+                data_dict={},
+                command=[],
+                config_directory="test",
             )
         with self.assertRaises(ValueError):
-            execute_in_subprocess(file_name=__file__, command=[], backend="flux")
+            execute_in_subprocess(
+                file_name=os.path.join(__file__, "..", "executorlib_cache", "test.h5"),
+                data_dict={},
+                command=[],
+                backend="flux",
+            )
 
     def tearDown(self):
         shutil.rmtree("executorlib_cache", ignore_errors=True)

tests/test_standalone_inputcheck.py

@@ -18,7 +18,6 @@
     check_hostname_localhost,
     check_pysqa_config_directory,
     check_file_exists,
-    check_terminate_tasks_on_shutdown,
     check_log_obj_size,
     validate_number_of_cores,
 )
@@ -125,7 +124,3 @@ def test_validate_number_of_cores(self):
     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)