Bug 1333255: and then there was only one Task class; r?jonasfj draft
authorDustin J. Mitchell <dustin@mozilla.com>
Thu, 09 Mar 2017 23:14:30 -0500
changeset 496943 34c580c283df2af7a7c3385456004b7119b762b7
parent 496942 aa552675d68c88a751370ba300384a8ab6803f00
child 496944 e97ea08fbb6e5738dd5a6874a622f2d5c62a379c
push id48749
push userdmitchell@mozilla.com
push dateFri, 10 Mar 2017 22:45:14 +0000
reviewersjonasfj
bugs1333255
milestone55.0a1
Bug 1333255: and then there was only one Task class; r?jonasfj Note that the to_json method prefers the taskgraph's dependencies information (edges) to that from the task.dependencies entries. At a few points in task-graph generation, these values differ, although that is expected (for example, the full task set contains no edges, but that information is still in task.dependencies). Unifying that representation leads to some difficulty with task transforms that reach into the dependency tree (beetmover), so the different representations are left as-is. MozReview-Commit-ID: GeW8HNwFA9Z
graphdiff.py
taskcluster/ci/docker-image/kind.yml
taskcluster/docs/taskgraph.rst
taskcluster/taskgraph/task/base.py
taskcluster/taskgraph/task/docker_image.py
taskcluster/taskgraph/task/test.py
taskcluster/taskgraph/task/transform.py
taskcluster/taskgraph/taskgraph.py
taskcluster/taskgraph/test/test_create.py
taskcluster/taskgraph/test/test_decision.py
taskcluster/taskgraph/test/test_optimize.py
taskcluster/taskgraph/test/test_target_tasks.py
taskcluster/taskgraph/test/test_taskgraph.py
taskcluster/taskgraph/test/test_try_option_syntax.py
taskcluster/taskgraph/test/util.py
taskcluster/taskgraph/transforms/docker_image.py
--- a/graphdiff.py
+++ b/graphdiff.py
@@ -24,16 +24,24 @@ def relabel(graph, before):
 
 def remove_stuff(graph):
     for v in graph.itervalues():
         # kind_implementation goes away in a few commits, so let's ignore for now
         try:
             del v['kind_implementation']
         except KeyError:
             pass
+        try:
+            del v['kind']
+        except KeyError:
+            pass
+        try:
+            del v['optimizations']
+        except KeyError:
+            pass
     return graph
 
 def update_docker_builds(graph1, graph2):
     # docker image build tasks are now produced via the task description transform,
     # so they are quite a bit differetj
     for label, task in graph1.iteritems():
         if not label.startswith('build-docker-image'):
             continue
@@ -58,16 +66,19 @@ def update_docker_builds(graph1, graph2)
         task['task']['extra']['treeherder']['groupName'] = "Docker Image Builds"
 
         # set an explicit tier
         task['task']['extra']['treeherder']['tier'] = 1
 
         # and of course this order changes, having no effect
         task['task']['extra']['treeherderEnv'].sort()
 
+        # imageMeta is no longer needed, so no longer included
+        del task['task']['extra']['imageMeta']
+
         # copy metadata from the new graph
         task['task']['metadata'] = graph2[label]['task']['metadata']
 
         # remove empty env var
         del task['task']['payload']['env']['CONTEXT_URL']
 
         # task.py adds this, and it doesn't hurt
         task['task']['payload']['env']['SCCACHE_DISABLE'] = "1"
--- a/taskcluster/ci/docker-image/kind.yml
+++ b/taskcluster/ci/docker-image/kind.yml
@@ -1,13 +1,13 @@
 # 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/.
 
-loader: taskgraph.task.docker_image:load_tasks
+loader: taskgraph.task.transform:load_tasks
 
 transforms:
   - taskgraph.transforms.docker_image:transforms
   - taskgraph.transforms.task:transforms
 
 # make a task for each docker-image we might want.  For the moment, since we
 # write artifacts for each, these are whitelisted, but ideally that will change
 # (to use subdirectory clones of the proper directory), at which point we can
