From 2dfb4140b259c5aede10d148b9cc30963333fde5 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 13:18:01 +0200 Subject: [PATCH 1/7] perf(create): schedule task creation from dependency counts create_tasks rescanned every not yet submitted task each time a createTask call completed, rebuilding its set of dependencies and checking whether their futures were done. This is quadratic in the number of tasks, and runs on the main thread while holding the GIL. It also recursed once per completed batch, so a dependency chain of a few thousand tasks raised RecursionError. Now the number of pending dependencies of each task is computed once, and decremented as dependencies are created. A task is submitted as soon as its count drops to zero. As before, tasks depending on a task that failed to be created are not submitted. With createTask mocked out, creating a 20,000 task graph takes 1.9s instead of 6.2s (binary tree) and 2.0s instead of 4.6s (fan-out/fan-in). --- src/taskgraph/create.py | 115 ++++++++++++++++++---------------------- test/test_create.py | 85 +++++++++++++++++++++++++++++ test/test_graph_perf.py | 50 +++++++++++++++++ 3 files changed, 187 insertions(+), 63 deletions(-) diff --git a/src/taskgraph/create.py b/src/taskgraph/create.py index dbcba1331..21030c768 100644 --- a/src/taskgraph/create.py +++ b/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 diff --git a/test/test_create.py b/test/test_create.py index ae0017ca0..67d0b2634 100644 --- a/test/test_create.py +++ b/test/test_create.py @@ -9,6 +9,7 @@ from concurrent import futures from unittest import mock +import pytest import responses from taskgraph import create @@ -244,3 +245,87 @@ def wrapper(fut): {"level": "4"}, decision_task_id="decisiontask", ) + + +def _chain_taskgraph(n): + """Build a taskgraph where tid-i depends on tid-(i-1).""" + tasks = {} + for i in range(n): + deps = [f"tid-{i - 1}"] if i else [] + tasks[f"tid-{i}"] = Task( + kind="test", label=f"t{i}", attributes={}, task={"dependencies": deps} + ) + edges = {(f"tid-{i}", f"tid-{i - 1}", "dep") for i in range(1, n)} + taskgraph = TaskGraph(tasks, Graph(nodes=set(tasks), edges=edges)) + label_to_taskid = {t.label: tid for tid, t in tasks.items()} + return taskgraph, label_to_taskid + + +def test_create_tasks_long_chain(mocker): + "tasks are created after their dependencies, even in very deep graphs" + mocker.patch.object(create, "get_session") + created = [] + mocker.patch.object( + create, + "create_task", + side_effect=lambda session, task_id, label, task_def: created.append(task_id), + ) + + n = 3000 + taskgraph, label_to_taskid = _chain_taskgraph(n) + create.create_tasks( + GRAPH_CONFIG, taskgraph, label_to_taskid, {"level": "4"}, "decisiontask" + ) + + assert created == [f"tid-{i}" for i in range(n)] + assert taskgraph.tasks["tid-0"].task["dependencies"] == ["decisiontask"] + assert taskgraph.tasks["tid-1"].task["dependencies"] == ["tid-0"] + + +def test_create_tasks_skips_dependents_of_failed_tasks(mocker): + "tasks depending on a task that failed to be created are not submitted" + mocker.patch.object(create, "get_session") + created = [] + + def fake_create_task(session, task_id, label, task_def): + if task_id == "tid-1": + raise Exception("oh no!") + created.append(task_id) + + mocker.patch.object(create, "create_task", side_effect=fake_create_task) + + taskgraph, label_to_taskid = _chain_taskgraph(4) + with pytest.raises(CreateTasksException) as excinfo: + create.create_tasks( + GRAPH_CONFIG, taskgraph, label_to_taskid, {"level": "4"}, "decisiontask" + ) + + assert created == ["tid-0"] + assert "Could not create 't1'" in str(excinfo.value) + assert "'t2'" not in str(excinfo.value) + + +def test_create_tasks_duplicates(mocker): + "tasks with the task_duplicates attribute are created multiple times" + mocker.patch.object(create, "get_session") + created = [] + mocker.patch.object( + create, + "create_task", + side_effect=lambda session, task_id, label, task_def: created.append( + (task_id, label) + ), + ) + + taskgraph, label_to_taskid = _chain_taskgraph(2) + taskgraph.tasks["tid-0"].attributes["task_duplicates"] = 3 + create.create_tasks( + GRAPH_CONFIG, taskgraph, label_to_taskid, {"level": "4"}, "decisiontask" + ) + + labels = [label for _, label in created] + assert labels.count("t0") == 3 + assert labels.count("t1") == 1 + assert len({task_id for task_id, _ in created}) == 4 + # t1 is only created once the primary t0 task was created. + assert labels.index("t1") > created.index(("tid-0", "t0")) diff --git a/test/test_graph_perf.py b/test/test_graph_perf.py index 83d400f1b..f6eb9cf4d 100644 --- a/test/test_graph_perf.py +++ b/test/test_graph_perf.py @@ -8,6 +8,7 @@ import pytest +from taskgraph import create from taskgraph.graph import Graph from taskgraph.task import Task from taskgraph.taskgraph import TaskGraph @@ -179,6 +180,55 @@ def test_taskgraph_to_json(geometry): assert len(data) == N +# --------------------------------------------------------------------------- +# Benchmarks – create_tasks scheduling +# --------------------------------------------------------------------------- + + +def _create_tasks_args(geometry): + """Copy a geometry into a fresh taskgraph (keyed by taskId) whose task + definitions list their dependencies, as `create_tasks` mutates them.""" + _, graph, tg = GEOMETRIES[geometry] + deps = graph.links_dict() + tasks = { + label: Task( + kind=task.kind, + label=label, + attributes={}, + task={"dependencies": sorted(deps[label])}, + ) + for label, task in tg.tasks.items() + } + taskgraph = TaskGraph(tasks, graph) + label_to_taskid = {label: label for label in tasks} + return ( + {"trust-domain": "domain"}, + taskgraph, + label_to_taskid, + {"level": "1"}, + "decision", + ), {} + + +@pytest.mark.benchmark +@pytest.mark.parametrize("geometry", ["fan", "btree"]) +def test_create_tasks(benchmark, mocker, geometry): + created = [] + mocker.patch.object(create, "get_session") + # list.append is atomic, unlike incrementing a Mock's call_count. + mocker.patch.object( + create, "create_task", side_effect=lambda *args: created.append(args[1]) + ) + + def setup(): + # The benchmark may run multiple rounds. + created.clear() + return _create_tasks_args(geometry) + + benchmark.pedantic(create.create_tasks, setup=setup) + assert len(created) == N + + # --------------------------------------------------------------------------- # Benchmarks – TransformSequence with a simple transform # --------------------------------------------------------------------------- From e0ba000b207e97313b82b799d47bf817ff8e4598 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 14:36:27 +0200 Subject: [PATCH 2/7] perf(parameterization): speed up resolving task definition references _recurse walks every task definition when resolving task references during optimization, and again when resolving timestamps during task creation. For each single-key dictionary it iterated over all the parameter functions, building a set of the dictionary keys to compare against each of them. Now the key is looked up directly in the parameter functions. Dictionaries are also checked before lists, as they are much more common in task definitions. On a typical task definition, resolve_task_references takes 14.5us instead of 18.3us, and resolve_timestamps 20.5us instead of 23.1us. --- src/taskgraph/util/parameterization.py | 22 ++++++------ test/test_graph_perf.py | 47 ++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 12 deletions(-) diff --git a/src/taskgraph/util/parameterization.py b/src/taskgraph/util/parameterization.py index 48e8f8966..ab9df23e3 100644 --- a/src/taskgraph/util/parameterization.py +++ b/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": {"": "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": {"": "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 diff --git a/test/test_graph_perf.py b/test/test_graph_perf.py index f6eb9cf4d..72845e343 100644 --- a/test/test_graph_perf.py +++ b/test/test_graph_perf.py @@ -13,6 +13,11 @@ from taskgraph.task import Task from taskgraph.taskgraph import TaskGraph from taskgraph.transforms.base import TransformSequence +from taskgraph.util.parameterization import ( + resolve_task_references, + resolve_timestamps, +) +from taskgraph.util.time import current_json_time # --------------------------------------------------------------------------- # Graph builders – each returns (tasks_dict, Graph, TaskGraph) for 1000 nodes @@ -229,6 +234,48 @@ def setup(): assert len(created) == N +# --------------------------------------------------------------------------- +# Benchmarks – task definition parameterization +# --------------------------------------------------------------------------- + + +def _make_task_def(i): + return { + "created": {"relative-datestamp": "0 seconds"}, + "deadline": {"relative-datestamp": "1 day"}, + "expires": {"relative-datestamp": "28 days"}, + "metadata": {"name": f"task-{i}", "description": "d", "owner": "o"}, + "routes": [f"index.domain.v2.project.task-{i}"], + "payload": { + "command": ["run-task", "--", "bash", "-c", f"echo {i}"], + "env": { + "BUILD": {"task-reference": ""}, + "ARTIFACT": {"artifact-reference": ""}, + **{f"VAR{j}": f"value-{j}" for j in range(20)}, + }, + "artifacts": [{"name": f"public/a{j}", "path": "/x"} for j in range(5)], + }, + } + + +TASK_DEFS = [_make_task_def(i) for i in range(N)] + + +@pytest.mark.benchmark +def test_resolve_task_references(): + for task_def in TASK_DEFS: + resolve_task_references( + "label", task_def, "task-id", "decision-id", {"build": "build-id"} + ) + + +@pytest.mark.benchmark +def test_resolve_timestamps(): + now = current_json_time(datetime_format=True) + for task_def in TASK_DEFS: + resolve_timestamps(now, task_def) + + # --------------------------------------------------------------------------- # Benchmarks – TransformSequence with a simple transform # --------------------------------------------------------------------------- From b2377657687b0b5a4aaf6aa3074fe5b4c8fe4d4c Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 14:36:27 +0200 Subject: [PATCH 3/7] perf(time): cache parsed relative time strings `value_of` parsed the same few relative time strings ("1 day", "28 days", ...) with a regex every time, although task definitions contain several of them and are resolved multiple times. Its results are timedeltas, which are immutable, so they are now cached. On a typical task definition, resolve_timestamps takes 16.4us instead of 20.5us. --- src/taskgraph/util/time.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/taskgraph/util/time.py b/src/taskgraph/util/time.py index a613f9e83..f2866f4f4 100644 --- a/src/taskgraph/util/time.py +++ b/src/taskgraph/util/time.py @@ -7,6 +7,7 @@ import datetime +import functools import re PATTERN = re.compile(r"((?:\d+)?\.?\d+) *([a-z]+)") @@ -57,6 +58,7 @@ class UnknownTimeMeasurement(Exception): pass +@functools.cache def value_of(input_str): """ Convert a string to a json date in the future From a12665b1bbc01b396d3297fda69e4acbb1ecd15e Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 13:27:32 +0200 Subject: [PATCH 4/7] 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. --- src/taskgraph/graph.py | 13 +++++++++++-- src/taskgraph/optimize/base.py | 4 ++-- test/test_graph_perf.py | 4 ++++ 3 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/taskgraph/graph.py b/src/taskgraph/graph.py index 576325514..f77d71ef3 100644 --- a/src/taskgraph/graph.py +++ b/src/taskgraph/graph.py @@ -107,6 +107,15 @@ 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 @@ -114,14 +123,14 @@ def visit_postorder(self): 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): diff --git a/src/taskgraph/optimize/base.py b/src/taskgraph/optimize/base.py index ba07ca276..aeb6cf5ff 100644 --- a/src/taskgraph/optimize/base.py +++ b/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() diff --git a/test/test_graph_perf.py b/test/test_graph_perf.py index 72845e343..362ac6ab8 100644 --- a/test/test_graph_perf.py +++ b/test/test_graph_perf.py @@ -130,6 +130,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 @@ -138,6 +140,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 From a2e2094dbdecc236bb58c1ff1f1306f41e613a8b Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 13:42:14 +0200 Subject: [PATCH 5/7] perf(optimize): resolve each dependent's deadline only once To decide whether a task can be replaced, replace_tasks computes the latest deadline of its dependents. It resolved the deadline of every dependent for every task, although tasks such as docker images or toolchains share thousands of dependents. Now each task's deadline is resolved at most once, relative to a single `now`. Similarly, IndexSearch parsed the same deadline (and the expiration of tasks used as replacement for multiple tasks) over and over, so parsed timestamps are now cached. On a synthetic graph of 20,210 tasks where the 210 build and docker tasks are replaced, replace_tasks takes 0.17s instead of 0.24s. --- src/taskgraph/optimize/base.py | 21 ++++++++++++++------- src/taskgraph/optimize/strategies.py | 12 ++++++++++-- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/src/taskgraph/optimize/base.py b/src/taskgraph/optimize/base.py index aeb6cf5ff..c425bb23e 100644 --- a/src/taskgraph/optimize/base.py +++ b/src/taskgraph/optimize/base.py @@ -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 diff --git a/src/taskgraph/optimize/strategies.py b/src/taskgraph/optimize/strategies.py index 8fed9e54a..459412c8a 100644 --- a/src/taskgraph/optimize/strategies.py +++ b/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}" ) From b90fe639ee1b9303726f3eb93ad17b88b4f3a951 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 13:42:29 +0200 Subject: [PATCH 6/7] perf(decision): stop re-parsing the full task graph The decision task rebuilt Task objects for the whole full task graph from its JSON representation, only to check that TaskGraph.from_json works. This is already covered by the TaskGraph tests, and costs time proportional to the size of the full task graph on every decision task. --- src/taskgraph/decision.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/taskgraph/decision.py b/src/taskgraph/decision.py index a27eda672..7d927ecd8 100644 --- a/src/taskgraph/decision.py +++ b/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())) From 55031b89c6fe848df00078da4e76f7558ab39629 Mon Sep 17 00:00:00 2001 From: Marco Castelluccio Date: Thu, 24 Sep 2026 13:43:59 +0200 Subject: [PATCH 7/7] perf(generator): look up kind dependency tasks by kind Before loading a kind, the generator filtered all the tasks loaded so far to find the ones belonging to the kind's dependencies (copying them first when loading kinds in parallel). This is proportional to the number of kinds times the number of tasks, and in parallel mode it happens on the main thread, delaying the submission of newly unblocked kinds. Now loaded tasks are also grouped by kind, so the tasks of each kind dependency are looked up directly. They are still passed in the order in which kinds were loaded. With 150 kinds of 270 tasks each (40,500 tasks), each depending on three other kinds, gathering the kind dependency tasks takes 0.01s in total instead of 0.42s. --- src/taskgraph/generator.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index bc8f9c400..1368ef01c 100644 --- a/src/taskgraph/generator.py +++ b/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]