fix cache directory in resource_dict

Author: jan-janssen

executorlib/task_scheduler/file/shared.py

@@ -52,7 +52,6 @@ def done(self) -> bool:
 
 def execute_tasks_h5(
     future_queue: queue.Queue,
-    cache_directory: str,
     execute_function: Callable,
     resource_dict: dict,
     terminate_function: Optional[Callable] = None,
@@ -65,7 +64,6 @@ def execute_tasks_h5(
 
     Args:
         future_queue (queue.Queue): The queue containing the tasks.
-        cache_directory (str): The directory to store the HDF5 files.
         resource_dict (dict): A dictionary of resources required by the task. With the following keys:
                               - cores (int): number of MPI cores to be used for each function call
                               - cwd (str/None): current working directory where the parallel python task is executed
@@ -104,6 +102,7 @@ def execute_tasks_h5(
                 {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)
+            cache_directory = task_resource_dict.pop("cache_directory")
             task_key, data_dict = serialize_funct_h5(
                 fn=task_dict["fn"],
                 fn_args=task_args,
@@ -146,15 +145,23 @@ def execute_tasks_h5(
                 file_name_dict[task_key] = os.path.join(
                     cache_directory, task_key + "_o.h5"
                 )
-                memory_dict[task_key] = task_dict["future"]
+                memory_dict[task_key] = {
+                    "future": task_dict["future"],
+                    "cache_directory": cache_directory,
+                }
             future_queue.task_done()
         else:
             memory_dict = {
-                key: _check_task_output(
-                    task_key=key, future_obj=value, cache_directory=cache_directory
-                )
+                key: {
+                    "future": _check_task_output(
+                        task_key=key,
+                        future_obj=value["future"],
+                        cache_directory=value["cache_directory"],
+                    ),
+                    "cache_directory": value["cache_directory"],
+                }
                 for key, value in memory_dict.items()
-                if not value.done()
+                if not value["future"].done()
             }
 
 

executorlib/task_scheduler/file/task_scheduler.py

@@ -27,7 +27,7 @@
 class FileTaskScheduler(TaskSchedulerBase):
     def __init__(
         self,
-        cache_directory: str = "executorlib_cache",
+        cache_directory: Optional[str] = None,
         resource_dict: Optional[dict] = None,
         execute_function: Callable = execute_with_pysqa,
         terminate_function: Optional[Callable] = None,
@@ -53,6 +53,7 @@ def __init__(
         default_resource_dict = {
             "cores": 1,
             "cwd": None,
+            "cache_directory": cache_directory,
         }
         if resource_dict is None:
             resource_dict = {}
@@ -61,12 +62,9 @@ def __init__(
         )
         if execute_function == execute_in_subprocess and terminate_function is None:
             terminate_function = terminate_subprocess
-        cache_directory_path = os.path.abspath(cache_directory)
-        os.makedirs(cache_directory_path, exist_ok=True)
         self._process_kwargs = {
             "future_queue": self._future_queue,
             "execute_function": execute_function,
-            "cache_directory": cache_directory_path,
             "resource_dict": resource_dict,
             "terminate_function": terminate_function,
             "pysqa_config_directory": pysqa_config_directory,
@@ -98,7 +96,9 @@ def create_file_executor(
     disable_dependencies: bool = False,
 ):
     if cache_directory is None:
-        cache_directory = "executorlib_cache"
+        resource_dict["cache_directory"] = os.path.abspath("executorlib_cache")
+    else:
+        resource_dict["cache_directory"] = os.path.abspath(cache_directory)
     if block_allocation:
         raise ValueError(
             "The option block_allocation is not available with the pysqa based backend."
@@ -114,7 +114,6 @@ def create_file_executor(
     check_nested_flux_executor(nested_flux_executor=flux_executor_nesting)
     check_flux_log_files(flux_log_files=flux_log_files)
     return FileTaskScheduler(
-        cache_directory=cache_directory,
         resource_dict=resource_dict,
         pysqa_config_directory=pysqa_config_directory,
         backend=backend.split("_submission")[0],