Minor improvements to DAG visualization (#348)

* Show if DAG was optimized on visualization * Left justify lines in graph caption in DAG
cubed-dev · Jan 16, 2024 · 6364895 · 6364895
1 parent f3c3b94
commit 6364895
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
@@ -265,9 +265,11 @@ def visualize(
         dag.graph["graph"] = {
             "rankdir": rankdir,
             "label": (
-                f"num tasks: {self.num_tasks(optimize_graph, optimize_function)}\n"
-                f"max projected memory: {memory_repr(self.max_projected_mem(optimize_graph, optimize_function))}\n"
-                f"total nbytes: {memory_repr(self.total_nbytes(optimize_graph, optimize_function))}"
+                # note that \l is used to left-justify each line (see https://www.graphviz.org/docs/attrs/nojustify/)
+                rf"num tasks: {self.num_tasks(optimize_graph, optimize_function)}\l"
+                rf"max projected memory: {memory_repr(self.max_projected_mem(optimize_graph, optimize_function))}\l"
+                rf"total nbytes: {memory_repr(self.total_nbytes(optimize_graph, optimize_function))}\l"
+                rf"optimized: {optimize_graph}\l"
             ),
             "labelloc": "bottom",
             "labeljust": "left",