fixes

Author: jan-janssen

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,16 +65,11 @@ def serialize_funct_h5(
         fn_kwargs = {}
     if resource_dict is None:
         resource_dict = {}
-    if "cache_key" in resource_dict:
-        task_key = resource_dict["cache_key"]
+    if cache_key is not None:
+        task_key = cache_key
     else:
         binary_all = cloudpickle.dumps(
-            {
-                "fn": fn,
-                "args": fn_args,
-                "kwargs": fn_kwargs,
-                "resource_dict": resource_dict,
-            }
+            {"fn": fn, "args": fn_args, "kwargs": fn_kwargs, "resource_dict": resource_dict}
         )
         task_key = fn.__name__ + _get_hash(binary=binary_all)
     data = {

executorlib/task_scheduler/file/shared.py

@@ -72,6 +72,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,11 +102,13 @@ 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 not (

executorlib/task_scheduler/interactive/shared.py

@@ -22,6 +22,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,
     **kwargs,
 ) -> None:
@@ -41,6 +42,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 "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.
     """
     interface = interface_bootup(
@@ -73,6 +76,7 @@ def execute_tasks(
                     task_dict=task_dict,
                     future_queue=future_queue,
                     cache_directory=cache_directory,
+                    cache_key=cache_key,
                 )
 
 
@@ -129,6 +133,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.
@@ -139,6 +144,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
 
@@ -147,6 +154,7 @@ 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(os.path.join(cache_directory, task_key), exist_ok=True)
     file_name = os.path.join(cache_directory, task_key, "cache.h5out")

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": "abc"})
+            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 = "./cache"
+        with SingleNodeExecutor(cache_directory=cache_directory) as exe:
+            self.assertTrue(exe)
+            future_lst = [exe.submit(sum, [i, i], resource_dict={"cache_key": "same_" + 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 = "./cache_error"
         with SingleNodeExecutor(cache_directory=cache_directory) as exe: