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/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())) 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] 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..c425bb23e 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() @@ -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}" ) 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/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 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..362ac6ab8 100644 --- a/test/test_graph_perf.py +++ b/test/test_graph_perf.py @@ -8,10 +8,16 @@ import pytest +from taskgraph import create from taskgraph.graph import Graph 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 @@ -124,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 @@ -132,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 @@ -179,6 +189,97 @@ 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 – 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 # ---------------------------------------------------------------------------