Merge branch 'main' into Py314u

Author: jan-janssen

src/executorlib/__init__.py

@@ -12,6 +12,7 @@
 pandas.DataFrame.
 """
 
+from concurrent.futures import Future
 from typing import Optional
 
 import executorlib._version
@@ -43,30 +44,82 @@ def get_cache_data(cache_directory: str) -> list[dict]:
     return get_cache_data(cache_directory=cache_directory)
 
 
+def get_future_from_cache(
+    cache_directory: str,
+    cache_key: str,
+) -> Future:
+    """
+    Reload future from HDF5 file in cache directory with the given cache key. The function checks if the output file
+    exists, if not it checks for the input file. If neither of them exist, it raises a FileNotFoundError. If the output
+    file exists, it loads the output and sets it as the result of the future. If only the input file exists, it checks
+    if the execution is finished and if there was an error. If there was no error, it sets the output as the result of
+    the future, otherwise it raises the error.
+
+    Args:
+        cache_directory (str): The directory to store cache files.
+        cache_key (str): The key of the cache file to be reloaded.
+
+    Returns:
+        Future: Future object containing the result of the execution of the python function.
+    """
+    from executorlib.standalone.hdf import get_future_from_cache
+
+    return get_future_from_cache(
+        cache_directory=cache_directory,
+        cache_key=cache_key,
+    )
+
+
 def terminate_tasks_in_cache(
     cache_directory: str,
-    config_directory: Optional[str] = None,
+    pysqa_config_directory: Optional[str] = None,
     backend: Optional[str] = None,
 ):
     """
     Delete all jobs stored in the cache directory from the queuing system
 
     Args:
         cache_directory (str): The directory to store cache files.
-        config_directory (str, optional): path to the config directory.
+        pysqa_config_directory (str, optional): path to the pysqa config directory.
         backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"].
     """
     from executorlib.task_scheduler.file.spawner_pysqa import terminate_tasks_in_cache
 
     return terminate_tasks_in_cache(
         cache_directory=cache_directory,
-        config_directory=config_directory,
+        pysqa_config_directory=pysqa_config_directory,
+        backend=backend,
+    )
+
+
+def terminate_task_in_cache(
+    cache_directory: str,
+    cache_key: str,
+    pysqa_config_directory: Optional[str] = None,
+    backend: Optional[str] = None,
+):
+    """
+    Delete a specific job stored in the cache directory from the queuing system
+
+    Args:
+        cache_directory (str): The directory to store cache files.
+        cache_key (str): The key of the cache file to be deleted.
+        pysqa_config_directory (str, optional): path to the pysqa config directory.
+        backend (str, optional): name of the backend used to spawn tasks ["slurm", "flux"].
+    """
+    from executorlib.task_scheduler.file.spawner_pysqa import terminate_task_in_cache
+
+    return terminate_task_in_cache(
+        cache_directory=cache_directory,
+        cache_key=cache_key,
+        pysqa_config_directory=pysqa_config_directory,
         backend=backend,
     )
 
 
 __all__: list[str] = [
     "get_cache_data",
+    "get_future_from_cache",
     "get_item_from_future",
     "split_future",
     "terminate_tasks_in_cache",

src/executorlib/backend/cache_parallel.py

@@ -39,11 +39,11 @@ def main() -> None:
 
     time_start = time.time()
     apply_dict = {}
-    if mpi_rank_zero:
-        apply_dict = backend_load_file(file_name=file_name)
-    apply_dict = MPI.COMM_WORLD.bcast(apply_dict, root=0)
-    output = apply_dict["fn"].__call__(*apply_dict["args"], **apply_dict["kwargs"])
     try:
+        if mpi_rank_zero:
+            apply_dict = backend_load_file(file_name=file_name)
+        apply_dict = MPI.COMM_WORLD.bcast(apply_dict, root=0)
+        output = apply_dict["fn"].__call__(*apply_dict["args"], **apply_dict["kwargs"])
         result = (
             MPI.COMM_WORLD.gather(output, root=0) if mpi_size_larger_one else output
         )

src/executorlib/executor/flux.py

@@ -415,6 +415,7 @@ def __init__(
                         init_function=init_function,
                         disable_dependencies=disable_dependencies,
                         wait=wait,
+                        refresh_rate=refresh_rate,
                     )
                 )
         else:

src/executorlib/executor/single.py

@@ -356,6 +356,7 @@ def __init__(
                     disable_dependencies=disable_dependencies,
                     execute_function=execute_in_subprocess,
                     wait=wait,
+                    refresh_rate=refresh_rate,
                 )
             )
         else:

src/executorlib/executor/slurm.py

@@ -215,6 +215,7 @@ def __init__(
                         init_function=init_function,
                         disable_dependencies=disable_dependencies,
                         wait=wait,
+                        refresh_rate=refresh_rate,
                     )
                 )
         else:

src/executorlib/standalone/hdf.py

@@ -1,4 +1,5 @@
 import os
+from concurrent.futures import Future
 from typing import Any, Optional
 
 import cloudpickle
@@ -161,6 +162,43 @@ def get_cache_files(cache_directory: str) -> list[str]:
     return file_lst
 
 
+def get_future_from_cache(
+    cache_directory: str,
+    cache_key: str,
+) -> Future:
+    """
+    Reload future from HDF5 file in cache directory with the given cache key. The function checks if the output file
+    exists, if not it checks for the input file. If neither of them exist, it raises a FileNotFoundError. If the output
+    file exists, it loads the output and sets it as the result of the future. If only the input file exists, it checks
+    if the execution is finished and if there was an error. If there was no error, it sets the output as the result of
+    the future, otherwise it raises the error.
+
+    Args:
+        cache_directory (str): The directory to store cache files.
+        cache_key (str): The key of the cache file to be reloaded.
+
+    Returns:
+        Future: Future object containing the result of the execution of the python function.
+    """
+    file_name_in = os.path.join(cache_directory, cache_key + "_i.h5")
+    file_name_out = os.path.join(cache_directory, cache_key + "_o.h5")
+    future: Future = Future()
+    if os.path.exists(file_name_out):
+        file_name = file_name_out
+    elif os.path.exists(file_name_in):
+        file_name = file_name_in
+    else:
+        raise FileNotFoundError(
+            f"Neither input nor output file for cache key {cache_key} found in cache directory {cache_directory}."
+        )
+    exec_flag, no_error_flag, result = get_output(file_name=file_name)
+    if exec_flag and no_error_flag:
+        future.set_result(result)
+    elif exec_flag:
+        raise result
+    return future
+
+
 def _get_content_of_file(file_name: str) -> dict:
     """
     Get content of an HDF5 file

src/executorlib/task_scheduler/base.py

@@ -201,7 +201,7 @@ def shutdown(self, wait: bool = True, *, cancel_futures: bool = False):
             self._future_queue.put(
                 {"shutdown": True, "wait": wait, "cancel_futures": cancel_futures}
             )
-            if wait and isinstance(self._process, Thread):
+            if isinstance(self._process, Thread):
                 self._process.join()
                 self._future_queue.join()
         self._process = None

src/executorlib/task_scheduler/file/backend.py

@@ -68,9 +68,9 @@ def backend_execute_task_in_file(file_name: str) -> None:
     Returns:
         None
     """
-    apply_dict = backend_load_file(file_name=file_name)
-    time_start = time.time()
     try:
+        apply_dict = backend_load_file(file_name=file_name)
+        time_start = time.time()
         result = {
             "result": apply_dict["fn"].__call__(
                 *apply_dict["args"], **apply_dict["kwargs"]