Perf/decision task

src/taskgraph/create.py

@@ -5,6 +5,7 @@
 
 import logging
 import sys
+from collections import defaultdict
 from concurrent import futures
 
 from slugid import nice as slugid
@@ -57,78 +58,66 @@ def create_tasks(graph_config, taskgraph, label_to_taskid, params, decision_task
         task_def["taskGroupId"] = decision_task_id
         task_def["schedulerId"] = scheduler_id
 
+    # We can't submit a task until its dependencies have been created. So
+    # track, for each task, how many of its dependencies within this graph are
+    # still pending, and submit it once that number drops to zero. Tasks
+    # depending (directly or not) on a task that failed to be created are never
+    # submitted, as their creation would fail anyway.
+    pending_deps = {}
+    dependents = defaultdict(list)
+    for task_id in taskgraph.graph.nodes:
+        # Some dependencies aren't in our graph, so make sure to filter those
+        # out.
+        deps = {
+            d
+            for d in taskgraph.tasks[task_id].task.get("dependencies", [])
+            if d in taskgraph.tasks
+        }
+        pending_deps[task_id] = len(deps)
+        for dep in deps:
+            dependents[dep].append(task_id)
+
     # If `testing` is True, then run without parallelization
     concurrency = CONCURRENCY if not testing else 1
     session = get_session()
     with futures.ThreadPoolExecutor(concurrency) as e:
-        fs = {}
+        # Maps each future to the task it is creating, and whether it is the
+        # primary creation of that task (as opposed to a duplicate).
         fs_to_task = {}
-        skipped = set()
-
-        # We can't submit a task until its dependencies have been submitted.
-        # So our strategy is to walk the graph and submit tasks once all
-        # their dependencies have been submitted.
-        tasklist = set(taskgraph.graph.visit_postorder())
-        alltasks = tasklist.copy()
-
-        def schedule_tasks():
-            to_remove = set()
-            new = set()
-
-            def submit(task_id, label, task_def):
-                fut = e.submit(create_task, session, task_id, label, task_def)
-                new.add(fut)
-                fs[task_id] = fut
-                fs_to_task[fut] = (task_id, label)
-
-                def mark_failed_as_skipped(fut):
-                    if fut.exception():
-                        task_id, _ = fs_to_task[fut]
-                        skipped.add(task_id)
-
-                fut.add_done_callback(mark_failed_as_skipped)
-
-            for task_id in tasklist:
-                task_def = taskgraph.tasks[task_id].task
-                # Some dependencies aren't in our graph, so make sure to filter
-                # those out
-                deps = set(task_def.get("dependencies", [])) & alltasks
-
-                # If one of the dependencies didn't get created, then
-                # don't attempt to submit as it would fail.
-                if any(d in skipped for d in deps):
-                    skipped.add(task_id)
-                    to_remove.add(task_id)
+        in_flight = set()
+
+        def submit(task_id):
+            task = taskgraph.tasks[task_id]
+            label = taskid_to_label[task_id]
+            # Schedule tasks as many times as task_duplicates indicates. We
+            # use slugid() for duplicates since we want a distinct task id.
+            for i in range(task.attributes.get("task_duplicates", 1)):
+                fut_task_id = task_id if i == 0 else slugid()
+                fut = e.submit(create_task, session, fut_task_id, label, task.task)
+                fs_to_task[fut] = (task_id, label, i == 0)
+                in_flight.add(fut)
+
+        for task_id, count in pending_deps.items():
+            if count == 0:
+                submit(task_id)
+
+        # As each of those futures complete, schedule the tasks that were
+        # waiting on them.
+        while in_flight:
+            done, _ = futures.wait(in_flight, return_when=futures.FIRST_COMPLETED)
+            for fut in done:
+                in_flight.remove(fut)
+                task_id, _, primary = fs_to_task[fut]
+                if not primary or fut.exception():
                     continue
-
-                # If we haven't finished submitting all our dependencies yet,
-                # come back to this later.
-                if any((d not in fs or not fs[d].done()) for d in deps):
-                    continue
-
-                submit(task_id, taskid_to_label[task_id], task_def)
-                to_remove.add(task_id)
-
-                # Schedule tasks as many times as task_duplicates indicates
-                attributes = taskgraph.tasks[task_id].attributes
-                for i in range(1, attributes.get("task_duplicates", 1)):
-                    # We use slugid() since we want a distinct task id
-                    submit(slugid(), taskid_to_label[task_id], task_def)
-            tasklist.difference_update(to_remove)
-
-            # As each of those futures complete, try to schedule more tasks.
-            for f in futures.as_completed(new):
-                schedule_tasks()
-
-        # Start scheduling tasks and run until everything is scheduled.
-        schedule_tasks()
-
-        # Wait for all futures to complete.
-        futures.wait(fs.values())
+                for dependent in dependents[task_id]:
+                    pending_deps[dependent] -= 1
+                    if pending_deps[dependent] == 0:
+                        submit(dependent)
 
         # Collect errors.
         errors = {}
-        for fut, (task_id, label) in fs_to_task.items():
+        for fut, (_, label, _) in fs_to_task.items():
             if exc := fut.exception():
                 errors[label] = exc
 

src/taskgraph/decision.py

@@ -17,7 +17,6 @@
 from taskgraph.create import create_tasks
 from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import Parameters, get_version
-from taskgraph.taskgraph import TaskGraph
 from taskgraph.util import json
 from taskgraph.util.python_path import find_object
 from taskgraph.util.schema import Schema, validate_schema
@@ -131,9 +130,6 @@ def taskgraph_decision(options, parameters=None):
         "runnable-jobs.json", full_task_graph_to_runnable_tasks(full_task_json)
     )
 
