feat(load-task): support using custom images

This implements `taskgraph load-task --image <image>`, where <image> can
be one of:
* task-id=<task id> - Downloads image.tar from specified task id
* image=<image> - Downloads image.tar from specified index path
* <image name> - Builds image that lives under
  `taskcluster/docker/<image name>`
* <image reference> - Uses specified image from a registry

Author: ahal

src/taskgraph/docker.py

@@ -6,11 +6,13 @@
 import os
 import shlex
 import subprocess
+import sys
 import tarfile
 import tempfile
 from io import BytesIO
+from pathlib import Path
 from textwrap import dedent
-from typing import Dict, Generator, List, Mapping, Optional
+from typing import Dict, Generator, List, Mapping, Optional, Union
 
 try:
     import zstandard as zstd
@@ -354,7 +356,48 @@ def _index(l: List, s: str) -> Optional[int]:
         pass
 
 
-def load_task(task_id: str, remove: bool = True, user: Optional[str] = None) -> int:
+def _resolve_image(image: Union[str, Dict[str, str]], graph_config: GraphConfig) -> str:
+    image_task_id = None
+
+    # Standard case, image comes from the task definition.
+    if isinstance(image, dict):
+        assert "type" in image
+
+        if image["type"] == "task-image":
+            image_task_id = image["taskId"]
+        elif image["type"] == "indexed-image":
+            image_task_id = find_task_id(image["namespace"])
+        else:
+            raise Exception(f"Tasks with {image['type']} images are not supported!")
+    else:
+        # Check if image refers to an in-tree image under taskcluster/docker,
+        # if so build it.
+        image_dir = docker.image_path(image, graph_config)
+        if Path(image_dir).is_dir():
+            tag = f"taskcluster/{image}:latest"
+            build_image(image, tag, graph_config, os.environ)
+            return tag
+
+        # Check if we're referencing a task or index.
+        if image.startswith("task-id="):
+            image_task_id = image.split("=", 1)[1]
+        elif image.startswith("index="):
+            index = image.split("=", 1)[1]
+            image_task_id = find_task_id(index)
+        else:
+            # Assume the string references an image from a registry.
+            return image
+
+    return load_image_by_task_id(image_task_id)
+
+
+def load_task(
+    graph_config: GraphConfig,
+    task_id: str,
+    remove: bool = True,
+    user: Optional[str] = None,
+    custom_image: Optional[str] = None,
+) -> int:
     """Load and run a task interactively in a Docker container.
 
     Downloads the docker image from a task's definition and runs it in an
@@ -363,9 +406,11 @@ def load_task(task_id: str, remove: bool = True, user: Optional[str] = None) ->
     the 'exec-task' function provided in the shell.
 
     Args:
+        graph_config: The graph configuration object.
         task_id: The ID of the task to load.
         remove: Whether to remove the container after exit (default True).
         user: The user to switch to in the container (default 'worker').
+        custom_image: A custom image to use instead of the task's image.
 
     Returns:
         int: The exit code from the Docker container.
@@ -387,6 +432,13 @@ def load_task(task_id: str, remove: bool = True, user: Optional[str] = None) ->
         print("Only tasks using `run-task` are supported!")
         return 1
 
+    try:
+        image = custom_image or image
+        image_tag = _resolve_image(image, graph_config)
+    except Exception as e:
+        print(e, file=sys.stderr)
+        return 1
+
     # Remove the payload section of the task's command. This way run-task will
     # set up the task (clone repos, download fetches, etc) but won't actually
     # start the core of the task. Instead we'll drop the user into an interactive
@@ -415,16 +467,6 @@ def load_task(task_id: str, remove: bool = True, user: Optional[str] = None) ->
         else:
             task_cwd = "$TASK_WORKDIR"
 
-    if image["type"] == "task-image":
-        image_task_id = image["taskId"]
-    elif image["type"] == "indexed-image":
-        image_task_id = find_task_id(image["namespace"])
-    else:
-        print(f"Tasks with {image['type']} images are not supported!")
-        return 1
-
-    image_tag = load_image_by_task_id(image_task_id)
-
     # Set some env vars the worker would normally set.
     env = {
         "RUN_ID": "0",

src/taskgraph/main.py

@@ -569,7 +569,7 @@ def show_taskgraph(options):
     "--root",
     "-r",
     default="taskcluster",
-    help="relative path for the root of the taskgraph definition",
+    help="Relative path to the root of the Taskgraph definition.",
 )
 @argument(
     "-t", "--tag", help="tag that the image should be built as.", metavar="name:tag"
@@ -683,11 +683,34 @@ def image_digest(args):
     help="Keep the docker container after exiting.",
 )
 @argument("--user", default=None, help="Container user to start shell with.")
+@argument(
+    "--image",
+    default=None,
+    help="Use a custom image instead of the task's image. Can be the name of "
+    "an image under `taskcluster/docker`, `task-id=<task id>`, or "
+    "`index=<index path>`.",
+)
+@argument(
+    "--root",
+    "-r",
+    default="taskcluster",
+    help="Relative path to the root of the Taskgraph definition.",
+)
 def load_task(args):
+    from taskgraph.config import load_graph_config  # noqa: PLC0415
     from taskgraph.docker import load_task  # noqa: PLC0415
 
     validate_docker()
-    return load_task(args["task_id"], remove=args["remove"], user=args["user"])
+
+    root = args["root"]
+    graph_config = load_graph_config(root)
+    return load_task(
+        graph_config,
+        args["task_id"],
+        remove=args["remove"],
+        user=args["user"],
+        custom_image=args["image"],
+    )
 
 
 @command("decision", help="Run the decision task")

test/test_docker.py

@@ -113,11 +113,22 @@ def mock_run(*popenargs, check=False, **kwargs):
 def run_load_task(mocker):
     task_id = "abc"
 
-    def inner(task, remove=False):
+    def inner(task, remove=False, custom_image=None):
         proc = mocker.MagicMock()
         proc.returncode = 0
 
+        graph_config = GraphConfig(
+            {
+                "trust-domain": "test-domain",
+                "docker-image-kind": "docker-image",
+            },
+            "test/data/taskcluster",
+        )
+
         mocks = {
+            "build_image": mocker.patch.object(
+                docker, "build_image", return_value=None
+            ),
             "get_task_definition": mocker.patch.object(
                 docker, "get_task_definition", return_value=task
             ),
@@ -129,7 +140,9 @@ def inner(task, remove=False):
             ),
         }
 
-        ret = docker.load_task(task_id, remove=remove)
+        ret = docker.load_task(
+            graph_config, task_id, remove=remove, custom_image=custom_image
+        )
         return ret, mocks
 
     return inner
@@ -330,8 +343,58 @@ def test_load_task_with_unsupported_image_type(capsys, run_load_task):
         },
     }
 
-    ret, mocks = run_load_task(task)
+    ret, _ = run_load_task(task)
     assert ret == 1
 
-    out, _ = capsys.readouterr()
-    assert "Tasks with unsupported-type images are not supported!" in out
+    _, err = capsys.readouterr()
+    assert "Tasks with unsupported-type images are not supported!" in err
+
+
+@pytest.fixture
+def task():
+    return {
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--task-cwd=/builds/worker",
+                "--",
+                "echo",
+                "test",
+            ],
+            "image": {"type": "task-image", "taskId": "abc"},
+        },
+    }
+
+
+def test_load_task_with_custom_image_in_tree(run_load_task, task):
+    image = "hello-world"
+    ret, mocks = run_load_task(task, custom_image=image)
+    assert ret == 0
+
+    mocks["build_image"].assert_called_once()
+    args = mocks["subprocess_run"].call_args[0][0]
+    tag = args[args.index("-it") + 1]
+    assert tag == f"taskcluster/{image}:latest"
+
+
+def test_load_task_with_custom_image_task_id(run_load_task, task):
+    image = "task-id=abc"
+    ret, mocks = run_load_task(task, custom_image=image)
+    assert ret == 0
+    mocks["load_image_by_task_id"].assert_called_once_with("abc")
+
+
+def test_load_task_with_custom_image_index(mocker, run_load_task, task):
+    image = "index=abc"
+    mocker.patch.object(docker, "find_task_id", return_value="abc")
+    ret, mocks = run_load_task(task, custom_image=image)
+    assert ret == 0
+    mocks["load_image_by_task_id"].assert_called_once_with("abc")
+
+
+def test_load_task_with_custom_image_registry(mocker, run_load_task, task):
+    image = "ubuntu:latest"
+    ret, mocks = run_load_task(task, custom_image=image)
+    assert ret == 0
+    assert not mocks["load_image_by_task_id"].called
+    assert not mocks["build_image"].called