[Feature] Add option to overwrite cache_key (#676)

* [Feature] Add option to overwrite cache_key

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

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

* fixes

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

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

* recursively iterate over folders

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

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

* fixes

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

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

* [Bug] Fix folder generation in dump() function

* fix

---------

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

Author: jan-janssen

executorlib/standalone/cache.py

@@ -13,6 +13,23 @@
 }
 
 
+def get_cache_files(cache_directory: str) -> list[str]:
+    """
+    Recursively find all HDF5 files in the cache_directory which contain outputs.
+
+    Args:
+        cache_directory (str): The directory to store cache files.
+
+    Returns:
+        list[str]: List of HDF5 file in the cache directory which contain outputs.
+    """
+    file_lst = []
+    cache_directory_abs = os.path.abspath(cache_directory)
+    for dirpath, _, filenames in os.walk(cache_directory_abs):
+        file_lst += [os.path.join(dirpath, f) for f in filenames if f.endswith("_o.h5")]
+    return file_lst
+
+
 def get_cache_data(cache_directory: str) -> list[dict]:
     """
     Collect all HDF5 files in the cache directory
@@ -27,15 +44,13 @@ def get_cache_data(cache_directory: str) -> list[dict]:
     import numpy as np
 
     file_lst = []
-    for task_key in os.listdir(cache_directory):
-        file_name = os.path.join(cache_directory, task_key)
-        if task_key[-5:] == "_o.h5":
-            with h5py.File(file_name, "r") as hdf:
-                file_content_dict = {
-                    key: cloudpickle.loads(np.void(hdf["/" + key]))
-                    for key in group_dict.values()
-                    if key in hdf
-                }
-            file_content_dict["filename"] = file_name
-            file_lst.append(file_content_dict)
+    for file_name in get_cache_files(cache_directory=cache_directory):
+        with h5py.File(file_name, "r") as hdf:
+            file_content_dict = {
+                key: cloudpickle.loads(np.void(hdf["/" + key]))
+                for key in group_dict.values()
+                if key in hdf
+            }
+        file_content_dict["filename"] = file_name
+        file_lst.append(file_content_dict)
     return file_lst

executorlib/standalone/serialize.py

@@ -33,6 +33,7 @@ def serialize_funct_h5(
     fn_args: Optional[list] = None,
     fn_kwargs: Optional[dict] = None,
     resource_dict: Optional[dict] = None,
+    cache_key: Optional[str] = None,
 ) -> tuple[str, dict]:
     """
     Serialize a function and its arguments and keyword arguments into an HDF5 file.
@@ -51,6 +52,8 @@ def serialize_funct_h5(
                                   executor: None,
                                   hostname_localhost: False,
                               }
+        cache_key (str, optional): By default the cache_key is generated based on the function hash, this can be
+                                   overwritten by setting the cache_key.
 
     Returns:
         Tuple[str, dict]: A tuple containing the task key and the serialized data.
