diff --git a/executorlib/standalone/plot.py b/executorlib/standalone/plot.py index 320df56d..610bfcff 100644 --- a/executorlib/standalone/plot.py +++ b/executorlib/standalone/plot.py @@ -5,7 +5,7 @@ import cloudpickle -def generate_nodes_and_edges( +def generate_nodes_and_edges_for_plotting( task_hash_dict: dict, future_hash_inverse_dict: dict ) -> tuple[list, list]: """ @@ -81,7 +81,9 @@ def add_element(arg, link_to, label=""): return node_lst, edge_lst -def generate_task_hash(task_dict: dict, future_hash_inverse_dict: dict) -> bytes: +def generate_task_hash_for_plotting( + task_dict: dict, future_hash_inverse_dict: dict +) -> bytes: """ Generate a hash for a task dictionary. @@ -132,7 +134,9 @@ def convert_arg(arg, future_hash_inverse_dict): ) -def draw(node_lst: list, edge_lst: list, filename: Optional[str] = None): +def plot_dependency_graph_function( + node_lst: list, edge_lst: list, filename: Optional[str] = None +): """ Draw the graph visualization of nodes and edges. diff --git a/executorlib/task_scheduler/interactive/dependency.py b/executorlib/task_scheduler/interactive/dependency.py index 66bf1531..cfb9e9f3 100644 --- a/executorlib/task_scheduler/interactive/dependency.py +++ b/executorlib/task_scheduler/interactive/dependency.py @@ -12,9 +12,9 @@ update_futures_in_input, ) from executorlib.standalone.plot import ( - draw, - generate_nodes_and_edges, - generate_task_hash, + generate_nodes_and_edges_for_plotting, + generate_task_hash_for_plotting, + plot_dependency_graph_function, ) from executorlib.task_scheduler.base import TaskSchedulerBase @@ -144,7 +144,7 @@ def submit( # type: ignore "future": f, "resource_dict": resource_dict, } - task_hash = generate_task_hash( + task_hash = generate_task_hash_for_plotting( task_dict=task_dict, future_hash_inverse_dict={ v: k for k, v in self._future_hash_dict.items() @@ -205,13 +205,13 @@ def __exit__( """ super().__exit__(exc_type=exc_type, exc_val=exc_val, exc_tb=exc_tb) # type: ignore if self._generate_dependency_graph: - node_lst, edge_lst = generate_nodes_and_edges( + node_lst, edge_lst = generate_nodes_and_edges_for_plotting( task_hash_dict=self._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in self._future_hash_dict.items() }, ) - return draw( + return plot_dependency_graph_function( node_lst=node_lst, edge_lst=edge_lst, filename=self._plot_dependency_graph_filename, diff --git a/tests/test_fluxjobexecutor_plot.py b/tests/test_fluxjobexecutor_plot.py index 1b71239e..2ee15dda 100644 --- a/tests/test_fluxjobexecutor_plot.py +++ b/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 +from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting from executorlib.standalone.serialize import cloudpickle_register @@ -55,7 +55,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -100,7 +100,7 @@ def test_many_to_one_plot(self): self.assertTrue(future_sum.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 7) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 7) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -126,7 +126,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -169,7 +169,7 @@ def test_many_to_one_plot(self): self.assertTrue(future_sum.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 7) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 7) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() diff --git a/tests/test_singlenodeexecutor_plot_dependency.py b/tests/test_singlenodeexecutor_plot_dependency.py index fd7c50ee..a3426ed4 100644 --- a/tests/test_singlenodeexecutor_plot_dependency.py +++ b/tests/test_singlenodeexecutor_plot_dependency.py @@ -7,7 +7,7 @@ SlurmJobExecutor, SlurmClusterExecutor, ) -from executorlib.standalone.plot import generate_nodes_and_edges +from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting from executorlib.standalone.serialize import cloudpickle_register @@ -61,7 +61,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -122,7 +122,7 @@ def test_many_to_one_plot(self): self.assertTrue(future_sum.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 7) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 7) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -139,7 +139,7 @@ def test_future_input_dict(self): ) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -167,7 +167,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -212,7 +212,7 @@ def test_many_to_one_plot(self): self.assertTrue(future_sum.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 7) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 7) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -238,7 +238,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() @@ -281,7 +281,7 @@ def test_many_to_one_plot(self): self.assertTrue(future_sum.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 7) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 7) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items() diff --git a/tests/test_testclusterexecutor.py b/tests/test_testclusterexecutor.py index 0a47aa5f..73e1d361 100644 --- a/tests/test_testclusterexecutor.py +++ b/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 +from executorlib.standalone.plot import generate_nodes_and_edges_for_plotting from executorlib.standalone.serialize import cloudpickle_register try: @@ -83,7 +83,7 @@ def test_executor_dependency_plot(self): self.assertTrue(future_2.done()) self.assertEqual(len(exe._task_scheduler._future_hash_dict), 2) self.assertEqual(len(exe._task_scheduler._task_hash_dict), 2) - nodes, edges = generate_nodes_and_edges( + nodes, edges = generate_nodes_and_edges_for_plotting( task_hash_dict=exe._task_scheduler._task_hash_dict, future_hash_inverse_dict={ v: k for k, v in exe._task_scheduler._future_hash_dict.items()