--- a/taskcluster/docs/taskgraph.rst
+++ b/taskcluster/docs/taskgraph.rst
@@ -203,36 +203,38 @@ Task graphs -- both the graph artifacts 
 output by the ``--json`` option to the ``mach taskgraph`` commands -- are JSON
 objects, keyed by label, or for optimized task graphs, by taskId.  For
 convenience, the decision task also writes out ``label-to-taskid.json``
 containing a mapping from label to taskId.  Each task in the graph is
 represented as a JSON object.
 
 Each task has the following properties:
 
+``kind``
+   The name of this task's kind
+
 ``task_id``
    The task's taskId (only for optimized task graphs)
 
 ``label``
    The task's label
 
 ``attributes``
    The task's attributes
 
 ``dependencies``
    The task's in-graph dependencies, represented as an object mapping
    dependency name to label (or to taskId for optimized task graphs)
 
+``optimizations``
+   The optimizations to be applied to this task
+
 ``task``
    The task's TaskCluster task definition.
 
-``kind_implementation``
-   The module and the class name which was used to implement this particular task.
-   It is always of the form ``<module-path>:<object-path>``
-
 The results from each command are in the same format, but with some differences
 in the content:
 
 * The ``tasks`` and ``target`` subcommands both return graphs with no edges.
   That is, just collections of tasks without any dependencies indicated.
 
 * The ``optimized`` subcommand returns tasks that have been assigned taskIds.
   The dependencies array, too, contains taskIds instead of labels, with
--- a/taskcluster/taskgraph/task/base.py
+++ b/taskcluster/taskgraph/task/base.py
@@ -1,16 +1,14 @@
 # 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 __future__ import absolute_import, print_function, unicode_literals
 
-import abc
-
 
 class Task(object):
     """
     Representation of a task in a TaskGraph.  Each Task has, at creation:
 
     - kind: the name of the task kind
     - label; the label for this task
     - attributes: a dictionary of attributes for this task (used for filtering)
@@ -19,23 +17,19 @@ class Task(object):
     - dependencies: tasks this one depends on, in the form {name: label}, for example
       {'build': 'build-linux64/opt', 'docker-image': 'build-docker-image-desktop-test'}
 
     And later, as the task-graph processing proceeds:
 
     - task_id -- TaskCluster taskId under which this task will be created
     - optimized -- true if this task need not be performed
 
-    A kind represents a collection of tasks that share common characteristics.
-    For example, all build jobs.  Each instance of a kind is intialized with a
-    path from which it draws its task configuration.  The instance is free to
-    store as much local state as it needs.
+    This class is just a convenience wraper for the data type and managing
+    display, comparison, serialization, etc. It has no functionality of its own.
     """
-    __metaclass__ = abc.ABCMeta
-
     def __init__(self, kind, label, attributes, task,
                  optimizations=None, dependencies=None):
         self.kind = kind
         self.label = label
         self.attributes = attributes
         self.task = task
 
         self.task_id = None
@@ -50,20 +44,43 @@ class Task(object):
         return self.kind == other.kind and \
             self.label == other.label and \
             self.attributes == other.attributes and \
             self.task == other.task and \
             self.task_id == other.task_id and \
             self.optimizations == other.optimizations and \
             self.dependencies == other.dependencies
 
+    def __repr__(self):
+        return ('Task({kind!r}, {label!r}, {attributes!r}, {task!r}, '
+                'optimizations={optimizations!r}, '
+                'dependencies={dependencies!r})'.format(**self.__dict__))
+
+    def to_json(self):
+        rv = {
+            'kind': self.kind,
+            'label': self.label,
+            'attributes': self.attributes,
+            'dependencies': self.dependencies,
+            'optimizations': self.optimizations,
+            'task': self.task,
+        }
+        if self.task_id:
+            rv['task_id'] = self.task_id
+        return rv
+
     @classmethod
     def from_json(cls, task_dict):
         """
         Given a data structure as produced by taskgraph.to_json, re-construct
         the original Task object.  This is used to "resume" the task-graph
         generation process, for example in Action tasks.
         """
