Add plot functionality for new future selector (#820)

* Implement splitting of future objects

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

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

* Update split.py

* Update split.py

* fix tests

* extend tests

* cover all lines

* Import from API

* fixes

* Add support for dictionaries

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

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

* use dict function

* move split to main

* Add docstrings

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

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

* Test response is None

* Extend plotting function to support splitfuture objects

* Rename splitfuture to futureselector

* sync with latest changes

* Rename module to select rather than split

* rename module

* fix plotting

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

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

* Add tests

* skip test

* remove redundant lines

* move plot module

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

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

---------

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

Author: jan-janssen

executorlib/task_scheduler/interactive/dependency.py

@@ -11,12 +11,12 @@
     get_future_objects_from_input,
     update_futures_in_input,
 )
-from executorlib.standalone.plot import (
+from executorlib.task_scheduler.base import TaskSchedulerBase
+from executorlib.task_scheduler.interactive.dependency_plot import (
     generate_nodes_and_edges_for_plotting,
     generate_task_hash_for_plotting,
     plot_dependency_graph_function,
 )
-from executorlib.task_scheduler.base import TaskSchedulerBase
 
 
 class DependencyTaskScheduler(TaskSchedulerBase):

executorlib/standalone/plot.py renamed to executorlib/task_scheduler/interactive/dependency_plot.py

@@ -4,6 +4,8 @@
 
 import cloudpickle
 
+from executorlib.standalone.select import FutureSelector
+
 
 def generate_nodes_and_edges_for_plotting(
     task_hash_dict: dict, future_hash_inverse_dict: dict
@@ -31,7 +33,15 @@ def add_element(arg, link_to, label=""):
             link_to: ID of the node to link the element to.
             label (str, optional): Label for the edge. Defaults to "".
         """
-        if isinstance(arg, Future):
+        if isinstance(arg, FutureSelector):
+            edge_lst.append(
+                {
+                    "start": hash_id_dict[future_hash_inverse_dict[arg._future]],
+                    "end": link_to,
+                    "label": label + str(arg._selector),
+                }
+            )
+        elif isinstance(arg, Future):
             edge_lst.append(
                 {
                     "start": hash_id_dict[future_hash_inverse_dict[arg]],
@@ -104,7 +114,24 @@ def convert_arg(arg, future_hash_inverse_dict):
         Returns:
             The hash representation of the argument.
         """
-        if isinstance(arg, Future):
+        if isinstance(arg, FutureSelector):
+            if arg not in future_hash_inverse_dict:
+                obj_dict = {
+                    "args": (),
+                    "kwargs": {
+                        "future": future_hash_inverse_dict[arg._future],
+                        "selector": arg._selector,
+                    },
+                }
+                if isinstance(arg._selector, str):
+                    obj_dict["fn"] = "get_item_from_future"
+                else:
+                    obj_dict["fn"] = "split_future"
+                arg_hash = cloudpickle.dumps(obj_dict)
+                future_hash_dict[arg_hash] = arg
+                future_hash_inverse_dict[arg] = arg_hash
+            return future_hash_inverse_dict[arg]
+        elif isinstance(arg, Future):
             return future_hash_inverse_dict[arg]
         elif isinstance(arg, list):
             return [

tests/test_fluxjobexecutor_plot.py

@@ -3,7 +3,7 @@
 from time import sleep
 
 from executorlib import FluxJobExecutor, FluxClusterExecutor
-from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting
+from executorlib.task_scheduler.interactive.dependency_plot import generate_nodes_and_edges_for_plotting
 from executorlib.standalone.serialize import cloudpickle_register
 
 

tests/test_singlenodeexecutor_plot_dependency.py

@@ -6,8 +6,10 @@
     SingleNodeExecutor,
     SlurmJobExecutor,
     SlurmClusterExecutor,
+    split_future,
+    get_item_from_future,
 )
-from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting
+from executorlib.task_scheduler.interactive.dependency_plot import generate_nodes_and_edges_for_plotting
 from executorlib.standalone.serialize import cloudpickle_register
 
 
@@ -43,6 +45,17 @@ def return_input_dict(input_dict):
     return input_dict
 
 
+def get_tuple(i):
+    return 1, 2, i
+
+
+def get_dict(i):
+    return {"a": 1, "b": i}
+
+
+def echo(i):
+    return i
+
 @unittest.skipIf(
     skip_graphviz_test,
     "graphviz is not installed, so the plot_dependency_graph tests are skipped.",
@@ -289,3 +302,28 @@ def test_many_to_one_plot(self):
             )
             self.assertEqual(len(nodes), 19)
             self.assertEqual(len(edges), 22)
+
+
+@unittest.skipIf(
+    skip_graphviz_test,
+    "graphviz is not installed, so the plot_dependency_graph tests are skipped.",
+)
+class TestSelectExecutorPlot(unittest.TestCase):
+    def test_split_future(self):
+        with SingleNodeExecutor(plot_dependency_graph=True) as exe:
+            f = exe.submit(get_tuple, 5)
+            f1, f2, f3 = split_future(future=f, n=3)
+            f12 = exe.submit(echo, f1)
+            f22 = exe.submit(echo, f2)
+            f32 = exe.submit(echo, f3)
+            self.assertIsNone(f12.result())
+            self.assertIsNone(f22.result())
+            self.assertIsNone(f32.result())
+
+    def test_get_item_from_future(self):
+        with SingleNodeExecutor(plot_dependency_graph=True) as exe:
+            f = exe.submit(get_dict, 5)
+            f1 = exe.submit(echo, get_item_from_future(future=f, key="a"))
+            f2 = exe.submit(echo, get_item_from_future(future=f, key="b"))
+            self.assertIsNone(f1.result())
+            self.assertIsNone(f2.result())

tests/test_testclusterexecutor.py

@@ -4,7 +4,7 @@
 
 from executorlib import get_cache_data
 from executorlib.api import TestClusterExecutor
-from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting
+from executorlib.task_scheduler.interactive.dependency_plot import generate_nodes_and_edges_for_plotting
 from executorlib.standalone.serialize import cloudpickle_register
 
 try: