feat: add --develop flag to taskgraph load-task

This flag will setup the current local checkout as a volume to the
task's checkout and avoid all VCS operations.

This allows developers to make changes to the source and immediately
see it in action in the task.

Author: ahal

docs/howto/load-task-locally.rst

@@ -105,4 +105,32 @@ by combining passing in a custom image locally, and piping a task definition via
 
    taskgraph morphed -J --tasks test-unit-py | jq -r 'to_entries | first | .value.task' | taskgraph load-task --image python -
 
+Developing in the Container
+---------------------------
+
+The ``taskgraph load-task`` command also accepts a ``--develop`` flag. This mounts your
+local repository as a volume in the task's expected checkout directory, allowing you to
+add print debugging or test fixes directly in a local container.
+
+Run:
+
+.. code-block:: shell
+
+   taskgraph load-task --develop <task-id>
+
+.. warning::
+
+   Tasks can do all sorts of strange things. Running in this mode has the
+   potential to modify your local source checkout in unexpected ways, up to and
+   including data loss. Use with caution and consider using a dedicated Git
+   worktree when running with ``--develop``.
+
+.. warning::
+
+   Only tasks that use the ``run-task`` script to bootstrap their commands
+   are currently supported.
+
+Using ``--develop`` in conjunction with ``--interactive`` can be a powerful way
+to iterate quickly in a task's environment.
+
 .. _docker-worker payload format: https://docs.taskcluster.net/docs/reference/workers/docker-worker/payload

src/taskgraph/docker.py

@@ -36,6 +36,7 @@
     get_task_definition,
     status_task,
 )
+from taskgraph.util.vcs import get_repository
 
 logger = logging.getLogger(__name__)
 RUN_TASK_RE = re.compile(r"run-task(-(git|hg))?$")
@@ -386,6 +387,27 @@ def _index(l: list, s: str) -> Optional[int]:
         pass
 
 
+def _extract_arg(cmd: list[str], arg: str) -> Optional[str]:
+    if index := _index(cmd, arg):
+        return cmd[index + 1]
+
+    for item in cmd:
+        if item.startswith(f"{arg}="):
+            return item.split("=", 1)[1]
+
+
+def _delete_arg(cmd: list[str], arg:str) -> bool:
+    if index := _index(cmd, arg):
+        del cmd[index: index + 2]
+        return True
+
+    for i, item in enumerate(cmd):
+        if item.startswith(f"{arg}="):
+            del cmd[i]
+            return True
+    return False
+
+
 def _resolve_image(image: Union[str, dict[str, str]], graph_config: GraphConfig) -> str:
     image_task_id = None
 
@@ -432,6 +454,7 @@ def load_task(
     custom_image: Optional[str] = None,
     interactive: Optional[bool] = False,
     volumes: Optional[list[tuple[str, str]]] = None,
+    develop: bool = False,
 ) -> int:
     """Load and run a task interactively in a Docker container.
 
@@ -449,6 +472,8 @@ def load_task(
         interactive: If True, execution of the task will be paused and user
           will be dropped into a shell. They can run `exec-task` to resume
           it (default: False).
+        develop: If True, the task will be configured to use the current
+          local checkout at the current revision (default: False).
 
     Returns:
         int: The exit code from the Docker container.
@@ -476,6 +501,10 @@ def load_task(
         logger.error("Only tasks using `run-task` are supported with --interactive!")
         return 1
 
+    if develop and not is_run_task:
+        logger.error("Only tasks using `run-task` are supported with --develop!")
+        return 1
+
     try:
         image = custom_image or image
         image_tag = _resolve_image(image, graph_config)
@@ -484,6 +513,37 @@ def load_task(
         return 1
 
     task_command = task_def["payload"].get("command")  # type: ignore
+    task_env = task_def["payload"].get("env", {})
+
+    if develop:
+        repositories = json.loads(task_env.get("REPOSITORIES", "{}"))
+        if not repositories:
+            logger.error("Can't use --develop with task that doesn't define any "
+                         "$REPOSITORIES!")
+            return 1
+
+        try:
+            repo = get_repository(os.getcwd())
+        except RuntimeError:
+            logger.error("Can't use --develop from outside a source repository!")
+            return 1
+
+        checkout_name = list(repositories.keys())[0]
+        checkout_arg = f"--{checkout_name}-checkout"
+        checkout_dir = _extract_arg(task_command, checkout_arg)
+        if not checkout_dir:
+            logger.error(f"Can't use --develop with task that doesn't use {checkout_arg}")
+            return 1
+        volumes = volumes or []
+        volumes.append((repo.path, checkout_dir))
+
+        # Delete environment and arguments related to this repo so that
+        # `run-task` doesn't attempt to fetch or checkout a new revision.
+        del repositories[checkout_name]
+        task_env["REPOSITORIES"] = json.dumps(repositories)
+        for arg in ("checkout", "sparse-profile", "shallow-clone"):
+            _delete_arg(task_command, f"--{checkout_name}-{arg}")
+
     exec_command = task_cwd = None
     if interactive:
         # Remove the payload section of the task's command. This way run-task will
@@ -503,16 +563,7 @@ def load_task(
             ]
 
         # Parse `--task-cwd` so we know where to execute the task's command later.
-        if index := _index(task_command, "--task-cwd"):
-            task_cwd = task_command[index + 1]
-        else:
-            for arg in task_command:
-                if arg.startswith("--task-cwd="):
-                    task_cwd = arg.split("=", 1)[1]
-                    break
-            else:
-                task_cwd = "$TASK_WORKDIR"
-
+        task_cwd = _extract_arg(task_command, "--task-cwd") or "$TASK_WORKDIR"
         task_command = [
             "bash",
             "-c",
@@ -527,7 +578,7 @@ def load_task(
         "TASKCLUSTER_ROOT_URL": get_root_url(),
     }
     # Add the task's environment variables.
-    env.update(task_def["payload"].get("env", {}))  # type: ignore
+    env.update(task_env)  # type: ignore
 
     # run-task expects the worker to mount a volume for each path defined in
     # TASKCLUSTER_CACHES; delete them to avoid needing to do the same, unless
@@ -545,6 +596,11 @@ def load_task(
         else:
             del env["TASKCLUSTER_CACHES"]
 
+    # run-task expects volumes listed under `TASKCLUSTER_VOLUMES` to be empty.
+    # This can interfere with load-task when using custom volumes.
+    if volumes and "TASKCLUSTER_VOLUMES" in env:
+        del env["TASKCLUSTER_VOLUMES"]
+
     envfile = None
     initfile = None
     isatty = os.isatty(sys.stdin.fileno())

src/taskgraph/main.py

@@ -767,10 +767,18 @@ def image_digest(args):
     action="store_true",
     default=False,
     help="Setup the task but pause execution before executing its command. "
-    "Repositories will be cloned, environment variables will be set and an"
+    "Repositories will be cloned, environment variables will be set and an "
     "executable script named `exec-task` will be provided to resume task "
     "execution. Only supported for `run-task` based tasks.",
 )
+@argument(
+    "--develop",
+    action="store_true",
+    default=False,
+    help="Configure the task to use the local source checkout at the current "
+    "revision instead of cloning and using the revision from CI."
+
+)
 @argument(
     "--keep",
     dest="remove",
@@ -837,6 +845,7 @@ def load_task(args):
         user=args["user"],
         custom_image=args["image"],
         volumes=volumes,
+        develop=args["develop"],
     )
 
 

src/taskgraph/util/vcs.py

@@ -582,7 +582,7 @@ def does_revision_exist_locally(self, revision):
             raise
 
 
-def get_repository(path):
+def get_repository(path: str):
     """Get a repository object for the repository at `path`.
     If `path` is not a known VCS repository, raise an exception.
     """

test/test_docker.py

@@ -8,6 +8,7 @@
 from taskgraph import docker
 from taskgraph.config import GraphConfig
 from taskgraph.transforms.docker_image import IMAGE_BUILDER_IMAGE
+from taskgraph.util.vcs import get_repository
 
 
 @pytest.fixture
@@ -477,6 +478,45 @@ def test_load_task_with_custom_image_registry(mocker, run_load_task, task):
     assert not mocks["build_image"].called
 
 
+def test_load_task_with_develop(mocker, run_load_task, task):
+    repo_name = "foo"
+    repo_path = "/workdir/vcs"
+    repo = get_repository(os.getcwd())
+
+    # No REPOSITORIES env
+    ret, _ = run_load_task(task, develop=True)
+    assert ret == 1
+
+    # No --checkout flag
+    task["payload"]["env"] = {"REPOSITORIES": f'{{"{repo_name}": "{repo_path}"}}'}
+    ret, mocks = run_load_task(task, develop=True)
+    assert ret == 1
+
+    env_file = None
+    task["payload"]["command"].insert(1, f"--{repo_name}-checkout={repo_path}")
+    m = mocker.patch("os.remove")
+    try:
+        ret, mocks = run_load_task(task, develop=True)
+        assert ret == 0
+        cmd = mocks["subprocess_run"].call_args[0][0]
+        cmdstr = " ".join(cmd)
+        assert f"-v {repo.path}:{repo_path}" in cmdstr
+        assert f"--{repo_name}-checkout" not in cmdstr
+
+        env_file = docker._extract_arg(cmd, "--env-file")
+        assert env_file
+        with open(env_file) as fh:
+            contents = fh.read()
+
+        assert "TASKCLUSTER_VOLUMES" not in contents
+        assert "REPOSITORIES" in contents
+        assert "foo" not in contents
+    finally:
+        if env_file:
+            m.reset_mock()
+            os.remove(env_file)
+
+
 @pytest.fixture
 def run_build_image(mocker):
     def inner(image_name, save_image=None, context_file=None, image_task=None):

test/test_main.py

@@ -467,6 +467,7 @@ def test_load_task_command(run_load_task):
         user=None,
         custom_image=None,
         volumes=[],
+        develop=False,
     )
 
     # Test with interactive flag
@@ -481,6 +482,22 @@ def test_load_task_command(run_load_task):
         user=None,
         custom_image=None,
         volumes=[],
+        develop=False,
+    )
+
+    # Test with develop flag
+    result, mocks = run_load_task(["load-task", "--develop", "task-id-456"])
+
+    assert result == 0
+    mocks["docker_load_task"].assert_called_once_with(
+        mocks["graph_config"],
+        "task-id-456",
+        interactive=False,
+        remove=True,
+        user=None,
+        custom_image=None,
+        volumes=[],
+        develop=True,
     )
 
 
@@ -501,6 +518,7 @@ def test_load_task_command_with_volume(run_load_task):
         user=None,
         custom_image=None,
         volumes=[("/host/path", "/builds/worker/checkouts")],
+        develop=False
     )
 
     # Test with no colon
@@ -542,6 +560,7 @@ def test_load_task_command_with_stdin(run_load_task):
         user=None,
         custom_image=None,
         volumes=[],
+        develop=False,
     )
 
 
@@ -560,6 +579,7 @@ def test_load_task_command_with_task_id(run_load_task):
         user=None,
         custom_image=None,
         volumes=[],
+        develop=False,
     )