-        return cls(
-            kind=task_dict['attributes']['kind'],
+        rv = cls(
+            kind=task_dict['kind'],
             label=task_dict['label'],
             attributes=task_dict['attributes'],
-            task=task_dict['task'])
+            task=task_dict['task'],
+            optimizations=task_dict['optimizations'],
+            dependencies=task_dict.get('dependencies'))
+        if 'task_id' in task_dict:
+            rv.task_id = task_dict['task_id']
+        return rv
deleted file mode 100644
--- a/taskcluster/taskgraph/task/docker_image.py
+++ /dev/null
@@ -1,43 +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 __future__ import absolute_import, print_function, unicode_literals
-
-import logging
-
-from . import transform
-from taskgraph.transforms.base import TransformSequence, TransformConfig
-from taskgraph.util.python_path import find_object
-
-logger = logging.getLogger(__name__)
-
-
-def transform_inputs(inputs, kind, path, config, params, loaded_tasks):
-    """
-    Transform a sequence of inputs according to the transform configuration.
-    """
-    transforms = TransformSequence()
-    for xform_path in config['transforms']:
-        transform = find_object(xform_path)
-        transforms.add(transform)
-
-    # perform the transformations
-    trans_config = TransformConfig(kind, path, config, params)
-    tasks = [DockerImageTask(kind, t)
-             for t in transforms(trans_config, inputs)]
-    return tasks
-
-
-def load_tasks(kind, path, config, params, loaded_tasks):
-    return transform_inputs(
-        transform.get_inputs(kind, path, config, params, loaded_tasks),
-        kind, path, config, params, loaded_tasks)
-
-
-class DockerImageTask(transform.TransformTask):
-
-    @classmethod
-    def from_json(cls, task_dict):
-        docker_image_task = cls(kind='docker-image', task=task_dict)
-        return docker_image_task
--- a/taskcluster/taskgraph/task/test.py
+++ b/taskcluster/taskgraph/task/test.py
@@ -16,17 +16,17 @@ logger = logging.getLogger(__name__)
 def get_inputs(kind, path, config, params, loaded_tasks):
     """
     Generate tasks implementing Gecko tests.
     """
 
     # the kind on which this one depends
     if len(config.get('kind-dependencies', [])) != 1:
         raise Exception(
-            "TestTask kinds must have exactly one item in kind-dependencies")
+            "Test kinds must have exactly one item in kind-dependencies")
     dep_kind = config['kind-dependencies'][0]
 
     # get build tasks, keyed by build platform
     builds_by_platform = get_builds_by_platform(dep_kind, loaded_tasks)
 
     # get the test platforms for those build tasks
     test_platforms_cfg = load_yaml(path, 'test-platforms.yml')
     test_platforms = get_test_platforms(test_platforms_cfg, builds_by_platform)
