fix: removed hardcoding from docker-image dir and kinds (#728)

* fix: removed hardcoding from docker-image dir and kinds

* fix: updated to review and added tests

* fix: updated docker tests

* fix: updated default case and add kinds dir to graph config

Author: abhishekmadan30

src/taskgraph/config.py

@@ -26,6 +26,10 @@
         # The trust-domain for this graph.
         # (See https://firefox-source-docs.mozilla.org/taskcluster/taskcluster/taskgraph.html#taskgraph-trust-domain)  # noqa
         Required("trust-domain"): str,
+        Optional(
+            "docker-image-kind",
+            description="Name of the docker image kind (default: docker-image)",
+        ): str,
         Required("task-priority"): optionally_keyed_by(
             "project",
             "level",
@@ -157,6 +161,14 @@ def vcs_root(self):
     def taskcluster_yml(self):
         return os.path.join(self.vcs_root, ".taskcluster.yml")
 
+    @property
+    def docker_dir(self):
+        return os.path.join(self.root_dir, "docker")
+
+    @property
+    def kinds_dir(self):
+        return os.path.join(self.root_dir, "kinds")
+
 
 def validate_graph_config(config):
     validate_schema(graph_config_schema, config, "Invalid graph configuration:")

src/taskgraph/docker.py

@@ -97,29 +97,29 @@ def load_image_by_task_id(task_id, tag=None):
     return tag
 
 
-def build_context(name, outputFile, args=None):
+def build_context(name, outputFile, graph_config, args=None):
     """Build a context.tar for image with specified name."""
     if not name:
         raise ValueError("must provide a Docker image name")
     if not outputFile:
         raise ValueError("must provide a outputFile")
 
-    image_dir = docker.image_path(name)
+    image_dir = docker.image_path(name, graph_config)
     if not os.path.isdir(image_dir):
         raise Exception(f"image directory does not exist: {image_dir}")
 
     docker.create_context_tar(".", image_dir, outputFile, args)
 
 
-def build_image(name, tag, args=None):
+def build_image(name, tag, graph_config, args=None):
     """Build a Docker image of specified name.
 
     Output from image building process will be printed to stdout.
     """
     if not name:
         raise ValueError("must provide a Docker image name")
 
-    image_dir = docker.image_path(name)
+    image_dir = docker.image_path(name, graph_config)
     if not os.path.isdir(image_dir):
         raise Exception(f"image directory does not exist: {image_dir}")
 

src/taskgraph/main.py

@@ -565,6 +565,12 @@ def show_taskgraph(options):
 
 @command("build-image", help="Build a Docker image")
 @argument("image_name", help="Name of the image to build")
+@argument(
+    "--root",
+    "-r",
+    default="taskcluster",
+    help="relative path for the root of the taskgraph definition",
+)
 @argument(
     "-t", "--tag", help="tag that the image should be built as.", metavar="name:tag"
 )
@@ -575,13 +581,20 @@ def show_taskgraph(options):
     metavar="context.tar",
 )
 def build_image(args):
+    from taskgraph.config import load_graph_config  # noqa: PLC0415
     from taskgraph.docker import build_context, build_image  # noqa: PLC0415
 
     validate_docker()
+
+    root = args["root"]
+    graph_config = load_graph_config(root)
+
     if args["context_only"] is None:
-        build_image(args["image_name"], args["tag"], os.environ)
+        build_image(args["image_name"], args["tag"], os.environ, graph_config)
     else:
-        build_context(args["image_name"], args["context_only"], os.environ)
+        build_context(
+            args["image_name"], args["context_only"], os.environ, graph_config
+        )
 
 
 @command(

src/taskgraph/transforms/task.py

@@ -575,7 +575,7 @@ def build_docker_worker_payload(config, task, task_def):
             }
 
             # Find VOLUME in Dockerfile.
-            volumes = dockerutil.parse_volumes(name)
+            volumes = dockerutil.parse_volumes(name, config.graph_config)
             for v in sorted(volumes):
                 if v in worker["volumes"]:
                     raise Exception(

src/taskgraph/util/docker.py

@@ -206,28 +206,34 @@ def stream_context_tar(topsrcdir, context_dir, out_file, args=None):
 
 
 @functools.lru_cache(maxsize=None)
-def image_paths():
+def image_paths(graph_config):
     """Return a map of image name to paths containing their Dockerfile."""
-    config = load_yaml("taskcluster", "kinds", "docker-image", "kind.yml")
+
+    config = load_yaml(
+        graph_config.kinds_dir,
+        graph_config.get("docker-image-kind", "docker-image"),
+        "kind.yml",
+    )
+
     return {
-        k: os.path.join(IMAGE_DIR, v.get("definition", k))
+        k: os.path.join(graph_config.docker_dir, v.get("definition", k))
         for k, v in config["tasks"].items()
     }
 
 
-def image_path(name):
-    paths = image_paths()
+def image_path(name, graph_config):
+    paths = image_paths(graph_config)
     if name in paths:
         return paths[name]
-    return os.path.join(IMAGE_DIR, name)
+    return os.path.join(graph_config.docker_dir, name)
 
 
 @functools.lru_cache(maxsize=None)
-def parse_volumes(image):
+def parse_volumes(image, graph_config):
     """Parse VOLUME entries from a Dockerfile for an image."""
     volumes = set()
 
-    path = image_path(image)
+    path = image_path(image, graph_config)
 
     with open(os.path.join(path, "Dockerfile"), "rb") as fh:
         for line in fh:

test/data/taskcluster/kinds/docker-image/kind.yml

@@ -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/.
+---
+meta:
+    - &uv_version 0.6.1
+
+loader: taskgraph.loader.transform:loader
+
+transforms:
+    - taskgraph.transforms.docker_image:transforms
+    - taskgraph.transforms.cached_tasks: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
+# generate tasks for every docker image in the directory, secure in the
+# knowledge that unnecessary images will be omitted from the target task graph
+tasks:
+    decision:
+        symbol: I(d)
+        parent: run-task
+    fetch:
+        symbol: I(fetch)
+    index-task:
+        symbol: I(idx)
+    python:
+        symbol: I(py)
+        args:
+            PYTHON_VERSIONS: "3.13 3.12 3.11 3.10 3.9 3.8"
+            UV_VERSION: *uv_version
+    run-task:
+        symbol: I(rt)
+        args:
+            UV_VERSION: *uv_version
+    skopeo:
+        symbol: I(skopeo)

test/test_docker.py

@@ -4,6 +4,7 @@
 import pytest
 
 from taskgraph import docker
+from taskgraph.config import GraphConfig
 
 
 @pytest.fixture(autouse=True, scope="module")
@@ -30,7 +31,15 @@ def test_build_image(capsys, mock_docker_build):
     image = "hello-world-tag"
     tag = f"test/{image}:1.0"
 
-    assert docker.build_image(image, None) is None
+    graph_config = GraphConfig(
+        {
+            "trust-domain": "test-domain",
+            "docker-image-kind": "docker-image",
+        },
+        "test/data/taskcluster",
+    )
+
+    assert docker.build_image(image, None, graph_config=graph_config) is None
     m_stream.assert_called_once()
     m_run.assert_called_once_with(
         ["docker", "image", "build", "--no-cache", f"-t={tag}", "-"],
@@ -47,7 +56,15 @@ def test_build_image_no_tag(capsys, mock_docker_build):
     m_stream, m_run = mock_docker_build
     image = "hello-world"
 
-    assert docker.build_image(image, None) is None
+    graph_config = GraphConfig(
+        {
+            "trust-domain": "test-domain",
+            "docker-image-kind": "docker-image",
+        },
+        "test/data/taskcluster",
+    )
+
+    assert docker.build_image(image, None, graph_config=graph_config) is None
     m_stream.assert_called_once()
     m_run.assert_called_once_with(
         ["docker", "image", "build", "--no-cache", "-"],
@@ -71,8 +88,16 @@ def mock_run(*popenargs, check=False, **kwargs):
     m_run.side_effect = mock_run
     image = "hello-world"
 
+    graph_config = GraphConfig(
+        {
+            "trust-domain": "test-domain",
+            "docker-image-kind": "docker-image",
+        },
+        "test/data/taskcluster",
+    )
+
     with pytest.raises(Exception):
-        docker.build_image(image, None)
+        docker.build_image(image, None, graph_config=graph_config)
     m_stream.assert_called_once()
     m_run.assert_called_once_with(
         ["docker", "image", "build", "--no-cache", "-"],

test/test_util_docker.py

@@ -14,6 +14,7 @@
 
 import taskcluster_urls as liburls
 
+from taskgraph.config import GraphConfig
 from taskgraph.util import docker
 
 from .mockedopen import MockedOpen
@@ -269,3 +270,81 @@ def test_stream_context_tar(self):
             )
         finally:
             shutil.rmtree(tmp)
+
+    def test_image_paths_with_custom_kind(self):
+        """Test image_paths function with graph_config parameter."""
+        temp_dir = tempfile.mkdtemp()
+        try:
+            # Create the kinds directory structure
+            kinds_dir = os.path.join(temp_dir, "kinds", "docker-test-image")
+            os.makedirs(kinds_dir)
+
+            # Create the kind.yml file with task definitions
+            kind_yml_path = os.path.join(kinds_dir, "kind.yml")
+            with open(kind_yml_path, "w") as f:
+                f.write("tasks:\n")
+                f.write("  test-image:\n")
+                f.write("    definition: test-image\n")
+                f.write("  another-image:\n")
+                f.write("    definition: custom-path\n")
+
+            # Create graph config pointing to our test directory
+            temp_graph_config = GraphConfig(
+                {
+                    "trust-domain": "test-domain",
+                    "docker-image-kind": "docker-test-image",
+                },
+                temp_dir,
+            )
+
+            paths = docker.image_paths(temp_graph_config)
+
+            expected_docker_dir = os.path.join(temp_graph_config.root_dir, "docker")
+            self.assertEqual(
+                paths["test-image"], os.path.join(expected_docker_dir, "test-image")
+            )
+            self.assertEqual(
+                paths["another-image"], os.path.join(expected_docker_dir, "custom-path")
+            )
+        finally:
+            shutil.rmtree(temp_dir)
+
+    def test_parse_volumes_with_graph_config(self):
+        """Test parse_volumes function with graph_config parameter."""
+        temp_dir = tempfile.mkdtemp()
+        try:
+            kinds_dir = os.path.join(temp_dir, "kinds", "docker-test-image")
+            os.makedirs(kinds_dir)
+
+            kind_yml_path = os.path.join(kinds_dir, "kind.yml")
+            with open(kind_yml_path, "w") as f:
+                f.write("tasks:\n")
+                f.write("  test-image:\n")
+                f.write("    definition: test-image\n")
+
+            docker_dir = os.path.join(temp_dir, "docker")
+            os.makedirs(docker_dir)
+
+            image_dir = os.path.join(docker_dir, "test-image")
+            os.makedirs(image_dir)
+
+            dockerfile_path = os.path.join(image_dir, "Dockerfile")
+            with open(dockerfile_path, "wb") as fh:
+                fh.write(b"VOLUME /foo/bar \n")
+                fh.write(b"VOLUME /hello  /world \n")
+
+            test_graph_config = GraphConfig(
+                {
+                    "trust-domain": "test-domain",
+                    "docker-image-kind": "docker-test-image",
+                },
+                temp_dir,
+            )
+
+            volumes = docker.parse_volumes("test-image", test_graph_config)
+
+            expected_volumes = {"/foo/bar", "/hello", "/world"}
+            self.assertEqual(volumes, expected_volumes)
+
+        finally:
+            shutil.rmtree(temp_dir)