diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py index 6206e9784..0bd7d8035 100644 --- a/src/taskgraph/docker.py +++ b/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 diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py index 5060f8fdf..d406e92ff 100644 --- a/src/taskgraph/main.py +++ b/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( diff --git a/test/test_docker.py b/test/test_docker.py index d662c1e7c..db010454d 100644 --- a/test/test_docker.py +++ b/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"