--- a/taskcluster/taskgraph/task/transform.py
+++ b/taskcluster/taskgraph/task/transform.py
@@ -58,35 +58,22 @@ def transform_inputs(inputs, kind, path,
     """
     transforms = TransformSequence()
     for xform_path in config['transforms']:
         transform = find_object(xform_path)
         transforms.add(transform)
 
     # perform the transformations
     trans_config = TransformConfig(kind, path, config, params)
-    tasks = [TransformTask(kind, t) for t in transforms(trans_config, inputs)]
+    tasks = [base.Task(kind,
+                       label=task_dict['label'],
+                       attributes=task_dict['attributes'],
+                       task=task_dict['task'],
+                       optimizations=task_dict.get('optimizations'),
+                       dependencies=task_dict.get('dependencies'))
+             for task_dict in transforms(trans_config, inputs)]
     return tasks
 
 
 def load_tasks(kind, path, config, params, loaded_tasks):
     return transform_inputs(
         get_inputs(kind, path, config, params, loaded_tasks),
         kind, path, config, params, loaded_tasks)
-
-
-class TransformTask(base.Task):
-    """
-    Tasks of this class are generated by applying transformations to a sequence
-    of input entities.  By default, it gets those inputs from YAML data in the
-    kind directory, but subclasses may override `get_inputs` to produce them in
-    some other way.
-    """
-
-    def __init__(self, kind, task):
-        super(TransformTask, self).__init__(kind, task['label'],
-                                            task['attributes'], task['task'],
-                                            optimizations=task.get('optimizations'),
-                                            dependencies=task.get('dependencies'))
-
-    @classmethod
-    def from_json(cls, task_dict):
-        return cls(task_dict['attributes']['kind'], task_dict)
--- a/taskcluster/taskgraph/taskgraph.py
+++ b/taskcluster/taskgraph/taskgraph.py
@@ -1,51 +1,31 @@
 # 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 __future__ import absolute_import, print_function, unicode_literals
 
 from .graph import Graph
-from .util.python_path import find_object
+from .task.base import Task
 
 
 class TaskGraph(object):
     """
     Representation of a task graph.
 
     A task graph is a combination of a Graph and a dictionary of tasks indexed
     by label.  TaskGraph instances should be treated as immutable.
     """
 
     def __init__(self, tasks, graph):
         assert set(tasks) == graph.nodes
         self.tasks = tasks
         self.graph = graph
 
-    def to_json(self):
-        "Return a JSON-able object representing the task graph, as documented"
-        named_links_dict = self.graph.named_links_dict()
-        # this dictionary may be keyed by label or by taskid, so let's just call it 'key'
-        tasks = {}
-        for key in self.graph.visit_postorder():
-            task = self.tasks[key]
-            implementation = task.__class__.__module__ + ":" + task.__class__.__name__
-            task_json = {
-                'label': task.label,
-                'attributes': task.attributes,
-                'dependencies': named_links_dict.get(key, {}),
-                'task': task.task,
-                'kind_implementation': implementation
-            }
-            if task.task_id:
-                task_json['task_id'] = task.task_id
-            tasks[key] = task_json
-        return tasks
-
     def for_each_task(self, f, *args, **kwargs):
         for task_label in self.graph.visit_postorder():
             task = self.tasks[task_label]
             f(task, self, *args, **kwargs)
 
     def __getitem__(self, label):
         "Get a task by label"
         return self.tasks[label]
@@ -55,28 +35,35 @@ class TaskGraph(object):
         return self.tasks.itervalues()
 
     def __repr__(self):
         return "<TaskGraph graph={!r} tasks={!r}>".format(self.graph, self.tasks)
 
     def __eq__(self, other):
         return self.tasks == other.tasks and self.graph == other.graph
 
+    def to_json(self):
+        "Return a JSON-able object representing the task graph, as documented"
+        named_links_dict = self.graph.named_links_dict()
+        # this dictionary may be keyed by label or by taskid, so let's just call it 'key'
+        tasks = {}
+        for key in self.graph.visit_postorder():
+            tasks[key] = self.tasks[key].to_json()
+            # overwrite dependencies with the information in the taskgraph's edges.
+            tasks[key]['dependencies'] = named_links_dict.get(key, {})
+        return tasks
+
     @classmethod
     def from_json(cls, tasks_dict):
         """
         This code is used to generate the a TaskGraph using a dictionary
         which is representative of the TaskGraph.
         """
         tasks = {}
         edges = set()
         for key, value in tasks_dict.iteritems():
-            # We get the implementation from JSON
-            implementation = value['kind_implementation']
-            # Loading the module and creating a Task from a dictionary
-            task_kind = find_object(implementation)
-            tasks[key] = task_kind.from_json(value)
+            tasks[key] = Task.from_json(value)
             if 'task_id' in value:
                 tasks[key].task_id = value['task_id']
             for depname, dep in value['dependencies'].iteritems():
                 edges.add((key, dep, depname))
         task_graph = cls(tasks, Graph(set(tasks), edges))
         return tasks, task_graph
--- a/taskcluster/taskgraph/test/test_create.py
+++ b/taskcluster/taskgraph/test/test_create.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 import os
 
 from .. import create
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 
 from mozunit import main
 
 
 class TestCreate(unittest.TestCase):
 
     def setUp(self):
         self.old_task_id = os.environ.get('TASK_ID')
@@ -32,18 +32,18 @@ class TestCreate(unittest.TestCase):
         elif 'TASK_ID' in os.environ:
             del os.environ['TASK_ID']
 
     def fake_create_task(self, session, task_id, label, task_def):
         self.created_tasks[task_id] = task_def
 
     def test_create_tasks(self):
         tasks = {
-            'tid-a': TestTask(label='a', task={'payload': 'hello world'}),
-            'tid-b': TestTask(label='b', task={'payload': 'hello world'}),
+            'tid-a': Task(kind='test', label='a', attributes={}, task={'payload': 'hello world'}),
+            'tid-b': Task(kind='test', label='b', attributes={}, task={'payload': 'hello world'}),
         }
         label_to_taskid = {'a': 'tid-a', 'b': 'tid-b'}
         graph = Graph(nodes={'tid-a', 'tid-b'}, edges={('tid-a', 'tid-b', 'edge')})
         taskgraph = TaskGraph(tasks, graph)
 
         create.create_tasks(taskgraph, label_to_taskid, {'level': '4'})
 
         for tid, task in self.created_tasks.iteritems():
@@ -55,17 +55,17 @@ class TestCreate(unittest.TestCase):
                     # Don't look for decisiontask here
                     continue
                 self.assertIn(depid, self.created_tasks)
 
     def test_create_task_without_dependencies(self):
         "a task with no dependencies depends on the decision task"
         os.environ['TASK_ID'] = 'decisiontask'
         tasks = {
-            'tid-a': TestTask(label='a', task={'payload': 'hello world'}),
+            'tid-a': Task(kind='test', label='a', attributes={}, task={'payload': 'hello world'}),
         }
         label_to_taskid = {'a': 'tid-a'}
         graph = Graph(nodes={'tid-a'}, edges=set())
         taskgraph = TaskGraph(tasks, graph)
 
         create.create_tasks(taskgraph, label_to_taskid, {'level': '4'})
 
         for tid, task in self.created_tasks.iteritems():
--- a/taskcluster/taskgraph/test/test_decision.py
+++ b/taskcluster/taskgraph/test/test_decision.py
@@ -7,51 +7,21 @@ from __future__ import absolute_import, 
 import os
 import json
 import yaml
 import shutil
 import unittest
 import tempfile
 
 from .. import decision
-from ..graph import Graph
-from ..taskgraph import TaskGraph
-from .util import TestTask
 from mozunit import main
 
 
 class TestDecision(unittest.TestCase):
 
-    def test_taskgraph_to_json(self):
-        tasks = {
-            'a': TestTask(label='a', attributes={'attr': 'a-task'}),
-            'b': TestTask(label='b', task={'task': 'def'}),
-        }
-        graph = Graph(nodes=set('ab'), edges={('a', 'b', 'edgelabel')})
-        taskgraph = TaskGraph(tasks, graph)
-
-        res = taskgraph.to_json()
-
-        self.assertEqual(res, {
-            'a': {
-                'label': 'a',
-                'attributes': {'attr': 'a-task', 'kind': 'test'},
-                'task': {},
-                'dependencies': {'edgelabel': 'b'},
-                'kind_implementation': 'taskgraph.test.util:TestTask',
-            },
-            'b': {
-                'label': 'b',
-                'attributes': {'kind': 'test'},
-                'task': {'task': 'def'},
-                'dependencies': {},
-                'kind_implementation': 'taskgraph.test.util:TestTask',
-            }
-        })
-
     def test_write_artifact_json(self):
         data = [{'some': 'data'}]
         tmpdir = tempfile.mkdtemp()
         try:
             decision.ARTIFACTS_DIR = os.path.join(tmpdir, "artifacts")
             decision.write_artifact("artifact.json", data)
             with open(os.path.join(decision.ARTIFACTS_DIR, "artifact.json")) as f:
                 self.assertEqual(json.load(f), data)
--- a/taskcluster/taskgraph/test/test_optimize.py
+++ b/taskcluster/taskgraph/test/test_optimize.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..optimize import optimize_task_graph, resolve_task_references, optimization
 from ..optimize import annotate_task_graph, get_subgraph
 from ..taskgraph import TaskGraph
 from .. import graph
-from .util import TestTask
+from ..task.base import Task
 
 
 class TestResolveTaskReferences(unittest.TestCase):
 
     def do(self, input, output):
         taskid_for_edge_name = {'edge%d' % n: 'tid%d' % n for n in range(1, 4)}
         self.assertEqual(resolve_task_references('subject', input, taskid_for_edge_name), output)
 
@@ -62,28 +62,28 @@ class TestOptimize(unittest.TestCase):
         # set up some simple optimization functions
         optimization('no-optimize')(lambda self, params: (False, None))
         optimization('optimize-away')(lambda self, params: (True, None))
         optimization('optimize-to-task')(lambda self, params, task: (True, task))
         optimization('false-with-taskid')(lambda self, params: (False, 'some-taskid'))
 
     def make_task(self, label, optimization=None, task_def=None, optimized=None, task_id=None):
         task_def = task_def or {'sample': 'task-def'}
-        task = TestTask(label=label, task=task_def)
+        task = Task(kind='test', label=label, attributes={}, task=task_def)
         task.optimized = optimized
         if optimization:
             task.optimizations = [optimization]
         else:
             task.optimizations = []
         task.task_id = task_id
         return task
 
     def make_graph(self, *tasks_and_edges):
-        tasks = {t.label: t for t in tasks_and_edges if isinstance(t, TestTask)}
-        edges = {e for e in tasks_and_edges if not isinstance(e, TestTask)}
+        tasks = {t.label: t for t in tasks_and_edges if isinstance(t, Task)}
+        edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
         return TaskGraph(tasks, graph.Graph(set(tasks), edges))
 
     def assert_annotations(self, graph, **annotations):
         def repl(task_id):
             return 'SLUGID' if task_id and len(task_id) == 22 else task_id
         got_annotations = {
             t.label: (t.optimized, repl(t.task_id)) for t in graph.tasks.itervalues()
         }
--- a/taskcluster/taskgraph/test/test_target_tasks.py
+++ b/taskcluster/taskgraph/test/test_target_tasks.py
@@ -5,17 +5,17 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from .. import target_tasks
 from .. import try_option_syntax
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 from mozunit import main
 
 
 class FakeTryOptionSyntax(object):
 
     def __init__(self, message, task_graph):
         self.trigger_tests = 0
         self.talos_trigger_tests = 0
@@ -29,18 +29,19 @@ class FakeTryOptionSyntax(object):
         return 'at-at' in attributes
 
 
 class TestTargetTasks(unittest.TestCase):
 
     def default_matches(self, run_on_projects, project):
         method = target_tasks.get_method('default')
         graph = TaskGraph(tasks={
-            'a': TestTask(kind='build', label='a',
-                          attributes={'run_on_projects': run_on_projects}),
+            'a': Task(kind='build', label='a',
+                      attributes={'run_on_projects': run_on_projects},
+                      task={}),
         }, graph=Graph(nodes={'a'}, edges=set()))
         parameters = {'project': project}
         return 'a' in method(graph, parameters)
 
     def test_default_all(self):
         """run_on_projects=[all] includes release, integration, and other projects"""
         self.assertTrue(self.default_matches(['all'], 'mozilla-central'))
         self.assertTrue(self.default_matches(['all'], 'mozilla-inbound'))
@@ -62,18 +63,18 @@ class TestTargetTasks(unittest.TestCase)
     def test_default_nothing(self):
         """run_on_projects=[] includes nothing"""
         self.assertFalse(self.default_matches([], 'mozilla-central'))
         self.assertFalse(self.default_matches([], 'mozilla-inbound'))
         self.assertFalse(self.default_matches([], 'baobab'))
 
     def test_try_option_syntax(self):
         tasks = {
-            'a': TestTask(kind=None, label='a'),
-            'b': TestTask(kind=None, label='b', attributes={'at-at': 'yep'}),
+            'a': Task(kind=None, label='a', attributes={}, task={}),
+            'b': Task(kind=None, label='b', attributes={'at-at': 'yep'}, task={}),
         }
         graph = Graph(nodes=set('ab'), edges=set())
         tg = TaskGraph(tasks, graph)
         params = {'message': 'try me'}
 
         orig_TryOptionSyntax = try_option_syntax.TryOptionSyntax
         try:
             try_option_syntax.TryOptionSyntax = FakeTryOptionSyntax
--- a/taskcluster/taskgraph/test/test_taskgraph.py
+++ b/taskcluster/taskgraph/test/test_taskgraph.py
@@ -2,51 +2,77 @@
 # 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 __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..graph import Graph
-from ..task.docker_image import DockerImageTask
-from ..task.transform import TransformTask
+from ..task.base import Task
 from ..taskgraph import TaskGraph
 from mozunit import main
 
 
 class TestTaskGraph(unittest.TestCase):
 
-    def test_from_json(self):
-        task = {
-            "routes": [],
-            "extra": {
-                "imageMeta": {
-                    "contextHash": "<hash>",
-                    "imageName": "<image>",
-                    "level": "1"
-                }
+    maxDiff = None
+
+    def test_taskgraph_to_json(self):
+        tasks = {
+            'a': Task(kind='test', label='a',
+                      attributes={'attr': 'a-task'},
+                      task={'taskdef': True}),
+            'b': Task(kind='test', label='b',
+                      attributes={},
+                      task={'task': 'def'},
+                      optimizations=[['seta']],
+                      # note that this dep is ignored, superseded by that
+                      # from the taskgraph's edges
+                      dependencies={'first': 'a'}),
+        }
+        graph = Graph(nodes=set('ab'), edges={('a', 'b', 'edgelabel')})
+        taskgraph = TaskGraph(tasks, graph)
+
+        res = taskgraph.to_json()
+
+        self.assertEqual(res, {
+            'a': {
+                'kind': 'test',
+                'label': 'a',
+                'attributes': {'attr': 'a-task', 'kind': 'test'},
+                'task': {'taskdef': True},
+                'dependencies': {'edgelabel': 'b'},
+                'optimizations': [],
+            },
+            'b': {
+                'kind': 'test',
+                'label': 'b',
+                'attributes': {'kind': 'test'},
+                'task': {'task': 'def'},
+                'dependencies': {},
+                'optimizations': [['seta']],
             }
-        }
+        })
+
+    def test_round_trip(self):
         graph = TaskGraph(tasks={
-            'a': TransformTask(
+            'a': Task(
                 kind='fancy',
-                task={
-                    'label': 'a',
-                    'attributes': {},
-                    'dependencies': {},
-                    'when': {},
-                    'task': {'task': 'def'},
-                }),
-            'b': DockerImageTask(kind='docker-image',
-                                 task={
-                                     'label': 'b',
-                                     'attributes': {},
-                                     'task': task,
-                                 }),
-        }, graph=Graph(nodes={'a', 'b'}, edges=set()))
+                label='a',
+                attributes={},
+                dependencies={'prereq': 'b'},  # must match edges, below
+                optimizations=[['seta']],
+                task={'task': 'def'}),
+            'b': Task(
+                kind='pre',
+                label='b',
+                attributes={},
+                dependencies={},
+                optimizations=[['seta']],
+                task={'task': 'def2'}),
+        }, graph=Graph(nodes={'a', 'b'}, edges={('a', 'b', 'prereq')}))
 
         tasks, new_graph = TaskGraph.from_json(graph.to_json())
-        self.assertEqual(graph.tasks['a'], new_graph.tasks['a'])
         self.assertEqual(graph, new_graph)
 
 if __name__ == '__main__':
     main()
--- a/taskcluster/taskgraph/test/test_try_option_syntax.py
+++ b/taskcluster/taskgraph/test/test_try_option_syntax.py
@@ -5,35 +5,35 @@
 from __future__ import absolute_import, print_function, unicode_literals
 
 import unittest
 
 from ..try_option_syntax import TryOptionSyntax
 from ..try_option_syntax import RIDEALONG_BUILDS
 from ..graph import Graph
 from ..taskgraph import TaskGraph
-from .util import TestTask
+from ..task.base import Task
 from mozunit import main
 
 # an empty graph, for things that don't look at it
 empty_graph = TaskGraph({}, Graph(set(), set()))
 
 
 def unittest_task(n, tp):
-    return (n, TestTask('test', n, {
+    return (n, Task('test', n, {
         'unittest_try_name': n,
         'test_platform': tp,
-    }))
+    }, {}))
 
 
 def talos_task(n, tp):
-    return (n, TestTask('test', n, {
+    return (n, Task('test', n, {
         'talos_try_name': n,
         'test_platform': tp,
-    }))
+    }, {}))
 
 tasks = {k: v for k, v in [
     unittest_task('mochitest-browser-chrome', 'linux'),
     unittest_task('mochitest-browser-chrome-e10s', 'linux64'),
     unittest_task('mochitest-chrome', 'linux'),
     unittest_task('mochitest-webgl', 'linux'),
     unittest_task('crashtest-e10s', 'linux'),
     unittest_task('gtest', 'linux64'),
deleted file mode 100644
--- a/taskcluster/taskgraph/test/util.py
+++ /dev/null
@@ -1,24 +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 __future__ import absolute_import, print_function, unicode_literals
-
-from ..task import base
-
-
-class TestTask(base.Task):
-
-    def __init__(self, kind=None, label=None, attributes=None, task=None):
-        super(TestTask, self).__init__(
-                kind or 'test',
-                label or 'test-label',
-                attributes or {},
-                task or {})
-
-    @classmethod
-    def load_tasks(cls, kind, path, config, parameters):
-        return []
-
-    def get_dependencies(self, taskgraph):
-        return []
--- a/taskcluster/taskgraph/transforms/docker_image.py
+++ b/taskcluster/taskgraph/transforms/docker_image.py
@@ -62,33 +62,24 @@ def fill_template(config, tasks):
         # for a given context hash, a worker within Taskcluster does not need to contain
         # the same image per branch.
         optimizations = [['index-search', '{}.level-{}.{}.hash.{}'.format(
             INDEX_PREFIX, level, image_name, context_hash)]
             for level in reversed(range(int(config.params['level']), 4))]
 
         # include some information that is useful in reconstructing this task
         # from JSON
-        extra = {
-            'imageMeta': {
-                'level': config.params['level'],
-                'contextHash': context_hash,
-                'imageName': image_name,
-            },
-        }
-
         taskdesc = {
             'label': 'build-docker-image-' + image_name,
             'description': description,
             'attributes': {'image_name': image_name},
             'expires-after': '1 year',
             'routes': routes,
             'optimizations': optimizations,
             'scopes': ['secrets:get:project/taskcluster/gecko/hgfingerprint'],
-            'extra': extra,
             'treeherder': {
                 'symbol': job_symbol,
                 'platform': 'taskcluster-images/opt',
                 'kind': 'other',
                 'tier': 1,
             },
             'run-on-projects': [],
             'worker-type': 'aws-provisioner-v1/gecko-images',