@@ -62,8 +65,17 @@ def serialize_funct_h5(
         fn_kwargs = {}
     if resource_dict is None:
         resource_dict = {}
-    binary_all = cloudpickle.dumps({"fn": fn, "args": fn_args, "kwargs": fn_kwargs})
-    task_key = fn.__name__ + _get_hash(binary=binary_all)
+    if cache_key is not None:
+        task_key = cache_key
+    else:
+        binary_all = cloudpickle.dumps(
+            {
+                "fn": fn,
+                "args": fn_args,
+                "kwargs": fn_kwargs,
+            }
+        )
+        task_key = fn.__name__ + _get_hash(binary=binary_all)
     data = {
         "fn": fn,
         "args": fn_args,

executorlib/task_scheduler/file/hdf.py

@@ -1,3 +1,4 @@
+import os
 from typing import Any, Optional
 
 import cloudpickle
@@ -16,7 +17,9 @@ def dump(file_name: Optional[str], data_dict: dict) -> None:
         data_dict (dict): dictionary containing the python function to be executed {"fn": ..., "args": (), "kwargs": {}}
     """
     if file_name is not None:
-        with h5py.File(file_name, "a") as fname:
+        file_name_abs = os.path.abspath(file_name)
+        os.makedirs(os.path.dirname(file_name_abs), exist_ok=True)
+        with h5py.File(file_name_abs, "a") as fname:
             for data_key, data_value in data_dict.items():
                 if data_key in group_dict:
                     fname.create_dataset(

executorlib/task_scheduler/file/shared.py

@@ -6,6 +6,7 @@
 from concurrent.futures import Future
 from typing import Any, Callable, Optional
 
+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
@@ -72,6 +73,7 @@ def execute_tasks_h5(
         terminate_function (Callable): The function to terminate the tasks.
         pysqa_config_directory (str, optional): path to the pysqa config directory (only for pysqa based backend).
         backend (str, optional): name of the backend used to spawn tasks.
+        disable_dependencies (boolean): Disable resolving future objects during the submission.
 
     Returns:
         None
@@ -101,15 +103,18 @@ def execute_tasks_h5(
             task_resource_dict.update(
                 {k: v for k, v in resource_dict.items() if k not in task_resource_dict}
             )
+            cache_key = task_resource_dict.pop("cache_key", None)
             task_key, data_dict = serialize_funct_h5(
                 fn=task_dict["fn"],
                 fn_args=task_args,
                 fn_kwargs=task_kwargs,
                 resource_dict=task_resource_dict,
+                cache_key=cache_key,
             )
             if task_key not in memory_dict:
-                if task_key + "_o.h5" not in os.listdir(cache_directory):
-                    os.makedirs(cache_directory, exist_ok=True)
+                if task_key + "_o.h5" not in get_cache_files(
+                    cache_directory=cache_directory
+                ):
                     file_name = os.path.join(cache_directory, task_key + "_i.h5")
                     dump(file_name=file_name, data_dict=data_dict)
                     if not disable_dependencies:

executorlib/task_scheduler/interactive/shared.py

@@ -6,6 +6,7 @@
 import time
 from typing import Callable, Optional
 
+from executorlib.standalone.cache import get_cache_files
 from executorlib.standalone.command import get_command_path
 from executorlib.standalone.interactive.communication import (
     SocketInterface,
@@ -22,6 +23,7 @@ def execute_tasks(
     hostname_localhost: Optional[bool] = None,
     init_function: Optional[Callable] = None,
     cache_directory: Optional[str] = None,
+    cache_key: Optional[str] = None,
     queue_join_on_shutdown: bool = True,
     log_obj_size: bool = False,
     **kwargs,
@@ -42,6 +44,8 @@ def execute_tasks(
                                      option to true
        init_function (Callable): optional function to preset arguments for functions which are submitted later
        cache_directory (str, optional): The directory to store cache files. Defaults to "executorlib_cache".
+       cache_key (str, optional): By default the cache_key is generated based on the function hash, this can be
+                                  overwritten by setting the cache_key.
        queue_join_on_shutdown (bool): Join communication queue when thread is closed. Defaults to True.
        log_obj_size (bool): Enable debug mode which reports the size of the communicated objects.
     """
@@ -76,6 +80,7 @@ def execute_tasks(
                     task_dict=task_dict,
                     future_queue=future_queue,
                     cache_directory=cache_directory,
+                    cache_key=cache_key,
                 )
 
 
@@ -132,6 +137,7 @@ def _execute_task_with_cache(
     task_dict: dict,
     future_queue: queue.Queue,
     cache_directory: str,
+    cache_key: Optional[str] = None,
 ):
     """
     Execute the task in the task_dict by communicating it via the interface using the cache in the cache directory.
@@ -142,6 +148,8 @@ def _execute_task_with_cache(
                           {"fn": Callable, "args": (), "kwargs": {}, "resource_dict": {}}
         future_queue (Queue): Queue for receiving new tasks.
         cache_directory (str): The directory to store cache files.
+        cache_key (str, optional): By default the cache_key is generated based on the function hash, this can be
+                                  overwritten by setting the cache_key.
     """
     from executorlib.task_scheduler.file.hdf import dump, get_output
 
@@ -150,10 +158,11 @@ def _execute_task_with_cache(
         fn_args=task_dict["args"],
         fn_kwargs=task_dict["kwargs"],
         resource_dict=task_dict.get("resource_dict", {}),
+        cache_key=cache_key,
     )
     os.makedirs(cache_directory, exist_ok=True)
-    file_name = os.path.join(cache_directory, task_key + "_o.h5")
-    if task_key + "_o.h5" not in os.listdir(cache_directory):
+    file_name = os.path.abspath(os.path.join(cache_directory, task_key + "_o.h5"))
+    if file_name not in get_cache_files(cache_directory=cache_directory):
         f = task_dict.pop("future")
         if f.set_running_or_notify_cancel():
             try:

tests/test_cache_fileexecutor_serial.py

@@ -42,6 +42,13 @@ def test_executor_mixed(self):
             self.assertEqual(fs1.result(), 3)
             self.assertTrue(fs1.done())
 
+    def test_executor_mixed_cache_key(self):
+        with FileTaskScheduler(execute_function=execute_in_subprocess) as exe:
+            fs1 = exe.submit(my_funct, 1, b=2, resource_dict={"cache_key": "a/b/c"})
+            self.assertFalse(fs1.done())
+            self.assertEqual(fs1.result(), 3)
+            self.assertTrue(fs1.done())
+
     def test_executor_dependence_mixed(self):
         with FileTaskScheduler(execute_function=execute_in_subprocess) as exe:
             fs1 = exe.submit(my_funct, 1, b=2)

tests/test_singlenodeexecutor_cache.py

@@ -34,6 +34,21 @@ def test_cache_data(self):
             sum([sum(c["input_args"][0]) for c in cache_lst]), sum(result_lst)
         )
 
+    def test_cache_key(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        with SingleNodeExecutor(cache_directory=cache_directory) as exe:
+            self.assertTrue(exe)
+            future_lst = [exe.submit(sum, [i, i], resource_dict={"cache_key": "same/j" + str(i)}) for i in range(1, 4)]
+            result_lst = [f.result() for f in future_lst]
+
+        cache_lst = get_cache_data(cache_directory=cache_directory)
+        for entry in cache_lst:
+            self.assertTrue("same" in entry['filename'])
+        self.assertEqual(sum([c["output"] for c in cache_lst]), sum(result_lst))
+        self.assertEqual(
+            sum([sum(c["input_args"][0]) for c in cache_lst]), sum(result_lst)
+        )
+
     def test_cache_error(self):
         cache_directory = os.path.abspath("cache_error")
         with SingleNodeExecutor(cache_directory=cache_directory) as exe: