Implement resource_dict for file executor

Author: jan-janssen

executorlib/__init__.py

@@ -138,9 +138,9 @@ def __new__(
             cache_directory (str, optional): The directory to store cache files. Defaults to "cache".
             max_cores (int): defines the number cores which can be used in parallel
             resource_dict (dict): A dictionary of resources required by the task. With the following keys:
-                                  - cores_per_worker (int): number of MPI cores to be used for each function call
+                                  - cores (int): number of MPI cores to be used for each function call
                                   - threads_per_core (int): number of OpenMP threads to be used for each function call
-                                  - gpus_per_worker (int): number of GPUs per worker - defaults to 0
+                                  - gpus_per_core (int): number of GPUs per worker - defaults to 0
                                   - cwd (str/None): current working directory where the parallel python task is executed
                                   - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI
                                                                   and SLURM only) - default False

executorlib/cache/executor.py

@@ -14,8 +14,7 @@ class FileExecutor(ExecutorBase):
     def __init__(
         self,
         cache_directory: str = "cache",
-        cores_per_worker: int = 1,
-        cwd: Optional[str] = None,
+        resource_dict: Optional[dict] = None,
         execute_function: callable = execute_in_subprocess,
         terminate_function: Optional[callable] = None,
         config_directory: Optional[str] = None,
@@ -26,6 +25,9 @@ def __init__(
 
         Args:
             cache_directory (str, optional): The directory to store cache files. Defaults to "cache".
+            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
             execute_function (callable, optional): The function to execute tasks. Defaults to execute_in_subprocess.
             cores_per_worker (int, optional): The number of CPU cores per worker. Defaults to 1.
             terminate_function (callable, optional): The function to terminate the tasks.
@@ -34,6 +36,15 @@ def __init__(
             backend (str, optional): name of the backend used to spawn tasks.
         """
         super().__init__()
+        default_resource_dict = {
+            "cores": 1,
+            "cwd": None,
+        }
+        if resource_dict is None:
+            resource_dict = {}
+        resource_dict.update(
+            {k: v for k, v in default_resource_dict.items() if k not in resource_dict}
+        )
         if execute_function == execute_in_subprocess and terminate_function is None:
             terminate_function = terminate_subprocess
         cache_directory_path = os.path.abspath(cache_directory)
@@ -45,8 +56,7 @@ def __init__(
                     "future_queue": self._future_queue,
                     "execute_function": execute_function,
                     "cache_directory": cache_directory_path,
-                    "cores_per_worker": cores_per_worker,
-                    "cwd": cwd,
+                    "resource_dict": resource_dict,
                     "terminate_function": terminate_function,
                     "config_directory": config_directory,
                     "backend": backend,

executorlib/cache/shared.py

@@ -50,8 +50,7 @@ def execute_tasks_h5(
     future_queue: queue.Queue,
     cache_directory: str,
     execute_function: callable,
-    cores_per_worker: int = 1,
-    cwd: Optional[str] = None,
+    resource_dict: dict,
     terminate_function: Optional[callable] = None,
     config_directory: Optional[str] = None,
     backend: Optional[str] = None,
@@ -62,9 +61,10 @@ def execute_tasks_h5(
     Args:
         future_queue (queue.Queue): The queue containing the tasks.
         cache_directory (str): The directory to store the HDF5 files.
-        cores_per_worker (int): The number of cores per worker.
+        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
         execute_function (callable): The function to execute the tasks.
-        cwd (str/None): current working directory where the parallel python task is executed
         terminate_function (callable): The function to terminate the tasks.
         config_directory (str, optional): path to the config directory.
         backend (str, optional): name of the backend used to spawn tasks.
@@ -97,16 +97,15 @@ def execute_tasks_h5(
                 memory_dict=memory_dict,
                 file_name_dict=file_name_dict,
             )
-            resource_dict = task_dict["resource_dict"].copy()
-            if "cores" not in resource_dict:
-                resource_dict["cores"] = cores_per_worker
-            if "cwd" not in resource_dict:
-                resource_dict["cwd"] = cwd
+            task_resource_dict = task_dict["resource_dict"].copy()
+            task_resource_dict.update(
+                {k: v for k, v in resource_dict.items() if k not in task_resource_dict}
+            )
             task_key, data_dict = serialize_funct_h5(
                 fn=task_dict["fn"],
                 fn_args=task_args,
                 fn_kwargs=task_kwargs,
-                resource_dict=resource_dict,
+                resource_dict=task_resource_dict,
             )
             if task_key not in memory_dict.keys():
                 if task_key + ".h5out" not in os.listdir(cache_directory):
@@ -115,12 +114,12 @@ def execute_tasks_h5(
                     process_dict[task_key] = execute_function(
                         command=_get_execute_command(
                             file_name=file_name,
-                            cores=cores_per_worker,
+                            cores=task_resource_dict["cores"],
                         ),
                         task_dependent_lst=[
                             process_dict[k] for k in future_wait_key_lst
                         ],
-                        resource_dict=resource_dict,
+                        resource_dict=task_resource_dict,
                         config_directory=config_directory,
                         backend=backend,
                     )

executorlib/interactive/executor.py

@@ -175,9 +175,9 @@ def create_executor(
         max_cores (int): defines the number cores which can be used in parallel
         cache_directory (str, optional): The directory to store cache files. Defaults to "cache".
         resource_dict (dict): A dictionary of resources required by the task. With the following keys:
-                              - cores_per_worker (int): number of MPI cores to be used for each function call
+                              - cores (int): number of MPI cores to be used for each function call
                               - threads_per_core (int): number of OpenMP threads to be used for each function call
-                              - gpus_per_worker (int): number of GPUs per worker - defaults to 0
+                              - gpus_per_core (int): number of GPUs per worker - defaults to 0
                               - cwd (str/None): current working directory where the parallel python task is executed
                               - openmpi_oversubscribe (bool): adds the `--oversubscribe` command line flag (OpenMPI and
                                                               SLURM only) - default False

tests/test_cache_executor_mpi.py

@@ -29,7 +29,7 @@ def mpi_funct(i):
 )
 class TestCacheExecutorMPI(unittest.TestCase):
     def test_executor(self):
-        with FileExecutor(cores_per_worker=2) as exe:
+        with FileExecutor(resource_dict={"cores": 2}) as exe:
             fs1 = exe.submit(mpi_funct, 1)
             self.assertFalse(fs1.done())
             self.assertEqual(fs1.result(), [(1, 2, 0), (1, 2, 1)])

tests/test_cache_executor_pysqa_flux.py

@@ -32,7 +32,7 @@ def mpi_funct(i):
 class TestCacheExecutorPysqa(unittest.TestCase):
     def test_executor(self):
         with FileExecutor(
-            cores_per_worker=2,
+            resource_dict={"cores": 2},
             execute_function=execute_with_pysqa,
             backend="flux",
         ) as exe:

tests/test_cache_executor_serial.py

@@ -48,7 +48,7 @@ def test_executor_dependence_mixed(self):
 
     def test_executor_working_directory(self):
         cwd = os.path.join(os.path.dirname(__file__), "executables")
-        with FileExecutor(cwd=cwd) as exe:
+        with FileExecutor(resource_dict={"cwd": cwd}) as exe:
             fs1 = exe.submit(list_files_in_working_directory)
             self.assertEqual(fs1.result(), os.listdir(cwd))
 
@@ -72,7 +72,7 @@ def test_executor_function(self):
                 "future_queue": q,
                 "cache_directory": cache_dir,
                 "execute_function": execute_in_subprocess,
-                "cores_per_worker": 1,
+                "resource_dict": {"cores": 1, "cwd": None},
                 "terminate_function": terminate_subprocess,
             },
         )
@@ -113,7 +113,7 @@ def test_executor_function_dependence_kwargs(self):
                 "future_queue": q,
                 "cache_directory": cache_dir,
                 "execute_function": execute_in_subprocess,
-                "cores_per_worker": 1,
+                "resource_dict": {"cores": 1, "cwd": None},
                 "terminate_function": terminate_subprocess,
             },
         )
@@ -154,7 +154,7 @@ def test_executor_function_dependence_args(self):
                 "future_queue": q,
                 "cache_directory": cache_dir,
                 "execute_function": execute_in_subprocess,
-                "cores_per_worker": 1,
+                "resource_dict": {"cores": 1, "cwd": None},
                 "terminate_function": terminate_subprocess,
             },
         )