-    # this is just a test to check whether the from_json() function is working
-    _, _ = TaskGraph.from_json(full_task_json)
-
     # write out the target task set to allow reproducing this as input
     write_artifact("target-tasks.json", list(tgg.target_task_set.tasks.keys()))
 

src/taskgraph/generator.py

@@ -282,8 +282,21 @@ def _load_kinds(self, graph_config, target_kinds=None):
                 except KindNotFound:
                     continue
 
+    @staticmethod
+    def _get_kind_dependencies_tasks(kind, tasks_by_kind):
+        """Return the tasks of the kinds `kind` depends on, in the order in
+        which they were loaded."""
+        kind_dependencies = kind.config.get("kind-dependencies", [])
+        return {
+            label: task
+            for kind_name, tasks in tasks_by_kind.items()
+            if kind_name in kind_dependencies
+            for label, task in tasks.items()
+        }
+
     def _load_tasks_serial(self, kinds, kind_graph, parameters):
         all_tasks = {}
+        tasks_by_kind = {}
         for kind_name in kind_graph.visit_postorder():
             logger.debug(f"Loading tasks for kind {kind_name}")
 
@@ -297,11 +310,7 @@ def _load_tasks_serial(self, kinds, kind_graph, parameters):
             try:
                 new_tasks = kind.load_tasks(
                     parameters,
-                    {
-                        k: t
-                        for k, t in all_tasks.items()
-                        if t.kind in kind.config.get("kind-dependencies", [])
-                    },
+                    self._get_kind_dependencies_tasks(kind, tasks_by_kind),
                     self._write_artifacts,
                 )
             except SchemaValidationError as exc:
@@ -310,15 +319,18 @@ def _load_tasks_serial(self, kinds, kind_graph, parameters):
             except Exception:
                 logger.exception(f"Error loading tasks for kind {kind_name}:")
                 raise
+            kind_tasks = tasks_by_kind.setdefault(kind_name, {})
             for task in new_tasks:
                 if task.label in all_tasks:
                     raise Exception("duplicate tasks with label " + task.label)
                 all_tasks[task.label] = task
+                kind_tasks[task.label] = task
 
         return all_tasks
 
     def _load_tasks_parallel(self, kinds, kind_graph, parameters, executor):
         all_tasks = {}
+        tasks_by_kind = {}
         futures_to_kind = {}
         futures = set()
         edges = set(kind_graph.edges)
@@ -328,7 +340,6 @@ def _load_tasks_parallel(self, kinds, kind_graph, parameters, executor):
             def submit_ready_kinds():
                 """Create the next batch of tasks for kinds without dependencies."""
                 nonlocal kinds, edges, futures
