fix: removed hardcoding from docker-image dir and kinds

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="What kind of docker image to use for the task.",
+        ): [str],
         Required("task-priority"): optionally_keyed_by(
             "project",
             "level",
@@ -157,6 +161,10 @@ 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")
+
 
 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, args=None, graph_config=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, args=None, graph_config=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,37 @@ def stream_context_tar(topsrcdir, context_dir, out_file, args=None):
 
 
 @functools.lru_cache(maxsize=None)
-def image_paths():
+def image_paths(graph_config=None):
     """Return a map of image name to paths containing their Dockerfile."""
-    config = load_yaml("taskcluster", "kinds", "docker-image", "kind.yml")
+    if graph_config:
+        config = load_yaml(
+            graph_config.docker_dir,
+            "kinds",
+            graph_config["docker_image_kind"],
+            "kind.yml",
+        )
+    else:
+        config = load_yaml("taskcluster", "kinds", "docker-image", "kind.yml")
+
     return {
         k: os.path.join(IMAGE_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=None):
+    paths = image_paths(graph_config)
     if name in paths:
         return paths[name]
     return os.path.join(IMAGE_DIR, name)
 
 
 @functools.lru_cache(maxsize=None)
-def parse_volumes(image):
+def parse_volumes(image, graph_config=None):
     """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: