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

@@ -5,10 +5,13 @@
 
 import json
 import os
+import shlex
 import subprocess
 import tarfile
+import tempfile
 from io import BytesIO
 from textwrap import dedent
+from typing import List, Optional
 
 try:
     import zstandard as zstd
@@ -19,6 +22,7 @@
 from taskgraph.util.taskcluster import (
     get_artifact_url,
     get_session,
+    get_task_definition,
 )
 
 DEPLOY_WARNING = """
@@ -90,7 +94,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 +241,109 @@ def download_and_modify_image():
         raise Exception("No repositories file found!")
 
     return info
+
+
+def _index(l: List, s: str) -> Optional[int]:
+    try:
+        return l.index(s)
+    except ValueError:
+        pass
+
+
+def load_task(task_id, remove=True):
+    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 an interactive
+    # shell and provide the ability to resume the task command.
+    task_command = None
+    if index := _index(command, "--"):
+        task_command = shlex.join(command[index + 1 :])
+        # I attempted to run the interactive bash shell here, but for some
+        # reason when executed through `run-task`, the interactive shell
+        # doesn't work well. There's no shell prompt on newlines and tab
+        # completion doesn't work. That's why it is executed outside of
+        # `run-task` below, and why we need to parse `--task-cwd`.
+        command[index + 1 :] = [
+            "echo",
+            "Task setup complete!\nRun `exec-task` to execute the task's command.",
+        ]
+
+    # Parse `--task-cwd` so we know where to execute the task's command later.
+    if index := _index(command, "--task-cwd"):
+        task_cwd = command[index + 1]
+    else:
+        for arg in command:
+            if arg.startswith("--task-cwd="):
+                task_cwd = arg.split("=", 1)[1]
+                break
+        else:
+            task_cwd = "$TASK_WORKDIR"
+
+    image_task_id = task_def["payload"]["image"]["taskId"]
+    image_tag = load_image_by_task_id(image_task_id)
+
+    env = task_def["payload"].get("env")
+
+    envfile = None
+    initfile = None
+    try:
+        command = [
+            "docker",
+            "run",
+            "-it",
+            image_tag,
+            "bash",
+            "-c",
+            f"{shlex.join(command)} && cd $TASK_WORKDIR && bash",
+        ]
+
+        if remove:
+            command.insert(2, "--rm")
+
+        if env:
+            envfile = tempfile.NamedTemporaryFile("w+", delete=False)
+            envfile.write("\n".join([f"{k}={v}" for k, v in env.items()]))
+            envfile.close()
+
+            command.insert(2, f"--env-file={envfile.name}")
+
+        if task_command:
+            initfile = tempfile.NamedTemporaryFile("w+", delete=False)
+            initfile.write(
+                dedent(
+                    f"""
+            function exec-task() {{
+                echo "Starting task: {task_command}";
+                pushd {task_cwd};
+                {task_command};
+                popd
+            }}
+            """
+                ).lstrip()
+            )
+            initfile.close()
+
+            command[2:2] = ["-v", f"{initfile.name}:/builds/worker/.bashrc"]
+
+        proc = subprocess.run(command)
+    finally:
+        if envfile:
+            os.remove(envfile.name)
+
+        if initfile:
+            os.remove(initfile.name)
+
+    return proc.returncode

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,28 @@ 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="The task id to load into a docker container.")
+@argument(
+    "--keep",
+    dest="remove",
+    action="store_false",
+    default=True,
+    help="Keep 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["remove"])
+
+
 @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,6 @@
+import re
+import tempfile
+
 import pytest
 
 from taskgraph import docker
@@ -79,3 +82,115 @@ 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):
+        proc = mocker.MagicMock()
+        proc.returncode = 0
+
+        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", return_value=proc
+            ),
+        }
+
+        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)
+
+    expected = [
+        "docker",
+        "run",
+        "-v",
+        re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
+        "-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!\n"
+        "Run `exec-task` to execute the task'\"'\"'s command.' && cd $TASK_WORKDIR && bash",
+    ]
+
+    mocks["subprocess_run"].assert_called_once()
+    actual = mocks["subprocess_run"].call_args[0][0]
+
+    assert len(expected) == len(actual)
+    for i, exp in enumerate(expected):
+        if isinstance(exp, re.Pattern):
+            assert exp.match(actual[i])
+        else:
+            assert exp == actual[i]
+
+
+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()
+    actual = mocks["subprocess_run"].call_args[0][0]
+    assert re.match(r"--env-file=/tmp/tmp.*", actual[4])
+    assert actual[5] == "--rm"