perf(graph): cache the visit order of graphs

Graphs are immutable, but visit_postorder and visit_preorder sorted them
topologically again every time they were called. The full task graph is
visited once per registered verification (11 times in taskgraph alone),
then again to serialize it. The target task graph is visited several
times during optimization. Now the order is computed once per graph and
direction, and cached like links_and_reverse_links_dict already is.

Also, during optimization:
- index paths are gathered by iterating over the tasks directly, as the
  order doesn't matter;
- remove_tasks uses the cached reverse links instead of building them
  again.

On a synthetic graph of 20,210 tasks and 40,200 edges, verifying the
full task graph takes 0.34s instead of 0.73s, and optimizing it takes
1.57s instead of 2.19s.

Author: marco-c

src/taskgraph/graph.py

@@ -107,21 +107,30 @@ def _visit(self, reverse):
                 f"Dependency loop detected involving the following nodes: {loopy_nodes}"
             )
 
+    @functools.cache
+    def _visit_order(self, reverse):
+        """
+        Return the order in which `_visit` yields nodes as a tuple. The graph
+        is immutable, so this is cached to avoid sorting it again every time
+        it is visited.
+        """
+        return tuple(self._visit(reverse))
+
     def visit_postorder(self):
         """
         Generate a sequence of nodes in postorder, such that every node is
         visited *after* any nodes it links to.
 
         Raises an exception if the graph contains a cycle.
         """
-        return self._visit(False)
+        return iter(self._visit_order(False))
 
     def visit_preorder(self):
         """
         Like visit_postorder, but in reverse: evrey node is visited *before*
         any nodes it links to.
         """
-        return self._visit(True)
+        return iter(self._visit_order(True))
 
     @functools.cache
     def links_and_reverse_links_dict(self):

src/taskgraph/optimize/base.py

@@ -78,7 +78,7 @@ def optimize_task_graph(
 
     # Gather each relevant task's index
     indexes = set()
-    for label in target_task_graph.graph.visit_postorder():
+    for label in target_task_graph.tasks:
         if label in do_not_optimize:
             continue
         _, strategy, arg = optimizations(label)
@@ -157,7 +157,7 @@ def remove_tasks(
     opt_counts = defaultdict(int)
     opt_reasons = {}
     removed = set()
-    dependents_of = target_task_graph.graph.reverse_links_dict()
+    _, dependents_of = target_task_graph.graph.links_and_reverse_links_dict()
     tasks = target_task_graph.tasks
     prune_candidates = set()
 

test/test_graph_perf.py

@@ -124,6 +124,8 @@ def test_transitive_closure(geometry):
 @pytest.mark.parametrize("geometry", ["linear", "fan", "btree", "diamond"])
 def test_visit_postorder(geometry):
     _, graph, _ = GEOMETRIES[geometry]
+    # Clear the functools.cache to measure actual computation each time
+    graph._visit_order.cache_clear()
     order = list(graph.visit_postorder())
     assert len(order) == N
 
@@ -132,6 +134,8 @@ def test_visit_postorder(geometry):
 @pytest.mark.parametrize("geometry", ["linear", "fan", "btree", "diamond"])
 def test_visit_preorder(geometry):
     _, graph, _ = GEOMETRIES[geometry]
+    # Clear the functools.cache to measure actual computation each time
+    graph._visit_order.cache_clear()
     order = list(graph.visit_preorder())
     assert len(order) == N