Fix TestClusterExecutor get_item_from_future (#910)

* Fix TestClusterExecutor get_item_from_future

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

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

* fix ruff

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

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

* fixes

* some more tests

* reset

* fix test

* fixes

---------

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

Author: jan-janssen

src/executorlib/task_scheduler/file/shared.py

@@ -11,7 +11,7 @@
 
 
 class FutureItem:
-    def __init__(self, file_name: str):
+    def __init__(self, file_name: str, selector: Optional[int | str] = None):
         """
         Initialize a FutureItem object.
 
@@ -20,6 +20,7 @@ def __init__(self, file_name: str):
 
         """
         self._file_name = file_name
+        self._selector = selector
 
     def result(self) -> Any:
         """
@@ -31,7 +32,10 @@ def result(self) -> Any:
         """
         exec_flag, no_error_flag, result = get_output(file_name=self._file_name)
         if exec_flag and no_error_flag:
-            return result
+            if self._selector is not None:
+                return result[self._selector]
+            else:
+                return result
         elif exec_flag:
             raise result
         else:
@@ -239,29 +243,45 @@ def _convert_args_and_kwargs(
     task_kwargs = {}
     future_wait_key_lst = []
     for arg in task_dict["args"]:
+        selector = None
         if isinstance(arg, Future):
+            if hasattr(arg, "_future") and hasattr(arg, "_selector"):
+                selector = arg._selector
+                future = arg._future
+            else:
+                future = arg
             match_found = False
             for k, v in memory_dict.items():
-                if arg == v:
-                    task_args.append(FutureItem(file_name=file_name_dict[k]))
+                if future == v:
+                    task_args.append(
+                        FutureItem(file_name=file_name_dict[k], selector=selector)
+                    )
                     future_wait_key_lst.append(k)
                     match_found = True
                     break
             if not match_found:
-                task_args.append(arg.result())
+                task_args.append(future.result())
         else:
             task_args.append(arg)
     for key, arg in task_dict["kwargs"].items():
+        selector = None
         if isinstance(arg, Future):
+            if hasattr(arg, "_future") and hasattr(arg, "_selector"):
+                selector = arg._selector
+                future = arg._future
+            else:
+                future = arg
             match_found = False
             for k, v in memory_dict.items():
-                if arg == v:
-                    task_kwargs[key] = FutureItem(file_name=file_name_dict[k])
+                if future == v:
+                    task_kwargs[key] = FutureItem(
+                        file_name=file_name_dict[k], selector=selector
+                    )
                     future_wait_key_lst.append(k)
                     match_found = True
                     break
             if not match_found:
-                task_kwargs[key] = arg.result()
+                task_kwargs[key] = future.result()
         else:
             task_kwargs[key] = arg
     return task_args, task_kwargs, future_wait_key_lst

tests/unit/task_scheduler/file/test_backend.py

@@ -3,10 +3,12 @@
 import shutil
 import unittest
 
+from executorlib.standalone.select import FutureSelector
+
 
 try:
     from executorlib.task_scheduler.file.backend import backend_execute_task_in_file
-    from executorlib.task_scheduler.file.shared import _check_task_output, FutureItem
+    from executorlib.task_scheduler.file.shared import _check_task_output, _convert_args_and_kwargs, FutureItem
     from executorlib.standalone.hdf import dump, get_runtime
     from executorlib.standalone.serialize import serialize_funct
 
@@ -19,6 +21,14 @@ def my_funct(a, b):
     return a + b
 
 
+def return_dict(a, b):
+    return {"a": a, "b": b}
+
+
+def return_list(a, b):
+    return [a, b]
+
+
 def get_error(a):
     raise ValueError(a)
 
@@ -36,7 +46,6 @@ def test_execute_function_mixed(self):
             fn_kwargs={"b": 2},
         )
         file_name = os.path.join(cache_directory, task_key + "_i.h5")
-        os.makedirs(cache_directory, exist_ok=True)
         dump(file_name=file_name, data_dict=data_dict)
         backend_execute_task_in_file(file_name=file_name)
         future_obj = Future()
@@ -55,6 +64,47 @@ def test_execute_function_mixed(self):
         self.assertTrue(future_file_obj.done())
         self.assertEqual(future_file_obj.result(), 3)
 
+    def test_execute_function_mixed_selector_convert(self):
+        cache_directory = os.path.abspath("executorlib_cache")
+        os.makedirs(cache_directory, exist_ok=True)
+        task_key_1, data_dict = serialize_funct(
+            fn=return_dict,
+            fn_args=[1],
+            fn_kwargs={"b": 2},
+        )
+        file_name_1 = os.path.join(cache_directory, task_key_1 + "_i.h5")
+        dump(file_name=file_name_1, data_dict=data_dict)
+        backend_execute_task_in_file(file_name=file_name_1)
+        f1 = Future()
+        _check_task_output(
+            task_key=task_key_1, future_obj=f1, cache_directory=cache_directory
+        )
+        task_key_2, data_dict = serialize_funct(
+            fn=return_list,
+            fn_args=[1],
+            fn_kwargs={"b": 2},
+        )
+        file_name_2 = os.path.join(cache_directory, task_key_2 + "_i.h5")
+        dump(file_name=file_name_2, data_dict=data_dict)
+        backend_execute_task_in_file(file_name=file_name_2)
+        f2 = Future()
+        _check_task_output(
+            task_key=task_key_2, future_obj=f2, cache_directory=cache_directory
+        )
+        fs1 = FutureSelector(future=f1, selector="a")
+        fs2 = FutureSelector(future=f2, selector=1)
+        task_args, task_kwargs, future_wait_key_lst = _convert_args_and_kwargs(
+            task_dict={"fn": 1, "args": (fs1,), "kwargs": {"b": fs2}},
+            memory_dict={task_key_1: f1, task_key_2: f2},
+            file_name_dict={
+                task_key_1: os.path.join(cache_directory, task_key_1 + "_o.h5"), 
+                task_key_2: os.path.join(cache_directory, task_key_2 + "_o.h5"),
+            },
+        )
+        self.assertEqual(task_args[0].result(), 1)
+        self.assertEqual(task_kwargs["b"].result(), 2)
+        self.assertTrue(len(future_wait_key_lst) == 2)
+
     def test_execute_function_args(self):
         cache_directory = os.path.abspath("executorlib_cache")
         os.makedirs(cache_directory, exist_ok=True)
@@ -92,7 +142,6 @@ def test_execute_function_kwargs(self):
             fn_kwargs={"a": 1, "b": 2},
         )
         file_name = os.path.join(cache_directory, task_key + "_i.h5")
-        os.makedirs(cache_directory, exist_ok=True)
         dump(file_name=file_name, data_dict=data_dict)
         backend_execute_task_in_file(file_name=file_name)
         future_obj = Future()
@@ -120,7 +169,6 @@ def test_execute_function_error(self):
             fn_kwargs={"a": 1},
         )
         file_name = os.path.join(cache_directory, task_key + "_i.h5")
-        os.makedirs(cache_directory, exist_ok=True)
         data_dict["error_log_file"] = os.path.join(cache_directory, "error.out")
         dump(file_name=file_name, data_dict=data_dict)
         backend_execute_task_in_file(file_name=file_name)