Bug 1759030 - [gecko_taskgraph] Use 'generator.py' from upstream Taskgraph, r=releng-reviewers,jlorenzo

Differential Revision: https://phabricator.services.mozilla.com/D161057

python/mozbuild/mozbuild/toolchains.py

@@ -8,7 +8,8 @@ import six
 
 
 def toolchain_task_definitions():
-    from gecko_taskgraph.generator import load_tasks_for_kind
+    import gecko_taskgraph  # noqa: triggers override of the `graph_config_schema`
+    from taskgraph.generator import load_tasks_for_kind
 
     # Don't import globally to allow this module being imported without
     # the taskgraph module being available (e.g. standalone js)

taskcluster/gecko_taskgraph/__init__.py

@@ -46,11 +46,11 @@ def register(graph_config):
     Args:
         graph_config: The graph configuration object.
     """
+    from taskgraph import generator
     from gecko_taskgraph.parameters import register_parameters
     from gecko_taskgraph import (  # noqa: trigger target task method registration
         target_tasks,
     )
-    from gecko_taskgraph import generator
     from gecko_taskgraph import morph  # noqa: trigger morph registration
     from gecko_taskgraph.util.verify import verifications
 

taskcluster/gecko_taskgraph/decision.py

@@ -22,6 +22,7 @@ from taskgraph.decision import (
     _determine_more_accurate_base_rev,
     _get_env_prefix,
 )
+from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import Parameters
 from taskgraph.taskgraph import TaskGraph
 from taskgraph.util.python_path import find_object
@@ -33,7 +34,6 @@ from voluptuous import Any, Optional, Required
 
 from . import GECKO
 from .actions import render_actions_json
-from .generator import TaskGraphGenerator
 from .parameters import get_app_version, get_version
 from .try_option_syntax import parse_message
 from .util.backstop import BACKSTOP_INDEX, is_backstop

taskcluster/gecko_taskgraph/docker.py

@@ -8,10 +8,10 @@ import os
 import tarfile
 from io import BytesIO
 
+from taskgraph.generator import load_tasks_for_kind
 from taskgraph.parameters import Parameters
 from taskgraph.util.taskcluster import get_session, get_artifact_url
 
-from gecko_taskgraph.generator import load_tasks_for_kind
 from gecko_taskgraph.optimize.strategies import IndexSearch
 from gecko_taskgraph.util import docker
 from . import GECKO

taskcluster/gecko_taskgraph/generator.py

@@ -1,447 +0,0 @@
-# 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 logging
-import os
-import copy
-
-import attr
-from taskgraph import filter_tasks
-from taskgraph.config import GraphConfig, load_graph_config
-from taskgraph.graph import Graph
-from taskgraph.morph import morph
-from taskgraph.optimize.base import optimize_task_graph
-from taskgraph.parameters import parameters_loader
-from taskgraph.task import Task
-from taskgraph.taskgraph import TaskGraph
-from taskgraph.transforms.base import TransformSequence, TransformConfig
-from taskgraph.util.python_path import find_object
-from taskgraph.util.yaml import load_yaml
-from taskgraph.util.verify import verifications
-
-logger = logging.getLogger(__name__)
-
-
-class KindNotFound(Exception):
-    """
-    Raised when trying to load kind from a directory without a kind.yml.
-    """
-
-
-@attr.s(frozen=True)
-class Kind:
-
-    name = attr.ib(type=str)
-    path = attr.ib(type=str)
-    config = attr.ib(type=dict)
-    graph_config = attr.ib(type=GraphConfig)
-
-    def _get_loader(self):
-        try:
-            loader = self.config["loader"]
-        except KeyError:
-            raise KeyError(f"{self.path!r} does not define `loader`")
-        return find_object(loader)
-
-    def load_tasks(self, parameters, loaded_tasks, write_artifacts):
-        loader = self._get_loader()
-        config = copy.deepcopy(self.config)
-
-        kind_dependencies = config.get("kind-dependencies", [])
-        kind_dependencies_tasks = {
-            task.label: task for task in loaded_tasks if task.kind in kind_dependencies
-        }
-
-        inputs = loader(self.name, self.path, config, parameters, loaded_tasks)
-
-        transforms = TransformSequence()
-        for xform_path in config["transforms"]:
-            transform = find_object(xform_path)
-            transforms.add(transform)
-
-        # perform the transformations on the loaded inputs
-        trans_config = TransformConfig(
-            self.name,
-            self.path,
-            config,
-            parameters,
-            kind_dependencies_tasks,
-            self.graph_config,
-            write_artifacts=write_artifacts,
-        )
-        tasks = [
-            Task(
-                self.name,
-                label=task_dict["label"],
-                description=task_dict["description"],
-                attributes=task_dict["attributes"],
-                task=task_dict["task"],
-                optimization=task_dict.get("optimization"),
-                dependencies=task_dict.get("dependencies"),
-                soft_dependencies=task_dict.get("soft-dependencies"),
-                if_dependencies=task_dict.get("if-dependencies"),
-            )
-            for task_dict in transforms(trans_config, inputs)
-        ]
-        return tasks
-
-    @classmethod
-    def load(cls, root_dir, graph_config, kind_name):
-        path = os.path.join(root_dir, kind_name)
-        kind_yml = os.path.join(path, "kind.yml")
-        if not os.path.exists(kind_yml):
-            raise KindNotFound(kind_yml)
-
-        logger.debug(f"loading kind `{kind_name}` from `{path}`")
-        config = load_yaml(kind_yml)
-
-        return cls(kind_name, path, config, graph_config)
-
-
-class TaskGraphGenerator:
-    """
-    The central controller for taskgraph.  This handles all phases of graph
-    generation.  The task is generated from all of the kinds defined in
-    subdirectories of the generator's root directory.
-
-    Access to the results of this generation, as well as intermediate values at
-    various phases of generation, is available via properties.  This encourages
-    the provision of all generation inputs at instance construction time.
-    """
-
-    # Task-graph generation is implemented as a Python generator that yields
-    # each "phase" of generation.  This allows some mach subcommands to short-
-    # circuit generation of the entire graph by never completing the generator.
-
-    def __init__(
-        self,
-        root_dir,
-        parameters,
-        decision_task_id="DECISION-TASK",
-        write_artifacts=False,
-    ):
-        """
-        @param root_dir: root directory, with subdirectories for each kind
-        @param paramaters: parameters for this task-graph generation, or callable
-            taking a `GraphConfig` and returning parameters
-        @type parameters: Union[Parameters, Callable[[GraphConfig], Parameters]]
-        """
-        if root_dir is None:
-            root_dir = "taskcluster/ci"
-        self.root_dir = root_dir
-        self._parameters = parameters
-        self._decision_task_id = decision_task_id
-        self._write_artifacts = write_artifacts
-
-        # start the generator
-        self._run = self._run()
-        self._run_results = {}
-
-    @property
-    def parameters(self):
-        """
-        The properties used for this graph.
-
-        @type: Properties
-        """
-        return self._run_until("parameters")
-
-    @property
-    def full_task_set(self):
-        """
-        The full task set: all tasks defined by any kind (a graph without edges)
-
-        @type: TaskGraph
-        """
-        return self._run_until("full_task_set")
-
-    @property
-    def full_task_graph(self):
-        """
-        The full task graph: the full task set, with edges representing
-        dependencies.
-
-        @type: TaskGraph
-        """
-        return self._run_until("full_task_graph")
-
-    @property
-    def target_task_set(self):
-        """
-        The set of targetted tasks (a graph without edges)
-
-        @type: TaskGraph
-        """
-        return self._run_until("target_task_set")
-
-    @property
-    def target_task_graph(self):
-        """
-        The set of targetted tasks and all of their dependencies
-
-        @type: TaskGraph
-        """
-        return self._run_until("target_task_graph")
-
-    @property
-    def optimized_task_graph(self):
-        """
-        The set of targetted tasks and all of their dependencies; tasks that
-        have been optimized out are either omitted or replaced with a Task
-        instance containing only a task_id.
-
-        @type: TaskGraph
-        """
-        return self._run_until("optimized_task_graph")
-
-    @property
-    def label_to_taskid(self):
-        """
-        A dictionary mapping task label to assigned taskId.  This property helps
-        in interpreting `optimized_task_graph`.
-
-        @type: dictionary
-        """
-        return self._run_until("label_to_taskid")
-
-    @property
-    def morphed_task_graph(self):
-        """
-        The optimized task graph, with any subsequent morphs applied. This graph
-        will have the same meaning as the optimized task graph, but be in a form
-        more palatable to TaskCluster.
-
-        @type: TaskGraph
-        """
-        return self._run_until("morphed_task_graph")
-
-    @property
-    def graph_config(self):
-        """
-        The configuration for this graph.
-
-        @type: TaskGraph
-        """
-        return self._run_until("graph_config")
-
-    def _load_kinds(self, graph_config, target_kind=None):
-        if target_kind:
-            # docker-image is an implicit dependency that never appears in
-            # kind-dependencies.
-            queue = [target_kind, "docker-image"]
-            seen_kinds = set()
-            while queue:
-                kind_name = queue.pop()
-                if kind_name in seen_kinds:
-                    continue
-                seen_kinds.add(kind_name)
-                kind = Kind.load(self.root_dir, graph_config, kind_name)
-                yield kind
-                queue.extend(kind.config.get("kind-dependencies", []))
-        else:
-            for kind_name in os.listdir(self.root_dir):
-                try:
-                    yield Kind.load(self.root_dir, graph_config, kind_name)
-                except KindNotFound:
-                    continue
-
-    def _run(self):
-        logger.info("Loading graph configuration.")
-        graph_config = load_graph_config(self.root_dir)
-
-        yield ("graph_config", graph_config)
-
-        graph_config.register()
-
-        # Initial verifications that don't depend on any generation state.
-        verifications("initial")
-
-        if callable(self._parameters):
-            parameters = self._parameters(graph_config)
-        else:
-            parameters = self._parameters
-
-        logger.info("Using {}".format(parameters))
-        logger.debug("Dumping parameters:\n{}".format(repr(parameters)))
-
-        filters = parameters.get("filters", [])
-        # Always add legacy target tasks method until we deprecate that API.
-        if "target_tasks_method" not in filters:
-            filters.insert(0, "target_tasks_method")
-        filters = [filter_tasks.filter_task_functions[f] for f in filters]
-
-        yield self.verify("parameters", parameters)
-
-        logger.info("Loading kinds")
-        # put the kinds into a graph and sort topologically so that kinds are loaded
-        # in post-order
-        if parameters.get("target-kind"):
-            target_kind = parameters["target-kind"]
-            logger.info(
-                "Limiting kinds to {target_kind} and dependencies".format(
-                    target_kind=target_kind
-                )
-            )
-        kinds = {
-            kind.name: kind
-            for kind in self._load_kinds(graph_config, parameters.get("target-kind"))
-        }
-        verifications("kinds", kinds)
-
-        edges = set()
-        for kind in kinds.values():
-            for dep in kind.config.get("kind-dependencies", []):
-                edges.add((kind.name, dep, "kind-dependency"))
-        kind_graph = Graph(set(kinds), edges)
-
-        if parameters.get("target-kind"):
-            kind_graph = kind_graph.transitive_closure({target_kind, "docker-image"})
-
-        logger.info("Generating full task set")
-        all_tasks = {}
-        for kind_name in kind_graph.visit_postorder():
-            logger.debug(f"Loading tasks for kind {kind_name}")
-            kind = kinds[kind_name]
-            try:
-                new_tasks = kind.load_tasks(
-                    parameters,
-                    list(all_tasks.values()),
-                    self._write_artifacts,
-                )
-            except Exception:
-                logger.exception(f"Error loading tasks for kind {kind_name}:")
-                raise
-            for task in new_tasks:
-                if task.label in all_tasks:
-                    raise Exception("duplicate tasks with label " + task.label)
-                all_tasks[task.label] = task
-            logger.info(f"Generated {len(new_tasks)} tasks for kind {kind_name}")
-        full_task_set = TaskGraph(all_tasks, Graph(set(all_tasks), set()))
-        yield self.verify("full_task_set", full_task_set, graph_config, parameters)
-
-        logger.info("Generating full task graph")
-        edges = set()
-        for t in full_task_set:
-            for depname, dep in t.dependencies.items():
-                edges.add((t.label, dep, depname))
-
-        full_task_graph = TaskGraph(all_tasks, Graph(full_task_set.graph.nodes, edges))
-        logger.info(
-            "Full task graph contains %d tasks and %d dependencies"
-            % (len(full_task_set.graph.nodes), len(edges))
-        )
-        yield self.verify("full_task_graph", full_task_graph, graph_config, parameters)
-
-        logger.info("Generating target task set")
-        target_task_set = TaskGraph(
-            dict(all_tasks), Graph(set(all_tasks.keys()), set())
-        )
-        for fltr in filters:
-            old_len = len(target_task_set.graph.nodes)
-            target_tasks = set(fltr(target_task_set, parameters, graph_config))
-            target_task_set = TaskGraph(
-                {l: all_tasks[l] for l in target_tasks}, Graph(target_tasks, set())
-            )
-            logger.info(
-                "Filter %s pruned %d tasks (%d remain)"
-                % (fltr.__name__, old_len - len(target_tasks), len(target_tasks))
-            )
-
-        yield self.verify("target_task_set", target_task_set, graph_config, parameters)
-
-        logger.info("Generating target task graph")
-        # include all docker-image build tasks here, in case they are needed for a graph morph
-        docker_image_tasks = {
-            t.label
-            for t in full_task_graph.tasks.values()
-            if t.attributes["kind"] == "docker-image"
-        }
-        # include all tasks with `always_target` set
-        if parameters["enable_always_target"]:
-            always_target_tasks = {
-                t.label
-                for t in full_task_graph.tasks.values()
-                if t.attributes.get("always_target")
-            }
-        else:
-            always_target_tasks = set()
-        logger.info(
-            "Adding %d tasks with `always_target` attribute"
-            % (len(always_target_tasks) - len(always_target_tasks & target_tasks))
-        )
-        requested_tasks = target_tasks | docker_image_tasks | always_target_tasks
-        target_graph = full_task_graph.graph.transitive_closure(requested_tasks)
-        target_task_graph = TaskGraph(
-            {l: all_tasks[l] for l in target_graph.nodes}, target_graph
-        )
-        yield self.verify(
-            "target_task_graph", target_task_graph, graph_config, parameters
-        )
-
-        logger.info("Generating optimized task graph")
-        existing_tasks = parameters.get("existing_tasks")
-        do_not_optimize = set(parameters.get("do_not_optimize", []))
-        if not parameters.get("optimize_target_tasks", True):
-            do_not_optimize = set(target_task_set.graph.nodes).union(do_not_optimize)
-
-        # this is used for testing experimental optimization strategies
-        strategies = os.environ.get(
-            "TASKGRAPH_OPTIMIZE_STRATEGIES", parameters.get("optimize_strategies")
-        )
-        if strategies:
-            strategies = find_object(strategies)
-
-        optimized_task_graph, label_to_taskid = optimize_task_graph(
-            target_task_graph,
-            requested_tasks,
-            parameters,
-            do_not_optimize,
-            self._decision_task_id,
-            existing_tasks=existing_tasks,
-            strategy_override=strategies,
-        )
-
-        yield self.verify(
-            "optimized_task_graph", optimized_task_graph, graph_config, parameters
-        )
-
-        morphed_task_graph, label_to_taskid = morph(
-            optimized_task_graph, label_to_taskid, parameters, graph_config
-        )
-
-        yield "label_to_taskid", label_to_taskid
-        yield self.verify(
-            "morphed_task_graph", morphed_task_graph, graph_config, parameters
-        )
-
-    def _run_until(self, name):
-        while name not in self._run_results:
-            try:
-                k, v = next(self._run)
-            except StopIteration:
-                raise AttributeError(f"No such run result {name}")
-            self._run_results[k] = v
-        return self._run_results[name]
-
-    def verify(self, name, obj, *args, **kwargs):
-        verifications(name, obj, *args, **kwargs)
-        return name, obj
-
-
-def load_tasks_for_kind(parameters, kind, root_dir=None):
-    """
-    Get all the tasks of a given kind.
-
-    This function is designed to be called from outside of taskgraph.
-    """
-    # make parameters read-write
-    parameters = dict(parameters)
-    parameters["target-kind"] = kind
-    parameters = parameters_loader(spec=None, strict=False, overrides=parameters)
-    tgg = TaskGraphGenerator(root_dir=root_dir, parameters=parameters)
-    return {
-        task.task["metadata"]["name"]: task
-        for task in tgg.full_task_set
-        if task.kind == kind
-    }

taskcluster/gecko_taskgraph/main.py

@@ -101,7 +101,7 @@ FORMAT_METHODS = {
 
 def get_taskgraph_generator(root, parameters):
     """Helper function to make testing a little easier."""
-    from gecko_taskgraph.generator import TaskGraphGenerator
+    from taskgraph.generator import TaskGraphGenerator
 
     return TaskGraphGenerator(root_dir=root, parameters=parameters)
 

taskcluster/gecko_taskgraph/test/conftest.py

@@ -6,18 +6,15 @@ import os
 import pytest
 from mach.logging import LoggingManager
 from responses import RequestsMock
-from taskgraph import target_tasks as target_tasks_mod
+from taskgraph import generator as generator_mod, target_tasks as target_tasks_mod
 from taskgraph.config import GraphConfig, load_graph_config
+from taskgraph.generator import TaskGraphGenerator, Kind
 from taskgraph.optimize import base as optimize_mod
 from taskgraph.optimize.base import OptimizationStrategy
 from taskgraph.parameters import Parameters
 
-from gecko_taskgraph import (
+from gecko_taskgraph import GECKO
-    GECKO,
-    generator,
-)
 from gecko_taskgraph.actions import render_actions_json
-from gecko_taskgraph.generator import TaskGraphGenerator, Kind
 from gecko_taskgraph.util.templates import merge
 
 
@@ -169,7 +166,7 @@ def maketgg(monkeypatch):
         )
         parameters.update(params)
 
-        monkeypatch.setattr(generator, "load_graph_config", fake_load_graph_config)
+        monkeypatch.setattr(generator_mod, "load_graph_config", fake_load_graph_config)
 
         tgg = WithFakeKind("/root", parameters)
         tgg.loaded_kinds = loaded_kinds

taskcluster/gecko_taskgraph/test/python.ini

@@ -4,7 +4,6 @@ subsuite = taskgraph
 [test_actions_util.py]
 [test_decision.py]
 [test_files_changed.py]
-[test_generator.py]
 [test_main.py]
 [test_morph.py]
 [test_optimize_strategies.py]

taskcluster/gecko_taskgraph/test/test_generator.py

@@ -1,137 +0,0 @@
-# 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/.
-
-from mozunit import main
-from taskgraph.graph import Graph
-
-from gecko_taskgraph.generator import load_tasks_for_kind
-from gecko_taskgraph import (
-    generator,
-)
-
-from conftest import (
-    FakeKind,
-    WithFakeKind,
-    fake_load_graph_config,
-)
-
-
-def test_kind_ordering(maketgg):
-    "When task kinds depend on each other, they are loaded in postorder"
-    tgg = maketgg(
-        kinds=[
-            ("_fake3", {"kind-dependencies": ["_fake2", "_fake1"]}),
-            ("_fake2", {"kind-dependencies": ["_fake1"]}),
-            ("_fake1", {"kind-dependencies": []}),
-        ]
-    )
-    tgg._run_until("full_task_set")
-    assert tgg.loaded_kinds == ["_fake1", "_fake2", "_fake3"]
-
-
-def test_full_task_set(maketgg):
-    "The full_task_set property has all tasks"
-    tgg = maketgg()
-    assert tgg.full_task_set.graph == Graph(
-        {"_fake-t-0", "_fake-t-1", "_fake-t-2"}, set()
-    )
-    assert sorted(tgg.full_task_set.tasks.keys()) == sorted(
-        ["_fake-t-0", "_fake-t-1", "_fake-t-2"]
-    )
-
-
-def test_full_task_graph(maketgg):
-    "The full_task_graph property has all tasks, and links"
-    tgg = maketgg()
-    assert tgg.full_task_graph.graph == Graph(
-        {"_fake-t-0", "_fake-t-1", "_fake-t-2"},
-        {
-            ("_fake-t-1", "_fake-t-0", "prev"),
-            ("_fake-t-2", "_fake-t-1", "prev"),
-        },
-    )
-    assert sorted(tgg.full_task_graph.tasks.keys()) == sorted(
-        ["_fake-t-0", "_fake-t-1", "_fake-t-2"]
-    )
-
-
-def test_target_task_set(maketgg):
-    "The target_task_set property has the targeted tasks"
-    tgg = maketgg(["_fake-t-1"])
-    assert tgg.target_task_set.graph == Graph({"_fake-t-1"}, set())
-    assert set(tgg.target_task_set.tasks.keys()) == {"_fake-t-1"}
-
-
-def test_target_task_graph(maketgg):
-    "The target_task_graph property has the targeted tasks and deps"
-    tgg = maketgg(["_fake-t-1"])
-    assert tgg.target_task_graph.graph == Graph(
-        {"_fake-t-0", "_fake-t-1"}, {("_fake-t-1", "_fake-t-0", "prev")}
-    )
-    assert sorted(tgg.target_task_graph.tasks.keys()) == sorted(
-        ["_fake-t-0", "_fake-t-1"]
-    )
-
-
-def test_always_target_tasks(maketgg):
-    "The target_task_graph includes tasks with 'always_target'"
-    tgg_args = {
-        "target_tasks": ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"],
-        "kinds": [
-            ("_fake", {"job-defaults": {"optimization": {"odd": None}}}),
-            (
-                "_ignore",
-                {
-                    "job-defaults": {
-                        "attributes": {"always_target": True},
-                        "optimization": {"even": None},
-                    }
-                },
-            ),
-        ],
-        "params": {"optimize_target_tasks": False, "enable_always_target": True},
-    }
-    tgg = maketgg(**tgg_args)
-    assert sorted(tgg.target_task_set.tasks.keys()) == sorted(
-        ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"]
-    )
-    assert sorted(tgg.target_task_graph.tasks.keys()) == sorted(
-        ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1", "_ignore-t-2"]
-    )
-    assert sorted(t.label for t in tgg.optimized_task_graph.tasks.values()) == sorted(
-        ["_fake-t-0", "_fake-t-1", "_ignore-t-0", "_ignore-t-1"]
-    )
-
-
-def test_optimized_task_graph(maketgg):
-    "The optimized task graph contains task ids"
-    tgg = maketgg(["_fake-t-2"])
-    tid = tgg.label_to_taskid
-    assert tgg.optimized_task_graph.graph == Graph(
-        {tid["_fake-t-0"], tid["_fake-t-1"], tid["_fake-t-2"]},
-        {
-            (tid["_fake-t-1"], tid["_fake-t-0"], "prev"),
-            (tid["_fake-t-2"], tid["_fake-t-1"], "prev"),
-        },
-    )
-
-
-def test_load_tasks_for_kind(monkeypatch):
-    """
-    `load_tasks_for_kinds` will load the tasks for the provided kind
-    """
-    FakeKind.loaded_kinds = []
-    monkeypatch.setattr(generator, "TaskGraphGenerator", WithFakeKind)
-    monkeypatch.setattr(generator, "load_graph_config", fake_load_graph_config)
-
-    tasks = load_tasks_for_kind(
-        {"_kinds": [("_example-kind", []), ("docker-image", [])]},
-        "_example-kind",
-        "/root",
-    )
-    assert "t-1" in tasks and tasks["t-1"].label == "_example-kind-t-1"
-
-
-if __name__ == "__main__":
-    main()

taskcluster/mach_commands.py

@@ -331,7 +331,7 @@ def setup_logging(command_context, quiet=False, verbose=True):
 def show_actions(command_context, options):
     import gecko_taskgraph
     import gecko_taskgraph.actions
-    import gecko_taskgraph.generator
+    from taskgraph.generator import TaskGraphGenerator
     from taskgraph.parameters import parameters_loader
 
     try:
@@ -340,7 +340,7 @@ def show_actions(command_context, options):
         )
         parameters = parameters_loader(options["parameters"])
 
-        tgg = gecko_taskgraph.generator.TaskGraphGenerator(
+        tgg = TaskGraphGenerator(
             root_dir=options.get("root"),
             parameters=parameters,
         )

taskcluster/test/conftest.py

@@ -9,9 +9,9 @@ import os
 
 import pytest
 from responses import RequestsMock, logger as rsps_logger
+from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import parameters_loader
 
-from gecko_taskgraph.generator import TaskGraphGenerator
 from gecko_taskgraph.util.hg import PUSHLOG_PUSHES_TMPL
 from gecko_taskgraph.util.bugbug import BUGBUG_BASE_URL
 

tools/tryselect/tasks.py

@@ -15,10 +15,10 @@ from mach.util import get_state_dir
 from mozbuild.base import MozbuildObject
 from mozpack.files import FileFinder
 from moztest.resolve import TestResolver, TestManifestLoader, get_suite_definition
+from taskgraph.generator import TaskGraphGenerator
 from taskgraph.parameters import ParameterMismatch, parameters_loader
 from taskgraph.taskgraph import TaskGraph
 
-from gecko_taskgraph.generator import TaskGraphGenerator
 
 here = os.path.abspath(os.path.dirname(__file__))
 build = MozbuildObject.from_environment(cwd=here)