-                loaded_tasks = all_tasks.copy()
                 kinds_with_deps = {edge[0] for edge in edges}
                 ready_kinds = (
                     set(kinds) - kinds_with_deps - set(futures_to_kind.values())
@@ -346,11 +357,7 @@ def submit_ready_kinds():
                     future = executor.submit(
                         kind.load_tasks,
                         dict(parameters),
-                        {
-                            k: t
-                            for k, t in loaded_tasks.items()
-                            if t.kind in kind.config.get("kind-dependencies", [])
-                        },
+                        self._get_kind_dependencies_tasks(kind, tasks_by_kind),
                         self._write_artifacts,
                     )
                     futures.add(future)
@@ -376,10 +383,12 @@ def submit_ready_kinds():
                     kind = futures_to_kind.pop(future)
                     futures.remove(future)
 
+                    kind_tasks = tasks_by_kind.setdefault(kind, {})
                     for task in future.result():
                         if task.label in all_tasks:
                             raise Exception("duplicate tasks with label " + task.label)
                         all_tasks[task.label] = task
+                        kind_tasks[task.label] = task
 
                     # Update state for next batch of futures.
                     del kinds[kind]

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()
 
@@ -301,6 +301,18 @@ def replace_tasks(
         target_task_graph.graph.links_and_reverse_links_dict()
     )
 
+    # Many tasks share dependents (e.g. docker images or toolchains), so
+    # resolve the deadline of each dependent only once.
+    now = datetime.datetime.now(datetime.timezone.utc)
+    deadlines = {}
+
+    def get_deadline(label):
+        if label not in deadlines:
+            deadlines[label] = resolve_timestamps(
+                now, target_task_graph.tasks[label].task["deadline"]
+            )
+        return deadlines[label]
+
     for label in target_task_graph.graph.visit_postorder():
         logger.debug(f"replace_tasks: {label}")
         # if we're not allowed to optimize, that's easy..
@@ -331,14 +343,9 @@ def replace_tasks(
         opt_by, opt, arg = optimizations(label)
 
         # compute latest deadline of dependents (if any)
-        dependents = [target_task_graph.tasks[l] for l in dependents_of[label]]
         deadline = None
-        if dependents:
-            now = datetime.datetime.now(datetime.timezone.utc)
-            deadline = max(
-                resolve_timestamps(now, task.task["deadline"])
-                for task in dependents  # type: ignore
-            )
+        if dependents_of[label]:
+            deadline = max(get_deadline(l) for l in dependents_of[label])
 
         if isinstance(opt, IndexSearch):
             arg = arg, index_to_taskid, taskid_to_status

src/taskgraph/optimize/strategies.py

@@ -1,3 +1,4 @@
+import functools
 import logging
 from datetime import datetime
 
@@ -10,6 +11,13 @@
 logger = logging.getLogger("optimization")
 
 
+@functools.cache
+def _parse_time(timestamp, fmt):
+    # Many tasks share the same deadline or replacement task, so avoid parsing
+    # the same timestamps over and over.
+    return datetime.strptime(timestamp, fmt)
+
+
 @register_strategy("index-search")
 class IndexSearch(OptimizationStrategy):
     # A task with no dependencies remaining after optimization will be replaced
@@ -56,10 +64,10 @@ def should_replace_task(self, task, params, deadline, arg):
                     )
                     continue
 
-                if deadline and datetime.strptime(
+                if deadline and _parse_time(
                     status["expires"],  # type: ignore
                     self.fmt,
-                ) < datetime.strptime(deadline, self.fmt):
+                ) < _parse_time(deadline, self.fmt):
                     logger.debug(
                         f"not replacing {task.label} with {task_id} because it expires before {deadline}"
                     )

src/taskgraph/util/parameterization.py

@@ -15,20 +15,18 @@
 
 def _recurse(val, param_fns):
     def recurse(val):
-        if isinstance(val, list):
-            return [recurse(v) for v in val]
-        elif isinstance(val, dict):
+        if isinstance(val, dict):
             if len(val) == 1:
-                for param_key, param_fn in param_fns.items():
-                    if set(val.keys()) == {param_key}:
-                        if isinstance(val[param_key], dict):
-                            # handle `{"task-reference": {"<foo>": "bar"}}`
-                            return {
-                                param_fn(key): recurse(v)
-                                for key, v in val[param_key].items()
-                            }
-                        return param_fn(val[param_key])
+                ((key, value),) = val.items()
+                param_fn = param_fns.get(key)
+                if param_fn is not None:
+                    if isinstance(value, dict):
+                        # handle `{"task-reference": {"<foo>": "bar"}}`
+                        return {param_fn(k): recurse(v) for k, v in value.items()}
+                    return param_fn(value)
             return {k: recurse(v) for k, v in val.items()}
+        elif isinstance(val, list):
+            return [recurse(v) for v in val]
         else:
             return val