diff --git a/docs/reference/source/taskgraph.util.rst b/docs/reference/source/taskgraph.util.rst index e9e363416..cd5d80920 100644 --- a/docs/reference/source/taskgraph.util.rst +++ b/docs/reference/source/taskgraph.util.rst @@ -60,6 +60,14 @@ taskgraph.util.keyed\_by module :undoc-members: :show-inheritance: +taskgraph.util.memory module +---------------------------- + +.. automodule:: taskgraph.util.memory + :members: + :undoc-members: + :show-inheritance: + taskgraph.util.parameterization module -------------------------------------- diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py index 45cdfc32b..41d24e26c 100644 --- a/src/taskgraph/generator.py +++ b/src/taskgraph/generator.py @@ -6,17 +6,21 @@ import inspect import logging import multiprocessing +import multiprocessing.connection import os +import pickle import platform +import signal +import sys import time +import traceback from concurrent.futures import ( FIRST_COMPLETED, - ProcessPoolExecutor, ThreadPoolExecutor, wait, ) from dataclasses import dataclass -from typing import Callable, Optional, Union +from typing import Callable, Optional, Union, cast from . import filter_tasks from .config import GraphConfig, load_graph_config @@ -27,6 +31,7 @@ from .task import Task from .taskgraph import TaskGraph from .transforms.base import TransformConfig, TransformSequence +from .util.memory import gc_disabled from .util.python_path import find_object from .util.schema import SchemaValidationError from .util.verify import verifications @@ -54,6 +59,12 @@ class KindNotFound(Exception): """ +def _max_workers(): + # os.process_cpu_count is new in Python 3.13. + cpu_count = getattr(os, "process_cpu_count", os.cpu_count) + return cpu_count() or 1 + + @dataclass(frozen=True) class Kind: name: str @@ -138,6 +149,16 @@ def load(cls, root_dir, graph_config, kind_name): return cls(kind_name, path, config, graph_config) +class _RemoteTraceback(Exception): + """The traceback of an exception raised in a child process.""" + + def __init__(self, tb): + self.tb = tb + + def __str__(self): + return f'\n"""\n{self.tb}"""' + + class TaskGraphGenerator: """ The central controller for taskgraph. This handles all phases of graph @@ -413,6 +434,157 @@ def submit_ready_kinds(): return all_tasks + def _load_tasks_forked(self, kinds, kind_graph, parameters): + """Load each kind in a child process forked once the kinds it depends + on are loaded. + + The child shares every task loaded so far with the parent, so the + tasks of the kinds it depends on don't need to be sent to it, and it + sends its own tasks back through a pipe of its own. + """ + all_tasks = {} + tasks_by_kind = {} + # Pipe to each running child -> (kind name, pid) + running = {} + remaining = kinds.copy() + edges = set(kind_graph.edges) + max_workers = _max_workers() + + try: + while remaining or running: + kinds_with_deps = {edge[0] for edge in edges} + ready_kinds = sorted(set(remaining) - kinds_with_deps) + for name in ready_kinds[: max_workers - len(running)]: + kind = remaining.pop(name) + reader, pid = self._fork_kind( + kind, + parameters, + self._get_kind_dependencies_tasks(kind, tasks_by_kind), + ) + running[reader] = (name, pid) + + # Waiting on nothing would block forever. _run made sure there + # is no dependency loop, so some kind is always running here. + assert running + for ready in multiprocessing.connection.wait(list(running)): + # wait returns the objects it was given, but is typed as + # also returning sockets and file descriptors. + reader = cast(multiprocessing.connection.Connection, ready) + name, pid = running[reader] + try: + result = reader.recv() + except EOFError: + # The child exited without sending anything. + result = None + del running[reader] + reader.close() + _, status = os.waitpid(pid, 0) + new_tasks = self._kind_result(name, result, status) + + kind_tasks = tasks_by_kind.setdefault(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 + + edges = {e for e in edges if e[1] != name} + finally: + # Only left with running children on error. + for reader, (_, pid) in running.items(): + reader.close() + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + + return all_tasks + + def _fork_kind(self, kind, parameters, kind_dependencies_tasks): + """Fork a child process loading the tasks of `kind`, and return the + pipe its result comes back through, and its pid.""" + reader, writer = multiprocessing.Pipe(duplex=False) + # The child inherits the disabled garbage collector, so it doesn't + # write to every page it shares with us traversing the objects in them. + with gc_disabled(): + pid = os.fork() + if pid == 0: + # Never leave this block, or return into the caller's code, in + # the child. + status = 1 + try: + reader.close() + status = self._load_kind_in_child( + kind, parameters, kind_dependencies_tasks, writer + ) + finally: + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + os._exit(status) + writer.close() + return reader, pid + + def _load_kind_in_child(self, kind, parameters, kind_dependencies_tasks, writer): + """Load the tasks of `kind` and send them, or the exception loading + them raised, through `writer`. Return the child's exit status.""" + try: + tasks = kind.load_tasks( + parameters, kind_dependencies_tasks, self._write_artifacts + ) + except BaseException as e: + tb = traceback.format_exc() + try: + # Make sure the parent will be able to unpickle it. + pickle.loads(pickle.dumps(e)) + except Exception: + e = Exception(f"{type(e).__name__}: {e}") + writer.send((e, tb, None)) + return 1 + + try: + writer.send((None, None, tasks)) + except Exception as e: + writer.send( + ( + Exception( + f"Could not send the tasks of kind {kind.name} " + f"to the parent process: {type(e).__name__}: {e}" + ), + traceback.format_exc(), + None, + ) + ) + return 1 + return 0 + + @staticmethod + def _kind_result(name, result, status): + """Return the tasks the child loading kind `name` sent, given what it + sent and its exit status. Raise the exception loading them raised in + the child, if any.""" + if result is None: + code = os.waitstatus_to_exitcode(status) + if code < 0: + how = f"was killed by signal {signal.Signals(-code).name}" + else: + how = f"exited with status {code}" + raise Exception( + f"Process loading tasks for kind {name} {how} without sending them" + ) + + exc, tb, tasks = result + if exc is not None: + exc.__cause__ = _RemoteTraceback(tb) + if isinstance(exc, SchemaValidationError): + logger.error(f"Error loading tasks for kind {name}:\n{exc}") + else: + logger.error( + f"Error loading tasks for kind {name}:", + exc_info=(type(exc), exc, exc.__traceback__), + ) + raise exc + return tasks + def _run(self): logger.info("Loading graph configuration.") graph_config = load_graph_config(self.root_dir) @@ -460,6 +632,12 @@ def _run(self): for kind in kinds.values(): for dep in kind.config.get("kind-dependencies", []): edges.add((kind.name, dep, "kind-dependency")) + for _, dep, _ in edges: + if dep not in kinds: + message = f'Could not find the kind "{dep}"\nAvailable kinds:\n' + for k in sorted(kinds): + message += f' - "{k}"\n' + raise Exception(message) kind_graph = Graph(frozenset(kinds), frozenset(edges)) if target_kinds: @@ -467,6 +645,11 @@ def _run(self): set(target_kinds) | {"docker-image"} ) + # The task loaders wait for each kind's dependencies to be loaded + # before loading it, so a dependency loop would hang them or silently + # leave kinds out. This raises if there is one. + kind_graph.visit_postorder() + yield "kind_graph", kind_graph logger.info("Generating full task set") @@ -474,7 +657,7 @@ def _run(self): # The next block deals with enabling parallel kind processing, which # currently has different support on different platforms. In summary: # * Parallel kind processing is supported and enabled by default on - # Linux. We use multiple processes by default, but experimental + # Linux. We fork a process per kind by default, but experimental # support for multiple threads can be enabled instead. # * On other platforms, we have experimental support for parallel # kind processing with multiple threads. @@ -494,21 +677,19 @@ def load_tasks(): if os.environ.get("TASKGRAPH_SERIAL"): return self._load_tasks_serial(kinds, kind_graph, parameters) elif os.environ.get("TASKGRAPH_USE_THREADS"): - executor = ThreadPoolExecutor(max_workers=os.process_cpu_count()) - else: - executor = ProcessPoolExecutor( - mp_context=multiprocessing.get_context("fork") + executor = ThreadPoolExecutor(max_workers=_max_workers()) + return self._load_tasks_parallel( + kinds, kind_graph, parameters, executor ) - return self._load_tasks_parallel( - kinds, kind_graph, parameters, executor - ) + else: + return self._load_tasks_forked(kinds, kind_graph, parameters) else: if os.environ.get("TASKGRAPH_SERIAL") or not os.environ.get( "TASKGRAPH_USE_THREADS" ): return self._load_tasks_serial(kinds, kind_graph, parameters) else: - executor = ThreadPoolExecutor(max_workers=os.process_cpu_count()) + executor = ThreadPoolExecutor(max_workers=_max_workers()) return self._load_tasks_parallel( kinds, kind_graph, parameters, executor ) @@ -619,7 +800,11 @@ def _run_until(self, name): while name not in self._run_results: start = time.monotonic() try: - k, v = next(self._run) # type: ignore + # Generation creates millions of objects that stay alive, so + # collecting garbage meanwhile only traverses them for nothing. + # The callers' code runs with the collector in its usual state. + with gc_disabled(): + k, v = next(self._run) # type: ignore except StopIteration: raise AttributeError(f"No such run result {name}") except BaseException: diff --git a/src/taskgraph/taskgraph.py b/src/taskgraph/taskgraph.py index aee37a178..fdfabf1b2 100644 --- a/src/taskgraph/taskgraph.py +++ b/src/taskgraph/taskgraph.py @@ -6,6 +6,7 @@ from .graph import Graph from .task import Task +from .util.memory import gc_disabled @dataclass(frozen=True) @@ -58,13 +59,14 @@ def from_json(cls, tasks_dict): """ tasks = {} edges = set() - for key, value in tasks_dict.items(): - tasks[key] = Task.from_json(value) - if "task_id" in value: - tasks[key].task_id = value["task_id"] - for depname, dep in value["dependencies"].items(): - # Task filtering can cause dependencies to be removed from the graph. - if dep in tasks_dict: - edges.add((key, dep, depname)) + with gc_disabled(): + for key, value in tasks_dict.items(): + tasks[key] = Task.from_json(value) + if "task_id" in value: + tasks[key].task_id = value["task_id"] + for depname, dep in value["dependencies"].items(): + # Task filtering can cause dependencies to be removed from the graph. + if dep in tasks_dict: + edges.add((key, dep, depname)) task_graph = cls(tasks, Graph(frozenset(tasks), edges)) return tasks, task_graph diff --git a/src/taskgraph/util/memory.py b/src/taskgraph/util/memory.py new file mode 100644 index 000000000..3dc4117c0 --- /dev/null +++ b/src/taskgraph/util/memory.py @@ -0,0 +1,24 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +import gc +from contextlib import contextmanager + + +@contextmanager +def gc_disabled(): + """Disable the garbage collector for the duration of the block, restoring + its previous state afterwards. + + Generating or loading a task graph creates millions of objects that stay + alive, so collections keep traversing a growing heap without freeing + anything. + """ + was_enabled = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable() diff --git a/test/test_generator.py b/test/test_generator.py index fa65f2302..2dda5c671 100644 --- a/test/test_generator.py +++ b/test/test_generator.py @@ -3,12 +3,16 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. +import gc import os import platform +import signal +import time from concurrent.futures import ProcessPoolExecutor import pytest from pytest_taskgraph import WithFakeKind, fake_load_graph_config +from pytest_taskgraph.fixtures.gen import FakeKind from taskgraph import generator, graph from taskgraph.generator import Kind, load_tasks_for_kind, load_tasks_for_kinds @@ -27,7 +31,7 @@ ) -class FakePPE(ProcessPoolExecutor): +class FakeTPE(ProcessPoolExecutor): loaded_kinds = [] def submit(self, kind_load_tasks, *args): @@ -35,18 +39,39 @@ def submit(self, kind_load_tasks, *args): return super().submit(kind_load_tasks, *args) -class FakeTPE(ProcessPoolExecutor): - loaded_kinds = [] +class RecordingKind(FakeKind): + """Records which process loaded it, and the kinds of the dependency + tasks it was given, in the file at `record_path`.""" - def submit(self, kind_load_tasks, *args): - self.loaded_kinds.append(kind_load_tasks.__self__.name) - return super().submit(kind_load_tasks, *args) + record_path = None + + def load_tasks(self, parameters, loaded_tasks, write_artifacts): + dep_kinds = sorted({t.kind for t in loaded_tasks.values()}) + line = f"{self.name} {os.getpid()} {','.join(dep_kinds)}\n" + fd = os.open(self.record_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT) + try: + os.write(fd, line.encode()) + finally: + os.close(fd) + return super().load_tasks(parameters, loaded_tasks, write_artifacts) + + +def use_kind_class(tgg, cls): + """Make `tgg` load its kinds as instances of `cls`.""" + load_kinds = tgg._load_kinds + + def _load_kinds(graph_config, target_kinds=None): + for kind in load_kinds(graph_config, target_kinds): + yield cls(kind.name, kind.path, kind.config, kind.graph_config) + + tgg._load_kinds = _load_kinds @linuxonly -def test_kind_ordering_multiprocess(mocker, maketgg): - "When task kinds depend on each other, they are loaded in postorder" - mocked_ppe = mocker.patch.object(generator, "ProcessPoolExecutor", new=FakePPE) +def test_kind_ordering_forked(monkeypatch, tmp_path, maketgg): + """Each kind is loaded in a child process, with its subclass' load_tasks, + after the kinds it depends on, and is given their tasks.""" + monkeypatch.setattr(RecordingKind, "record_path", str(tmp_path / "record")) tgg = maketgg( kinds=[ ("_fake3", {"kind-dependencies": ["_fake2", "_fake1"]}), @@ -54,8 +79,124 @@ def test_kind_ordering_multiprocess(mocker, maketgg): ("_fake1", {"kind-dependencies": []}), ] ) + use_kind_class(tgg, RecordingKind) tgg._run_until("full_task_set") - assert mocked_ppe.loaded_kinds == ["_fake1", "_fake2", "_fake3"] + assert len(tgg.full_task_set.tasks) == 9 + + records = [ + line.split(" ") for line in (tmp_path / "record").read_text().splitlines() + ] + assert [(name, deps) for name, _, deps in records] == [ + ("_fake1", ""), + ("_fake2", "_fake1"), + ("_fake3", "_fake1,_fake2"), + ] + pids = {int(pid) for _, pid, _ in records} + assert len(pids) == 3 + assert os.getpid() not in pids + + +@linuxonly +def test_forked_max_workers(mocker, monkeypatch, tmp_path, maketgg): + "Independent kinds are all loaded when fewer can be loaded at once." + mocker.patch.object(generator, "_max_workers", return_value=1) + tgg = maketgg(kinds=[(f"_fake{i}", {}) for i in range(4)]) + assert len(tgg.full_task_set.tasks) == 12 + + +class UnpicklableError(Exception): + def __init__(self, a, b): + super().__init__(f"{a} {b}") + + +@linuxonly +def test_forked_exception_traceback(mocker, caplog, maketgg): + """An exception raised loading a kind in a child is raised in the parent, + with the traceback from the child.""" + + def fail_in_child(self, *args, **kwargs): + raise UnpicklableError("not", "picklable") + + mocker.patch.object(Kind, "load_tasks", fail_in_child) + tgg = maketgg() + + with caplog.at_level("ERROR"): + with pytest.raises(Exception, match="UnpicklableError: not picklable") as e: + tgg._run_until("full_task_set") + + assert "in fail_in_child" in str(e.value.__cause__) + records = [r for r in caplog.records if "Error loading tasks" in r.message] + assert len(records) == 1 + assert records[0].message == "Error loading tasks for kind _fake:" + assert "in fail_in_child" in caplog.text + + +@linuxonly +@pytest.mark.parametrize( + "exit,expected", + ( + pytest.param(lambda: os._exit(3), "exited with status 3", id="exit"), + pytest.param( + lambda: os.kill(os.getpid(), signal.SIGKILL), + "was killed by signal SIGKILL", + id="killed", + ), + ), +) +def test_forked_child_dies(mocker, maketgg, exit, expected): + "A child exiting without sending its tasks is reported as an error." + + def die(self, *args, **kwargs): + exit() + + mocker.patch.object(Kind, "load_tasks", die) + tgg = maketgg() + with pytest.raises( + Exception, match=f"Process loading tasks for kind _fake {expected}" + ): + tgg._run_until("full_task_set") + + +@linuxonly +def test_forked_unpicklable_result(mocker, maketgg): + "Tasks that can't be sent back are reported as an error loading the kind." + + def unpicklable(self, *args, **kwargs): + return [lambda: None] + + mocker.patch.object(Kind, "load_tasks", unpicklable) + tgg = maketgg() + with pytest.raises( + Exception, match="Could not send the tasks of kind _fake to the parent" + ): + tgg._run_until("full_task_set") + + +@linuxonly +def test_forked_children_killed_on_error(mocker, tmp_path, maketgg): + "When loading a kind fails, the children loading other kinds are killed." + pid_file = tmp_path / "pid" + load_tasks = Kind.load_tasks + + def fail_or_hang(self, *args, **kwargs): + if self.name == "_fail": + # Wait for the other child to be running. + while not pid_file.exists(): + time.sleep(0.01) + raise RuntimeError("failed") + if self.name == "_hang": + pid_file.write_text(str(os.getpid())) + time.sleep(60) + return load_tasks(self, *args, **kwargs) + + mocker.patch.object(Kind, "load_tasks", fail_or_hang) + tgg = maketgg(kinds=[("_fail", {}), ("_hang", {})]) + start = time.monotonic() + with pytest.raises(RuntimeError, match="failed"): + tgg._run_until("full_task_set") + assert time.monotonic() - start < 30 + with pytest.raises(ProcessLookupError): + os.kill(int(pid_file.read_text()), 0) @threadsonly @@ -261,6 +402,50 @@ def test_verifications(mocker, maketgg): m.assert_not_called() +@pytest.mark.parametrize("enabled", (True, False)) +def test_gc_disabled_during_generation(mocker, maketgg, enabled): + """The garbage collector is disabled while generating each phase, and + restored to its previous state between phases.""" + states = [] + mocker.patch.object( + generator, "verifications", side_effect=lambda *a: states.append(gc.isenabled()) + ) + tgg = maketgg(["_fake-t-2"]) + was_enabled = gc.isenabled() + if not enabled: + gc.disable() + try: + tgg.full_task_set + assert gc.isenabled() == enabled + tgg.morphed_task_graph + assert gc.isenabled() == enabled + finally: + if was_enabled: + gc.enable() + assert states and not any(states) + + +@linuxonly +def test_gc_disabled_in_forked_children(mocker, tmp_path, maketgg): + "Children loading kinds inherit the disabled garbage collector." + record = tmp_path / "gc" + load_tasks = Kind.load_tasks + + def record_gc(self, *args, **kwargs): + record.write_text(str(gc.isenabled())) + return load_tasks(self, *args, **kwargs) + + mocker.patch.object(Kind, "load_tasks", record_gc) + tgg = maketgg() + tgg._load_tasks_forked( + {k.name: k for k in tgg._load_kinds(tgg.graph_config)}, + tgg.kind_graph, + tgg.parameters, + ) + assert gc.isenabled() + assert record.read_text() == "False" + + def test_load_tasks_for_kind(monkeypatch): """ `load_tasks_for_kinds` will load the tasks for the provided kind @@ -437,6 +622,31 @@ def test_kind_graph(maketgg): } +def test_kind_graph_missing_kind(maketgg): + "A dependency on a kind that doesn't exist is an error" + tgg = maketgg( + kinds=[ + ("_fake2", {"kind-dependencies": ["_missing"]}), + ("_fake1", {"kind-dependencies": []}), + ] + ) + with pytest.raises(Exception, match='Could not find the kind "_missing"'): + tgg.kind_graph + + +def test_kind_graph_dependency_loop(maketgg): + "A dependency loop between kinds is an error" + tgg = maketgg( + kinds=[ + ("_fake3", {"kind-dependencies": ["_fake2"]}), + ("_fake2", {"kind-dependencies": ["_fake3"]}), + ("_fake1", {"kind-dependencies": []}), + ] + ) + with pytest.raises(Exception, match="Dependency loop detected"): + tgg.kind_graph + + def test_kind_graph_with_target_kinds(maketgg): "The kind_graph property respects target_kinds parameter" tgg = maketgg( diff --git a/test/test_util_memory.py b/test/test_util_memory.py new file mode 100644 index 000000000..8b5495d98 --- /dev/null +++ b/test/test_util_memory.py @@ -0,0 +1,38 @@ +# This Source Code Form is subject to the terms of the Mozilla Public +# License, v. 2.0. If a copy of the MPL was not distributed with this +# file, You can obtain one at http://mozilla.org/MPL/2.0/. + +import gc + +import pytest + +from taskgraph.util.memory import gc_disabled + + +@pytest.fixture +def restore_gc(): + was_enabled = gc.isenabled() + yield + if was_enabled: + gc.enable() + else: + gc.disable() + + +def test_gc_disabled(restore_gc): + gc.enable() + with gc_disabled(): + assert not gc.isenabled() + assert gc.isenabled() + + with pytest.raises(RuntimeError): + with gc_disabled(): + raise RuntimeError() + assert gc.isenabled() + + +def test_gc_disabled_keeps_disabled(restore_gc): + gc.disable() + with gc_disabled(): + assert not gc.isenabled() + assert not gc.isenabled()