feat: implement new 'load-task' command to debug tasks locally

This command is an extension of the `load-image` command. Given a
taskId, it will:

1. Find and download the docker image from the parent task.
2. Spin up a docker container using this image.
3. Ensure environment variables in the task's definition are set.
4. Run the "setup" portion of `run-task` (checkouts, fetches, etc)
5. Drop the user into a bash shell.

This will make it easier to debug tasks locally. Note for now only tasks
using run-task with a docker-worker (or D2G) payload are supported.

Author: ahal

src/taskgraph/docker.py

@@ -7,6 +7,7 @@
 import os
 import subprocess
 import tarfile
+import tempfile
 from io import BytesIO
 from textwrap import dedent
 
@@ -19,6 +20,7 @@
 from taskgraph.util.taskcluster import (
     get_artifact_url,
     get_session,
+    get_task_definition,
 )
 
 DEPLOY_WARNING = """
@@ -90,7 +92,7 @@ def load_image_by_task_id(task_id, tag=None):
     else:
         tag = "{}:{}".format(result["image"], result["tag"])
     print(f"Try: docker run -ti --rm {tag} bash")
-    return True
+    return tag
 
 
 def build_context(name, outputFile, args=None):
@@ -237,3 +239,59 @@ def download_and_modify_image():
         raise Exception("No repositories file found!")
 
     return info
+
+
+def load_task(task_id, remove=False):
+    task_def = get_task_definition(task_id)
+
+    if (
+        impl := task_def.get("tags", {}).get("worker-implementation")
+    ) != "docker-worker":
+        print(f"Tasks with worker-implementation '{impl}' are not supported!")
+        return 1
+
+    command = task_def["payload"].get("command")
+    if not command or not command[0].endswith("run-task"):
+        print("Only tasks using `run-task` are supported!")
+        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 a command
+    # shell.
+    if (index := command.index("--")) > -1:
+        command[index + 1 :] = ["echo 'Task setup complete!'"]
+
+    image_task_id = task_def["payload"]["image"]["taskId"]
+    image_tag = load_image_by_task_id(image_task_id, None)
+
+    env = task_def["payload"].get("env")
+
+    tmpfile = None
+    try:
+        command = [
+            "docker",
+            "run",
+            "-it",
+            image_tag,
+            "bash",
+            "-c",
+            f"{' '.join(command)} && bash",
+        ]
+
+        if remove:
+            command.insert(2, "--rm")
+
+        if env:
+            tmpfile = tempfile.NamedTemporaryFile("w+", delete=False)
+            tmpfile.write("\n".join([f"{k}={v}" for k, v in env.items()]))
+            tmpfile.close()
+
+            command.insert(2, f"--env-file={tmpfile.name}")
+
+        subprocess.run(command)
+    finally:
+        if tmpfile:
+            os.remove(tmpfile.name)
+
+    return 0

src/taskgraph/main.py

@@ -618,10 +618,10 @@ def load_image(args):
     validate_docker()
     try:
         if args["task_id"]:
-            ok = load_image_by_task_id(args["task_id"], args.get("tag"))
+            tag = load_image_by_task_id(args["task_id"], args.get("tag"))
         else:
-            ok = load_image_by_name(args["image_name"], args.get("tag"))
-        if not ok:
+            tag = load_image_by_name(args["image_name"], args.get("tag"))
+        if not tag:
             sys.exit(1)
     except Exception:
         traceback.print_exc()
@@ -652,6 +652,31 @@ def image_digest(args):
         sys.exit(1)
 
 
+@command(
+    "load-task",
+    help="Loads a pre-built Docker image and drops you into a container with "
+    "the same environment variables and run-task setup as the specified task. "
+    "The task's payload.command will be replaced with 'bash'. You need to have "
+    "docker installed and running for this to work.",
+)
+@argument(
+    "task_id",
+    help="Load the image at public/image.tar.zst in this task, "
+    "rather than searching the index",
+)
+@argument(
+    "--rm",
+    action="store_true",
+    default=False,
+    help="Delete the docker container after exiting.",
+)
+def load_task(args):
+    from taskgraph.docker import load_task
+
+    validate_docker()
+    return load_task(args["task_id"], remove=args["rm"])
+
+
 @command("decision", help="Run the decision task")
 @argument("--root", "-r", help="root of the taskgraph definition relative to topsrcdir")
 @argument(

test/test_docker.py

@@ -1,3 +1,5 @@
+import re
+
 import pytest
 
 from taskgraph import docker
@@ -79,3 +81,99 @@ def mock_run(*popenargs, check=False, **kwargs):
 
     out, _ = capsys.readouterr()
     assert f"Successfully built {image}" not in out
+
+
+@pytest.fixture
+def run_load_task(mocker):
+    task_id = "abc"
+
+    def inner(task, remove=False):
+        mocks = {
+            "get_task_definition": mocker.patch.object(
+                docker, "get_task_definition", return_value=task
+            ),
+            "load_image_by_task_id": mocker.patch.object(
+                docker, "load_image_by_task_id", return_value="image/tag"
+            ),
+            "subprocess_run": mocker.patch.object(docker.subprocess, "run"),
+        }
+
+        ret = docker.load_task(task_id, remove=remove)
+        return ret, mocks
+
+    return inner
+
+
+def test_load_task_invalid_task(run_load_task):
+    task = {}
+    assert run_load_task(task)[0] == 1
+
+    task["tags"] = {"worker-implementation": "generic-worker"}
+    assert run_load_task(task)[0] == 1
+
+    task["tags"]["worker-implementation"] = "docker-worker"
+    task["payload"] = {"command": []}
+    assert run_load_task(task)[0] == 1
+
+    task["payload"]["command"] = ["echo", "foo"]
+    assert run_load_task(task)[0] == 1
+
+
+def test_load_task(run_load_task):
+    image_task_id = "def"
+    task = {
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--repo-checkout=/builds/worker/vcs/repo",
+                "--task-cwd=/builds/worker/vcs/repo",
+                "--",
+                "echo foo",
+            ],
+            "image": {"taskId": image_task_id},
+        },
+        "tags": {"worker-implementation": "docker-worker"},
+    }
+    ret, mocks = run_load_task(task)
+    assert ret == 0
+
+    mocks["get_task_definition"].assert_called_once_with("abc")
+    mocks["load_image_by_task_id"].assert_called_once_with(image_task_id, None)
+    mocks["subprocess_run"].assert_called_once_with(
+        [
+            "docker",
+            "run",
+            "-it",
+            "image/tag",
+            "bash",
+            "-c",
+            "/usr/bin/run-task --repo-checkout=/builds/worker/vcs/repo "
+            "--task-cwd=/builds/worker/vcs/repo -- echo 'Task setup complete!' && "
+            "bash",
+        ]
+    )
+
+
+def test_load_task_env_and_remove(run_load_task):
+    image_task_id = "def"
+    task = {
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--repo-checkout=/builds/worker/vcs/repo",
+                "--task-cwd=/builds/worker/vcs/repo",
+                "--",
+                "echo foo",
+            ],
+            "env": {"FOO": "BAR", "BAZ": 1},
+            "image": {"taskId": image_task_id},
+        },
+        "tags": {"worker-implementation": "docker-worker"},
+    }
+    ret, mocks = run_load_task(task, remove=True)
+    assert ret == 0
+
+    mocks["subprocess_run"].assert_called_once()
+    args = mocks["subprocess_run"].call_args[0][0]
+    assert re.match(r"--env-file=/tmp/tmp.*", args[2])
+    assert args[3] == "--rm"